import hdnntools as hdt
import nmstools as nmt
import pyanitools as pyt
import pyaniasetools as aat
import pyanitrainer as atr
import pymolfrag as pmf

from pyNeuroChem import cachegenerator as cg

import numpy as np

from time import sleep
import subprocess
import random
#import pyssh
import re
import os

import pyaniasetools as aat

from multiprocessing import Process
import shutil

import matplotlib.pyplot as plt

class alconformationalsampler():

    # Constructor
    def __init__(self, ldtdir, datdir, optlfile, fpatoms, netdict):
        self.ldtdir = ldtdir # local working dir
        self.datdir = datdir # working data dir
        self.cdir = ldtdir+datdir+'/confs/' # confs store dir (the data gen code looks here for conformations to run QM on)

        self.fpatoms = fpatoms # atomic species being sampled

        self.optlfile = optlfile # Optimized molecules store path file

        self.idir = [f for f in open(optlfile).read().split('\n') if f != ''] # read and store the paths to the opt files

        # create cdir if it does not exist
        if not os.path.exists(self.cdir):
            os.mkdir(self.cdir)

        # store network parameters dictionary
        self.netdict = netdict

    # Runs NMS sampling (single GPU only currently)
    def run_sampling_nms(self, nmsparams, gpus=[0]):
        print('Running NMS sampling...')
        p = Process(target=self.normal_mode_sampling, args=(nmsparams['T'],
                                                            nmsparams['Ngen'],
                                                            nmsparams['Nkep'],
                                                            nmsparams['maxd'],
                                                            nmsparams['sig'],
                                                            gpus[0]))
        p.start()
        p.join()

    # Run MD sampling on N GPUs. This code will automatically run 2 mds per GPU for better utilization
    def run_sampling_md(self, mdparams, perc=1.0, gpus=[0]):
        md_work = []
        for di, id in enumerate(self.idir):
            files = os.listdir(id)
            for f in files:
                if len(f) > 4:
                    if ".ipt" in f[-4:]:
                        md_work.append(id+f)
                    else:
                        print('Incorrect extension:',id+f)

        gpus2 = gpus

        md_work = np.array(md_work)
        np.random.shuffle(md_work)
        md_work = md_work[0:int(perc*md_work.size)]
        md_work = np.array_split(md_work,len(gpus2))

        proc = []
        for i,(md,g) in enumerate(zip(md_work,gpus2)):
            proc.append(Process(target=self.mol_dyn_sampling, args=(md,i,
                                                                    mdparams['N'],
                                                                    mdparams['T1'],
                                                                    mdparams['T2'],
                                                                    mdparams['dt'],
                                                                    mdparams['Nc'],
                                                                    mdparams['Ns'],
                                                                    mdparams['sig'],
                                                                    g)))
        print('Running MD Sampling...')
        for i,p in enumerate(proc):
            p.start()

        for p in proc:
            p.join()
        print('Finished sampling.')

    # Run the dimer sampling code on N gpus
    def run_sampling_dimer(self, dmparams, gpus=[0]):

        proc = []
        for i,g in enumerate(gpus):
            proc.append(Process(target=self.dimer_sampling, args=(i, int(dmparams['Nr']/len(gpus)),
                                                                  dmparams,
                                                                  g)))
        print('Running Dimer-MD Sampling...')
        for i,p in enumerate(proc):
            p.start()

        for p in proc:
            p.join()
        print('Finished sampling.')

    # Run cluster sampling on N gpus
    def run_sampling_cluster(self, gcmddict, gpus=[0]):

        Nmols = np.random.randint(low=gcmddict['MolLow'],
                                  high=gcmddict['MolHigh'],
                                  size=gcmddict['Nr'])

        Ntmps = np.random.randint(low=gcmddict['T']-200,
                                  high=gcmddict['T']+200,
                                  size=gcmddict['Nr'])

        print('Box Sizes:',Nmols)
        print('Sim Temps:',Ntmps)

        seeds = np.random.randint(0,1000000,len(gpus))

        mol_sizes = np.array_split(Nmols, len(gpus))
        mol_temps = np.array_split(Ntmps, len(gpus))

        proc = []
        for i,g in enumerate(gpus):
            proc.append(Process(target=self.cluster_sampling, args=(i, int(gcmddict['Nr']/len(gpus)),
                                                                    mol_sizes[i],
                                                                    mol_temps[i],
                                                                    seeds[i],
                                                                    gcmddict,
                                                                    g)))
        print('Running Cluster-MD Sampling...')
        for i,p in enumerate(proc):
            p.start()

        for p in proc:
            p.join()
        print('Finished sampling.')

    # Normal mode sampler function
    def normal_mode_sampling(self, T, Ngen, Nkep, maxd, sig, gpuid):
        of = open(self.ldtdir + self.datdir + '/info_data_nms.nfo', 'w')

        aevsize = self.netdict['aevsize']

        anicv = aat.anicrossvalidationconformer(self.netdict['cnstfile'],
                                                self.netdict['saefile'],
                                                self.netdict['nnfprefix'],
                                                self.netdict['num_nets'],
                                                [gpuid], False)

        dc = aat.diverseconformers(self.netdict['cnstfile'],
                                   self.netdict['saefile'],
                                   self.netdict['nnfprefix']+'0/networks/',
                                   aevsize,
                                   gpuid, False)

        Nkp = 0
        Nkt = 0
        Ntt = 0
        idx = 0
        for di,id in enumerate(self.idir):
            of.write(str(di)+' of '+str(len(self.idir))+') dir: '+ str(id) +'\n')
            #print(di,'of',len(self.idir),') dir:', id)
            files = os.listdir(id)
            files.sort()

            Nk = 0
            Nt = 0
            for fi,f in enumerate(files):
                print(f)
                data = hdt.read_rcdb_coordsandnm(id+f)

                #print(id+f)
                spc = data["species"]
                xyz = data["coordinates"]
                nmc = data["nmdisplacements"]
                frc = data["forceconstant"]

                if "charge" in data and "multip" in data:
                    chg = data["charge"]
                    mlt = data["multip"]
                else:
                    chg = "0"
                    mlt = "1"

                nms = nmt.nmsgenerator(xyz,nmc,frc,spc,T,minfc=5.0E-2,maxd=maxd)
                conformers = nms.get_Nrandom_structures(Ngen)

                if conformers.shape[0] > 0:
                    if conformers.shape[0] > Nkep:
                        ids = dc.get_divconfs_ids(conformers, spc, Ngen, Nkep, [])
                        conformers = conformers[ids]

                    sigma = anicv.compute_stddev_conformations(conformers,spc)
                    sid = np.where( sigma >  sig )[0]

                    Nt += sigma.size
                    Nk += sid.size
                    if 100.0*sid.size/float(Ngen) > 0:
                        Nkp += sid.size
                        cfn = f.split('.')[0].split('-')[0]+'_'+str(idx).zfill(5)+'-'+f.split('.')[0].split('-')[1]+'_2.xyz'
                        cmts = [' '+chg+' '+mlt for c in range(Nk)]
                        hdt.writexyzfilewc(self.cdir+cfn,conformers[sid],spc,cmts)
                idx += 1

            Nkt += Nk
            Ntt += Nt
            of.write('    -Total: '+str(Nk)+' of '+str(Nt)+' percent: '+"{:.2f}".format(100.0*Nk/Nt)+'\n')
            of.flush()
            #print('    -Total:',Nk,'of',Nt,'percent:',"{:.2f}".format(100.0*Nk/Nt))

        del anicv
        del dc

        of.write('\nGrand Total: '+ str(Nkt)+ ' of '+ str(Ntt)+' percent: '+"{:.2f}".format(100.0*Nkt/Ntt)+ ' Kept '+str(Nkp)+'\n')
        #print('\nGrand Total:', Nkt, 'of', Ntt,'percent:',"{:.2f}".format(100.0*Nkt/Ntt), 'Kept',Nkp)
        of.close()

    # MD Sampling function
    def mol_dyn_sampling(self,md_work, i, N, T1, T2, dt, Nc, Ns, sig, gpuid):
        activ = aat.moldynactivelearning(self.netdict['cnstfile'],
                                         self.netdict['saefile'],
                                         self.netdict['nnfprefix'],
                                         self.netdict['num_nets'],
                                         gpuid)

        difo = open(self.ldtdir + self.datdir + '/info_data_mdso-'+str(i)+'.nfo', 'w')
        Nmol = 0
        dnfo = 'MD Sampler running: ' + str(md_work.size)
        difo.write(dnfo + '\n')
        Nmol = md_work.size
        ftme_t = 0.0
        for di, id in enumerate(md_work):
            data = hdt.read_rcdb_coordsandnm(id)
            #print(di, ') Working on', id, '...')
            S = data["species"]

            # Set mols
            activ.setmol(data["coordinates"], S)

            # Generate conformations
            X = activ.generate_conformations(N, T1, T2, dt, Nc, Ns, dS=sig)

            ftme_t += activ.failtime

            nfo = activ._infostr_
            m = id.rsplit('/',1)[1].rsplit('.',1)[0]
            difo.write('  -' + m + ': ' + nfo + '\n')
            difo.flush()
            #print(nfo)

            if X.size > 0:
                hdt.writexyzfile(self.cdir + 'mds_' + m.split('.')[0] + '_' + str(i).zfill(2) + str(di).zfill(4) + '.xyz', X, S)
        difo.write('Complete mean fail time: ' + "{:.2f}".format(ftme_t / float(Nmol)) + '\n')
        print(Nmol)
        del activ
        difo.close()

    # Dimer sampling function
    def dimer_sampling(self, tid, Nr, dparam, gpuid):
        mds_select = dparam['mdselect']
        #N = dparam['N']
        T = dparam['T']
        L = dparam['L']
        V = dparam['V']
        maxNa = dparam['maxNa']
        dt = dparam['dt']
        sig = dparam['sig']
        Nm = dparam['Nm']
        Ni = dparam['Ni']
        Ns = dparam['Ns']

        mols = []
        difo = open(self.ldtdir + self.datdir + '/info_data_mddimer-'+str(tid)+'.nfo', 'w')
        for di,id in enumerate(dparam['mdselect']):
            files = os.listdir(self.idir[id[1]])
            random.shuffle(files)

            dnfo = str(di) + ' of ' + str(len(dparam['mdselect'])) + ') dir: ' + str(self.idir[id[1]]) + ' Selecting: '+str(id[0]*len(files))
            #print(dnfo)
            difo.write(dnfo+'\n')

            for i in range(id[0]):
                for n,m in enumerate(files):
                        data = hdt.read_rcdb_coordsandnm(self.idir[id[1]]+m)
                        if len(data['species']) < maxNa:
                            mols.append(data)

        dgen = pmf.dimergenerator(self.netdict['cnstfile'],
                                  self.netdict['saefile'],
                                  self.netdict['nnfprefix'],
                                  self.netdict['num_nets'],
                                  mols, gpuid)

        difo.write('Beginning dimer generation...\n')

        Nt = 0
        Nd = 0
        for i in range(Nr):
            dgen.init_dynamics(Nm, V, L, dt, T)

            for j in range(Ns):
                if j != 0:
                    dgen.run_dynamics(Ni)

                fname = self.cdir + 'dimer-'+str(tid).zfill(2)+str(i).zfill(2)+'-'+str(j).zfill(2)+'_'
                max_sig = dgen.__fragmentbox__(fname,sig)
                print('MaxSig:',max_sig)
                #difo.write('Step ('+str(i)+',',+str(j)+') ['+ str(dgen.Nd)+ '/'+ str(dgen.Nt)+']\n')
                difo.write('Step ('+str(i)+','+str(j)+') ['+ str(dgen.Nd)+ '/'+ str(dgen.Nt)+'] max sigma: ' + "{:.2f}".format(max_sig) + ' generated '+str(len(dgen.frag_list))+' dimers...\n')

                Nt += dgen.Nt
                Nd += dgen.Nd

                #print('Step (',tid,',',i,') [', str(dgen.Nd), '/', str(dgen.Nt),'] generated ',len(dgen.frag_list), 'dimers...')
                #difo.write('Step ('+str(i)+') ['+ str(dgen.Nd)+ '/'+ str(dgen.Nt)+'] generated '+str(len(dgen.frag_list))+'dimers...\n')
                if max_sig > 3.0*sig:
                    difo.write('Terminating dynamics -- max sigma: '+"{:.2f}".format(max_sig)+' Ran for: '+"{:.2f}".format(j*Ni*dt)+'fs\n')
                    break

        difo.write('Generated '+str(Nd)+' of '+str(Nt)+' tested dimers. Percent: ' + "{:.2f}".format(100.0*Nd/float(Nt)))
        difo.close()

    # Cluster sampling function
    def cluster_sampling(self, tid, Nr, mol_sizes, mol_temps, seed, gcmddict, gpuid):
        os.environ["OMP_NUM_THREADS"] = "2"

        dictc = gcmddict.copy()
        solv_file = dictc['solv_file']
        solu_dirs = dictc['solu_dirs']

        np.random.seed(seed)

        dictc['Nr'] = Nr
        dictc['molfile'] = self.cdir + 'clst'
        dictc['dstore'] = self.ldtdir + self.datdir + '/'

        solv = [hdt.read_rcdb_coordsandnm(solv_file)]

        if solu_dirs:
            solu = [hdt.read_rcdb_coordsandnm(solu_dirs+f) for f in os.listdir(solu_dirs)]
        else:
            solu = []

        dgen = pmf.clustergenerator(self.netdict['cnstfile'],
                                    self.netdict['saefile'],
                                    self.netdict['nnfprefix'],
                                    self.netdict['num_nets'],
                                    solv, solu, gpuid)

        dgen.generate_clusters(dictc, mol_sizes, mol_temps, tid)

    # Run the TS sampler
    def run_sampling_TS(self, tsparams, gpus=[0], perc=1.0):
        TS_infiles = []
        for di, id in enumerate(tsparams['tsfiles']):
            files = [fl for fl in os.listdir(id) if '.xyz' in fl]
            for f in files:
                TS_infiles.append(id+f)

        gpus2 = gpus

        TS_infiles = np.array(TS_infiles)
        np.random.shuffle(TS_infiles)
        TS_infiles = TS_infiles[0:int(perc*len(TS_infiles))]
        TS_infiles = np.array_split(TS_infiles,len(gpus2))

        proc = []
        for i,g in enumerate(gpus2):
            proc.append(Process(target=self.TS_sampling, args=(i, TS_infiles[i], tsparams, g)))
        print('Running MD Sampling...')
        for p in proc:
            p.start()

        for p in proc:
            p.join()

        print('Finished sampling.')

    # TS sampler function
    def TS_sampling(self, tid, TS_infiles, tsparams, gpuid):
        activ = aat.MD_Sampler(TS_infiles,
                               self.netdict['cnstfile'],
                               self.netdict['saefile'],
                               self.netdict['nnfprefix'],
                               self.netdict['num_nets'],
                               gpuid)
        T=tsparams['T']
        sig=tsparams['sig']
        Ns=tsparams['n_samples']
        n_steps=tsparams['n_steps']
        steps=tsparams['steps']
        min_steps=tsparams['min_steps']
        nm=tsparams['normalmode']
        displacement=tsparams['displacement']
        difo = open(self.ldtdir + self.datdir + '/info_tssampler-'+str(tid)+'.nfo', 'w')
        for f in TS_infiles:
            X = []
            ftme_t = 0.0
            fail_count=0
            sumsig = 0.0
            for i in range(Ns):
                #x, S, t, stddev, fail, temp = activ.run_md(f, T, steps, n_steps, nmfile=f.rsplit(".",1)[0]+'.log', displacement=displacement, min_steps=min_steps, sig=sig, nm=nm)
                x, S, t, stddev, fail, temp = activ.run_md(f, T, steps, n_steps, min_steps=min_steps, sig=sig, nm=nm)
                sumsig += stddev
                if fail:
                    #print('Job '+str(i)+' failed in '+"{:.2f}".format(t)+' Sigma: ' + "{:.2f}".format(stddev)+' SetTemp: '+"{:.2f}".format(temp))
                    difo.write('Job '+str(i)+' failed in '+"{:.2f}".format(t)+'fs Sigma: ' + "{:.2f}".format(stddev) + ' SetTemp: ' + "{:.2f}".format(temp) + '\n')
                    X.append(x[np.newaxis,:,:])
                    fail_count+=1
                else:
                    #print('Job '+str(i)+' succeeded.')
                    difo.write('Job '+str(i)+' succeeded.\n')
                ftme_t += t
            print('Complete mean fail time: ' + "{:.2f}".format(ftme_t / float(Ns)) + ' failed ' + str(fail_count) + '/' + str(Ns) + '\n')
            difo.write('Complete mean fail time: ' + "{:.2f}".format(ftme_t / float(Ns)) + ' failed ' + str(fail_count) + '/' + str(Ns) + ' MeanSig: ' + "{:.2f}".format(sumsig / float(Ns)) + '\n')
            X = np.vstack(X)
            hdt.writexyzfile(self.cdir + os.path.basename(f), X, S)
        del activ
        difo.close()

    # Run the dihedral sampler
    def run_sampling_dhl(self, dhlparams, gpus):
        dhlparams['Nmol'] = int(np.ceil(dhlparams['Nmol']/len(gpus)))

        seeds = np.random.randint(0,1000000,len(gpus))

        proc = []
        for i,g in enumerate(gpus):
            proc.append(Process(target=self.DHL_sampling, args=(i, dhlparams, self.fpatoms, g, seeds[i])))

        print('Running DHL Sampling...')
        for p in proc:
            p.start()

        for p in proc:
            p.join()

        print('Finished sampling.')

    # Dihedral sampling function
    def DHL_sampling(self, i, dhlparams, fpatoms, gpuid, seed):
        activ = aat.aniTortionSampler(self.netdict,
                                      self.cdir,
                                      dhlparams['smilefile'],
                                      dhlparams['Nmol'],
                                      dhlparams['Nsamp'],
                                      dhlparams['sig'],
                                      dhlparams['rng'],
                                      fpatoms,
                                      seed,
                                      gpuid)

        activ.generate_dhl_samples(MaxNa=dhlparams['MaxNa'], fpref='dhl_scan-'+str(i).zfill(2), freqname='vib'+str(i)+'.')

    def run_sampling_pDynTS(self, pdynparams, gpus=0):

        gpus2 = gpus
        proc = []
        for g in enumerate(gpus2):
            proc.append(Process(target=self.pDyn_QMsampling, args=(pdynparams, g)))
        print('Running pDynamo Sampling...')
        for p in proc:
            p.start()

        for p in proc:
            p.join()

        print('Finished pDynamo sampling.')

    def pDyn_QMsampling(self, pdynparams, gpuid):
                                                                  #Call subproc_pDyn class in pyaniasetools as activ
        activ = aat.subproc_pDyn(self.netdict['cnstfile'],
                                         self.netdict['saefile'],
                                         self.netdict['nnfprefix'],
                                         self.netdict['num_nets'],
                                         gpuid)
        pDyn_dir=pdynparams['pDyn_dir']                         #Folder to write pDynamo input file
        num_rxn=pdynparams['num_rxn']                           #Number of input rxn
        logfile_OPT=pdynparams['logfile_OPT']                   #logfile for FIRE OPT output
        logfile_TS=pdynparams['logfile_TS']                     #logfile for ANI TS output
        logfile_IRC=pdynparams['logfile_IRC']                   #logfile for ANI IRC output
        sbproc_cmdOPT=pdynparams['sbproc_cmdOPT']               #Subprocess commands to run pDyanmo
        sbproc_cmdTS=pdynparams['sbproc_cmdTS']
        sbproc_cmdIRC=pdynparams['sbproc_cmdIRC']
        IRCdir=pdynparams['IRCdir']                             #path to get pDynamo saved IRC points
        indir=pdynparams['indir']                               #path to save XYZ files of IRC points to check stddev
        XYZfile=pdynparams['XYZfile']                           #XYZ file with high standard deviations structures
        l_val=pdynparams['l_val']                               #Ri --> randomly perturb in the interval [+x,-x]
        h_val=pdynparams['h_val']
        n_points=pdynparams['n_points']                         #Number of points along IRC (forward+backward+1 for TS)
        sig=pdynparams['sig']
        N=pdynparams['N']
        wkdir=pdynparams['wkdir']
        cnstfilecv=pdynparams['cnstfilecv']
        saefilecv=pdynparams['saefilecv']
        Nnt=pdynparams['Nnt']

        # --------------------------------- Run pDynamo ---------------------------
        # auto-TS ---> FIRE constraint OPT of core C atoms ---> ANI TS ---> ANI IRC

        activ.write_pDynOPT(num_rxn, pDyn_dir, wkdir, cnstfilecv, saefilecv, Nnt)                              #Write pDynamo input file in pDyndir
        activ.write_pDynTS(num_rxn, pDyn_dir, wkdir, cnstfilecv, saefilecv, Nnt)
        activ.write_pDynIRC(num_rxn, pDyn_dir, wkdir, cnstfilecv, saefilecv, Nnt)

        chk_OPT = activ.subprocess_cmd(sbproc_cmdOPT, False, logfile_OPT)
        if chk_OPT == 0:                                                                                #Wait until previous subproc is done!!
            chk_TS = activ.subprocess_cmd(sbproc_cmdTS, False, logfile_TS)
            if chk_TS == 0:
                chk_IRC = activ.subprocess_cmd(sbproc_cmdIRC, False, logfile_IRC)

        # ----------------------- Save points along ANI IRC ------------------------
        IRCfils=os.listdir(IRCdir)
        IRCfils.sort()

        for f in IRCfils:
            activ.getIRCpoints_toXYZ(n_points, IRCdir+f, f, indir)
        infils=os.listdir(indir)
        infils.sort()

        # ------ Check for high standard deviation structures and get vib modes -----
        for f in infils:
            stdev = activ.check_stddev(indir+f, sig)
            if stdev > sig:                      #if stddev is high then get modes for that point
                nmc = activ.get_nm(indir+f)      #save modes in memory for later use

            activ.write_nm_xyz(XYZfile)          #writes all the structures with high standard deviations to xyz file

        # ----------------------------- Read XYZ for NM -----------------------------
        X, spc, Na, C = hdt.readxyz3(XYZfile)

        # --------- NMS for XYZs with high stddev --------
        for i in range(len(X)):
            for j in range (len(nmc)):
                gen = nmt.nmsgenerator_RXN(X[i],nmc[j],spc[i],l_val,h_val)      # xyz,nmo,fcc,spc,T,Ri_-x,Ri_+x,minfc = 1.0E-3

                N = 500
                gen_crd = np.zeros((N, len(spc[i]),3),dtype=np.float32)
                for k in range(N):
                    gen_crd[k] = gen.get_random_structure()

                hdt.writexyzfile(self.cdir + 'nms_TS%i.xyz' %N, gen_crd, spc[i])

        del activ

