######## XY STMap v1.1.b - python code written by Martin Viel www.martinviel.com ########
import nuke

import re
import random
import array

import time

def log(func):
    def wrapper(node):
        start=time.time()
        func(node)
        end=time.time()
        wrapper.kronos=end-start
    return wrapper


@log
def convert(ST):
    
    all_tracks=re.findall(r'"(.*)"',ST['tracks'].toScript())
    
    #test
    if all_tracks==[]:
        nuke.message('No tracks to convert!')
        return
    if ST['stmap_file'].value()=='':
        nuke.message('1. Choose a STMap file<br>2. Convert')
        return
        
    key_display=int(ST['keyframe_display'].value())
    ST['keyframe_display'].setValue(3)
    
    
    
    #### GET X,Y VALUES FOR EACH KEY FRAME ####
    
    Tracks=[]
    TRS=[]
    
    inputWidth=ST.format().width()
    inputHeight=ST.format().height()
    
    
    for t in range(len(all_tracks)):
        
        TRS.append((ST['tracks'].getValue(t*31+6),ST['tracks'].getValue(t*31+7),ST['tracks'].getValue(t*31+8)))
        
        framesX=set()
        for key in range(ST['tracks'].getNumKeys(t*31+2)):
            framesX.add(ST['tracks'].getKeyTime(key,t*31+2))
            
        framesY=set()
        for key in range(ST['tracks'].getNumKeys(t*31+3)):
            framesY.add(ST['tracks'].getKeyTime(key,t*31+3))
            
        all_frames=framesX|framesY
        all_frames=sorted(all_frames)
        
        K=[]
        trackX=[]
        trackY=[]
        
        for key in all_frames:
            x=ST['tracks'].getValueAt(key,t*31+2)
            y=ST['tracks'].getValueAt(key,t*31+3)
            
            k=0
            if key in framesX:
                k+=1
            if key in framesY:
                k+=2
                
            K.append(k)
            trackX.append(x)
            trackY.append(y)
              
        Tracks.append((all_frames,K,trackX,trackY))
    
    
    
    #### SAMPLE STMAP ####
    
    #run progress bar
    pt=nuke.ProgressTask("Loading STMap")
    
    #create read
    read=nuke.nodes.Read()
    read.hideControlPanel()
    read.setXYpos(ST.xpos()-100,ST.ypos())
    read['autolabel'].setValue("'XY STMap'")
    read['file'].setValue(f"{ST['stmap_file'].value()}")
    
    Width=read.width()
    Height=read.height()
    
    #test stmap
    r=0
    g=0
    for test in range(20):
        w=random.randint(1,Width)
        h=random.randint(1,Height)
        r=r+read.sample('r',w+0.5,h+0.5)
        g=g+read.sample('g',w+0.5,h+0.5)
        
    if r==0 or g==0:
        nuke.delete(read)
        nuke.message('No STMap')
        return
    
    
    
    # SAMPLE
    mapRG=[]
    
    i=1
    for h in range(Height):
        for w in range(Width):
            
            r=read.sample('r',w+0.5,h+0.5)
            g=read.sample('g',w+0.5,h+0.5)
            
            #input r,g / stmap i
            mapRG.append((r,g,i))
            i+=1
            
        #progress bar#
        pt.setProgress(int(h/(Height)*100))
        #cancel#
        if pt.isCancelled():
            nuke.delete(read)
            ST['keyframe_display'].setValue(key_display)
            del pt
            return
            
    nuke.delete(read)
    
    
    
    #### SORTING STMAP ####
    
    #progress bar
    pt=nuke.ProgressTask("Sorting STMap")
    pt.setProgress(100)
    
    mapR=sorted(mapRG)
    arrayR=array.array('f',[map[0] for map in mapR])
    
    mapG=sorted(mapRG, key=lambda y:y[1])
    arrayG=array.array('f',[map[1] for map in mapG])
    
    lenMap=len(mapR)
    
    del mapRG
    
    
    
    #### CONVERT EACH TRACK THROUGH STMAP ####
    
    #run progress bar
    pt=nuke.ProgressTask("Converting values")
    all_steps=sum([len(elem[0]) for elem in Tracks])
    step=0
    
    for all_frames,K,trackX,trackY in Tracks:
        for e,(x,y) in enumerate(zip(trackX,trackY)):
            
            #progress bar
            step+=1
            pt.setProgress(int(step/(all_steps)*100))
            #cancel
            if pt.isCancelled():
                ST['keyframe_display'].setValue(key_display)
                del pt
                return
                
                
                
            # CONVERT x,y into r,g
            r=x/inputWidth
            g=y/inputHeight
            
            if not (arrayR[0] <= r <= arrayR[-1]) or not (arrayG[0] <= g <= arrayG[-1]) :
                K[e]=0
                
            else:
                # GUESS i in mapR and mapG
                iR=int(r*Width*Height)
                iG=int(g*Width*Height)
                
                
                # ADJUST i in mapR and mapG
                def adjust(map,i,v):
                    """
                    Args:
                        map (list): sorted sample stmap.
                        i (int): estimated index in map.
                        v (float): value of the track to compare against (r or g).
                        
                    Returns:
                        int: adjusted i.
                    """
                    if map[i]<=v:
                        for j,m in enumerate(map[i:]):
                            if m>=v:
                                break
                        i+=j
                        
                    elif map[i]>v:
                        for j,m in enumerate(map[i::-1]):
                            if m<=v:
                                break
                        i-=j
                        
                    return i
                    
                iR=adjust(arrayR,iR,r) # iR in mapR
                iG=adjust(arrayG,iG,g) # iG in mapG
                
                
                # DEFINE LIMITS OF EXPANSION
                IR=(iR//Height)*Height
                IG=(iG//Width)*Width
                
                E=5 #Edge thickness(px)
                onEdge=False
                
                if iR < Height*E:
                    minR=0
                    maxR=IR+IR
                    onEdge=True
                elif iR > lenMap-Height*E:
                    minR=IR-(lenMap-IR)
                    maxR=lenMap
                    onEdge=True
                else:
                    minR=0
                    maxR=lenMap
                    
                if iG < Width*E:
                    minG=0
                    maxG=IG+IG
                    onEdge=True
                elif iG > lenMap-Width*E:
                    minG=IG-(lenMap-IG)
                    maxG=lenMap
                    onEdge=True
                else:
                    minG=0
                    maxG=lenMap
                    
                    
                # WIDDEN RANGE OF SEARCH
                closestRG=set()
                
                firstR=(iR//Height)*Height
                lastR=(iR//Height)*Height
                firstG=(iG//Width)*Width
                lastG=(iG//Width)*Width
                
                while len(closestRG)<15:
                    
                    if minR <= firstR-Height:
                        firstR-=Height
                    if lastR+Height <= maxR:
                        lastR+=Height
                                
                    if minG <= firstG-Width:
                        firstG-=Width
                    if lastG+Width <= maxG:
                        lastG+=Width
                    
                    closestRG = set(mapR[firstR:lastR]) & set(mapG[firstG:lastG])
                    
                    if onEdge==True:
                        if (firstR==minR and lastR==maxR) or (firstG==minG and lastG==maxG):
                            break
                            
                            
                if len(closestRG)==0:
                    K[e]=0
                else:
                    # NEW X, NEW Y
                    bestW=[]
                    
                    for R,G,i in closestRG:
                        
                        #Weight: 1/hypothenuse
                        dR=abs(R-r)
                        dG=abs(G-g)*(inputWidth/inputHeight)
                        w=1/(dR**2+dG**2)**0.5
                        
                        bestW.append((w,i))
                        
                    last=None
                    if onEdge==True:
                        bestW.sort(reverse=True)
                        last=E
                        
                    div=sum(w for w,i in bestW[:last])
                    
                    trackX[e]=round(sum(w*(i%Width-0.5) for w,i in bestW[:last])/div,3)
                    trackY[e]=round(sum(w*(i//Width+0.5) for w,i in bestW[:last])/div,3)
                
                
                
    #### CREATE TRACKER WITH NEW TRACKS VALUES ####
    
    #new tracker
    TT=nuke.nodes.Tracker4()
    TT.setName(ST.name()+'_STMap')
    panel_onCreate(TT)
    TT.showControlPanel()
    TT.setXYpos(ST.xpos()+100,ST.ypos())
    TT.setInput(0,ST.input(0))
    TT.setTab(0)
    
    #create and rename tracks
    for t in range(len(all_tracks)):
        TT['add_track'].execute()
        
    all_names=iter(all_tracks)
    
    def next_name(m):
        return f'"{next(all_names)}"'
        
    new_toScript=re.sub(r'"(.*)"',next_name,TT['tracks'].toScript())
    TT['tracks'].fromScript(new_toScript)
    
    
    #replace old coordinates with new coordinates
    for t,(all_frames,K,trackX,trackY) in enumerate(Tracks):
        
        TT['tracks'].setValue(TRS[t][0],t*31+6)
        TT['tracks'].setValue(TRS[t][1],t*31+7)
        TT['tracks'].setValue(TRS[t][2],t*31+8)
        
        for f,k,x,y in zip(all_frames,K,trackX,trackY):
            
            if k==1:
                TT['tracks'].setValueAt(x,f,t*31+2)
            elif k==2:
                TT['tracks'].setValueAt(y,f,t*31+3)
            elif k==3:
                TT['tracks'].setValueAt(x,f,t*31+2)
                TT['tracks'].setValueAt(y,f,t*31+3)
                
    if nuke.frame() not in TT['tracks'].getKeyList():
        for t in range(len(all_tracks)):
            for c in range(31):
                TT['tracks'].removeKey(t*31+c)
        
    ST['keyframe_display'].setValue(key_display)
    
    del pt
    
    
    
def panel_onCreate(XY):
    
    if 'XY_STMaptab' not in XY.knobs():
        #tab
        XY.addKnob(nuke.Tab_Knob('XY_STMaptab','XY STMap'))
        
        #stmap
        XY.addKnob(nuke.Text_Knob('howto_stmap','','1. Choose STMap file. Must be RGB with values from 0 to 1.'))
        XY.addKnob(nuke.File_Knob('stmap_file',''))
        
        #convert
        XY.addKnob(nuke.Text_Knob('howto_convert','','2. Convert all <b>tracks</b> through STMap. Input format:'))
        XY.addKnob(nuke.PyScript_Knob('convert','Convert and Create new Tracker','xy_stmap.convert(nuke.thisNode())'))
        XY['convert'].setFlag(nuke.STARTLINE) #STARTLINE
        
        #witness
        XY.addKnob(nuke.Int_Knob('witness',''))
        XY['witness'].setExpression("0[python nuke.thisNode().knob('howto_convert').setValue('2. Convert all <b>tracks</b> through STMap. Input format: '+str(nuke.thisNode().format().width())+'x'+str(nuke.thisNode().format().height()))]")
        XY['witness'].setVisible(False) #HIDE
        
        XY.addKnob(nuke.Text_Knob('about','',"<p style='color:gray'><br><b>XY STMap v1.1.b</b> - More on <a href='https://www.martinviel.com/python' style='color:gray'>martinviel.com/python</a>"))
        
    XY['witness'].setFlag(0x20000000) #NO CURVE EDITOR



######## XY STMap v1.1.b - python code written by Martin Viel www.martinviel.com ########