import numpy as np
import os
from researchUtils import Constants

class DMC:
    def __init__(self,
                 simName = "DMC_Sim",
                 outputFolder = "DMCResults/",
                 weighting = 'discrete',
                 initialWalkers = 1000,
                 nTimeSteps = 10000,
                 equilTime = 2000,
                 wfnSpacing = 1000,
                 DwSteps = 50,
                 atoms = [],
                 dimensions = 3,
                 deltaT = 5,
                 D = 0.5,
                 potential=None,
                 masses = None,
                 startStructure = None,
                 branch_every = 1
                 ):
        """
        :param simName:Simulation name for saving wavefunctions
        :type simName:str
        :param outputFolder:The folder where the results will be stored, including wavefunctions and energies
        :type outputFolder:str
        :param weighting:Discrete or Continuous weighting DMC.  Continuous means that there are fixed number of walkers
        :type weighting:str
        :param initialWalkers:Number of walkers we will start the simulation with
        :type initialWalkers:int
        :param nTimeSteps:Total time steps we will be propagating the walkers.  nTimeSteps*deltaT = total time in A.U.
        :type nTimeSteps:int
        :param equilTime: Time before we start collecting wavefunctions
        :type equilTime:int
        :param wfnSpacing:How many time steps in between we will propagate before collecting another wavefunction
        :type wfnSpacing:int
        :param DwSteps:Number of time steps for descendant weighting.
        :type DwSteps: int
        :param atoms:List of atoms for the simulation
        :type atoms:list
        :param dimensions: 3 leads to a 3N dimensional simulation. This should always be 3 for real systems.
        :type dimensions:int
        :param deltaT: The length of the time step; how many atomic units of time are you going in one time step.
        :type deltaT: int
        :param D: Diffusion Coefficient.  Usually set at 0.5
        :type D:float
        :param potential: Takes in coordinates, gives back energies
        :type potential: function
        :param masses:For feeding in artificial masses in atomic units.  If not, then the atoms param will designate masses
        :type masses: list
        :param startStructure:An initial structure to initialize all your walkers
        :type startStructure:np.ndarray
        """
        self.atoms=atoms
        self.simName = simName
        self.outputFolder = outputFolder
        self.initialWalkers = initialWalkers
        self.nTimeSteps = nTimeSteps
        self.potential = potential
        self.weighting = weighting
        self.DwSteps=DwSteps
        self.branch_every = branch_every
        self.branch_step = np.arange(0,nTimeSteps,self.branch_every)
        self.WfnSaveStep = np.arange(equilTime,nTimeSteps,wfnSpacing)
        self.DwSaveStep = self.WfnSaveStep+self.DwSteps
        self.whoFrom = None #Not descendant weighting yet
        self.walkerV = np.zeros(self.initialWalkers)
        self.vrefAr = np.zeros(self.nTimeSteps)
        self.popAr = np.zeros(self.nTimeSteps)
        self.deltaT = deltaT
        self.alpha = 1.0 / (2.0 * deltaT)  # simulation parameter - adjustable
        if startStructure is None:
            self.walkerC = np.zeros((self.initialWalkers,len(atoms),dimensions))
        else:
            self.walkerC = np.repeat(np.expand_dims(startStructure, axis=0), self.initialWalkers, axis=0)
        if masses is None:
            masses = np.array([ Constants.mass(a) for a in self.atoms ])
        self.sigmas = np.sqrt((2 * D * deltaT) / masses)
        if not os.path.isdir(self.outputFolder):
            os.makedirs(self.outputFolder)
        if self.weighting == 'continuous':
            self.contWts = np.ones(self.initialWalkers)
        else:
            self.contWts = None

    def birthOrDeath_vec(self,vref, Desc):
        """
        Chooses whether or not the walker made a bad enough random walk to be removed from the simulation.
        For discrete weighting, this leads to removal or duplication of the walkers.  For continuous, this leads
         to an update of the weights and a potential branching of a large weight walker to the smallest one
         """
        if self.weighting == 'discrete':
            randNums = np.random.random(len(self.walkerC))
            deathMask = np.logical_or((1 - np.exp(-1. * (self.walkerV - vref) * self.deltaT)) < randNums, self.walkerV < vref)
            self.walkerC = self.walkerC[deathMask]
            self.walkerV = self.walkerV[deathMask]
            randNums = randNums[deathMask]
            if Desc:
                self.whoFrom = self.whoFrom[deathMask]

            birthMask = np.logical_and((np.exp(-1. * (self.walkerV - vref) * self.deltaT) - 1) > randNums, self.walkerV < vref)
            self.walkerC = np.concatenate((self.walkerC, self.walkerC[birthMask]))
            self.walkerV = np.concatenate((self.walkerV, self.walkerV[birthMask]))
            if Desc:
                self.whoFrom = np.concatenate((self.whoFrom, self.whoFrom[birthMask]))
            return self.whoFrom,self.walkerC,self.walkerV
        else:
            self.contWts = self.contWts*np.exp(-1.0*(self.walkerV - vref) * self.deltaT)
            thresh = 1.0/self.initialWalkers
            killMark = np.where(self.contWts < thresh)[0]
            for walker in killMark:
                maxWalker = np.argmax(self.contWts)
                self.walkerC[walker] = np.copy(self.walkerC[maxWalker])
                self.walkerV[walker] = np.copy(self.walkerV[maxWalker])
                if Desc:
                    self.whoFrom[walker]=self.whoFrom[maxWalker]
                self.contWts[maxWalker] /= 2.0
                self.contWts[walker] = np.copy(self.contWts[maxWalker])
            return self.contWts,self.whoFrom,self.walkerC,self.walkerV

    def moveRandomly(self,walkerC):
        disps = np.random.normal(0.0, self.sigmas, size=np.shape(walkerC.transpose(0,2,1))).transpose(0,2,1)
        return walkerC + disps

    def getVref(self):  # Use potential of all walkers to calculate vref
        """
             Use the energy of all walkers to calculate vref with a correction for the fluctuation in the population
             or weight.
         """
        if self.weighting == 'discrete':
            Vbar = np.average(self.walkerV)
            correction = (len(self.walkerV) - self.initialWalkers) / self.initialWalkers
        else:
            Vbar = np.average(self.walkerV,weights=self.contWts)
            correction = (np.sum(self.contWts - np.ones(self.initialWalkers))) / self.initialWalkers
        vref = Vbar - (self.alpha * correction)
        return vref

    def propagate(self):
        """
             The main DMC loop.
             1. Move Randomly
             2. Calculate the Potential Energy
             3. Birth/Death
             4. Update Vref
             Additionally, checks when the wavefunction has hit a point where it should save / start descendent
             weighting.
         """
        DW=False
        for prop in range(self.nTimeSteps):
            if prop % 100 == 0:
                print(f'propagation step {prop}')
                if self.weighting == 'discrete':
                    print(f'num walkers : {len(self.walkerC)}')
            self.walkerC = self.moveRandomly(self.walkerC)
            self.walkerV = self.potential(self.walkerC)
            if prop == 0:
                Vref = self.getVref()
            if prop in self.WfnSaveStep:
                dwts = np.zeros(len(self.walkerC))
                parent = np.copy(self.walkerC)
                self.whoFrom = np.arange(len(self.walkerC))
                DW = True
            if prop in self.DwSaveStep:
                DW = False
                if self.weighting == 'discrete':
                    unique, counts = np.unique(self.whoFrom, return_counts=True)
                    dwts[unique]=counts
                else:
                    for q in range(len(self.contWts)):
                        dwts[q] = np.sum(self.contWts[self.whoFrom == q])
                np.savez(self.outputFolder+"/"+self.simName+"_wfn_"+str(prop-self.DwSteps)+"ts",
                         coords=parent,
                         weights=dwts,
                         nDw = self.DwSteps,
                         atms = self.atoms,
                         vref=self.vrefAr
                         )
            if prop in self.branch_step:
                print(f"branching at step {prop}")
                if self.weighting=='discrete':
                    self.whoFrom, self.walkerC, self.walkerV = self.birthOrDeath_vec(Vref,DW)
                else:
                    self.contWts,self.whoFrom, self.walkerC, self.walkerV = self.birthOrDeath_vec(Vref, DW)
            else:
                if self.weighting=='continuous':
                    self.contWts = self.contWts*np.exp(-1.0*(self.walkerV - Vref) * self.deltaT)

            Vref = self.getVref()
            self.vrefAr[prop] = Vref
            self.popAr[prop] = len(self.walkerC)
    def run(self):
        self.propagate()
        np.save(self.outputFolder+"/"+self.simName+"_energies.npy",Constants.convert(self.vrefAr,"wavenumbers",to_AU=False))
        if self.weighting == 'discrete':
            np.save(self.outputFolder + "/" + self.simName + "_population" + ".npy",self.popAr)