def interval(v,S):
    ps = 0.0
    ds = 1.0 / float(S)
    for s in range(S):
        if v > ps and v <= ps+ds:
            return s
        ps = ps + ds

class anitrainerinputdesigner:
    def __init__(self):
        self.params = {"sflparamsfile":None, # AEV parameters file
                       "ntwkStoreDir":"networks/", # Store network dir
                       "atomEnergyFile":None, # Atomic energy shift file
                       "nmax": 0, # Max training iterations
                       "tolr": 50, # Annealing tolerance (patience)
                       "emult": 0.5, # Annealing multiplier
                       "eta": 0.001, # Learning rate
                       "tcrit": 1.0e-5, # eta termination crit.
                       "tmax": 0, # Maximum time (0 = inf)
                       "tbtchsz": 2048, # training batch size
                       "vbtchsz": 2048, # validation batch size
                       "gpuid": 0, # Default GPU id (is overridden by -g flag for HDAtomNNP-Trainer exe)
                       "ntwshr": 0, # Use a single network for all types... (THIS IS BROKEN, DO NOT USE)
                       "nkde": 2, # Energy delta regularization
                       "energy": 1, # Enable/disable energy training
                       "force": 0, # Enable/disable force training
                       "fmult": 1.0, # Multiplier of force cost
                       "pbc": 0, # Use PBC in training (Warning, this only works for data with a single rect. box size)
                       "cmult": 1.0, # Charge cost multiplier (CHARGE TRAINING BROKEN IN CURRENT VERSION)
                       "runtype" : "ANNP_CREATE_HDNN_AND_TRAIN", # DO NOT CHANGE - For NeuroChem backend
                       "adptlrn" : "OFF",
                       "decrate" : 0.9,
                       "moment" : "ADAM",
                       "mu" : 0.99
                       }

        self.layers = dict()

    def add_layer(self, atomtype, layer_dict):
        layer_dict.update({"type":0})
        if atomtype not in self.layers:
            self.layers[atomtype]=[layer_dict]
        else:
            self.layers[atomtype].append(layer_dict)

    def set_parameter(self,key,value):
        self.params[key]=value

    def print_layer_parameters(self):
        for ak in self.layers.keys():
            print('Species:',ak)
            for l in self.layers[ak]:
                print('  -',l)

    def print_training_parameters(self):
        print(self.params)

    def __get_value_string__(self,value):

        if type(value)==float:
            string="{0:10.7e}".format(value)
        else:
            string=str(value)

        return string

    def __build_network_str__(self, iptsize):

        network  = "network_setup {\n"
        network += "    inputsize="+str(iptsize)+";\n"

        for ak in self.layers.keys():
            network += "    atom_net " + ak + " $\n"
            self.layers[ak].append({"nodes":1,"activation":6,"type":0})
            for l in self.layers[ak]:
                network += "        layer [\n"
                for key in l.keys():
                    network += "            "+key+"="+self.__get_value_string__(l[key])+";\n"
                network += "        ]\n"
            network += "    $\n"
        network += "}\n"
        return network

    def write_input_file(self, file, iptsize):
        f = open(file,'w')
        for key in self.params.keys():
            f.write(key+'='+self.__get_value_string__(self.params[key])+'\n')
        f.write(self.__build_network_str__(iptsize))
        f.close()

