##| Copyright: (C) 2019-2020 Kevin Larke <contact AT larke DOT org> 
##| License: GNU GPL version 3.0 or above. See the accompanying LICENSE file.
import os,sys,pickle,json
import numpy as np

import matplotlib.pyplot as plt
import matplotlib._color_data as mcd
from matplotlib.pyplot import figure

from rms_analysis import rms_analysis_main 
from rms_analysis import note_stats
from rms_analysis import key_info_dictionary
from rms_analysis import select_first_stable_note_by_delta_db
from rms_analysis import select_first_stable_note_by_dur


def do_plot(r, statsL ):

    fig,ax = plt.subplots()

    x = [ i / r.rms_srate for i in range(len(r.tdRmsDbV)) ]
    ax.plot( x,r.tdRmsDbV, color="blue" )
    
    x = [ i / r.rms_srate for i in range(len(r.rmsDbV)) ]
    ax.plot( x,r.rmsDbV, color="green")

    ymx = np.max(r.tdRmsDbV)
    ymn = np.min(r.tdRmsDbV)
    
    for r in statsL:
        x = r.pkSmpSec
        ax.axvline(x,ymax=ymx,ymin=ymn)
        ax.text(x,r.pkDb+1,"%i ms" % r.durMs)
        ax.text(x,r.pkDb+2,"%4.1f dB" % r.pkDb)
        ax.text(x,r.pkDb+3,"%i us" % r.pulse_us)

        if hasattr(r,"MIN"):
            ax.plot(x,r.pkDb,marker="*",color="red")
        
    plt.show()


def select_min_note( rr, statsL, minDurMs=600, minDb=8, contextSecs=10 ):

    sel_note_r = None
    
    for r in statsL:

        if r.pkDb > minDb and r.durMs > minDurMs:
            sel_note_r = r
            break


    #print(rr.tdRmsDbV.shape,rr.rmsDbV.shape)

    sL = []
    
    if sel_note_r is None:
        print("ERROR: No min note found.")
    else:
        #print(sel_note_r)
    
        bi = max(0,                   int(round(sel_note_r.pkSmpSec * rr.rms_srate - contextSecs*rr.rms_srate)))
        ei = min(rr.tdRmsDbV.shape[0],int(round(sel_note_r.pkSmpSec * rr.rms_srate + contextSecs*rr.rms_srate)))

        
        
        rr.tdRmsDbV = rr.tdRmsDbV[bi:ei]    
        rr.rmsDbV   = rr.rmsDbV[bi:ei]    
        offsSec     = bi / rr.rms_srate
        
        for r in statsL:
            begSmpIdx = int(round(r.begSmpSec * rr.rms_srate))
            endSmpIdx = int(round(r.endSmpSec * rr.rms_srate))

            if begSmpIdx > bi and endSmpIdx < ei:
                r0 = r
                r0.begSmpSec = r.begSmpSec - offsSec
                r0.endSmpSec = r.endSmpSec - offsSec
                r0.pkSmpSec  = r.pkSmpSec  - offsSec

                if r.begSmpSec == sel_note_r.begSmpSec:
                    setattr(r0,"MIN",True)
                sL.append(r0)


    return rr,sL
        

def plot_note_audio_reanalysis( inDir ):

    rmsWndMs=300
    rmsHopMs=30
    dbLinRef=0.001
    harmCandN=5
    harmN=3
    durDecayPct = 50
    
    path = os.path.normpath(inDir)
    
    pathL = inDir.split(os.sep)
    take_id    = int(pathL[-1])
    midi_pitch = int(pathL[-2])

    r = rms_analysis_main( inDir, midi_pitch, rmsWndMs=rmsWndMs, rmsHopMs=rmsHopMs, dbLinRef=dbLinRef, harmCandN=harmCandN, harmN=harmN, durDecayPct=durDecayPct )

    r,statsL = select_min_note(r,r.statsL)
    
    do_plot(r,statsL)
    

def plot_note_audio_reanalysis_dir( inDir, dirL ):


    for folder in dirL:

        path = os.path.join(inDir,str(folder),"0")

        plot_note_audio_reanalysis( path )
        
        

def get_all_note_durations( inDir, cacheFn ):

    folderL = os.listdir( inDir )

    yL = []
    for folder in folderL:

        takeId = 0
        rmsWndMs=300
        rmsHopMs=30
        dbLinRef=0.001
        harmCandN=5
        harmN=3
        durDecayPct = 40
    
        path = os.path.normpath(inDir)
    
        midi_pitch = int( folder )
        takePath = os.path.join(inDir,folder,str(takeId))

        if os.path.isfile(os.path.join(takePath,'seq.json')):
            print(midi_pitch)
            r = rms_analysis_main( takePath, midi_pitch, rmsWndMs=rmsWndMs, rmsHopMs=rmsHopMs, dbLinRef=dbLinRef, harmCandN=harmCandN, harmN=harmN, durDecayPct=durDecayPct )

            xL = []
            for i,sr in enumerate(r.statsL):
                xL.append((r.pkUsL[i],sr.durMs,sr.pkDb,sr.quality))

            yL.append((midi_pitch,xL))

    with open(cacheFn,"wb") as f:
        pickle.dump(yL,f)


    