if __name__ == "__main__":
    def PatrickShingle(cds):
        import subprocess as sub
        np.savetxt("PES/PES0/hoh_coord.dat", cds.reshape(cds.shape[0]*cds.shape[1],cds.shape[-1]), header=str(len(cds)), comments="")
        sub.run('./calc_h2o_pot', cwd='PES/PES0')
        return np.loadtxt('PES/PES0/hoh_pot.dat')

    # def protCluster(cds):
    #     atms = ['H','H','H','O','H','H','O','H','H','O']
    #     import subprocess as sub
    #     import multiprocessing as mp
    #     splt = np.array_split(cds,mp.cpu_count())
    #     for k in range(mp.cpu_count()):
    #         tmm2 = time.time()
    #         fllK = open('big'+str(k)+'coord.dat','w+')
    #         fllK.write("10\n")
    #         fllK.write("%d\n" % len(splt[k]))
    #         for walk in range(len(splt[k])):
    #             for atm in range(len(atms)):
    #                 fllK.write("%0.18f %0.18f %0.18f %s\n" % (splt[k][walk,atm,0],splt[k][walk,atm,1],splt[k][walk,atm,2],atms[atm]))
    #         fllK.close()
    #     print(f'THAT took {time.time() - tmm2} seconds.')
    #     sub.call('runPots.sh')
    #     vprime = np.loadtxt("big1/eng_dip.dat")[:,0]
    #     for k in range(2,mp.cpu_count+1):
    #         v = np.concatenate(vprime,np.loadtxt("big"+str(k)+"/eng_dip.dat")[:,0])
    #     return v
    def HODMC(cds):
        omega = Constants.convert(3000.,'wavenumbers',to_AU=True)
        mass = Constants.mass('H',to_AU=True)
        return np.squeeze(0.5*mass*omega**2*cds**2)

    dmc_HO = DMC(simName = "DMC_con_test",
                   outputFolder="~/HODMC/",
                   weighting='discrete',
                   initialWalkers=10000,
                   nTimeSteps=10000+1,
                   equilTime=1000,
                   wfnSpacing=5000,
                   DwSteps=50,
                   atoms=['H'],
                   dimensions=1,
                   deltaT=5,
                   D=0.5,
                   potential=HODMC,
                   masses=None,
                   startStructure = Constants.convert(
                       np.array([[0.00000]]),"angstroms",to_AU=True))
    dmc_HO.run()
    # dmcTrimer = DMC(simName = "DMC_con_test",
    #                outputFolder="DMCResults/",
    #                weighting='discrete',
    #                initialWalkers=1000,
    #                nTimeSteps=1000+1,
    #                equilTime=500,
    #                wfnSpacing=100,
    #                DwSteps=50,
    #                atoms=['H','H','H','O','H','H','O','H','H','O'],
    #                dimensions=3,
    #                deltaT=5,
    #                D=0.5,
    #                potential=protCluster,
    #                masses=None,
    #                startStructure = Constants.convert(
    #                    np.array([[0.00000, 0.91527, -0.05817],
    #                        [0.00000, 1.67720, 0.53729],
    #                        [-0.87302, 0.35992, 0.01707],
    #                        [2.56267, -0.75858, 0.76451],
    #                        [2.70113, -0.40578, -0.73813],
    #                        [2.07091, -0.46191, -0.00993],
    #                        [0.87302, 0.35993, 0.01707],
    #                        [-2.70115, -0.40575, -0.73811],
    #                        [-2.56265, -0.75862, 0.76451],
    #                        [-2.07092, -0.46190, -0.00993]]
    #                             )* 1.01
    #                    ,"angstroms",to_AU=True))
    # dmcTrimer.run()


import numpy as np
from fractions import Fraction
from collections import namedtuple

# | prefer name | other names | intervals |
DIATONIC_SCALES = [
    ['ionian mode', ['mjaor scale'], [0, 2, 4, 5, 7, 9, 11]],
    ['dorian mode', [], [0, 2, 3, 5, 7, 9, 10]],
    ['phrygian mode', [], [0, 1, 3, 5, 7, 8, 10]],
    ['lydian mode', [], [0, 2, 4, 6, 7, 9, 11]],
    ['mixolydian scale', [], [0, 2, 4, 5, 7, 9, 10]],
    ['aeolian mode', ['nature minor mode'], [0, 2, 3, 5, 7, 8, 10]],
    ['locrian mode', [], [0, 1, 3, 5, 6, 8, 10]]
]

# | prefer name | other names | intervals | index to diatonic scale | operations |
MELODIC_MINOR_SCALES = [
    ['ascending melodic minor', [], [0, 2, 3, 5, 7, 9, 11], DIATONIC_SCALES[0], ['3-']],
    ['phrygidorian', ['phrygian ♮6', 'dorian ♭2', 'assyrian'], [0, 1, 3, 5, 7, 9, 10],  DIATONIC_SCALES[1], ['2-']],
    ['lydian augmented', ['lydian ♯5'], [0, 2, 4, 6, 8, 9, 11],  DIATONIC_SCALES[3], ['5+']],
    ['lydian dominant', ['lydian ♭7', 'acoustic scale', 'mixolydian ♯4', 'overtone', 'lydomyxian'], [0, 2, 4, 6, 7, 9, 10],  DIATONIC_SCALES[3], ['7-']],
    ['melodic major', ['mixolydian ♭6', 'fifth mode of melodic minor', 'hindu', 'myxaeolian'], [0, 2, 4, 5, 7, 8, 10],  DIATONIC_SCALES[4], ['6-']],
    ['aeolocrian', ['locrian ♮2', 'half-diminished'], [0, 2, 3, 5, 6, 8, 10],  DIATONIC_SCALES[6], ['2+']],
    ['altered scale', ['super locrian', 'altered dominant scale'], [0, 1, 3, 4, 6, 8, 10],  DIATONIC_SCALES[6], ['4-']]
]

BEBOP_SCALES = [
    ['bebop dominant scale', [], [0, 2, 4, 5, 7, 9, 10, 11], DIATONIC_SCALES[4], ['add7+']],
    ['bebop major scale', [], [0, 2, 4, 5, 7, 8, 9, 11], DIATONIC_SCALES[0], ['add5+']]
]

PENTATONE_SCALES = [
    ['major pentatonic scale', ['gong'], [0, 2, 4, 7, 9], DIATONIC_SCALES[0], ['omit4', 'omit7']],
    ['egyptian', ['suspended', 'shang'], [0, 2, 5, 7, 10], DIATONIC_SCALES[1], ['omit3', 'omit6']],
    ['blues minor', ['man gong', 'jue'], [0, 3, 5, 8, 10], DIATONIC_SCALES[2], ['omit2', 'omit5']],
    ['blues major', ['Ritsusen', 'yo scale', 'zhi'], [0, 2, 5, 7, 9], DIATONIC_SCALES[4], ['omit3', 'omit7']],
    ['minor pentatonic', ['yu'], [0, 3, 5, 7, 10], DIATONIC_SCALES[5], ['omit2', 'omit6']]
]