class alaniensembletrainer():
    def __init__(self, train_root, netdict, h5dir, Nn):
        self.train_root = train_root
        #self.train_pref = train_pref
        self.h5dir = h5dir
        self.Nn = Nn
        self.netdict = netdict

        self.h5file = [f for f in os.listdir(self.h5dir) if f.rsplit('.',1)[1] == 'h5']
        #print(self.h5dir,self.h5file)

    def build_training_cache(self,forces=True):
        store_dir = self.train_root + "cache-data-"
        N = self.Nn

        for i in range(N):
            if not os.path.exists(store_dir + str(i)):
                os.mkdir(store_dir + str(i))

            if os.path.exists(store_dir + str(i) + '/../testset/testset' + str(i) + '.h5'):
                os.remove(store_dir + str(i) + '/../testset/testset' + str(i) + '.h5')

            if not os.path.exists(store_dir + str(i) + '/../testset'):
                os.mkdir(store_dir + str(i) + '/../testset')

        cachet = [cg('_train', self.netdict['saefile'], store_dir + str(r) + '/', False) for r in range(N)]
        cachev = [cg('_valid', self.netdict['saefile'], store_dir + str(r) + '/', False) for r in range(N)]
        testh5 = [pyt.datapacker(store_dir + str(r) + '/../testset/testset' + str(r) + '.h5') for r in range(N)]

        Nd = np.zeros(N, dtype=np.int32)
        Nbf = 0
        for f, fn in enumerate(self.h5file):
            print('Processing file(' + str(f + 1) + ' of ' + str(len(self.h5file)) + '):', fn)
            adl = pyt.anidataloader(self.h5dir+fn)

            To = adl.size()
            Ndc = 0
            Fmt = []
            Emt = []
            for c, data in enumerate(adl):
                Pn = data['path'] + '_' + str(f).zfill(6) + '_' + str(c).zfill(6)

                # Progress indicator
                #sys.stdout.write("\r%d%% %s" % (int(100 * c / float(To)), Pn))
                #sys.stdout.flush()

                # print(data.keys())

                # Extract the data
                X = data['coordinates']
                E = data['energies']
                S = data['species']

                # 0.0 forces if key doesnt exist
                if forces:
                    F = data['forces']
                else:
                    F = 0.0*X

                Fmt.append(np.max(np.linalg.norm(F, axis=2), axis=1))
                Emt.append(E)
                Mv = np.max(np.linalg.norm(F, axis=2), axis=1)

                index = np.where(Mv > 10.5)[0]
                indexk = np.where(Mv <= 10.5)[0]

                Nbf += index.size

                # CLear forces
                X = X[indexk]
                F = F[indexk]
                E = E[indexk]

                Esae = hdt.compute_sae(self.netdict['saefile'],S)

                hidx = np.where(np.abs(E-Esae) > 10.0)
                lidx = np.where(np.abs(E-Esae) <= 10.0)
                if hidx[0].size > 0:
                    print('  -('+str(c).zfill(3)+')High energies detected:\n    ',E[hidx])

                X = X[lidx]
                E = E[lidx]
                F = F[lidx]

                Ndc += E.size

                if (set(S).issubset(self.netdict['atomtyp'])):
                #if (set(S).issubset(['C', 'N', 'O', 'H', 'F', 'S', 'Cl'])):

                    # Random mask
                    R = np.random.uniform(0.0, 1.0, E.shape[0])
                    idx = np.array([interval(r, N) for r in R])

                    # Build random split lists
                    split = []
                    for j in range(N):
                        split.append([i for i, s in enumerate(idx) if s == j])
                        nd = len([i for i, s in enumerate(idx) if s == j])
                        Nd[j] = Nd[j] + nd

                    # Store data
                    for i, t, v, te in zip(range(N), cachet, cachev, testh5):
                        ## Store training data
                        X_t = np.array(np.concatenate([X[s] for j, s in enumerate(split) if j != i]), order='C',
                                       dtype=np.float32)
                        F_t = np.array(np.concatenate([F[s] for j, s in enumerate(split) if j != i]), order='C',
                                       dtype=np.float32)
                        E_t = np.array(np.concatenate([E[s] for j, s in enumerate(split) if j != i]), order='C',
                                       dtype=np.float64)

                        if E_t.shape[0] != 0:
                            t.insertdata(X_t, F_t, E_t, list(S))

                        ## Split test/valid data and store\
                        #tv_split = np.array_split(split[i], 2)

                        ## Store Validation
                        if np.array(split[i]).size > 0:
                            X_v = np.array(X[split[i]], order='C', dtype=np.float32)
                            F_v = np.array(F[split[i]], order='C', dtype=np.float32)
                            E_v = np.array(E[split[i]], order='C', dtype=np.float64)
                            if E_v.shape[0] != 0:
                                v.insertdata(X_v, F_v, E_v, list(S))

                        ## Store testset
                        #if tv_split[1].size > 0:
                            #X_te = np.array(X[split[i]], order='C', dtype=np.float32)
                            #F_te = np.array(F[split[i]], order='C', dtype=np.float32)
                            #E_te = np.array(E[split[i]], order='C', dtype=np.float64)
                            #if E_te.shape[0] != 0:
                            #    te.store_data(Pn, coordinates=X_te, forces=F_te, energies=E_te, species=list(S))

            #sys.stdout.write("\r%d%%" % int(100))
            #print(" Data Kept: ", Ndc, 'High Force: ', Nbf)
            #sys.stdout.flush()
            #print("")

        # Print some stats
        print('Data count:', Nd)
        print('Data split:', 100.0 * Nd / np.sum(Nd), '%')

        # Save train and valid meta file and cleanup testh5
        for t, v, th in zip(cachet, cachev, testh5):
            t.makemetadata()
            v.makemetadata()
            th.cleanup()

    def sae_linear_fitting(self, Ekey='energies', energy_unit=1.0, Eax0sum=False):
        from sklearn import linear_model
        print('Performing linear fitting...')

        datadir = self.h5dir
        sae_out = self.netdict['saefile']

        smap = dict()
        for i,Z in enumerate(self.netdict['atomtyp']):
            smap.update({Z:i})

        Na = len(smap)
        files = os.listdir(datadir)

        X = []
        y = []
        for f in files[0:20]:
            print(f)
            adl = pyt.anidataloader(datadir + f)
            for data in adl:
                # print(data['path'])
                S = data['species']

                if data[Ekey].size > 0:
                    if Eax0sum:
                        E = energy_unit*np.sum(np.array(data[Ekey], order='C', dtype=np.float64), axis=1)
                    else:
                        E = energy_unit*np.array(data[Ekey], order='C', dtype=np.float64)

                    S = S[0:data['coordinates'].shape[1]]
                    unique, counts = np.unique(S, return_counts=True)
                    x = np.zeros(Na, dtype=np.float64)
                    for u, c in zip(unique, counts):
                        x[smap[u]] = c

                    for e in E:
                        X.append(np.array(x))
                        y.append(np.array(e))

        X = np.array(X)
        y = np.array(y).reshape(-1, 1)

        lin = linear_model.LinearRegression(fit_intercept=False)
        lin.fit(X, y)

        coef = lin.coef_
        print(coef)

        sae = open(sae_out, 'w')
        for i, c in enumerate(coef[0]):
            sae.write(next(key for key, value in smap.items() if value == i) + ',' + str(i) + '=' + str(c) + '\n')

        sae.close()

        print('Linear fitting complete.')

    def build_strided_training_cache(self,Nblocks,Nvalid,Ntest,build_test=True, build_valid=False, forces=True, grad=False, Fkey='forces', forces_unit=1.0, Ekey='energies', energy_unit=1.0, Eax0sum=False, rmhighe=True):
        if not os.path.isfile(self.netdict['saefile']):
            self.sae_linear_fitting(Ekey=Ekey, energy_unit=energy_unit, Eax0sum=Eax0sum)

        h5d = self.h5dir

        store_dir = self.train_root + "cache-data-"
        N = self.Nn
        Ntrain = Nblocks - Nvalid - Ntest

        if Nblocks % N != 0:
            raise ValueError('Error: number of networks must evenly divide number of blocks.')

        Nstride = Nblocks/N

        for i in range(N):
            if not os.path.exists(store_dir + str(i)):
                os.mkdir(store_dir + str(i))

            if build_test:
                if os.path.exists(store_dir + str(i) + '/../testset/testset' + str(i) + '.h5'):
                    os.remove(store_dir + str(i) + '/../testset/testset' + str(i) + '.h5')

                if not os.path.exists(store_dir + str(i) + '/../testset'):
                    os.mkdir(store_dir + str(i) + '/../testset')

        cachet = [cg('_train', self.netdict['saefile'], store_dir + str(r) + '/', False) for r in range(N)]
        cachev = [cg('_valid', self.netdict['saefile'], store_dir + str(r) + '/', False) for r in range(N)]

        if build_test:
            testh5 = [pyt.datapacker(store_dir + str(r) + '/../testset/testset' + str(r) + '.h5') for r in range(N)]

        if build_valid:
            valdh5 = [pyt.datapacker(store_dir + str(r) + '/../testset/valdset' + str(r) + '.h5') for r in range(N)]

        if rmhighe:
            dE = []
            for f in self.h5file:
                adl = pyt.anidataloader(h5d+f)
                for data in adl:
                    S = data['species']
                    E = data['energies']
                    X = data['coordinates']

                    Esae = hdt.compute_sae(self.netdict['saefile'], S)

                    dE.append((E-Esae)/np.sqrt(len(S)))

            dE = np.concatenate(dE)
            cidx = np.where(np.abs(dE) < 15.0)
            std = np.abs(dE[cidx]).std()
            men = np.mean(dE[cidx])

            print(men,std,men+std)
            idx = np.intersect1d(np.where(dE>=-np.abs(15*std+men))[0],np.where(dE<=np.abs(11*std+men))[0])
            cnt = idx.size
            print('DATADIST: ',dE.size,cnt,(dE.size-cnt),100.0*((dE.size-cnt)/dE.size))

        E = []
        data_count = np.zeros((N,3),dtype=np.int32)
        for f in self.h5file:
            print('Reading data file:',h5d+f)
            adl = pyt.anidataloader(h5d+f)
            for data in adl:
                #print(data['path'],data['energies'].size)

                S = data['species']

                if data[Ekey].size > 0 and (set(S).issubset(self.netdict['atomtyp'])):

                    X = np.array(data['coordinates'], order='C',dtype=np.float32)

                    #print(np.array(data[Ekey].shape),np.sum(np.array(data[Ekey], order='C', dtype=np.float64),axis=1).shape,data[Fkey].shape)

                    if Eax0sum:
                        E = energy_unit*np.sum(np.array(data[Ekey], order='C', dtype=np.float64),axis=1)
                    else:
                        E = energy_unit*np.array(data[Ekey], order='C',dtype=np.float64)

                    if forces and not grad:
                        F = forces_unit*np.array(data[Fkey], order='C', dtype=np.float32)
                    elif forces and grad:
                        F = -forces_unit*np.array(data[Fkey], order='C', dtype=np.float32)
                    else:
                        F = 0.0*X

                    if rmhighe:
                        Esae = hdt.compute_sae(self.netdict['saefile'], S)

                        ind_dE = (E - Esae)/np.sqrt(len(S))

                        hidx = np.union1d(np.where(ind_dE<-(15.0*std+men))[0],np.where(ind_dE>(11.0*std+men))[0])
                        lidx = np.intersect1d(np.where(ind_dE>=-(15.0*std+men))[0],np.where(ind_dE<=(11.0*std+men))[0])

                        if hidx.size > 0:
                            print('  -(' + f + ':' + data['path'] + ')High energies detected:\n    ', (E[hidx]-Esae)/np.sqrt(len(S)))

                        X = X[lidx]
                        E = E[lidx]
                        F = F[lidx]

                    # Build random split index
                    ridx = np.random.randint(0,Nblocks,size=E.size)
                    Didx = [np.argsort(ridx)[np.where(ridx == i)] for i in range(Nblocks)]

                    # Build training cache
                    for nid,cache in enumerate(cachet):
                        set_idx = np.concatenate([Didx[((bid+nid*int(Nstride)) % Nblocks)] for bid in range(Ntrain)])
                        if set_idx.size != 0:
                            data_count[nid,0]+=set_idx.size
                            cache.insertdata(X[set_idx], F[set_idx], E[set_idx], list(S))

                    # for nid,cache in enumerate(cachev):
                    #     set_idx = np.concatenate([Didx[((1+bid+nid*int(Nstride)) % Nblocks)] for bid in range(Ntrain)])
                    #     if set_idx.size != 0:
                    #         data_count[nid,0]+=set_idx.size
                    #         cache.insertdata(X[set_idx], F[set_idx], E[set_idx], list(S))

                    for nid,cache in enumerate(cachev):
                        set_idx = np.concatenate([Didx[(Ntrain+bid+nid*int(Nstride)) % Nblocks] for bid in range(Nvalid)])
                        if set_idx.size != 0:
                            data_count[nid, 1] += set_idx.size
                            cache.insertdata(X[set_idx], F[set_idx], E[set_idx], list(S))
                            if build_valid:
                                valdh5[nid].store_data(f+data['path'], coordinates=X[set_idx], forces=F[set_idx], energies=E[set_idx], species=list(S))

                    if build_test:
                        for nid,th5 in enumerate(testh5):
                            set_idx = np.concatenate([Didx[(Ntrain+Nvalid+bid+nid*int(Nstride)) % Nblocks] for bid in range(Ntest)])
                            if set_idx.size != 0:
                                data_count[nid, 2] += set_idx.size
                                th5.store_data(f+data['path'], coordinates=X[set_idx], forces=F[set_idx], energies=E[set_idx], species=list(S))

        # Save train and valid meta file and cleanup testh5
        for t, v in zip(cachet, cachev):
            t.makemetadata()
            v.makemetadata()

        if build_test:
            for th in testh5:
                th.cleanup()

        if build_valid:
            for vh in valdh5:
                vh.cleanup()

        print(' Train ',' Valid ',' Test ')
        print(data_count)
        print('Training set built.')

    def train_ensemble(self, GPUList, remove_existing=False):
        print('Training Ensemble...')
        processes = []
        indicies = np.array_split(np.arange(self.Nn), len(GPUList))

        for gpu,idc in enumerate(indicies):
            processes.append(Process(target=self.train_network, args=(GPUList[gpu], idc, remove_existing)))
            processes[-1].start()
            #self.train_network(pyncdict, trdict, layers, id, i)

        for p in processes:
            p.join()
        print('Training Complete.')

    def train_network(self, gpuid, indicies, remove_existing=False):
        for index in indicies:
            pyncdict = dict()
            pyncdict['wkdir'] = self.train_root + 'train' + str(index) + '/'
            pyncdict['ntwkStoreDir'] = self.train_root + 'train' + str(index) + '/' + 'networks/'
            pyncdict['datadir'] = self.train_root + "cache-data-" + str(index) + '/'
            pyncdict['gpuid'] = str(gpuid)

            if not os.path.exists(pyncdict['wkdir']):
                os.mkdir(pyncdict['wkdir'])

            if remove_existing:
                shutil.rmtree(pyncdict['ntwkStoreDir'])

            if not os.path.exists(pyncdict['ntwkStoreDir']):
                os.mkdir(pyncdict['ntwkStoreDir'])

            outputfile = pyncdict['wkdir']+'output.opt'

            shutil.copy2(self.netdict['iptfile'], pyncdict['wkdir'])
            shutil.copy2(self.netdict['cnstfile'], pyncdict['wkdir'])
            shutil.copy2(self.netdict['saefile'], pyncdict['wkdir'])

            if "/" in self.netdict['iptfile']:
                nfile = self.netdict['iptfile'].rsplit("/",1)[1]
            else:
                nfile = self.netdict['iptfile']

            command = "cd " + pyncdict['wkdir'] + " && HDAtomNNP-Trainer -i " + nfile + " -d " + pyncdict['datadir'] + " -p 1.0 -m -g " + pyncdict['gpuid'] + " > output.opt"
            proc = subprocess.Popen (command, shell=True)
            proc.communicate()

            print('  -Model',index,'complete')

    def get_train_stats(self):
        #rerr = re.compile('EPOCH\s+?(\d+?)\n[\s\S]+?E \(kcal\/mol\)\s+?(\d+?\.\d+?)\s+?(\d+?\.\d+?)\s+?(\d+?\.\d+?)\n\s+?dE \(kcal\/mol\)\s+?(\d+?\.\d+?)\s+?(\d+?\.\d+?)\s+?(\d+?\.\d+?)\n[\s\S]+?Current best:\s+?(\d+?)\n[\s\S]+?Learning Rate:\s+?(\S+?)\n[\s\S]+?TotalEpoch:\s+([\s\S]+?)\n')
        #rerr = re.compile('EPOCH\s+?(\d+?)\s+?\n[\s\S]+?E \(kcal\/mol\)\s+?(\S+?)\s+?(\S+?)\s+?(\S+?)\n\s+?dE \(kcal\/mol\)\s+?(\S+?)\s+?(\S+?)\s+?(\S+?)\n')
        rblk = re.compile('=+?\n([\s\S]+?=+?\n[\s\S]+?(?:=|Deleting))')
        repo = re.compile('EPOCH\s+?(\d+?)\s+?\n')
        rerr = re.compile('\s+?(\S+?\s+?\(\S+?)\s+?((?:\d|inf)\S*?)\s+?((?:\d|inf)\S*?)\s+?((?:\d|inf)\S*?)\n')
        rtme = re.compile('TotalEpoch:\s+?(\d+?)\s+?dy\.\s+?(\d+?)\s+?hr\.\s+?(\d+?)\s+?mn\.\s+?(\d+?\.\d+?)\s+?sc\.')

        allnets = []
        for index in range(self.Nn):
            print('reading:',self.train_root + 'train' + str(index) + '/' + 'output.opt')
            optfile = open(self.train_root + 'train' + str(index) + '/' + 'output.opt','r').read()
            matches = re.findall(rblk, optfile)

            run = dict({'EPOCH':[],'RTIME':[],'ERROR':dict()})
            for i,data in enumerate(matches):
                run['EPOCH'].append(int(re.search(repo,data).group(1)))

                m = re.search(rtme, data)
                run['RTIME'].append(86400.0*float(m.group(1))+
                                     3600.0*float(m.group(2))+
                                       60.0*float(m.group(3))+
                                            float(m.group(4)))

                err = re.findall(rerr,data)
                for e in err:
                    if e[0] in run['ERROR']:
                        run['ERROR'][e[0]].append(np.array([float(e[1]),float(e[2]),float(e[3])],dtype=np.float64))
                    else:
                        run['ERROR'].update({e[0]:[np.array([float(e[1]), float(e[2]), float(e[3])], dtype=np.float64)]})

            for key in run['ERROR'].keys():
                run['ERROR'][key] = np.vstack(run['ERROR'][key])

            allnets.append(run)
        return allnets


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2021 The TARTRL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