def plot_all_note_durations( cacheFn, pitchL=None, axN=12, yamlCfgFn=None, minDurMs=800, maxPulseUs=None ):

    keyMapD = None
    if yamlCfgFn is not None:
        keyMapD = key_info_dictionary( keyMapL=None, yamlCfgFn=yamlCfgFn)
        
    fig,axL = plt.subplots(axN,1)
    fig.set_size_inches(18.5, 10.5*axN)
    
    #cL =  list(mcd.CSS4_COLORS.values())
    cL =  ['black','brown','orangered','saddlebrown','peru','olivedrab','lightgreen','springgreen','cadetblue','slategray','royalblue','navy','darkviolet','deeppink','crimson']

    
    yD=[]
    with open(cacheFn,"rb") as f:
        yD = dict(pickle.load(f))
        

    cn = 3 #min(1,len(cL)//len(yD))

    ci = 0
    i = 0
    for midi_pitch in pitchL:

        if (pitchL is not None and midi_pitch not in pitchL) or midi_pitch not in yD:
            continue

        
        xL = yD[midi_pitch]

        pkUsL,durMsL,pkDbL,qualityL = tuple(zip(*xL))

        if maxPulseUs is not None:
            pkUsL = np.array(pkUsL)
            pkUsL    = pkUsL[ pkUsL < maxPulseUs ]
            durMsL   = durMsL[0:len(pkUsL)]
            pkDbL    = pkDbL[0:len(pkUsL)]
            qualityL = qualityL[0:len(pkUsL)]

        axi = i//(len(pitchL)//axN)

        if keyMapD is None:
            legendLabel = str(midi_pitch)
        else:
            legendLabel = getattr(keyMapD[midi_pitch],'type') + " " + getattr(keyMapD[midi_pitch],'class') + str(midi_pitch)
            
        axL[axi].plot(pkUsL,durMsL,color=cL[ci],label=legendLabel)

        # plot the quietest stable note
        if minDurMs is not None:
            sni = select_first_stable_note_by_dur( durMsL, minDurMs )

            if sni is not None:
                axL[axi].plot(pkUsL[sni],durMsL[sni],marker=".",color='red')
                axL[axi].text(pkUsL[sni],durMsL[sni] + 50,"%4.1f" % pkDbL[sni])

        sni = select_first_stable_note_by_delta_db( pkDbL, pkUsL )
        
        if sni is not None:
            axL[axi].plot(pkUsL[sni],durMsL[sni],marker=".",color='blue')
            axL[axi].text(pkUsL[sni],durMsL[sni] + 50,"%4.1f" % pkDbL[sni])
                
        ci += cn
        if ci >= len(cL):
            ci = ci - len(cL)

        axL[axi].legend()
        i+=1

        
    plt.show()

def plot_quiet_note_db( cacheFn, yamlCfgFn, minDurMs=700 ):
    
    keyMapD = key_info_dictionary( keyMapL=None, yamlCfgFn=yamlCfgFn)
    
    yD=[]
    with open(cacheFn,"rb") as f:
        yD = dict(pickle.load(f))

    dbL = []
        
    for midi_pitch in range(24,108):

        pk0Db = 0
        pk1Db = 0
        minDb = 0
        if midi_pitch in yD:

            xL = yD[midi_pitch]

            pkUsL,durMsL,pkDbL,qualityL = tuple(zip(*xL))

            # plot the quietest stable note
            sni = select_first_stable_note_by_dur( durMsL, minDurMs )
            if sni is not None:
                pk0Db = pkDbL[sni]

            sni = select_first_stable_note_by_delta_db( pkDbL, pkUsL )
            if sni is not None:
                pk1Db = pkDbL[sni]
            
            minDb = min(pk0Db,pk1Db)
            
        dbL.append( (midi_pitch, minDb, pk0Db,pk1Db) )


    fig,ax = plt.subplots()

    pitchL,minDbL,pk0DbL,pk1DbL = tuple(zip(*dbL))

        
    
    ax.plot( pitchL, pk0DbL, label="dur" )
    ax.plot( pitchL, pk1DbL, label="delta" )
    #ax.plot( pitchL, minDbL, label="min" )

    for i,pitch in enumerate(pitchL):
        ax.text( pitch, pk0DbL[i]+1, "%i %s" % (pitch,getattr(keyMapD[pitch],'type')))

    ax.axhline( np.mean(minDbL), label="mean", color="blue" )
    ax.axhline( np.median(minDbL), label="median", color="green" )
    
    ax.legend()
    plt.show()
    
    
def dump_hold_duty_pct( inDir ):

    pitchL = []
    folderL = os.listdir(inDir)

    for folder in folderL:

        midi_pitch = int(folder)
        
        fn = os.path.join( inDir,folder,"0","seq.json")
        
        if not os.path.isfile(fn):
            print("No sequence file:%s" % (fn))
        else:
            with open(fn,"r") as f:
                d = json.load(f)

                if 'holdDutyPct' in d:
                    holdDutyPctL = [ [0,d['holdDutyPct']] ]
                else:
                    holdDutyPctL = d['holdDutyPctL']

                pitchL.append( {'pitch':midi_pitch, 'holdDutyPctL':holdDutyPctL} )
                #print(midi_pitch, holdDutyPctL)
        
    pitchL = sorted(pitchL, key=lambda x:  x['pitch'])

    for d in pitchL:
        print("{",d['pitch'],":",d['holdDutyPctL'],"},")
                    
                
if __name__ == "__main__":

    #inDir = sys.argv[1]
    
    #plot_note_analysis( inDir )

    pitchL = [ 30,31,32,33,34,35 ]
    pitchL = [ 70,71,72,73,74,75 ]
    
    plot_note_audio_reanalysis_dir( "/home/kevin/temp/p_ac_3c",pitchL)

    durFn = "/home/kevin/temp/cache_note_dur.pickle"
    #get_all_note_durations("/home/kevin/temp/p_ac_3c",durFn)
    #plot_all_note_durations(durFn, np.arange(45,55),2,"p_ac.yml",800,20000)

    #plot_quiet_note_db(durFn,"p_ac.yml")

    #dump_hold_duty_pct( "/home/kevin/temp/p_ac_3c" )