HEPTATONIC_SCALES = [
    # diatonic scales
    ['ionian mode (major scale)', [0, 2, 4, 5, 7, 9, 11]],
    ['dorian mode', [0, 2, 3, 5, 7, 9, 10]],
    ['phrygian mode', [0, 1, 3, 5, 7, 8, 10]],
    ['lydian mode', [0, 2, 4, 6, 7, 9, 11]],
    ['mixolydian scale', [0, 2, 4, 5, 7, 9, 10]],
    ['aeolian mode (nature minor mode)', [0, 2, 3, 5, 7, 8, 10]],
    ['locrian mode', [0, 1, 3, 5, 6, 8, 10]],
    # TODO
    ['adonai malakh scale', [0, 2, 4, 5, 7, 8, 10]],
    ['algerian scale', [0, 2, 3, 6, 7, 8, 11]],
    ['altered scale', [0, 1, 3, 4, 6, 8, 10]],
    ['double harmonic scale', [0, 1, 4, 5, 7, 8, 11]],
    ['enigmatic scale', [0, 1, 4, 6, 8, 10, 11]],
    ['flamenco mode', [0, 1, 4, 5, 7, 8, 11]],
    ['gypsy scale', [0, 2, 3, 6, 7, 8, 10]],
    ['half diminished scale', [0, 2, 3, 5, 6, 8, 10]],
    ['harmonic major scale', [0, 2, 4, 5, 7, 8, 11]],
    ['harmonic minor scale', [0, 2, 3, 5, 7, 8, 11]],
    ['hungarian gypsy scale', [0, 2, 3, 6, 7, 8, 11]],
    ['hungarian minor scale', [0, 2, 3, 6, 7, 8, 11]],
    ['lydian augmented scale', [0, 2, 4, 6, 8, 9, 11]],
    ['major locran scale', [0, 2, 4, 5, 6, 8, 10]],
    ['melodic minor scale (ascending)', [0, 2, 3, 5, 7, 9, 11]],
    ['neapolitan major scale', [0, 1, 3, 5, 7, 9, 11]],
    ['neapolitan minor scale', [0, 1, 3, 5, 7, 8, 11]],
    ['persian scale', [0, 1, 4, 5, 6, 8, 11]],
    ['phrygian dominant scale', [0, 1, 4, 5, 7, 8, 10]],
    ['ukranian dorian scale', [0, 2, 3, 6, 7, 9, 10]]
]

SCALES_COLLECTION = [
    ['acoustic scale', [0, 2, 4, 6, 7, 9]],
    ['adonai malakh scale', [0, 2, 4, 5, 7, 8, 10]],
    ['aeolian mode', [0, 2, 3, 5, 7, 8, 10]],
    ['nature minor mode', [0, 2, 3, 5, 7, 8, 10]],
    ['algerian scale', [0, 2, 3, 6, 7, 8, 11]],
    ['altered scale', [0, 1, 3, 4, 6, 8, 10]],
    ['augmented scale', [0, 3, 4, 7, 8, 11]],
    ['bebop dominant scale', [0, 2, 4, 5, 7, 9, 10, 11]],
    ['blues scale', [0, 3, 5, 6, 7, 10]],
    ['chromatic scale', [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]],
    ['dorian mode', [0, 2, 3, 5, 7, 9, 10]],
    ['double harmonic scale', [0, 1, 4, 5, 7, 8, 11]],
    ['enigmatic scale', [0, 1, 4, 6, 8, 10, 11]],
    ['flamenco mode', [0, 1, 4, 5, 7, 8, 11]],
    ['gypsy scale', [0, 2, 3, 6, 7, 8, 10]],
    ['half diminished scale', [0, 2, 3, 5, 6, 8, 10]],
    ['harmonic major scale', [0, 2, 4, 5, 7, 8, 11]],
    ['harmonic minor scale', [0, 2, 3, 5, 7, 8, 11]],
    ['hirajoshi scale', [0, 2, 3, 7, 8]],
    ['hungarian gypsy scale', [0, 2, 3, 6, 7, 8, 11]],
    ['hungarian minor scale', [0, 2, 3, 6, 7, 8, 11]],
    ['insen scale', [0, 1, 5, 7, 10]],
    ['ionian mode', [0, 2, 4, 5, 7, 9, 11]],
    ['major scale', [0, 2, 4, 5, 7, 9, 11]],
    ['istrian scale', [0, 1, 3, 4, 6, 7]],
    ['iwato scale', [0, 1, 5, 6, 10]],
    ['locrian mode', [0, 1, 3, 5, 6, 8, 10]],
    ['lydian augmented scale', [0, 2, 4, 6, 8, 9, 11]],
    ['lydian mode', [0, 2, 4, 6, 7, 9, 11]],
    ['major bepop scale', [0, 2, 4, 5, 7, 8, 9, 11]],
    ['major locran scale', [0, 2, 4, 5, 6, 8, 10]],
    ['major pentatonic scale', [0, 2, 4, 7, 9]],
    ['melodic minor scale', [0, 2, 3, 5, 7, 8, 9, 10, 11]],
    ['melodic minor scale (ascending)', [0, 2, 3, 5, 7, 9, 11]],
    ['minor pentatonic scale', [0, 3, 5, 7, 10]],
    ['mixolydian scale', [0, 2, 4, 5, 7, 9, 10]],
    ['neapolitan major scale', [0, 1, 3, 5, 7, 9, 11]],
    ['neapolitan minor scale', [0, 1, 3, 5, 7, 8, 11]],
    ['octatonic scale', [0, 2, 3, 5, 6, 8, 9, 11]],
    ['persian scale', [0, 1, 4, 5, 6, 8, 11]],
    ['phrygian dominant scale', [0, 1, 4, 5, 7, 8, 10]],
    ['phrygian mode', [0, 1, 3, 5, 7, 8, 10]],
    ['prometheus scale', [0, 2, 4, 6, 9, 10]],
    ['tritone scale', [0, 1, 4, 6, 7, 10]],
    ['ukranian dorian scale', [0, 2, 3, 6, 7, 9, 10]],
    ['whole tone scale', [0, 2, 4, 6, 8, 10]]
]

# | full name | short names | intervals | scale names |
CHORDS_TABLE = [
    # triad chords
    ['major triad', ['M'], [0, 4, 7], ['ionian']],
    ['minor triad', ['m'], [0, 3, 7], ['aeolian']],
    ['augmented triad', ['aug'], [0, 4, 8], ['lydian augmented']],
    ['diminished triad', ['dim'], [0, 3, 6], ['locrian']],
    # seventh chords
    ['diminished seventh', ['°'], [0, 3, 6, 9], [1, 3, 5, 7]],
    ['half-diminished seventh', ['φ'], [0, 3, 6, 10], ['locrian']],
    ['minor seventh', ['-'], [0, 3, 7, 10], ['dorian']],
    ['minor major seventh', ['-Δ'], [0, 3, 7, 11], ['melodic minor']],
    ['dominant seventh', ['7'], [0, 4, 7, 10], ['mixolydian']],
    ['major seventh', ['Δ'], [0, 4, 7, 11], ['ionian']],
    ['augmented seventh', ['7♯5'], [0, 4, 8, 10], ['melodic major']],
    ['augmented major seventh', ['Δ♯5'], [0, 4, 8, 11], ['lydian augmented']],
    # extented chords
    ['dominant ninth', ['9'], [0, 2, 4, 7, 10], ['mixolydian']],
    ['dominant eleventh', ['11'], [0, 2, 4, 5, 7, 10], ['mixolydian']],
    ['dominant thirteenth', ['13'], [0, 2, 4, 5, 7, 9, 10], ['mixolydian']],
    # atered chords
    ['seventh augmented fifth', ['7♯5'], [0, 4, 8, 10], ['melodic major']],
    ['seventh minor ninth', ['7♭9'], [0, 1, 4, 7, 10], ['phrygian dominant']],
    ['seventh sharp ninth', ['7♯9'], [0, 3, 4, 7, 10], [1, 2, 3, 5, 7]],
    ['seventh augmented eleventh', ['7♯11'], [0, 2, 4, 6, 7, 10]],
    ['seventh diminished thirteenth', ['7♭13'], [0, 2, 4, 5, 7, 8, 10]],
    ['half-diminished seventh', ['φ'], [0, 3, 6, 10], [1, 3, 5, 7]],
    # added tone chords
    ['add nine', ['add9'], [0, 2, 4, 7], [1, 2, 3, 5]],
    ['add fourth', ['add11'], [0, 4, 5, 7], [1, 3, 4, 5]],
    ['add sixth', ['6'], [0, 4, 7, 9], [1, 3, 5, 6]],
    ['six-nine', ['6/9'], [0, 2, 4, 7, 9], [1, 2, 3, 5, 6]],
    ['seven-six', ['7/6'], [0, 4, 7, 9, 10], [1, 3, 5, 6, 7]],
    ['mixed-third', ['--'], [0, 3, 4, 7], [1, 2, 3, 5]],
    # suspended chords
    ['suspended second', ['sus2'], [0, 2, 7], [1, 2, 5]],
    ['suspended fourth', ['sus4'], [0, 5, 7], [1, 4, 5]],
    ['jazz sus', ['9sus4'], [0, 4, 5, 7, 10], [1, 3, 4, 5, 7]]
]