""""""

import sys
from PIL import Image
import numpy as np
import torch

from dfd.utils import check_file
from dfd.params import padding_image, img_mean, img_std, resize, DeepFakeModel
from dfd.timm.models import create_deepfake_model_v4

def test_img(model_path, img_files):
    assert all(check_file(img_file) for img_file in img_files), 'file not exist!'

    use_cuda = True
    use_half = True
    print('To load model from {}'.format(model_path))
    model = create_deepfake_model_v4(
        'efficientnet_deepfake_v4',
        num_classes=2,
        in_chans=12,
        checkpoint_path=model_path,
        strict=False)
    print('Model loaded!')
    model = DeepFakeModel(model)
    if use_cuda:
        model.cuda()
        if use_half:
            model.half()
    model.eval()
    for img_file in img_files:
        img = np.transpose(padding_image(resize(np.array(Image.open(img_file).convert('RGB'), np.uint8))), (2, 0, 1))
        img = torch.from_numpy(img).float()
        img = img.sub_(img_mean).div_(img_std)
        if use_cuda:
            img = img.cuda()
            if use_half:
                img = img.half()
        img = [torch.cat([img, img.clone(), img.clone(), img.clone()], dim=0)]
        with torch.no_grad():
            scores = model(torch.stack(img, dim=0))
        scores = scores.cpu().numpy()[:, 0].tolist()
        print('{}\'s fake score:{}'.format(img_file, scores[0]))