class FreqRatio():
    FREQ_RATIOS = ['1/1', '16/15', '9/8', '6/5', '5/4', '4/3',
                   '45/32', '3/2', '8/5', '5/3', '9/5', '15/8']

    def __init__(self):
        pass

    def chromatic_freq_sizes(self):
        l = []
        for fr in self.FREQ_RATIOS:
            f = Fraction(fr)
            n = f.numerator
            d = f.denominator
            l.append(n * d)
        return l

    def calc_freq_radio_size(self, intervals):
        l = list(intervals)
        l.sort()
        denominators = []
        numerators = []
        for i in l:
            f = Fraction(self.FREQ_RATIOS[i])
            denominators.append(f.denominator)
            numerators.append(f.numerator)
        lcm = np.lcm.reduce(denominators)
        total = []
        for d,n in zip(denominators, numerators):
            total.append(n * (lcm / d))
        prod = np.prod(total)
        return prod

SCALES_QUERY_TABLE = {
    'ionian': ['diatonic', DIATONIC_SCALES[0]],
    'dorian': ['diatonic', DIATONIC_SCALES[1]],
    'phrygian': ['diatonic', DIATONIC_SCALES[2]],
    'lydian': ['diatonic', DIATONIC_SCALES[3]],
    'mixolydian': ['diatonic', DIATONIC_SCALES[4]],
    'aeolian': ['diatonic', DIATONIC_SCALES[5]],
    'locrian': ['diatonic', DIATONIC_SCALES[6]],
    'melodic minor': ['melodic', MELODIC_MINOR_SCALES[0]],
    'lydian dominant': ['melodic', MELODIC_MINOR_SCALES[3]],
    'melodic major': ['melodic', MELODIC_MINOR_SCALES[4]],
    'aeolocrian': ['melodic', MELODIC_MINOR_SCALES[5]],
    'altered scale': ['melodic', MELODIC_MINOR_SCALES[6]],
    'gong': ['pentatone', PENTATONE_SCALES[0]],
    'shang': ['pentatone', PENTATONE_SCALES[1]],
    'jue': ['pentatone', PENTATONE_SCALES[2]],
    'zhi': ['pentatone', PENTATONE_SCALES[3]],
    'yu': ['pentatone', PENTATONE_SCALES[4]],
    'bebop dominant': ['bebop', BEBOP_SCALES[0]],
    'bebop major': ['bebop', BEBOP_SCALES[1]],
}

class Scales:
    diatonic_scales = ['C', 0, 'D', 0, 'E', 'F', 0, 'G', 0, 'A', 0, 'B']
    scales_query_table = SCALES_QUERY_TABLE

    def __init__(self):
        pass

    def rotate(self, n, l, right=False):
        llen = len(l)
        if right:
            return l[-n:] + l[:-n]
        n1 = (llen - n) % llen
        return l[-n1:] + l[:-n1]

    def get_accidentals(self, key):
        alter = 0
        if len(key) > 1:
            ct = key.count('-')
            if ct > 0:
                alter = -ct
            ct = key.count('#')
            if ct > 0:
                alter = ct
        return alter

    def __query(self, mode):
        return self.scales_query_table[mode]

    def flat_list(self, l, output_list):
        for i in l:
            if type(i) == list:
                self.flat_list(i, output_list)
            else:
                output_list.append(i)

    def to_diatonic_scales(self, key, query):
        intervals = query[2]
        alter = self.get_accidentals(key)
        step = key[0]
        scale_count = len(self.diatonic_scales)
        idx = self.diatonic_scales.index(step)
        rotated_scales = self.rotate(idx, self.diatonic_scales)
        index_list = []
        scale_steps = []
        for i in range(scale_count):
            if type(rotated_scales[i]) == str:
                index_list.append(i)
                scale_steps.append(rotated_scales[i])
        a1 = np.array(index_list)
        a2 = np.array(intervals)
        dist = a2 - a1 + alter
        notes = []
        for n,d in zip(scale_steps, dist):
            if d > 0:
                note = "{}{}".format(n, '#' * d)
            elif d < 0:
                note = "{}{}".format(n, '-' * -d)
            else:
                note = n
            notes.append(note)
        return notes

    def filter_digit(self, s):
        num = int(''.join(filter(str.isdigit, s)))
        return num

    def update_accidental(self, note, alter):
        step = note[0]
        curr_alter = self.get_accidentals(note)
        final_alter = curr_alter + alter
        if final_alter > 0:
            nt = "{}{}".format(step, '#' * final_alter)
        elif final_alter < 0:
            nt = "{}{}".format(step, '-' * -final_alter)
        else:
            nt = step
        return nt

    def notes_operate(self, notes, opts):
        for opt in opts:
            if opt.count('omit') > 0:
                idx = self.filter_digit(opt)
                notes[idx - 1] = 0
            elif opt.count('add') > 0:
                idx = self.filter_digit(opt)
                alter = self.get_accidentals(opt)
                n = notes[idx - 1]
                nt = self.update_accidental(n, alter)
                if alter > 0:
                    notes[idx - 1] = [n, nt]
                else:
                    notes[idx - 1] = [nt, n]
            else:
                idx = self.filter_digit(opt)
                alter = self.get_accidentals(opt)
                n = notes[idx - 1]
                nt = self.update_accidental(n, alter)
                notes[idx - 1] = nt
        notes = list(filter(lambda a: a != 0, notes))
        final_notes = []
        self.flat_list(notes, final_notes)
        return final_notes

    def to_pentatone_scales(self, key, query):
        notes = []
        opts = query[4]
        notes = self.to_diatonic_scales(key, query[3])
        notes = self.notes_operate(notes, opts)
        return notes

    def to_bebop_scales(self, key, query):
        notes = []
        opts = query[4]
        notes = self.to_diatonic_scales(key, query[3])
        notes = self.notes_operate(notes, opts)
        return notes

    def to_scales(self, key, mode):
        query = self.__query(mode)
        notes = []
        if query[0] == 'diatonic':
            notes = self.to_diatonic_scales(key, query[1])
        elif query[0] == 'pentatone':
            notes = self.to_pentatone_scales(key, query[1])
        elif query[0] == 'melodic':
            notes = self.to_diatonic_scales(key, query[1])
        elif query[0] == 'bebop':
            notes = self.to_bebop_scales(key, query[1])
        return notes

CHORDS_QUERY_TABLE = {
    'maj': ['triad', CHORDS_TABLE[0]],
    'min': ['triad', CHORDS_TABLE[1]],
    'maj7': ['seventh', CHORDS_TABLE[9]],
    'min7': ['seventh', CHORDS_TABLE[6]],
}

class Chord(Scales):
    chords_query_table = CHORDS_QUERY_TABLE
    ChordsTable = namedtuple('ChordsTable', ['common_name', 'short_names',
                                             'intervals', 'scale_names'])

    def __init__(self):
        pass

    def __query(self, mode):
        return self.chords_query_table[mode]

    def to_triad(self, scales):
        chord = [scales[0], scales[2], scales[4]]
        return chord

    def to_seventh(self, scales):
        chord = [scales[0], scales[2], scales[4], scales[6]]
        return chord

    def to_chord(self, key, mode):
        query = self.__query(mode)
        chords_table = self.ChordsTable._make(query[1])
        scales = self.to_scales(key, chords_table.scale_names[0])
        chord = []
        if query[0] == 'triad':
            chord = self.to_triad(scales)
        elif query[0] == 'seventh':
            chord = self.to_seventh(scales)
        return chord