if __name__ == '__main__':
    model_path = '../models/model_half.pth.tar'
    if len(sys.argv) <= 1:
        print('Please input your images. e.g. python test_images.py image_path1 image_path2')
        exit()
    img_files = sys.argv[1:]
    test_img(model_path, img_files)


import copy

from .models import User, Friend, Recommendation, LogEntry
from collections import defaultdict
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
import datetime
from . import utils
from . import models
import logging
import numpy as np
from django.db.models import Q

logger = logging.getLogger(__name__)

IN_PERSON_REC_DESC = (
    "In-person interactions allow for more control "
    "over the interaction, like body language "
    "and other visual cues. Online and over-the-phone "
    "interactions should be used once a good relationship "
    "is already built."
)

GENERALLY_POSITIVE_REC_DESC = (
    "Looks like your interactions with "
    "this friend are typically positive - "
    "try to have more of them!"
)

DIFFERENT_TIME_REC_DESC = (
    "It's okay to be tired! "
    "Talk to this friend when you feel more awake, "
    "like after your morning coffee."
)

GENERALLY_NEGATIVE_REC_DESC = (
    "If this friend makes you feel angry, "
    "it's a good idea to re-evaluate your "
    "relationship or the context around "
    "interactions that make you angry."
)

AVOID_REC_DESC = (
    "People that make you feel sad "
    "shouldn't be in your life, "
    "it's that simple!"
)

SPREAD_INTERACTIONS_REC_DESC = (
    "Concentrating all of your interactions "
    "in one part of the day makes it easier "
    "for negative ones to ruin your day "
    "and harder ones to positively influence your "
    "day. Remember: some people need their "
    "coffee in the morning!"
)

AVOID_SMALL_REC_DESC = (
    "Studies have shown that having more small talk "
    "than not has a negative impact on happiness "
    "and interaction quality. Focus on having more "
    "substantive conversations, especially with this "
    "friend."
)

LAUGH_REC_DESC = (
    "Starting a conversation with an authentic friendly "
    "smile creates a friendly and open look that invites "
    "contact. If the other person smiles back, "
    "then you've already started a form of contact!"
)

COMPL_REC_DESC = (
    "Starting a conversation with a compliment "
    "is usually more interesting than  a remark "
    "about the weather. "
    "Pick any feature that actually interests you, "
    "otherwise it won't come off as genuine. "
    "Feel free to continue on the topic if they show "
    "interest!"
)

BODY_LANGUAGE_REC_DESC = (
    "You can create an 'understanding' "
    "without using words by mirroring "
    "the other person's body attitude. "
    "This typically makes relating much "
    "easier."
)

EVALUATE_REC_DESC = (
    "Remember to evaluate reactions from "
    "the other person and regarding verbal "
    "or non-verbal communication. "
    "You should also express what "
    "you have understood from their behavior."
)

HUMOR_REC_DESC = (
    "Humor creates an open and conversational "
    "atmosphere, but remember not to joke at "
    "the expense of other people and respect "
    "any other sentiments. "
)

OPEN_UP_REC_DESC = (
    "You're not a reporter! The more you show "
    "yourself, the more interested they become. "
    "A good conversation involves being interested "
    "in the other person as well as opening up. "
    "Opening up makes it easier for others to "
    "relate to you, but make sure you're authentic."
)

ML_POS_REC_DESC = (
    "You should have more positive interactions "
    "with this person. Your past successes have "
    "shown that this person is someone who you may "
    "get along with. Factors that we looked at include {}"
)

ML_NEG_REC_DESC = (
    "Your interaction experience with this person is "
    "mostly negative, and we really don't think this will "
    "improve. We recommend interacting with this person on "
    "a need to know and surface level basis to avoid discomfort. "
    "The factors that we looked into while making this decision "
    "include {}"
)

ML_AVOID_REC_DESC = (
    "This person looks a lot like people whom you avoid in "
    "in your past. Our recommendation is to minimize interactions "
    "interactions entirely with this person. Get on a need "
    "to know basis with this person and stay there. "
    "Factors that we looked into while making this decision "
    "include {}"
)

def _create_and_save_recommendation(rec, friend_id=None):
    r = Recommendation()
    r.rec_typ = rec['rec_type']
    r.recommendation = rec['recommendation']
    r.rec_description = rec['rec_description']
    r.recommend_person_id = rec['recommend_person']

    if friend_id:
        r.about_person = Friend.objects.get(id=friend_id)
    r.save()
    return r

def _check_in_person_recommendations(logs, user_id):
    in_person = len([None for log in logs if log.interaction_medium == 'IP'])
    if len(logs) != 0:
        in_person_frac = in_person / len(logs) * 100
    else:
        in_person_frac = 0

    if in_person_frac < 50:
        rec = dict(
            rec_type = 'PO',
            recommendation = "Try to have more in-person interactions.",
            rec_description = IN_PERSON_REC_DESC,
            recommend_person = user_id,
        )

        return _create_and_save_recommendation(rec)
    return None

def _check_time_fairness(logs, user_id):
    morning = 0
    afternoon = 0
    evening = 0

    for log in logs:
        if log.time_of_day == 'MO':
            morning += 1
        elif log.time_of_day == 'AF':
            afternoon += 1
        elif log.time_of_day == 'EV':
            evening += 1

    if len(logs) != 0:
        morning_frac = morning / len(logs) * 100
        afternoon_frac = afternoon / len(logs) * 100
        evening_frac = evening / len(logs) *  100
    else:
        morning_frac = 0
        afternoon_frac = 0
        evening_frac = 0

    if morning_frac > 50 or afternoon_frac > 50 or evening_frac > 50:
        rec = dict(
            rec_type = 'PO',
            recommendation = "Spread your interactions throughout the day.",
            rec_description = SPREAD_INTERACTIONS_REC_DESC,
            recommend_person = user_id,
        )
        return _create_and_save_recommendation(rec)
    return None

def _count_all(logs):
    # Takes an arr and a list of selectors and counts the keys
    # with a special "total" key that aggreages the row
    reactions = defaultdict(lambda: defaultdict(lambda: 0))

    for log in logs:
        reactions[log.loggee][log.reaction] += 1
        reactions[log.loggee][log.time_of_day] += 1
        reactions[log.loggee][log.social_context] += 1
        reactions[log.loggee][log.interaction_medium] += 1

    return reactions

def _count_reactions(logs):
    # id -> reaction -> count
    reactions = defaultdict(lambda: defaultdict(lambda: 0))
    reacc_dict = dict(LogEntry.REACTION_CHOICES)

    for log in logs:
        reactions[log.loggee][reacc_dict[log.reaction]] += 1

    return dict(reactions)