class RomanNumeral(Chord):
    roman_numerals = ['I', 'II', 'III', 'IV', 'V', 'VI', 'VII']

    def __init__(self, key):
        self.__notes = self.to_scales(key, 'diatonic')

    def parser(self, rn):
        offset = rn[0]
        pass

    def to_note(self, rn):
        pass

    def to_chord(self, rn):
        pass

class ChordProgression():
    # | genre | chord progression | expression |
    CHORD_PROGRESSION_TABLE = [
        ["50' progression", ['I', 'vi', 'IV', 'V'], 0],
        ["50' progression", ['I', 'vi', 'ii', 'V'], 1],
        ["50' progression", ['I', 'V', 'vi', 'IV'], 2]
    ]

    def __init__(self):
        pass

class Style():
    STYLES = ['Folk song', 'Hip hop', 'Jazz', 'Latin', 'Pop', 'R&B and soul', 'Rock',
              'Classical music', 'Country']

class Instruments():
    pass

class EvalSystem():
    pass


"""Probabilistic autoregressive model."""

import copy
import logging
import pickle
import uuid

import numpy as np
import pandas as pd
import rdt
import torch
from deepecho.models.base import DeepEcho
from deepecho.sequences import assemble_sequences
from sdv.metadata import Table
from sdv.tabular.copulas import GaussianCopula
from tqdm import tqdm

LOGGER = logging.getLogger(__name__)

def batch_to_df(sequences, columns, entity):
    df = pd.DataFrame()
    for k, v in sequences.items():
        if isinstance(v, torch.Tensor):
            list_ = list(v.numpy())
        else:
            list_ = list(v)
        if len(list_) < len(sequences["t"]):  # removes all the next_* and *_to_now features
            continue
        else:
            df[k] = list_

    df = df[columns + entity]

    return df

class PARNet(torch.nn.Module):
    """PARModel ANN model."""

    def __init__(self, data_size, context_size, hidden_size=32):
        super(PARNet, self).__init__()
        self.context_size = context_size
        self.down = torch.nn.Linear(data_size + context_size, hidden_size)
        self.rnn = torch.nn.GRU(hidden_size, hidden_size)
        self.up = torch.nn.Linear(hidden_size, data_size)

    def forward(self, x, c):
        """Forward passing computation."""
        if isinstance(x, torch.nn.utils.rnn.PackedSequence):
            x, lengths = torch.nn.utils.rnn.pad_packed_sequence(x)
            if self.context_size:
                x = torch.cat(
                    [x, c.unsqueeze(0).expand(x.shape[0], c.shape[0], c.shape[1])],
                    dim=2,
                )

            x = self.down(x)
            x = torch.nn.utils.rnn.pack_padded_sequence(x, lengths, enforce_sorted=False)
            x, _ = self.rnn(x)
            x, lengths = torch.nn.utils.rnn.pad_packed_sequence(x)
            x = self.up(x)
            x = torch.nn.utils.rnn.pack_padded_sequence(x, lengths, enforce_sorted=False)

        else:
            if self.context_size:
                x = torch.cat(
                    [x, c.unsqueeze(0).expand(x.shape[0], c.shape[0], c.shape[1])],
                    dim=2,
                )

            x = self.down(x)
            x, _ = self.rnn(x)
            x = self.up(x)

        return x

class PARModel(DeepEcho):
    def __init__(self, epochs=128, sample_size=1, cuda=True, verbose=True):
        self.epochs = epochs
        self.sample_size = sample_size

        if not cuda or not torch.cuda.is_available():
            device = "cpu"
        elif isinstance(cuda, str):
            device = cuda
        else:
            device = "cuda"

        self.device = torch.device(device)
        self.verbose = verbose

        LOGGER.info("%s instance created", self)
        if verbose:
            print(self, "instance created")

    def __repr__(self):
        return "{}(epochs={}, sample_size={}, cuda='{}', verbose={})".format(
            self.__class__.__name__,
            self.epochs,
            self.sample_size,
            self.device,
            self.verbose,
        )

    def _idx_map(self, x, t):
        idx = 0
        idx_map = {}
        for i, t in enumerate(t):
            if t == "continuous" or t == "datetime":
                try:
                    if i == 14:
                        x[i] = [hash(k) for k in x[i]]
                    idx_map[i] = {
                        "type": t,
                        "mu": np.nanmean(x[i]),
                        "std": np.nanstd(x[i]),
                        "nulls": pd.isnull(x[i]).any(),
                        "indices": (idx, idx + 1, idx + 2),
                    }
                    idx += 3
                except Exception as e:
                    print(e)

            elif t == "count":
                idx_map[i] = {
                    "type": t,
                    "min": np.nanmin(x[i]),
                    "range": np.nanmax(x[i]) - np.nanmin(x[i]),
                    "nulls": pd.isnull(x[i]).any(),
                    "indices": (idx, idx + 1, idx + 2),
                }
                idx += 3

            elif t == "categorical" or t == "ordinal":
                idx_map[i] = {"type": t, "indices": {}}
                idx += 1
                for v in set(x[i]):
                    if pd.isnull(v):
                        v = None

                    idx_map[i]["indices"][str(v)] = idx
                    idx += 1

            else:
                raise ValueError("Unsupported type: {}".format(t))

        return idx_map, idx

    def _build(self, examples, types):

        # TODO: define better suited values
        min_length = 0
        max_length = int(1e3)

        self._fixed_length = min_length == max_length
        self._min_length = min_length
        self._max_length = max_length

        self._ctx_map, self._ctx_dims = self._idx_map(
            examples["context"], types["context"]
        )  # all possible values => per batch!
        self._data_map, self._data_dims = self._idx_map(examples["data"], types["data"])
        self._data_map["<TOKEN>"] = {
            "type": "categorical",
            "indices": {
                "<START>": self._data_dims,
                "<END>": self._data_dims + 1,
                "<BODY>": self._data_dims + 2,
            },
        }
        self._data_dims += 3

    def _data_to_tensor(self, data):
        seq_len = len(data[0]) if data != [] else 0
        X = []

        x = torch.zeros(self._data_dims)
        x[self._data_map["<TOKEN>"]["indices"]["<START>"]] = 1.0
        X.append(x)

        for i in range(seq_len):
            x = torch.zeros(self._data_dims)
            for key, props in self._data_map.items():
                if key == "<TOKEN>":
                    x[self._data_map["<TOKEN>"]["indices"]["<BODY>"]] = 1.0

                elif props["type"] in ["continuous", "timestamp"]:
                    mu_idx, sigma_idx, missing_idx = props["indices"]
                    if data[key][i].isalpha() and not (props["std"] == 0):
                        x[mu_idx] = (
                            (float(data[key][i]) - props["mu"]) / props["std"]
                            if isinstance(data[key][i], float)
                            else 0.0
                        )
                    else:
                        x[mu_idx] = 0.0

                    x[sigma_idx] = 0.0
                    x[missing_idx] = 1.0 if data[key][i] is None else 0.0

                elif props["type"] in ["count"]:
                    r_idx, p_idx, missing_idx = props["indices"]
                    x[r_idx] = (
                        0.0
                        if (data[key][i] is None or props["range"] == 0)
                        else (data[key][i] - props["min"]) / props["range"]
                    )
                    x[p_idx] = 0.0
                    x[missing_idx] = 1.0 if data[key][i] is None else 0.0

                elif props["type"] in ["categorical", "ordinal"]:  # categorical
                    value = data[key][i]
                    if pd.isnull(value):
                        value = None
                    x[props["indices"][value]] = 1.0

                else:
                    raise ValueError()

            X.append(x)

        x = torch.zeros(self._data_dims)
        x[self._data_map["<TOKEN>"]["indices"]["<END>"]] = 1.0
        X.append(x)

        return torch.stack(X, dim=0).to(self.device)

    def _context_to_tensor(self, context):
        if not self._ctx_dims:
            return None

        x = torch.zeros(self._ctx_dims)
        for key, props in self._ctx_map.items():
            if props["type"] in ["continuous", "datetime"]:
                mu_idx, sigma_idx, missing_idx = props["indices"]
                x[mu_idx] = (
                    0.0
                    if (pd.isnull(context[key]) or props["std"] == 0)
                    else (context[key] - props["mu"]) / props["std"]
                )
                x[sigma_idx] = 0.0
                x[missing_idx] = 1.0 if pd.isnull(context[key]) else 0.0

            elif props["type"] in ["count"]:
                r_idx, p_idx, missing_idx = props["indices"]
                x[r_idx] = (
                    0.0
                    if (pd.isnull(context[key]) or props["range"] == 0)
                    else (context[key] - props["min"]) / props["range"]
                )
                x[p_idx] = 0.0
                x[missing_idx] = 1.0 if pd.isnull(context[key]) else 0.0

            elif props["type"] in ["categorical", "ordinal"]:
                value = context[key]
                if pd.isnull(value):
                    value = None
                x[props["indices"][value]] = 1.0

            else:
                raise ValueError()

        return x.to(self.device)

    def _transform_sequence_index(self, sequences, dem):
        sequence_index_idx = dem._data_columns.index(dem._sequence_index)
        for sequence in sequences:
            data = sequence["data"]
            sequence_index = data[sequence_index_idx]
            diffs = np.diff(sequence_index).tolist()
            data[sequence_index_idx] = diffs[0:1] + diffs
            data.append(sequence_index[0:1] * len(sequence_index))

    def fit_sequences(self, sequence_loader, columns, types, examples, dem):
        """Fit a model to the specified sequences.

        Args:
            sequences (list):
                List of sequences. Each sequence is a single training example
                (i.e. an example of a multivariate time series with some context).
                For example, a sequence might look something like::

                    {
                        "context": [1],
                        "data": [
                            [1, 3, 4, 5, 11, 3, 4],
                            [2, 2, 3, 4,  5, 1, 2],
                            [1, 3, 4, 5,  2, 3, 1]
                        ]
                    }

                The "context" attribute maps to a list of variables which
                should be used for conditioning. These are variables which
                do not change over time.

                The "data" attribute contains a list of lists corrsponding
                to the actual time series data such that `data[i][j]` contains
                the value at the jth time step of the ith channel of the
                multivariate time series.
            context_types (list):
                List of strings indicating the type of each value in context.
                he value at `context[i]` must match the type specified by
                `context_types[i]`. Valid types include the following: `categorical`,
                `continuous`, `ordinal`, `count`, and `datetime`.
            data_types (list):
                List of strings indicating the type of each channel in data.
                Each value in the list at data[i] must match the type specified by
                `data_types[i]`. The valid types are the same as for `context_types`.
        """
        # build variable2id maps
        self._build(examples, types)

        # Initialize model with data dimension & context dimension
        self._model = PARNet(self._data_dims, self._ctx_dims).to(self.device)

        # Specify optimizer
        optimizer = torch.optim.Adam(self._model.parameters(), lr=1e-3)

        iterator = range(self.epochs)
        if self.verbose:
            iterator = tqdm(iterator)

        # Training
        for epoch in iterator:
            for sequences in sequence_loader:

                # process sequence as in _fit
                timeseries_data = batch_to_df(sequences, columns, dem._entity_columns)
                sequences = assemble_sequences(
                    timeseries_data,
                    dem._entity_columns,
                    dem._context_columns,
                    dem._segment_size,
                    dem._sequence_index,
                    drop_sequence_index=False,
                )
                if dem._sequence_index:
                    self._transform_sequence_index(sequences, dem)

                X, C = [], []
                for sequence in sequences:
                    dataFeed = np.array(sequence["data"]).tolist()
                    # if dataFeed!=[]:
                    X.append(self._data_to_tensor(dataFeed))
                    C.append(self._context_to_tensor(sequence["context"]))

                X = torch.nn.utils.rnn.pack_sequence(X, enforce_sorted=False).to(self.device)
                if self._ctx_dims:
                    C = torch.stack(C, dim=0).to(self.device)

                X_padded, seq_len = torch.nn.utils.rnn.pad_packed_sequence(X)

                Y = self._model(X, C)
                Y_padded, _ = torch.nn.utils.rnn.pad_packed_sequence(Y)

                optimizer.zero_grad()
                loss = self._compute_loss(X_padded[1:, :, :], Y_padded[:-1, :, :], seq_len)
                loss.backward()
                if self.verbose:
                    iterator.set_description("Epoch {} | Loss {}".format(epoch + 1, loss.item()))

                optimizer.step()

    def _compute_loss(self, X_padded, Y_padded, seq_len):
        """Compute the loss between X and Y.

        Given X[i,:,:], the neural network predicts the value at the next
        timestep (i+1); this prediction is provided in Y[i,:,:]. This function
        returns the loss between the predicted and actual sequence.

        .. note::
            The `i`th time series (with padding removed) can be indexed with
            `X[:seq_len[i], i, :]`.

        Args:
            X_padded (tensor):
                This contains the input to the model.
            Y_padded (tensor):
                This contains the output of the model.
            seq_len (list):
                This list contains the length of each sequence.
        """
        log_likelihood = 0.0
        _, batch_size, input_size = X_padded.shape

        for key, props in self._data_map.items():
            if props["type"] in ["continuous", "timestamp"]:
                mu_idx, sigma_idx, missing_idx = props["indices"]
                mu = Y_padded[:, :, mu_idx]
                sigma = torch.nn.functional.softplus(Y_padded[:, :, sigma_idx])
                missing = torch.nn.LogSigmoid()(Y_padded[:, :, missing_idx])

                for i in range(batch_size):
                    dist = torch.distributions.normal.Normal(mu[: seq_len[i], i], sigma[: seq_len[i], i])
                    log_likelihood += torch.sum(dist.log_prob(X_padded[-seq_len[i] :, i, mu_idx]))

                    p_true = X_padded[: seq_len[i], i, missing_idx]
                    p_pred = missing[: seq_len[i], i]
                    log_likelihood += torch.sum(p_true * p_pred)
                    log_likelihood += torch.sum((1.0 - p_true) * torch.log(1.0 - torch.exp(p_pred)))

            elif props["type"] in ["count"]:
                r_idx, p_idx, missing_idx = props["indices"]
                r = torch.nn.functional.softplus(Y_padded[:, :, r_idx]) * props["range"]
                p = torch.sigmoid(Y_padded[:, :, p_idx])
                x = X_padded[:, :, r_idx] * props["range"]
                missing = torch.nn.LogSigmoid()(Y_padded[:, :, missing_idx])

                for i in range(batch_size):
                    dist = torch.distributions.negative_binomial.NegativeBinomial(
                        r[: seq_len[i], i], p[: seq_len[i], i]
                    )
                    log_likelihood += torch.sum(dist.log_prob(x[: seq_len[i], i]))

                    p_true = X_padded[: seq_len[i], i, missing_idx]
                    p_pred = missing[: seq_len[i], i]
                    log_likelihood += torch.sum(p_true * p_pred)
                    log_likelihood += torch.sum((1.0 - p_true) * torch.log(1.0 - torch.exp(p_pred)))

            elif props["type"] in ["categorical", "ordinal"]:
                idx = list(props["indices"].values())
                log_softmax = torch.nn.functional.log_softmax(Y_padded[:, :, idx], dim=2)

                for i in range(batch_size):
                    target = X_padded[: seq_len[i], i, idx]
                    predicted = log_softmax[: seq_len[i], i]
                    target = torch.argmax(target, dim=1).unsqueeze(dim=1)
                    log_likelihood += torch.sum(predicted.gather(dim=1, index=target))

            else:
                raise ValueError()

        return -log_likelihood / (batch_size * len(self._data_map) * batch_size)

    def _tensor_to_data(self, x):
        # Force CPU on x
        x = x.to(torch.device("cpu"))

        seq_len, batch_size, _ = x.shape
        assert batch_size == 1

        data = [None] * (len(self._data_map) - 1)
        for key, props in self._data_map.items():
            if key == "<TOKEN>":
                continue

            data[key] = []
            for i in range(seq_len):
                if props["type"] in ["continuous", "datetime"]:
                    mu_idx, sigma_idx, missing_idx = props["indices"]
                    if (x[i, 0, missing_idx] > 0) and props["nulls"]:
                        data[key].append(None)
                    else:
                        data[key].append(x[i, 0, mu_idx].item() * props["std"] + props["mu"])

                elif props["type"] in ["count"]:
                    r_idx, p_idx, missing_idx = props["indices"]
                    if x[i, 0, missing_idx] > 0 and props["nulls"]:
                        data[key].append(None)
                    else:
                        sample = x[i, 0, r_idx].item() * props["range"] + props["min"]
                        data[key].append(int(sample))

                elif props["type"] in ["categorical", "ordinal"]:
                    ml_value, max_x = None, float("-inf")
                    for value, idx in props["indices"].items():
                        if x[i, 0, idx] > max_x:
                            max_x = x[i, 0, idx]
                            ml_value = value

                    data[key].append(ml_value)

                else:
                    raise ValueError()

        return data

    def _sample_state(self, x):
        log_likelihood = 0.0
        seq_len, batch_size, input_size = x.shape
        assert seq_len == 1 and batch_size == 1

        for key, props in self._data_map.items():
            if props["type"] in ["continuous", "timestamp"]:
                mu_idx, sigma_idx, missing_idx = props["indices"]
                mu = x[0, 0, mu_idx]
                sigma = torch.nn.functional.softplus(x[0, 0, sigma_idx])
                dist = torch.distributions.normal.Normal(mu, sigma)
                x[0, 0, mu_idx] = dist.sample()
                x[0, 0, sigma_idx] = 0.0
                log_likelihood += torch.sum(dist.log_prob(x[0, 0, mu_idx]))

                dist = torch.distributions.Bernoulli(torch.sigmoid(x[0, 0, missing_idx]))
                x[0, 0, missing_idx] = dist.sample()
                x[0, 0, mu_idx] = x[0, 0, mu_idx] * (1.0 - x[0, 0, missing_idx])
                log_likelihood += torch.sum(dist.log_prob(x[0, 0, missing_idx]))

            elif props["type"] in ["count"]:
                r_idx, p_idx, missing_idx = props["indices"]
                r = torch.nn.functional.softplus(x[0, 0, r_idx]) * props["range"]
                p = torch.sigmoid(x[0, 0, p_idx])
                dist = torch.distributions.negative_binomial.NegativeBinomial(r, p)
                x[0, 0, r_idx] = dist.sample()
                x[0, 0, p_idx] = 0.0
                log_likelihood += torch.sum(dist.log_prob(x[0, 0, r_idx]))
                x[0, 0, r_idx] /= props["range"]

                dist = torch.distributions.Bernoulli(torch.sigmoid(x[0, 0, missing_idx]))
                x[0, 0, missing_idx] = dist.sample()
                x[0, 0, r_idx] = x[0, 0, r_idx] * (1.0 - x[0, 0, missing_idx])
                log_likelihood += torch.sum(dist.log_prob(x[0, 0, missing_idx]))

            elif props["type"] in ["categorical", "ordinal"]:
                idx = list(props["indices"].values())
                p = torch.nn.functional.softmax(x[0, 0, idx], dim=0)
                x_new = torch.zeros(p.size()).to(self.device)
                x_new.scatter_(dim=0, index=torch.multinomial(p, 1), value=1)
                x[0, 0, idx] = x_new
                log_likelihood += torch.sum(torch.log(p) * x_new)

            else:
                raise ValueError()

        return x, log_likelihood

    def _sample_sequence(self, context, min_length, max_length):
        log_likelihood = 0.0

        x = torch.zeros(self._data_dims).to(self.device)
        x[self._data_map["<TOKEN>"]["indices"]["<START>"]] = 1.0
        x = x.unsqueeze(0).unsqueeze(0)

        for step in range(max_length):
            next_x, ll = self._sample_state(self._model(x, context)[-1:, :, :])
            x = torch.cat([x, next_x], dim=0)
            log_likelihood += ll
            if next_x[0, 0, self._data_map["<TOKEN>"]["indices"]["<END>"]] > 0.0:
                if min_length <= step + 1 <= max_length:
                    break  # received end token

                next_x[0, 0, self._data_map["<TOKEN>"]["indices"]["<BODY>"]] = 1.0
                next_x[0, 0, self._data_map["<TOKEN>"]["indices"]["<END>"]] = 0.0

        return x[1:, :, :], log_likelihood

    def sample_sequence(self, context, sequence_length=None):

        if sequence_length is not None:
            min_length = max_length = sequence_length
        else:
            min_length = self._min_length
            max_length = self._max_length

        if self._ctx_dims:
            context = self._context_to_tensor(context).unsqueeze(0)
        else:
            context = None

        best_x, best_ll = None, float("-inf")
        for _ in range(self.sample_size):
            with torch.no_grad():
                x, log_likelihood = self._sample_sequence(context, min_length, max_length)

            if log_likelihood > best_ll:
                best_x = x
                best_ll = log_likelihood

        return self._tensor_to_data(best_x)