def _save_generic_recommendations(logs, user_id):
    rec_type = 'GE'

    user = User.objects.get(id=user_id)
    generic_recs = []

    laugh_rec = dict(
        recommendation = "Laugh and smile when starting an interaction.",
        rec_description = LAUGH_REC_DESC,
        recommend_person = user_id,
        rec_type = rec_type
        )

    generic_recs.append(_create_and_save_recommendation(laugh_rec))

    compl_rec = dict(
        recommendation = "Start with a compliment",
        rec_description = COMPL_REC_DESC,
        recommend_person = user_id,
        rec_type = rec_type
        )

    generic_recs.append(_create_and_save_recommendation(compl_rec))

    body_rec = dict(
        recommendation = "Use body language.",
        rec_description = BODY_LANGUAGE_REC_DESC,
        recommend_person = user_id,
        rec_type = rec_type
        )

    generic_recs.append(_create_and_save_recommendation(body_rec))

    eval_rec = dict(
        recommendation = "Evaluate and Understand",
        rec_description = EVALUATE_REC_DESC,
        recommend_person = user_id,
        rec_type = rec_type
        )

    generic_recs.append(_create_and_save_recommendation(eval_rec))

    humor_rec = dict(
        recommendation = "Use humor!",
        rec_description = HUMOR_REC_DESC,
        recommend_person = user_id,
        rec_type = rec_type
        )

    generic_recs.append(_create_and_save_recommendation(humor_rec))

    open_rec = dict(
        recommendation = "Open up.",
        rec_description = OPEN_UP_REC_DESC,
        recommend_person = user_id,
        rec_type = rec_type
        )

    generic_recs.append(_create_and_save_recommendation(open_rec))

    user.has_generic_recs = True
    user.save()
    return generic_recs

def _count_small_talk(logs):
    small_talk = defaultdict(lambda: {'small': 0, 'long': 0}) # key is friend ID
    count_small_talk = 0

    for log in logs:
        if log.content_class == 'SM':
            count_small_talk += 1
            # TODO: need to handle null case
            small_talk[log.loggee]['small'] += 1
        else:
            small_talk[log.loggee]['long'] += 1
    return small_talk, count_small_talk

def _recommended_this_week(user_id):
    base = Recommendation.objects.filter(recommend_person_id=user_id)
    sun = utils.last_sunday()
    base = base.filter(created_at__gt=sun)
    return base.all()

def recommendations_from_logs(logs, user_id):
    # Don't regenerate recommendations if we already have recommendations
    has_generated = _recommended_this_week(user_id)
    if len(has_generated) != 0:
        return has_generated

    rec_list = list()

    user = User.objects.get(id=user_id)
    # First store some generic recommendations
    if not user.has_generic_recs:
        recs = _save_generic_recommendations(logs, user_id)
        rec_list.extend(recs)

    # Check if they have the right amount of in person
    # interactions
    rec = _check_in_person_recommendations(logs, user_id)
    if rec:
        rec_list.append(rec)

    # Make sure interactions are not a majority in any
    # category
    rec = _check_time_fairness(logs, user_id)
    if rec:
        rec_list.append(rec)

    reactions = _count_reactions(logs)

    small_talk, count_small_talk = _count_small_talk(logs)

    for friend in reactions:
        if friend is None:
            continue
        name = friend.name
        total = sum(reactions[friend].values())

        r = dict(recommend_person= user_id,)

        if (reactions[friend]['Happy']/total * 100) > 50:
            r['rec_type'] = 'PO'
            r["recommendation"] = "Spend more time with {}.".format(name)
            r["rec_description"] = GENERALLY_POSITIVE_REC_DESC
            rec = _create_and_save_recommendation(r, friend_id=friend.id)
            rec_list.append(rec)
        elif (reactions[friend]['Tired']/total * 100) > 10:
            r['rec_type'] = 'PO'
            r["recommendation"] = "Talk to {} at different times.".format(name)
            r["rec_description"] = DIFFERENT_TIME_REC_DESC
            rec = _create_and_save_recommendation(r, friend_id=friend.id)
            rec_list.append(rec)
        elif (reactions[friend]['Angry']/total * 100) > 10:
            r['rec_type'] = 'NE'
            r["recommendation"] = "Talk to {} less.".format(name)
            r["rec_description"] = GENERALLY_NEGATIVE_REC_DESC
            rec = _create_and_save_recommendation(r, friend_id=friend.id)
            rec_list.append(rec)
        elif (reactions[friend]['Sad']/total * 100) > 10:
            r['rec_type'] = 'AV'
            r["recommendation"] = "Avoid talking to {}.".format(name)
            r["rec_description"] = AVOID_REC_DESC
            rec = _create_and_save_recommendation(r, friend_id=friend.id)
            rec_list.append(rec)

    for friend in small_talk:
        if friend is None:
            continue
        name = friend.name
        small_frac = small_talk[friend]['small']/(small_talk[friend]['small'] + small_talk[friend]['long']) * 100

        if small_frac > 50:
            small_talk_rec = dict(
                recommend_person=user_id,
                rec_type='AV',
                recommendation="Avoid small talk with {}.".format(name),
                rec_description=AVOID_SMALL_REC_DESC,
                )

            rec = _create_and_save_recommendation(small_talk_rec, friend_id=friend.id)
            rec_list.append(rec)
    return rec_list

def group_list_by_sel(lis, sel):
    ret = defaultdict(lambda: [])
    for li in lis:
        key = sel(li)
        ret[key].append(li)
    return dict(ret)

standard_date_format = "%y-%m-%d"

def logs_by_week(logs):
    selector = lambda x: utils.round_to_sun(x.created_at).strftime(standard_date_format)
    return group_list_by_sel(logs, selector)

def recs_by_week(user_id):
    recs = Recommendation.objects.filter(recommend_person_id=user_id).filter(~Q(rec_typ = 'GE'))
    selector = lambda x: utils.round_to_sun(x.created_at).strftime(standard_date_format)
    return group_list_by_sel(recs, selector)

def check_recs_now(recs_week):
    ls = utils.last_sunday()
    most_recent_key = ls.strftime(standard_date_format)
    return recs_week.get(most_recent_key)

def _aggregate_and_normalize(diction, keys):
    ret = []
    for key in keys:
        ret.append(diction[key])

    tot = sum(ret)
    if tot == 0:
        return ret
    for i in range(len(ret)):
        ret[i] /= tot
    return ret

def rolling_disposition_by_friend(logs_week, recs_week):
    reacc_keys = ['HA', 'NE', 'AN', "SA",]
    social_keys = ['AC', 'SO', 'WO',]
    medium_keys = ['IP', 'ON', 'PH']

    english_keys = [
        'Happy Interaction Percentage',
        'Negative Interaction Percentage',
        'Angry Interaction Percentage',
        'Sad Interaction Percentage',
        'Academic Interaction Percentage',
        'Social Interaction Percentage',
        'Work Interaction Percentage',
        'In-Person Interaction Percentage',
        'Online Interaction Percentage',
        'Phone Interaction Percentage',
    ]

    rec_arr = [x[0] for x in models.Recommendation.RECOMMENDATION_CHOICES]
    feed_arr = [x[0] for x in models.RecommendationFeedback.FEEDBACK_CHOICES]
    Xs = []
    Xs_neg = []
    ys = []
    ys_neg = []
    ordered_keys = list(sorted(logs_week.keys()))
    for i in range(len(ordered_keys) - 1):
        key = ordered_keys[i]
        rec = recs_week.get(key)
        if not rec:
            continue

        # Get all the logs that are earlier
        preceding_logs = []
        for k, v in logs_week.items():
            if k > key:
                continue
            preceding_logs.extend(v)

        counts = _count_all(preceding_logs)
        for person, count in counts.items():
            rec_obj = next((x for x in rec if person != None and x.about_person_id == person.id), None)
            if not rec_obj:
                continue

            feeds = list(rec_obj.tied_recommendation.all())
            feed_typ = next((x for x in feeds if x.feedback_typ != ''), None)
            if not feed_typ:
                continue

            row = []
            row.extend(_aggregate_and_normalize(count, reacc_keys))
            row.extend(_aggregate_and_normalize(count, social_keys))
            row.extend(_aggregate_and_normalize(count, medium_keys))
            if feed_typ.feedback_typ == 'WO':
                Xs.append(row)
                ys.append(rec_arr.index(rec_obj.rec_typ))
            else:
                Xs_neg.append(row)
                ys_neg.append(rec_arr.index(rec_obj.rec_typ))

        last_week = []
        all_logs = []
        for v in logs_week.values():
            all_logs.extend(v)
        counts = _count_all(all_logs)

        for person, count in counts.items():
            if person == None:
                continue
            row = []
            # Row and then date
            row.extend(_aggregate_and_normalize(count, reacc_keys))
            row.extend(_aggregate_and_normalize(count, social_keys))
            row.extend(_aggregate_and_normalize(count, medium_keys))
            last_week.append((person.id, row))

        return Xs, ys, Xs_neg, ys_neg, last_week, english_keys

def generate_ml_recommendations(logs_week, recs_week):

    # Create a cumulative rolling sum of the data by person
    pos_classifier = DecisionTreeClassifier(criterion='entropy', splitter='random', random_state=104)
    Xs, ys, _, _, lw, ek = rolling_disposition_by_friend(logs_week, recs_week)
    pos_classifier.fit(Xs, ys)

    feats = [x[-1] for x in lw]

    pos_preds = pos_classifier.predict(feats)
    paths = pos_classifier.decision_path(feats)

    ret = []
    for pos_pred, path, (person_id, _) in zip(pos_preds, paths, lw):
        idxes = path.nonzero()[-1]
        english_labels = [ek[i] for i in idxes][:3]
        if len(english_labels) == 1:
            english_label = english_label[0]
        elif len(english_labels) == 2:
            english_label = "{} and {}".format(*english_labels)
        else:
            english_label = "{}, {}, and {}".format(*english_labels)
        rec_enum = Recommendation.RECOMMENDATION_CHOICES[pos_pred][0]
        ret.append((rec_enum, english_label, person_id))

    return ret

def recommendations_from_ml(logs, user_id, from_dt=None, to_dt=None):
    # All the setup recommendations should be good
    logs_by = logs_by_week(logs)
    recs_week = recs_by_week(user_id)

    has_recs = check_recs_now(recs_week)
    rec_list = []
    if has_recs:
        return has_recs
    ml_recs = generate_ml_recommendations(logs_by, recs_week)
    for rec_enum, english_label, person_id in ml_recs:
        person = models.Friend.objects.filter(id=person_id).first()

        # We don't give generic recommendation
        if rec_enum == 'PO':
            positive_ml_rec = dict(
                recommend_person=user_id,
                rec_type=rec_enum,
                recommendation="Try acting more positive with {}".format(person.name),
                rec_description=ML_POS_REC_DESC.format(english_label),
                )
            rec = _create_and_save_recommendation(positive_ml_rec, friend_id=person_id)
            rec_list.append(rec)

        elif rec_enum == 'NE':
            neg_ml_rec = dict(
                recommend_person=user_id,
                rec_type=rec_enum,
                recommendation="Try a surface-level relationship with {}".format(person.name),
                rec_description=ML_NEG_REC_DESC.format(english_label),
                )
            rec = _create_and_save_recommendation(neg_ml_rec, friend_id=person_id)
            rec_list.append(rec)
        else:
            avoid_ml_rec = dict(
                recommend_person=user_id,
                rec_type=rec_enum,
                recommendation="Try avoiding {}".format(person.name),
                rec_description=ML_AVOID_REC_DESC.format(english_label),
                )
            rec = _create_and_save_recommendation(avoid_ml_rec, friend_id=person_id)
            rec_list.append(rec)

    return rec_list