class BaseTimeseriesModel:
    """Base class for timeseries models.

    Args:
        field_names (list[str]):
            List of names of the fields that need to be modeled
            and included in the generated output data. Any additional
            fields found in the data will be ignored and will not be
            included in the generated output.
            If ``None``, all the fields found in the data are used.
        field_types (dict[str, dict]):
            Dictinary specifying the data types and subtypes
            of the fields that will be modeled. Field types and subtypes
            combinations must be compatible with the SDV Metadata Schema.
        anonymize_fields (dict[str, str]):
            Dict specifying which fields to anonymize and what faker
            category they belong to.
        primary_key (str):
            Name of the field which is the primary key of the table.
        entity_columns (list[str]):
            Names of the columns which identify different time series
            sequences. These will be used to group the data in separated
            training examples.
        context_columns (list[str]):
            The columns in the dataframe which are constant within each
            group/entity. These columns will be provided at sampling time
            (i.e. the samples will be conditioned on the context variables).
        segment_size (int, pd.Timedelta or str):
            If specified, cut each training sequence in several segments of
            the indicated size. The size can either can passed as an integer
            value, which will interpreted as the number of data points to
            put on each segment, or as a pd.Timedelta (or equivalent str
            representation), which will be interpreted as the segment length
            in time. Timedelta segment sizes can only be used with sequence
            indexes of type datetime.
        sequence_index (str):
            Name of the column that acts as the order index of each
            sequence. The sequence index column can be of any type that can
            be sorted, such as integer values or datetimes.
        context_model (str or sdv.tabular.BaseTabularModel):
            Model to use to sample the context rows. It can be passed as a
            a string, which must be one of the following:

            * `gaussian_copula` (default): Use a GaussianCopula model.

            Alternatively, a preconfigured Tabular model instance can be
            passed.

        table_metadata (dict or metadata.Table):
            Table metadata instance or dict representation.
            If given alongside any other metadata-related arguments, an
            exception will be raised.
            If not given at all, it will be built using the other
            arguments or learned from the data.
    """

    _DTYPE_TRANSFORMERS = {
        "i": None,
        "f": None,
        "M": rdt.transformers.DatetimeTransformer(strip_constant=True),
        "b": None,
        "O": None,
    }
    _CONTEXT_MODELS = {
        "gaussian_copula": (
            GaussianCopula,
            {"categorical_transformer": "categorical_fuzzy"},
        )
    }

    _metadata = None

    def __init__(
        self,
        field_names=None,
        field_types=None,
        anonymize_fields=None,
        primary_key=None,
        entity_columns=None,
        context_columns=None,
        sequence_index=None,
        segment_size=None,
        context_model=None,
        table_metadata=None,
    ):
        if table_metadata is None:
            self._metadata = Table(
                field_names=field_names,
                primary_key=primary_key,
                field_types=field_types,
                anonymize_fields=anonymize_fields,
                dtype_transformers=self._DTYPE_TRANSFORMERS,
                sequence_index=sequence_index,
                entity_columns=entity_columns,
                context_columns=context_columns,
            )
            self._metadata_fitted = False
        else:
            null_args = (
                field_names,
                primary_key,
                field_types,
                anonymize_fields,
                sequence_index,
                entity_columns,
                context_columns,
            )
            for arg in null_args:
                if arg:
                    raise ValueError("If table_metadata is given {} must be None".format(arg.__name__))

            if isinstance(table_metadata, dict):
                table_metadata = Table.from_dict(
                    table_metadata,
                    dtype_transformers=self._DTYPE_TRANSFORMERS,
                )

            self._metadata = table_metadata
            self._metadata_fitted = table_metadata.fitted

        # Validate arguments
        if segment_size is not None and not isinstance(segment_size, int):
            if sequence_index is None:
                raise TypeError("`segment_size` must be of type `int` if " "no `sequence_index` is given.")

            segment_size = pd.to_timedelta(segment_size)

        self._context_columns = self._metadata._context_columns
        self._entity_columns = self._metadata._entity_columns
        self._sequence_index = self._metadata._sequence_index
        self._segment_size = segment_size

        context_model = context_model or "gaussian_copula"
        if isinstance(context_model, str):
            context_model = self._CONTEXT_MODELS[context_model]

        self._context_model_template = context_model

    def _fit(self, sequence_loader, timeseries_data_example):
        raise NotImplementedError()

    def _fit_context_model(self, transformed):
        template = self._context_model_template
        default_kwargs = {
            "primary_key": self._entity_columns,
            "field_types": {
                name: meta for name, meta in self._metadata.get_fields().items() if name in self._entity_columns
            },
        }
        if isinstance(template, tuple):
            context_model_class, context_model_kwargs = copy.deepcopy(template)
            if "primary_key" not in context_model_kwargs:
                context_model_kwargs["primary_key"] = self._entity_columns
                for keyword, argument in default_kwargs.items():
                    if keyword not in context_model_kwargs:
                        context_model_kwargs[keyword] = argument

            self._context_model = context_model_class(**context_model_kwargs)
        elif isinstance(template, type):
            self._context_model = template(**default_kwargs)
        else:
            self._context_model = copy.deepcopy(template)

        LOGGER.debug("Fitting context model %s", self._context_model.__class__.__name__)
        if self._context_columns:
            context = transformed[self._entity_columns + self._context_columns]
        else:
            context = transformed[self._entity_columns].copy()
            # Add constant column to allow modeling
            context[str(uuid.uuid4())] = 0

        context = context.groupby(self._entity_columns).first().reset_index()
        self._context_model.fit(context)

    def fit(self, sequence_loader, columns, types, examples):
        self._fit(sequence_loader, columns, types, examples)

    def get_metadata(self):
        """Get metadata about the table.

        This will return an ``sdv.metadata.Table`` object containing
        the information about the data that this model has learned.

        This Table metadata will contain some common information,
        such as field names and data types, as well as additional
        information that each Sub-class might add, such as the
        observed data field distributions and their parameters.

        Returns:
            sdv.metadata.Table:
                Table metadata.
        """
        return self._metadata

    def _sample(self, context=None, sequence_length=None):
        raise NotImplementedError()

    def sample(self, context=None, sequence_length=None):
        """Sample new sequences.

        Args:
            num_sequences (int):
                Number of sequences to sample. If context is
                passed, this is ignored. If not given, the
                same number of sequences as in the original
                timeseries_data is sampled.
            context (pandas.DataFrame):
                Context values to use when generating the sequences.
                If not passed, the context values will be sampled
                using the specified tabular model.
            sequence_length (int):
                If passed, sample sequences of this length. If not
                given, the sequence length will be sampled from
                the model.

        Returns:
            pandas.DataFrame:
                Table containing the sampled sequences in the same
                format as that he training data had.
        """
        num_sequences = 1  # TODO this was missing. Should it be an argument?
        if not self._entity_columns:
            if context is not None:
                raise TypeError("If there are no entity_columns, context must be None")

            context = pd.DataFrame(index=range(num_sequences or 1))
        elif context is None:
            context = self._context_model.sample(num_sequences)
            for column in self._entity_columns or []:
                if column not in context:
                    context[column] = range(len(context))

        sampled = self._sample(context, sequence_length)
        return sampled

    def save(self, path):
        """Save this model instance to the given path using pickle.

        Args:
            path (str):
                Path where the SDV instance will be serialized.
        """
        with open(path, "wb") as output:
            pickle.dump(self, output)

    @classmethod
    def load(cls, path):
        """Load a TabularModel instance from a given path.

        Args:
            path (str):
                Path from which to load the instance.

        Returns:
            TabularModel:
                The loaded tabular model.
        """
        with open(path, "rb") as f:
            return pickle.load(f)

class DeepEchoModel(BaseTimeseriesModel):
    """Base class for all the SDV Time series models based on DeepEcho."""

    _MODEL_CLASS = None
    _model_kwargs = None

    _DATA_TYPES = {
        "numerical": "continuous",
        "categorical": "categorical",
        "boolean": "categorical",
        "datetime": "datetime",
    }

    _verbose = True

    def _build_model(self):
        return self._MODEL_CLASS(**self._model_kwargs)

    def _transform_sequence_index(self, sequences):
        sequence_index_idx = self._data_columns.index(self._sequence_index)
        for sequence in sequences:
            data = sequence["data"]
            sequence_index = data[sequence_index_idx]
            diffs = np.diff(sequence_index).tolist()
            data[sequence_index_idx] = diffs[0:1] + diffs
            data.append(sequence_index[0:1] * len(sequence_index))

    def _fit(self, sequence_loader, columns, types, examples):

        self._model = self._build_model()

        # Specify output columns in self._data_columns
        self._output_columns = columns
        self._data_columns = [
            column for column in columns if column not in self._entity_columns + self._context_columns
        ]

        # Validate and fit
        dem = self
        self._model.fit_sequences(sequence_loader, columns, types, examples, dem)

    def _sample(self, context=None, sequence_length=None):
        """Sample new sequences.

        Args:
            context (pandas.DataFrame):
                Context values to use when generating the sequences.
                If not passed, the context values will be sampled
                using the specified tabular model.
            sequence_length (int):
                If passed, sample sequences of this length. If not
                given, the sequence length will be sampled from
                the model.

        Returns:
            pandas.DataFrame:
                Table containing the sampled sequences in the same
                format as that he training data had.
        """
        # Set the entity_columns as index to properly iterate over them
        if self._entity_columns:
            context = context.set_index(self._entity_columns)

        iterator = tqdm(context.iterrows(), disable=not self._verbose, total=len(context))

        output = list()
        for entity_values, context_values in iterator:
            context_values = context_values.tolist()
            sequence = self._model.sample_sequence(context_values, sequence_length)
            if self._sequence_index:
                sequence_index_idx = self._data_columns.index(self._sequence_index)
                diffs = sequence[sequence_index_idx]
                start = sequence[-1]
                sequence[sequence_index_idx] = np.cumsum(diffs) - diffs[0] + start

            # Reformat as a DataFrame
            group = pd.DataFrame(dict(zip(self._data_columns, sequence)), columns=self._data_columns)
            group[self._entity_columns] = entity_values
            for column, value in zip(self._context_columns, context_values):
                if column == self._sequence_index:
                    sequence_index = group[column]
                    group[column] = sequence_index.cumsum() - sequence_index.iloc[0] + value
                else:
                    group[column] = value

            output.append(group)

        output = pd.concat(output)
        return output[self._output_columns].reset_index(drop=True)

class PAREdit(DeepEchoModel):
    _MODEL_CLASS = PARModel

    def __init__(
        self,
        field_names=None,
        field_types=None,
        anonymize_fields=None,
        primary_key=None,
        entity_columns=None,
        context_columns=None,
        sequence_index=None,
        segment_size=None,
        context_model=None,
        table_metadata=None,
        epochs=128,
        sample_size=1,
        cuda=True,
        verbose=True,
    ):
        super().__init__(
            field_names=field_names,
            field_types=field_types,
            anonymize_fields=anonymize_fields,
            primary_key=primary_key,
            entity_columns=entity_columns,
            context_columns=context_columns,
            sequence_index=sequence_index,
            segment_size=segment_size,
            context_model=context_model,
            table_metadata=table_metadata,
        )

        self._model_kwargs = {
            "epochs": epochs,
            "sample_size": sample_size,
            "cuda": cuda,
            "verbose": verbose,
        }
        self._verbose = verbose