"""
Drawing commands that use vertex buffer objects (VBOs).

This module contains commands for basic graphics drawing commands,
but uses Vertex Buffer Objects. This keeps the vertices loaded on
the graphics card for much faster render times.
"""

import math
import itertools
from collections import defaultdict
import pyglet.gl as gl
import numpy as np

from typing import List, Iterable, Sequence
from typing import TypeVar
from typing import Generic
from typing import cast

from arcade import Color
from arcade import rotate_point
from arcade import Point, PointList
from arcade import get_four_byte_color
from arcade import get_window
from arcade import get_points_for_thick_line
from arcade.gl import BufferDescription

class Shape:
    """
    Primitive drawing shape. This can be part of a ShapeElementList so
    shapes can be drawn faster in batch.
    """
    def __init__(self):
        self.vao = None
        self.vbo = None
        self.program = None
        self.mode = None
        self.line_width = 1

    def draw(self):
        """
        Draw this shape. Drawing this way isn't as fast as drawing multiple
        shapes batched together in a ShapeElementList.
        """
        assert(self.line_width == 1)
        gl.glLineWidth(self.line_width)

        gl.glEnable(gl.GL_BLEND)
        gl.glBlendFunc(gl.GL_SRC_ALPHA, gl.GL_ONE_MINUS_SRC_ALPHA)
        gl.glEnable(gl.GL_LINE_SMOOTH)
        gl.glHint(gl.GL_LINE_SMOOTH_HINT, gl.GL_NICEST)
        gl.glHint(gl.GL_POLYGON_SMOOTH_HINT, gl.GL_NICEST)
        gl.glEnable(gl.GL_PRIMITIVE_RESTART)
        gl.glPrimitiveRestartIndex(2 ** 32 - 1)

        self.vao.render(self.program, mode=self.mode)

def create_line(start_x: float, start_y: float, end_x: float, end_y: float,
                color: Color, line_width: float = 1) -> Shape:
    """
    Create a line to be rendered later. This works faster than draw_line because
    the vertexes are only loaded to the graphics card once, rather than each frame.

    :param float start_x:
    :param float start_y:
    :param float end_x:
    :param float end_y:
    :param Color color:
    :param float line_width:

    :Returns Shape:

    """

    points = get_points_for_thick_line(start_x, start_y, end_x, end_y, line_width)
    color_list = [color, color, color, color]
    triangle_point_list = points[1], points[0], points[2], points[3]
    shape = create_triangles_filled_with_colors(triangle_point_list, color_list)
    return shape

def create_line_generic_with_colors(point_list: PointList,
                                    color_list: Iterable[Color],
                                    shape_mode: int,
                                    line_width: float = 1) -> Shape:
    """
    This function is used by ``create_line_strip`` and ``create_line_loop``,
    just changing the OpenGL type for the line drawing.

    :param PointList point_list:
    :param Iterable[Color] color_list:
    :param float shape_mode:
    :param float line_width:

    :Returns Shape:
    """
    window = get_window()
    ctx = window.ctx
    program = ctx.line_generic_with_colors_program

    buffer_type = np.dtype([('vertex', '2f4'), ('color', '4B')])
    data = np.zeros(len(point_list), dtype=buffer_type)
    data['vertex'] = point_list
    data['color'] = [get_four_byte_color(color) for color in color_list]

    vbo = ctx.buffer(data=data.tobytes())
    vao_content = [
        BufferDescription(
            vbo,
            '2f 4f1',
            ('in_vert', 'in_color'),
            normalized=['in_color']
        )
    ]

    vao = ctx.geometry(vao_content)

    shape = Shape()
    shape.vao = vao
    shape.vbo = vbo
    shape.program = program
    shape.mode = shape_mode
    shape.line_width = line_width

    return shape

def create_line_generic(point_list: PointList,
                        color: Color,
                        shape_mode: int, line_width: float = 1) -> Shape:
    """
    This function is used by ``create_line_strip`` and ``create_line_loop``,
    just changing the OpenGL type for the line drawing.
    """
    colors = [get_four_byte_color(color)] * len(point_list)
    shape = create_line_generic_with_colors(
        point_list,
        colors,
        shape_mode,
        line_width)

    return shape

def create_line_strip(point_list: PointList,
                      color: Color, line_width: float = 1):
    """
    Create a multi-point line to be rendered later. This works faster than draw_line because
    the vertexes are only loaded to the graphics card once, rather than each frame.

    Internally, thick lines are created by two triangles.

    :param PointList point_list:
    :param Color color:
    :param PointList line_width:

    :Returns Shape:

    """
    if line_width == 1:
        return create_line_generic(point_list, color, gl.GL_LINE_STRIP, line_width)
    else:
        triangle_point_list: List[Point] = []
        new_color_list: List[Color] = []
        for i in range(1, len(point_list)):
            start_x = point_list[i - 1][0]
            start_y = point_list[i - 1][1]
            end_x = point_list[i][0]
            end_y = point_list[i][1]
            color1 = color
            color2 = color
            points = get_points_for_thick_line(start_x, start_y, end_x, end_y, line_width)
            new_color_list += color1, color2, color1, color2
            triangle_point_list += points[1], points[0], points[2], points[3]

        shape = create_triangles_filled_with_colors(triangle_point_list, new_color_list)
        return shape

def create_line_loop(point_list: PointList,
                     color: Color, line_width: float = 1):
    """
    Create a multi-point line loop to be rendered later. This works faster than draw_line because
    the vertexes are only loaded to the graphics card once, rather than each frame.

    :param PointList point_list:
    :param Color color:
    :param float line_width:

    :Returns Shape:

    """
    point_list = list(point_list) + [point_list[0]]
    return create_line_strip(point_list, color, line_width)

def create_lines(point_list: PointList,
                 color: Color, line_width: float = 1):
    """
    Create a multi-point line loop to be rendered later. This works faster than draw_line because
    the vertexes are only loaded to the graphics card once, rather than each frame.

    :param PointList point_list:
    :param Color color:
    :param float line_width:

    :Returns Shape:

    """
    return create_line_generic(point_list, color, gl.GL_LINES, line_width)

def create_lines_with_colors(point_list: PointList,
                             color_list: Sequence[Color],
                             line_width: float = 1):

    if line_width == 1:
        return create_line_generic_with_colors(point_list, color_list, gl.GL_LINES, line_width)
    else:

        triangle_point_list: List[Point] = []
        new_color_list: List[Color] = []
        for i in range(1, len(point_list), 2):
            start_x = point_list[i-1][0]
            start_y = point_list[i-1][1]
            end_x = point_list[i][0]
            end_y = point_list[i][1]
            color1 = color_list[i-1]
            color2 = color_list[i]
            points = get_points_for_thick_line(start_x, start_y, end_x, end_y, line_width)
            new_color_list += color1, color1, color2, color2
            triangle_point_list += points[1], points[0], points[2], points[3]

            shape = create_triangles_filled_with_colors(triangle_point_list, new_color_list)
            return shape

def create_polygon(point_list: PointList,
                   color: Color):
    """
    Draw a convex polygon. This will NOT draw a concave polygon.
    Because of this, you might not want to use this function.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    :param PointList point_list:
    :param color:

    :Returns Shape:

    """
    # We assume points were given in order, either clockwise or counter clockwise.
    # Polygon is assumed to be monotone.
    # To fill the polygon, we start by one vertex, and we chain triangle strips
    # alternating with vertices to the left and vertices to the right of the
    # initial vertex.
    half = len(point_list) // 2
    interleaved = itertools.chain.from_iterable(
        itertools.zip_longest(point_list[:half], reversed(point_list[half:]))
    )
    point_list = [p for p in interleaved if p is not None]
    return create_line_generic(point_list, color, gl.GL_TRIANGLE_STRIP, 1)

def create_rectangle_filled(center_x: float, center_y: float, width: float,
                            height: float, color: Color,
                            tilt_angle: float = 0) -> Shape:
    """
    Create a filled rectangle.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    :param float center_x:
    :param float center_y:
    :param float width:
    :param float height:
    :param Color color:
    :param float tilt_angle:

    :Returns Shape:

    """
    return create_rectangle(center_x, center_y, width, height,
                            color, tilt_angle=tilt_angle)

def create_rectangle_outline(center_x: float, center_y: float, width: float,
                             height: float, color: Color,
                             border_width: float = 1, tilt_angle: float = 0) -> Shape:
    """
    Create a rectangle outline.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    Args:
        center_x:
        center_y:
        width:
        height:
        color:
        border_width:
        tilt_angle:

    Returns:

    """
    return create_rectangle(center_x, center_y, width, height,
                            color, border_width, tilt_angle, filled=False)

def get_rectangle_points(center_x: float, center_y: float, width: float,
                         height: float, tilt_angle: float = 0) -> PointList:
    """
    Utility function that will return all four coordinate points of a
    rectangle given the x, y center, width, height, and rotation.

    Args:
        center_x:
        center_y:
        width:
        height:
        tilt_angle:

    Returns:

    """
    x1 = -width / 2 + center_x
    y1 = -height / 2 + center_y

    x2 = -width / 2 + center_x
    y2 = height / 2 + center_y

    x3 = width / 2 + center_x
    y3 = height / 2 + center_y

    x4 = width / 2 + center_x
    y4 = -height / 2 + center_y

    if tilt_angle:
        x1, y1 = rotate_point(x1, y1, center_x, center_y, tilt_angle)
        x2, y2 = rotate_point(x2, y2, center_x, center_y, tilt_angle)
        x3, y3 = rotate_point(x3, y3, center_x, center_y, tilt_angle)
        x4, y4 = rotate_point(x4, y4, center_x, center_y, tilt_angle)

    data = [(x1, y1),
            (x2, y2),
            (x3, y3),
            (x4, y4)]

    return data

def create_rectangle(center_x: float, center_y: float, width: float,
                     height: float, color: Color,
                     border_width: float = 1, tilt_angle: float = 0,
                     filled=True) -> Shape:
    """
    This function creates a rectangle using a vertex buffer object.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    :param float center_x:
    :param float center_y:
    :param float width:
    :param float height:
    :param Color color:
    :param float border_width:
    :param float tilt_angle:
    :param bool filled:

    """
    data: List[Point] = cast(List[Point], get_rectangle_points(center_x, center_y, width, height, tilt_angle))

    if filled:
        shape_mode = gl.GL_TRIANGLE_STRIP
        data[-2:] = reversed(data[-2:])
    else:

        i_lb = center_x - width / 2 + border_width / 2, center_y - height / 2 + border_width / 2
        i_rb = center_x + width / 2 - border_width / 2, center_y - height / 2 + border_width / 2
        i_rt = center_x + width / 2 - border_width / 2, center_y + height / 2 - border_width / 2
        i_lt = center_x - width / 2 + border_width / 2, center_y + height / 2 - border_width / 2

        o_lb = center_x - width / 2 - border_width / 2, center_y - height / 2 - border_width / 2
        o_rb = center_x + width / 2 + border_width / 2, center_y - height / 2 - border_width / 2
        o_rt = center_x + width / 2 + border_width / 2, center_y + height / 2 + border_width / 2
        o_lt = center_x - width / 2 - border_width / 2, center_y + height / 2 + border_width / 2

        data = [o_lt, i_lt, o_rt, i_rt, o_rb, i_rb, o_lb, i_lb, o_lt, i_lt]

        if tilt_angle != 0:
            point_list_2: List[Point] = []
            for point in data:
                new_point = rotate_point(point[0], point[1], center_x, center_y, tilt_angle)
                point_list_2.append(new_point)
            data = point_list_2

        border_width = 1
        shape_mode = gl.GL_TRIANGLE_STRIP

        # _generic_draw_line_strip(point_list, color, gl.GL_TRIANGLE_STRIP)

        # shape_mode = gl.GL_LINE_STRIP
        # data.append(data[0])

    shape = create_line_generic(data, color, shape_mode, border_width)
    return shape

def create_rectangle_filled_with_colors(point_list, color_list) -> Shape:
    """
    This function creates one rectangle/quad using a vertex buffer object.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    """

    shape_mode = gl.GL_TRIANGLE_STRIP
    new_point_list = [point_list[0], point_list[1], point_list[3], point_list[2]]
    new_color_list = [color_list[0], color_list[1], color_list[3], color_list[2]]
    return create_line_generic_with_colors(new_point_list, new_color_list, shape_mode)

def create_rectangles_filled_with_colors(point_list, color_list) -> Shape:
    """
    This function creates multiple rectangle/quads using a vertex buffer object.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    """

    shape_mode = gl.GL_TRIANGLES
    new_point_list: List[Point] = []
    new_color_list: List[Color] = []
    for i in range(0, len(point_list), 4):
        new_point_list += [point_list[0 + i], point_list[1 + i], point_list[3 + i]]
        new_point_list += [point_list[1 + i], point_list[3 + i], point_list[2 + i]]

        new_color_list += [color_list[0 + i], color_list[1 + i], color_list[3 + i]]
        new_color_list += [color_list[1 + i], color_list[3 + i], color_list[2 + i]]

    return create_line_generic_with_colors(new_point_list, new_color_list, shape_mode)

def create_triangles_filled_with_colors(point_list, color_list) -> Shape:
    """
    This function creates multiple rectangle/quads using a vertex buffer object.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    """

    shape_mode = gl.GL_TRIANGLE_STRIP
    return create_line_generic_with_colors(point_list, color_list, shape_mode)

def create_ellipse_filled(center_x: float, center_y: float,
                          width: float, height: float, color: Color,
                          tilt_angle: float = 0, num_segments: int = 128) -> Shape:
    """
    Create a filled ellipse. Or circle if you use the same width and height.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    """

    border_width = 1
    return create_ellipse(center_x, center_y, width, height, color,
                          border_width, tilt_angle, num_segments, filled=True)

def create_ellipse_outline(center_x: float, center_y: float,
                           width: float, height: float, color: Color,
                           border_width: float = 1,
                           tilt_angle: float = 0, num_segments: int = 128) -> Shape:
    """
    Create an outline of an ellipse.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    """

    return create_ellipse(center_x, center_y, width, height, color,
                          border_width, tilt_angle, num_segments, filled=False)

def create_ellipse(center_x: float, center_y: float,
                   width: float, height: float, color: Color,
                   border_width: float = 1,
                   tilt_angle: float = 0, num_segments: int = 32,
                   filled=True) -> Shape:

    """
    This creates an ellipse vertex buffer object (VBO).

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    """
    # Create an array with the vertex point_list
    point_list = []

    for segment in range(num_segments):
        theta = 2.0 * 3.1415926 * segment / num_segments

        x = width * math.cos(theta) + center_x
        y = height * math.sin(theta) + center_y

        if tilt_angle:
            x, y = rotate_point(x, y, center_x, center_y, tilt_angle)

        point_list.append((x, y))

    if filled:
        half = len(point_list) // 2
        interleaved = itertools.chain.from_iterable(
            itertools.zip_longest(point_list[:half], reversed(point_list[half:]))
        )
        point_list = [p for p in interleaved if p is not None]
        shape_mode = gl.GL_TRIANGLE_STRIP
    else:
        point_list.append(point_list[0])
        shape_mode = gl.GL_LINE_STRIP

    return create_line_generic(point_list, color, shape_mode, border_width)

def create_ellipse_filled_with_colors(center_x: float, center_y: float,
                                      width: float, height: float,
                                      outside_color: Color, inside_color: Color,
                                      tilt_angle: float = 0, num_segments: int = 32) -> Shape:
    """
    Draw an ellipse, and specify inside/outside color. Used for doing gradients.

    The function returns a Shape object that can be drawn with ``my_shape.draw()``.
    Don't create the shape in the draw method, create it in the setup method and then
    draw it in ``on_draw``.

    For even faster performance, add multiple shapes into a ShapeElementList and
    draw that list. This allows nearly unlimited shapes to be drawn just as fast
    as one.

    :param float center_x:
    :param float center_y:
    :param float width:
    :param float height:
    :param Color outside_color:
    :param float inside_color:
    :param float tilt_angle:
    :param int num_segments:

    :Returns Shape:

    """

    # Create an array with the vertex data
    # Create an array with the vertex point_list
    point_list = [(center_x, center_y)]

    for segment in range(num_segments):
        theta = 2.0 * 3.1415926 * segment / num_segments

        x = width * math.cos(theta) + center_x
        y = height * math.sin(theta) + center_y

        if tilt_angle:
            x, y = rotate_point(x, y, center_x, center_y, tilt_angle)

        point_list.append((x, y))
    point_list.append(point_list[1])

    color_list = [inside_color] + [outside_color] * (num_segments + 1)
    return create_line_generic_with_colors(point_list, color_list, gl.GL_TRIANGLE_FAN)

TShape = TypeVar('TShape', bound=Shape)

class ShapeElementList(Generic[TShape]):
    """
    A program can put multiple drawing primitives in a ShapeElementList, and then
    move and draw them as one. Do this when you want to create a more complex object
    out of simpler primitives. This also speeds rendering as all objects are drawn
    in one operation.
    """
    def __init__(self):
        """
        Initialize the sprite list
        """
        # The context this shape list belongs to
        self.ctx = get_window().ctx
        # List of sprites in the sprite list
        self.shape_list = []
        self.change_x = 0
        self.change_y = 0
        self._center_x = 0
        self._center_y = 0
        self._angle = 0
        self.program = self.ctx.shape_element_list_program
        # Could do much better using just one vbo and glDrawElementsBaseVertex
        self.batches = defaultdict(_Batch)
        self.dirties = set()

    def append(self, item: TShape):
        """
        Add a new shape to the list.
        """
        self.shape_list.append(item)
        group = (item.mode, item.line_width)
        self.batches[group].items.append(item)
        self.dirties.add(group)

    def remove(self, item: TShape):
        """
        Remove a specific shape from the list.
        """
        self.shape_list.remove(item)
        group = (item.mode, item.line_width)
        self.batches[group].items.remove(item)
        self.dirties.add(group)

    def _refresh_shape(self, group):
        # Create a buffer large enough to hold all the shapes buffers
        batch = self.batches[group]

        total_vbo_bytes = sum(s.vbo.size for s in batch.items)
        vbo = self.ctx.buffer(reserve=total_vbo_bytes)
        offset = 0
        # Copy all the shapes buffer in our own vbo
        for shape in batch.items:
            vbo.copy_from_buffer(shape.vbo, offset=offset)
            offset += shape.vbo.size

        # Create an index buffer object. It should count starting from 0. We need to
        # use a reset_idx to indicate that a new shape will start.
        reset_idx = 2 ** 32 - 1
        indices = []
        counter = itertools.count()
        for shape in batch.items:
            indices.extend(itertools.islice(counter, shape.vao.num_vertices))
            indices.append(reset_idx)
        del indices[-1]
        indices = np.array(indices)
        ibo = self.ctx.buffer(data=indices.astype('i4').tobytes())

        vao_content = [
            BufferDescription(
                vbo,
                '2f 4f1',
                ('in_vert', 'in_color'),
                normalized=['in_color']
            )
        ]
        vao = self.ctx.geometry(vao_content, ibo)
        self.program['Position'] = [self.center_x, self.center_y]
        self.program['Angle'] = self.angle

        batch.shape.vao = vao
        batch.shape.vbo = vbo
        batch.shape.ibo = ibo
        batch.shape.program = self.program
        mode, line_width = group
        batch.shape.mode = mode
        batch.shape.line_width = line_width

    def move(self, change_x: float, change_y: float):
        """
        Move all the shapes ion the list
        :param change_x: Amount to move on the x axis
        :param change_y: Amount to move on the y axis
        """
        self.center_x += change_x
        self.center_y += change_y

    def __len__(self) -> int:
        """ Return the length of the sprite list. """
        return len(self.shape_list)

    def __iter__(self) -> Iterable[TShape]:
        """ Return an iterable object of sprites. """
        return iter(self.shape_list)

    def __getitem__(self, i):
        return self.shape_list[i]

    def draw(self):
        """
        Draw everything in the list.
        """
        self.program['Position'] = [self._center_x, self._center_y]

        for group in self.dirties:
            self._refresh_shape(group)
        self.dirties.clear()
        for batch in self.batches.values():
            batch.shape.draw()

    def _get_center_x(self) -> float:
        """Get the center x coordinate of the ShapeElementList."""
        return self._center_x

    def _set_center_x(self, value: float):
        """Set the center x coordinate of the ShapeElementList."""
        self._center_x = value

    center_x = property(_get_center_x, _set_center_x)

    def _get_center_y(self) -> float:
        """Get the center y coordinate of the ShapeElementList."""
        return self._center_y

    def _set_center_y(self, value: float):
        """Set the center y coordinate of the ShapeElementList."""
        self._center_y = value
        self.program['Position'] = [self._center_x, self._center_y]

    center_y = property(_get_center_y, _set_center_y)

    def _get_angle(self) -> float:
        """Get the angle of the ShapeElementList in degrees."""
        return self._angle

    def _set_angle(self, value: float):
        """Set the angle of the ShapeElementList in degrees."""
        self._angle = value
        self.program['Angle'] = self._angle

    angle = property(_get_angle, _set_angle)

class _Batch(Generic[TShape]):
    def __init__(self):
        self.shape = Shape()
        self.items = []