import gvar as gv
import corrfitter as cf
import numpy as np
import collections
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.ticker import MultipleLocator
matplotlib.use('Agg')
plt.rc("font",**{"size":18})
import datetime
import os
import pickle
import copy
#from plotting import *
import lsqfit
lsqfit.nonlinear_fit.set(fitter='gsl_multifit',alg='subspace2D',scaler='more',solver='cholesky')#,solver='cholesky')
####################################
maxiter=5000
#######################################################################################################

def read_setup(setup):
    #Reads in setups, and strips out currents, parents and daughters, as well as which is which
    daughters = []
    currents = []
    parents = []
    for element in setup:
        lab = element.split('-')
        daughters.append(lab[0])
        currents.append(lab[1])
        parents.append(lab[2])
    return(daughters,currents,parents)
######################################################################################################

def strip_list(l): #Strips elemenst from list l
    stripped = ''
    for element in l:
        stripped = '{0}{1}'.format(stripped,element)
    return(stripped)

######################################################################################################

def make_params(Fit,FitMasses,FitTwists,FitTs,daughters,currents,parents):
    #Removes things we do not want to fit, specified by FitMasses, FitTwists, FitTs assumes parents have varing mass and daughters varing twist
    j = 0
    for i in range(len(Fit['masses'])):
        if i not in FitMasses:
            del Fit['masses'][i-j]
            for element in set(parents):
                del Fit['tmaxes{0}'.format(element)][i-j]
            j += 1
    j = 0
    for i in range(len(Fit['twists'])):
        if i not in FitTwists:
            del Fit['twists'][i-j]
            for element in set(daughters):
                del Fit['tmaxes{0}'.format(element)][i-j]
            j += 1
    j = 0
    for i in range(len(Fit['Ts'])):
        if i not in FitTs:
            del Fit['Ts'][i-j]
            j += 1
    return()

#######################################################################################################

def make_data(filename,binsize):
    # Reads in filename.gpl, checks all keys have same configuration numbers, returns averaged data
    print('Reading data, binsize = ', binsize)
    dset = cf.read_dataset(filename,binsize=binsize)
    sizes = []
    for key in dset:
        #print(key,np.shape(dset[key]))
        sizes.append(np.shape(dset[key]))
    if len(set(sizes)) != 1:
        print('Not all elements of gpl the same size')
        for key in dset:
            print(key,np.shape(dset[key]))
    return(gv.dataset.avg_data(dset))

######################################################################################################

def make_pdata(filename,models,binsize):
    # Reads in filename.gpl, checks all keys have same configuration numbers, returns averaged data
    print('Reading processed data, binsize = ', binsize)
    dset = cf.read_dataset(filename,binsize=binsize)
    sizes = []
    for key in dset:
        #print(key,np.shape(dset[key]))
        sizes.append(np.shape(dset[key]))
    if len(set(sizes)) != 1:
        print('Not all elements of gpl the same size')
        for key in dset:
            print(key,np.shape(dset[key]))
    return(cf.process_dataset(dset, models))

#######################################################################################################

def effective_mass_calc(tag,correlator,tp):
    #finds the effective mass and amplitude of a two point correlator
    M_effs = []
    for t in range(2,len(correlator)-2):
        thing  = (correlator[t-2] + correlator[t+2])/(2*correlator[t])
        if thing >= 1:
            M_effs.append(gv.arccosh(thing)/2)
    #M_effs is all positive masses, we now take a rolling average of 4, and find where this changes the least
    rav = []
    for i in range(len(M_effs)-4):
        rav.append((M_effs[i] + M_effs[i+1] + M_effs[i+2] + M_effs[i+3])/4)
    M_eff = rav[0]
    diff = abs((rav[1] - rav[0]).mean)
    for i in range(1,len(rav)-1):
        if abs((rav[i+1]-rav[i]).mean) < diff:
            diff = abs((rav[i+1]-rav[i]).mean)
            M_eff = (rav[i] + rav[i+1])/2
    return(M_eff)

######################################################################################################

def effective_amplitude_calc(tag,correlator,tp,M_eff,Fit,corr):
    #finds the effective mass and amplitude of a two point correlator
    tmin = Fit['tmin{0}'.format(corr)]
    A_effs = []
    if len(correlator) == tp:
        tmin = 0
    for t in range(tmin,tmin+len(correlator)):
        numerator = correlator[t-tmin]
        if numerator >= 0:
            A_effs.append( gv.sqrt(numerator/(gv.exp(-M_eff*t)+gv.exp(-M_eff*(tp-t)))))
    rav = []
    for i in range(len(A_effs)-4):
        rav.append((A_effs[i] + A_effs[i+1] + A_effs[i+2] + A_effs[i+3])/4)
    A_eff = rav[0]
    diff = abs((rav[1] - rav[0]).mean)
    for i in range(1,len(rav)-1):
        if abs((rav[i+1]-rav[i]).mean) < diff:
            diff = abs((rav[i+1]-rav[i]).mean)
            A_eff = (rav[i] + rav[i+1])/2
    an = gv.gvar(Fit['an'])
    if A_eff.sdev/A_eff.mean > 0.5:
        print('Replaced A_eff for {0} {1} -> {2}'.format(tag,A_eff,an))
        A_eff = an
    return(A_eff)

########################################################################################

def effective_V_calc(corr,daughter,parent,correlator,dcorr,pcorr,Fit,mass,twist,pA_eff,dA_eff):
    #finds the effective V_nn[0][0]
    tp = Fit['tp']
    T = Fit['Ts'][-1]
    dtmin = Fit['tmin{0}'.format(daughter)]
    ptmin = Fit['tmin{0}'.format(parent)]
    Vtmin = Fit['{0}tmin'.format(corr)]
    dcorr2 = []
    pcorr2 = []
    Vcorr2 = []
    V_effs = []

    #print(corr,daughter,parent,mass,twist)
    if len(dcorr) == int(tp):
        dcorr2 = dcorr
    else:
        for i in range(dtmin):
            dcorr2.append(0)
        dcorr2.extend(dcorr)
        for i in range(int(tp/2)-len(dcorr2)+1):
            dcorr2.append(0)
    #print(dcorr2)
    if len(pcorr) == int(tp):
        pcorr2 = pcorr
    else:
        for i in range(ptmin):
            pcorr2.append(0)
        pcorr2.extend(pcorr)
        for i in range(int(tp/2)-len(pcorr2)+1):
            pcorr2.append(0)
    #print(pcorr2)
    if len(correlator) == int(tp):
        Vcorr2 = correlator
    else:
        for i in range(Vtmin):
            Vcorr2.append(0)
        Vcorr2.extend(correlator)
        for i in range(T-len(Vcorr2)+1):
            Vcorr2.append(0)
    #print(Vcorr2)

    for t in range(T):
        numerator = Vcorr2[t]*pA_eff*dA_eff
        denominator = dcorr2[t]*pcorr2[T-t]
        if numerator != 0 and denominator !=0:
            V_effs.append(numerator/denominator)
    rav = []
    for i in range(len(V_effs)-4):
        rav.append((V_effs[i] + V_effs[i+1] + V_effs[i+2] + V_effs[i+3])/4)
    V_eff = rav[0]
    diff = abs((rav[1] - rav[0]).mean)
    for i in range(1,len(rav)-1):
        if abs((rav[i+1]-rav[i]).mean) < diff:
            diff = abs((rav[i+1]-rav[i]).mean)
            if (rav[i] + rav[i+1]) > 0:
                V_eff = (rav[i] + rav[i+1])/2
    V = gv.gvar(Fit['{0}Vnn0'.format(corr)])
    if abs((V_eff.mean-V).mean/(V_eff.mean-V).sdev) > 1:
        print('Replaced V_eff for {0} m {1} tw {2}: {3} --> {4}'.format(corr,mass,twist,V_eff,V))
        V_eff = V
    return(V_eff)

#######################################################################################################

def SVD_diagnosis(Fit,models,corrs,svdfac,currents,SepMass):
    binsize = Fit['binsize']
    #Feed models and corrs (list of corrs in this SVD cut)
    if list(set(corrs).intersection(currents)) ==[]:
        filename = 'SVD/{0}{1}{2}{3}{4}{5}{6}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(corrs),binsize,SepMass)
    else:
        filename = 'SVD/{0}{1}{2}{3}{4}{5}{6}{7}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(corrs),strip_list(Fit['Ts']),binsize,SepMass)
    #print(filename)
    for corr in corrs:
       if 'tmin{0}'.format(corr) in Fit:
           filename += '{0}'.format(Fit['tmin{0}'.format(corr)])
           for element in Fit['tmaxes{0}'.format(corr)]:
               filename += '{0}'.format(element)
       if '{0}tmin'.format(corr) in Fit:
           filename += '{0}'.format(Fit['{0}tmin'.format(corr)])

    #print(filename)
    if os.path.isfile(filename) and os.path.getsize(filename) > 0:
        pickle_off = open(filename,"rb")
        svd = pickle.load(pickle_off)
        print('Loaded SVD for {0} : {1:.2g} x {2} = {3:.2g}'.format(corrs,svd,svdfac,svd*svdfac))
        pickle_off.close()
    else:
        print('Calculating SVD for {0}'.format(corrs))
        s = gv.dataset.svd_diagnosis(cf.read_dataset('{0}{1}.gpl'.format(Fit['file_location'],Fit['filename']),binsize=binsize), models=models, nbstrap=20)
        svd = s.svdcut
        ######## save plot ##########################
        plt.figure()
        x = s.val / s.val[-1]
        ratio = s.bsval / s.val
        idx = x > s.mincut
        ratio = ratio[idx]
        x = x[idx]
        y = gv.mean(ratio)
        yerr = gv.sdev(ratio)
        plt.errorbar(x=x, y=y, yerr=yerr, fmt='+', color='b')
        sig = (2. / len(s.val)) ** 0.5
        plt.plot([x[0], x[-1]], [1. - sig, 1. - sig], 'k:')
        plt.axhline(1,ls='--',color='k')
        plt.axvline(s.svdcut,ls=':',color='g')
        #plt.axvline(0.013,ls='--',color='g')
        plt.xscale('log')
        plt.savefig('svd_plots/{0}.pdf'.format(filename.split('/')[1]))
        ###############################################
        pickle_on = open(filename,"wb")
        print('Calculated SVD for {0} : {1:.2g} x {2} = {3:.2g}'.format(corrs,svd,svdfac,svd*svdfac))
        pickle.dump(svd,pickle_on)
    return(svd*svdfac)

#######################################################################################################

def make_models(Fit,FitCorrs,notwist0,non_oscillating,daughters,currents,parents,svdfac,Chained,allcorrs,links,parrlinks,SepMass,NoSVD=False):
    #several forms [(A,B,C,D)],[(A,B),(C),(D)],[(A,B),[(C),(D)]]
    #First make all models and then stick them into the correct chain
    models = collections.OrderedDict()
    tp = Fit['tp']
    for corr in set(parents):
        if corr in allcorrs:
            models['{0}'.format(corr)] = []
            for i,mass in enumerate(Fit['masses']):
                tag = Fit['{0}-Tag'.format(corr)].format(mass)
                models['{0}'.format(corr)].append(cf.Corr2(datatag=tag, tp=tp, tmin=Fit['tmin{0}'.format(corr)], tmax=Fit['tmaxes{0}'.format(corr)][i], a=('{0}:a'.format(tag), 'o{0}:a'.format(tag)), b=('{0}:a'.format(tag), 'o{0}:a'.format(tag)), dE=('dE:{0}'.format(tag), 'dE:o{0}'.format(tag)),s=(1,-1)))

    for corr in set(daughters):
        if corr in allcorrs:
            models['{0}'.format(corr)] = []
            for i,twist in enumerate(Fit['twists']):
                tag = Fit['{0}-Tag'.format(corr)].format(twist)
                if twist == '0' and corr in notwist0:
                    pass
                elif twist == '0' and corr in non_oscillating:
                    models['{0}'.format(corr)].append(cf.Corr2(datatag=tag, tp=tp, tmin=Fit['tmin{0}'.format(corr)], tmax=Fit['tmaxes{0}'.format(corr)][i], a=('{0}:a'.format(tag)), b=('{0}:a'.format(tag)), dE=('dE:{0}'.format(tag))))
                else:
                    models['{0}'.format(corr)].append(cf.Corr2(datatag=tag, tp=tp, tmin=Fit['tmin{0}'.format(corr)], tmax=Fit['tmaxes{0}'.format(corr)][i], a=('{0}:a'.format(tag), 'o{0}:a'.format(tag)), b=('{0}:a'.format(tag), 'o{0}:a'.format(tag)), dE=('dE:{0}'.format(tag), 'dE:o{0}'.format(tag)),s=(1,-1)))

    for i,corr in enumerate(currents):
        if corr in allcorrs:
            models['{0}'.format(corr)] = []
            for  mass in Fit['masses']:
                for twist in Fit['twists']:
                    for T in Fit['Ts']:
                        tag = Fit['threePtTag{0}'.format(corr)].format(T,Fit['m_s'],mass,Fit['m_l'],twist)
                        ptag = Fit['{0}-Tag'.format(parents[i])].format(mass)
                        dtag = Fit['{0}-Tag'.format(daughters[i])].format(twist)
                        if twist == '0' and corr in notwist0:
                            pass
                        elif twist == '0' and daughters[i] in non_oscillating:
                            models['{0}'.format(corr)].append(cf.Corr3(datatag=tag, T=T, tmin=Fit['{0}tmin'.format(corr)], a=('{0}:a'.format(dtag)), dEa=('dE:{0}'.format(dtag)), b=('{0}:a'.format(ptag), 'o{0}:a'.format(ptag)), dEb=('dE:{0}'.format(ptag), 'dE:o{0}'.format(ptag)), sb=(1,-1), Vnn='{0}Vnn_m{1}_tw{2}'.format(corr,mass,twist), Vno='{0}Vno_m{1}_tw{2}'.format(corr,mass,twist)))

                        else:
                            models['{0}'.format(corr)].append(cf.Corr3(datatag=tag, T=T, tmin=Fit['{0}tmin'.format(corr)], a=('{0}:a'.format(dtag), 'o{0}:a'.format(dtag)), dEa=('dE:{0}'.format(dtag), 'dE:o{0}'.format(dtag)), sa=(1,-1), b=('{0}:a'.format(ptag), 'o{0}:a'.format(ptag)), dEb=('dE:{0}'.format(ptag), 'dE:o{0}'.format(ptag)), sb=(1,-1), Vnn='{0}Vnn_m{1}_tw{2}'.format(corr,mass,twist), Vno='{0}Vno_m{1}_tw{2}'.format(corr,mass,twist),Von='{0}Von_m{1}_tw{2}'.format(corr,mass,twist),Voo='{0}Voo_m{1}_tw{2}'.format(corr,mass,twist)))
    #Now we make these models into our chain calculating an svd cut for each. We make them in two halves so we can sndwich a marginalisation term if we like later

    if Chained:
        finalmodelsA = []
        finalmodelsB = []
        intermediate = []
        for key in links:
            link = [] #link is models in link
            for corr in links[key]:
                link.extend(models['{0}'.format(corr)])
            svd = SVD_diagnosis(Fit,link,links[key],svdfac,currents,SepMass)
            finalmodelsA.append({'svdcut':svd})
            finalmodelsA.append(tuple(link))
        for key in parrlinks:
            link = [] #link is models in link
            for corr in parrlinks[key]:
                link.extend(models['{0}'.format(corr)])
            svd = SVD_diagnosis(Fit,link,parrlinks[key],svdfac,currents,SepMass)
            intermediate.append({'svdcut':svd})
            intermediate.append(tuple(link))
        finalmodelsB.append(intermediate)
        return(finalmodelsA,finalmodelsB)
    else:
        finalmodels = []
        for corr in allcorrs:
            finalmodels.extend(models['{0}'.format(corr)])
        if NoSVD == False:
            svd = SVD_diagnosis(Fit,finalmodels,allcorrs,svdfac,currents,SepMass)
            return(tuple(finalmodels),svd)
        else:
            return(tuple(finalmodels))

#######################################################################################################

def elements_in_FitCorrs(a):
    # reads [A,[B,C],[[D,E],F]] and interprets which elements will be chained and how. Returns alphabetical list of all elements, links in chain and links in parallell chain
    allcorrs = []
    links = collections.OrderedDict()
    parrlinks = collections.OrderedDict()
    for i in range(np.shape(a)[0]):
        links[i] =[]
        if len(np.shape(a[i])) == 0: #deals with one corr in chain
            #print(a[i],i,'fit alone in chain')
            allcorrs.append(a[i])
            links[i].append(a[i])
        elif len(np.shape(a[i][0])) == 0 : #deals with multiple elements in chain
            for j in range(len(a[i])):
                #print(a[i][j],i,'fit together in chain')
                allcorrs.append(a[i][j])
                links[i].append(a[i][j])
        else:
            del links[i]  #don't need thi key if it is in paralell
            for j in range(np.shape(a[i])[0]):
                parrlinks[j] = []
                if len(np.shape(a[i][j])) == 0: #deals with one corr in parr chain
                    allcorrs.append(a[i][j])
                    parrlinks[j].append(a[i][j])
                else:                           # deals with multiple elements in parralell chain
                    for k in range(len(a[i][j])):
                        allcorrs.append(a[i][j][k])
                        parrlinks[j].append(a[i][j][k])
    return(sorted(allcorrs),links,parrlinks)

######################################################################################################

def make_prior(Fit,N,allcorrs,currents,daughters,parents,loosener,data,notwist0,non_oscillating):
    No = N  # number of oscillating exponentials
    prior =  gv.BufferDict()
    tw_corr = True
    otw_corr = True
    if len(daughters) != 0 and '0' in Fit['twists'] and tw_corr:
        for corr in set(daughters).intersection(allcorrs):
            prior['d2_{0}'.format(corr)] = gv.gvar('0.0(1.0)')
            prior['c2_{0}'.format(corr)] = gv.gvar('0.0(1.0)')
        print('Daughter twists correlated')
    if len(daughters) != 0 and '0' in Fit['twists'] and otw_corr:
        for corr in set(daughters).intersection(allcorrs):
            prior['oc2_{0}'.format(corr)] = gv.gvar('0.0(1.0)')
        print('Daughter oscillating twists correlated')

    tp = Fit['tp']
    En = '{0}({1})'.format(0.5*Fit['a'],0.25*Fit['a']*loosener) #Lambda with error of half
    an = '{0}({1})'.format(gv.gvar(Fit['an']).mean,gv.gvar(Fit['an']).sdev*loosener)
    aon = '{0}({1})'.format(gv.gvar(Fit['aon']).mean,gv.gvar(Fit['aon']).sdev*loosener)
    for corr in allcorrs:
        if corr in parents:
            for mass in Fit['masses']:
                tag = Fit['{0}-Tag'.format(corr)].format(mass)
                M_eff = effective_mass_calc(tag,data[tag],tp)
                a_eff = effective_amplitude_calc(tag,data[tag],tp,M_eff,Fit,corr)
                # Parent
                prior['log({0}:a)'.format(tag)] = gv.log(gv.gvar(N * [an]))
                prior['log(dE:{0})'.format(tag)] = gv.log(gv.gvar(N * [En]))
                prior['log({0}:a)'.format(tag)][0] = gv.log(gv.gvar(a_eff.mean,loosener*Fit['loosener']*a_eff.mean))
                prior['log(dE:{0})'.format(tag)][0] = gv.log(gv.gvar(M_eff.mean,loosener*Fit['Mloosener']*M_eff.mean))
                # Parent -- oscillating part
                prior['log(o{0}:a)'.format(tag)] = gv.log(gv.gvar(No * [an]))
                prior['log(dE:o{0})'.format(tag)] = gv.log(gv.gvar(No * [En]))
                prior['log(dE:o{0})'.format(tag)][0] = gv.log(gv.gvar((M_eff+gv.gvar(En)*(4/5)).mean,loosener*Fit['oMloosener']*((M_eff+gv.gvar(En)*(4/5)).mean)))

        if corr in daughters:
            for twist in Fit['twists']:
                if twist =='0' and corr in notwist0:
                    pass
                else:
                    ap2 = 3*(np.pi*float(twist)/Fit['L'])**2
                    #print(twist,ap2)
                    tag0 = Fit['{0}-Tag'.format(corr)].format('0')
                    M_eff = np.sqrt(effective_mass_calc(tag0,data[tag0],tp)**2 +  ap2)   #from dispersion relation
                    tag = Fit['{0}-Tag'.format(corr)].format(twist)
                    a_eff = effective_amplitude_calc(tag,data[tag],tp,M_eff,Fit,corr)
                    # Daughter
                    prior['log({0}:a)'.format(tag)] = gv.log(gv.gvar(N * [an]))
                    prior['log(dE:{0})'.format(tag)] = gv.log(gv.gvar(N * [En]))
                    #prior['log(dE:{0})'.format(tag)][1] = gv.log(gv.gvar(gv.gvar(En).mean,0.01*gv.gvar(En).mean))
                    if twist !='0' and '0' in Fit['twists'] and 'log(dE:{0})'.format(tag0) in prior and tw_corr:
                        prior['log(dE:{0})'.format(tag)][0] = gv.log(gv.sqrt(prior['dE:{0}'.format(tag0)][0]**2 + ap2) * (1 + prior['c2_{0}'.format(corr)]*ap2/(np.pi)**2) )
                        prior['log({0}:a)'.format(tag)][0] = gv.log((prior['{0}:a'.format(tag0)][0]/gv.sqrt(gv.sqrt(1 + ap2/(prior['dE:{0}'.format(tag0)][0])**2))) * (1 + prior['d2_{0}'.format(corr)]*ap2/(np.pi)**2) )

                    else:
                        prior['log(dE:{0})'.format(tag)][0] = gv.log(gv.gvar(M_eff.mean,loosener*Fit['Mloosener']*M_eff.mean))
                        prior['log({0}:a)'.format(tag)][0] = gv.log(gv.gvar(a_eff.mean,loosener*Fit['loosener']*a_eff.mean))
                    # Daughter -- oscillating part
                    if twist =='0' and corr in non_oscillating:
                        pass
                    else:
                        newaon = aon
                        if twist == '0':
                            newaon = '{0}({1})'.format(gv.gvar(aon).mean/4,gv.gvar(aon).mean/2) #v small in the case of tw0
                        prior['log(o{0}:a)'.format(tag)] = gv.log(gv.gvar(No * [newaon]))
                        prior['log(dE:o{0})'.format(tag)] = gv.log(gv.gvar(No * [En]))
                        if twist !='0' and '0' in Fit['twists'] and 'log(dE:o{0})'.format(tag0) in prior and otw_corr:
                            prior['log(dE:o{0})'.format(tag)][0] = gv.log(gv.sqrt(prior['dE:o{0}'.format(tag0)][0]**2 + ap2) * (1 + prior['oc2_{0}'.format(corr)]*ap2/(np.pi)**2) )
                            #prior['log(o{0}:a)'.format(tag)][0] = gv.log((prior['o{0}:a'.format(tag0)][0]/gv.sqrt(1 + ap2/(prior['dE:o{0}'.format(tag0)][0])**2)) * (1 + prior['od2']*ap2/(np.pi)**2) )
                            prior['log(o{0}:a)'.format(tag)][0] = gv.log(gv.gvar(gv.gvar(newaon).mean,loosener*Fit['oloosener']*gv.gvar(newaon).mean))
                        else:
                            prior['log(dE:o{0})'.format(tag)][0] = gv.log(gv.gvar((M_eff+gv.gvar(En)/2).mean,loosener*Fit['oMloosener']*((M_eff+gv.gvar(En)/2).mean))) # kaon splitting
                            #prior['log(dE:o{0})'.format(tag)][0] = gv.log(prior['dE:{0}'.format(tag)][0] + gv.gvar(En))
                            prior['log(o{0}:a)'.format(tag)][0] = gv.log(gv.gvar(gv.gvar(newaon).mean,loosener*Fit['oloosener']*gv.gvar(newaon).mean))

        if corr in currents:
            for mass in Fit['masses']:
                for twist in Fit['twists']:
                    if twist =='0' and corr in notwist0:
                        pass
                    else:
                        daughter=daughters[currents.index(corr)]
                        parent=parents[currents.index(corr)]
                        dcorr = data[Fit['{0}-Tag'.format(daughter)].format(twist)]
                        pcorr = data[Fit['{0}-Tag'.format(parent)].format(mass)]
                        correlator = data[Fit['threePtTag{0}'.format(corr)].format(Fit['Ts'][-1],Fit['m_s'],mass,Fit['m_l'],twist)]
                        ptag = Fit['{0}-Tag'.format(parent)].format(mass)
                        pM_eff = effective_mass_calc(ptag,data[ptag],tp)
                        pa_eff = effective_amplitude_calc(ptag,data[ptag],tp,pM_eff,Fit,parent)

                        dtag = Fit['{0}-Tag'.format(daughter)].format(twist)
                        dM_eff = effective_mass_calc(dtag,data[dtag],tp)
                        da_eff = effective_amplitude_calc(dtag,data[dtag],tp,dM_eff,Fit,daughter)

                        V_eff = effective_V_calc(corr,daughter,parent,correlator,dcorr,pcorr,Fit,mass,twist,da_eff,pa_eff)
                        if V_eff.mean != gv.gvar(Fit['{0}Vnn0'.format(corr)]).mean:
                            Vnn0 = '{0}({1})'.format(V_eff.mean,loosener*V_eff.mean*Fit['Vloosener'])
                        else:
                            Vnn0 = '{0}({1})'.format(V_eff.mean,loosener*V_eff.sdev)
                        Vn = '{0}({1})'.format(gv.gvar(Fit['{0}Vn'.format(corr)]).mean,loosener*gv.gvar(Fit['{0}Vn'.format(corr)]).sdev)
                        V0 = '{0}({1})'.format(gv.gvar(Fit['{0}V0'.format(corr)]).mean,loosener*gv.gvar(Fit['{0}V0'.format(corr)]).sdev)
                    if twist =='0' and corr in notwist0:
                        pass
                    elif twist =='0' and daughters[currents.index(corr)] in non_oscillating :
                        prior['{0}Vnn_m{1}_tw{2}'.format(corr,mass,twist)] = gv.gvar(N * [N * [Vn]])
                        prior['{0}Vnn_m{1}_tw{2}'.format(corr,mass,twist)][0][0] = gv.gvar(Vnn0)
                        prior['{0}Vno_m{1}_tw{2}'.format(corr,mass,twist)] = gv.gvar(N * [No* [Vn]])
                        prior['{0}Vno_m{1}_tw{2}'.format(corr,mass,twist)][0][0] = gv.gvar(V0)
                    else:
                        prior['{0}Vnn_m{1}_tw{2}'.format(corr,mass,twist)] = gv.gvar(N * [N * [Vn]])
                        prior['{0}Vnn_m{1}_tw{2}'.format(corr,mass,twist)][0][0] = gv.gvar(Vnn0)
                        prior['{0}Vno_m{1}_tw{2}'.format(corr,mass,twist)] = gv.gvar(N * [No * [Vn]])
                        prior['{0}Vno_m{1}_tw{2}'.format(corr,mass,twist)][0][0] = gv.gvar(V0)
                        prior['{0}Voo_m{1}_tw{2}'.format(corr,mass,twist)] = gv.gvar(No * [No * [Vn]])
                        prior['{0}Voo_m{1}_tw{2}'.format(corr,mass,twist)][0][0] = gv.gvar(V0)
                        prior['{0}Von_m{1}_tw{2}'.format(corr,mass,twist)] = gv.gvar(No * [N * [Vn]])
                        prior['{0}Von_m{1}_tw{2}'.format(corr,mass,twist)][0][0] = gv.gvar(V0)
           # for key in prior:
           #     if key[0] == corr:
           #         for i in range(1,N):
           #             for j in range(1,N):
           #                 prior[key][i][j] = gv.gvar('0.0(5)')

    return(prior)

######################################################################################################

def get_p0(Fit,fittype,Nexp,allcorrs,prior,FitCorrs):
    # We want to take in several scenarios in this order, choosing the highest in preference.
    # 1) This exact fit has been done before, modulo priors, svds t0s etc
    # 2) Same but different type of fit, eg marginalised
    # 3) This fit has been done before with Nexp+1
    # 4) This fit has been done beofore with Nexp-1
    # 5a) Some elemnts have bene fitted to Nexp before,
    # 5b) Some elements of the fit have been fitted in other combinations before

    filename1 = 'p0/{0}{1}{2}{3}{4}{5}{6}{7}{8}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),FitCorrs,strip_list(Fit['Ts']),fittype,Nexp)
    filename2 = 'p0/{0}{1}{2}{3}{4}{5}{6}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),strip_list(Fit['Ts']),Nexp)
    filename3 = 'p0/{0}{1}{2}{3}{4}{5}{6}{7}{8}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),FitCorrs,strip_list(Fit['Ts']),fittype,Nexp+1)
    filename4 = 'p0/{0}{1}{2}{3}{4}{5}{6}{7}{8}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),FitCorrs,strip_list(Fit['Ts']),fittype,Nexp-1)
    filename5a = 'p0/{0}{1}{2}'.format(Fit['conf'],Fit['filename'],Nexp)
    filename5b = 'p0/{0}{1}'.format(Fit['conf'],Fit['filename'])
    #case 1
    if os.path.isfile(filename1):
        p0 = gv.load(filename1)
        print('Loaded p0 from exact fit')
    #case 2

    elif os.path.isfile(filename2):
        p0 = gv.load(filename2)
        print('Loaded p0 from exact fit of different type')
    #case 3
    elif os.path.isfile(filename3):
        p0 = gv.load(filename3)
        print('Loaded p0 from exact fit Nexp+1')

    #case 4
    elif os.path.isfile(filename4):
        p0 = gv.load(filename4)
        print('Loaded p0 from exact fit Nexp-1')

    #case 5
    elif os.path.isfile(filename5b):
        p0 = gv.load(filename5b)
        print('Loaded global p0')
        if os.path.isfile(filename5a):
            pnexp = gv.load(filename5a)
            for key in pnexp:
                if key in prior:
                    if key not in p0:
                        print('Error: {0} in global Nexp but not in global fit'.format(key))
                        p0[key] = pnexp[key]
                    del p0[key]
                    p0[key] = pnexp[key]
                    print('Loaded {0} p0 from global Nexp'.format(key))

    else:
        p0 = None
    return(p0)
######################################################################################################

def update_p0(p,finalp,Fit,fittype,Nexp,allcorrs,FitCorrs,Q,marg=False):
    # We want to take in several scenarios in this order
    # 1) This exact fit has been done before, modulo priors, svds t0s etc
    # 2) Same but different type of fit, eg marginalised
    # 3) Global Nexp
    # 4) Global
    # 5) if Marg is True, we don't want to save anything but filename 1 as Nexp = nmarg and is not similar to if we do other fits
    filename1 = 'p0/{0}{1}{2}{3}{4}{5}{6}{7}{8}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),FitCorrs,strip_list(Fit['Ts']),fittype,Nexp)
    filename2 = 'p0/{0}{1}{2}{3}{4}{5}{6}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),strip_list(Fit['Ts']),Nexp)
    filename3 = 'p0/{0}{1}{2}'.format(Fit['conf'],Fit['filename'],Nexp)
    filename4 = 'p0/{0}{1}'.format(Fit['conf'],Fit['filename'])

    #case 1

    for element in ['c2','d2','oc2']:
        for corr in allcorrs:
            if '{0}_{1}'.format(element,corr) in p:
                del p['{0}_{1}'.format(element,corr)]

    for element in ['c2','d2','oc2']:
        for corr in allcorrs:
            if '{0}_{1}'.format(element,corr) in finalp:
                del finalp['{0}_{1}'.format(element,corr)]
    gv.dump(p,filename1)
    if marg == False:
        #case 2
        gv.dump(finalp,filename2)

        #case 3
        if os.path.isfile(filename3) and Q > 0.05:
            p0 = gv.load(filename3) #load exisiting global Nexp
            for key in finalp:  # key in this output
                p0[key] =  finalp[key]  #Update exisiting and add new
            gv.dump(p0,filename3)

        else:
            gv.dump(finalp,filename3)

        if os.path.isfile(filename4) and Q > 0.05:
            p0 = gv.load(filename4) # load existing, could be any length
            for key in finalp:  # key in new
                if key in p0: # if
                    if len(np.shape(p0[key])) == 1 and len(p0[key]) <= Nexp:
                        #print('shape p0[key]',np.shape(p0[key]),key)
                        del p0[key]
                        p0[key] = finalp[key]
                        print('Updated global p0 {0}'.format(key))
                    elif np.shape(p0[key])[0] <= Nexp:
                        #print('shape p0[key]',np.shape(p0[key]),key)
                        del p0[key]
                        p0[key] = finalp[key]
                        print('Updated global p0 {0}'.format(key))
                else:
                    p0[key] =  finalp[key]
                    print('Added new element to global p0 {0}'.format(key))
            gv.dump(p0,filename4)
        else:
            gv.dump(finalp,filename4)
    return()

######################################################################################################

def save_fit(fit,Fit,allcorrs,fittype,Nexp,SvdFactor,PriorLoosener,currents,smallsave):
    filename = 'Fits/{0}{1}{2}{3}{4}{5}{6}_Nexp{7}_sfac{8}_pfac{9}_Q{10:.2f}_chi{11:.3f}_sm{12}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),strip_list(Fit['Ts']),fittype,Nexp,SvdFactor,PriorLoosener,fit.Q,fit.chi2/fit.dof,smallsave)
    for corr in allcorrs:
        if corr in currents:
            filename += '_{0}tmin{1}'.format(corr,Fit['{0}tmin'.format(corr)])
    savedict = gv.BufferDict()
    if smallsave:
        for key in fit.p:
            if key[0] == 'l':
                key2 = key.split('(')[1].split(')')[0]
                if key2.split(':')[0] =='dE' and key2.split(':')[1][0] != 'o':
                    savedict[key] = [fit.p[key][0]] #was palt
            elif key[2] =='n' and key[3] == 'n':
                savedict[key] = [[fit.p[key][0][0]]] #was palt
    elif smallsave == False:
        savedict = fit.p
    print('Started gv.gdump to {1}, smallsave = {0}'.format(smallsave,'{0}.pickle'.format(filename)),datetime.datetime.now())
    gv.gdump(savedict,'{0}.pickle'.format(filename))
    print('Finished gv.gdump fit, starting save fit output',datetime.datetime.now())
    f = open('{0}.txt'.format(filename),'w')
    f.write(fit.format(pstyle='v'))
    f.close()
    print('Finished save fit output',datetime.datetime.now())
    return()

######################################################################################################

def do_chained_fit(data,prior,Nexp,modelsA,modelsB,Fit,noise,currents,allcorrs,SvdFactor,PriorLoosener,FitCorrs,save,smallsave,GBF):#if GBF = None doesn't pass GBF, else passed GBF
    #do chained fit with no marginalisation Nexp = NMax
    models = copy.deepcopy(modelsA)
    if len(modelsB[0]) !=0:
        models.extend(modelsB)
    print('Models',models)
    fitter = cf.CorrFitter(models=models, maxit=maxiter, fast=False, tol=(1e-6,0.0,0.0))
    p0 = get_p0(Fit,'chained',Nexp,allcorrs,prior,FitCorrs)
    print(30 * '=','Chained-Unmarginalised','Nexp =',Nexp,'Date',datetime.datetime.now())
    fit = fitter.chained_lsqfit(data=data, prior=prior, p0=p0, noise=noise,debug=True)
    update_p0([f.pmean for f in fit.chained_fits.values()],fit.pmean,Fit,'chained',Nexp,allcorrs,FitCorrs,fit.Q) #fittype=chained, for marg,includeN
    if GBF == None:
        print(fit)
        print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
        print_results(fit.p,prior)
        print_Z_V(fit.p,Fit,allcorrs)
        if fit.Q > 0.05 and save: #threshold for a 'good' fit
            save_fit(fit,Fit,allcorrs,'chained',Nexp,SvdFactor,PriorLoosener,currents,smallsave)
            #print_fit_results(fit) do this later
        return()
    elif fit.logGBF - GBF < 1 and fit.logGBF - GBF > 0:
        print('log(GBF) went up by less than 1: {0:.2f}'.format(fit.logGBF - GBF))
        return(fit.logGBF)
    elif fit.logGBF - GBF < 0:
        print('log(GBF) went down {0:.2f}'.format(fit.logGBF - GBF))
        return(fit.logGBF)
    else:
        print(fit)
        print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
        print_results(fit.p,prior)
        print_Z_V(fit.p,Fit,allcorrs)
        print('log(GBF) went up {0:.2f}'.format(fit.logGBF - GBF))
        if fit.Q > 0.05 and save: #threshold for a 'good' fit
            save_fit(fit,Fit,allcorrs,'chained',Nexp,SvdFactor,PriorLoosener,currents,smallsave)
            #print_fit_results(fit) do this later
        return(fit.logGBF)

######################################################################################################

def do_chained_marginalised_fit(data,prior,Nexp,modelsA,modelsB,Fit,noise,currents,allcorrs,SvdFactor,PriorLoosener,FitCorrs,save,smallsave,GBF,Marginalised):#if GBF = None doesn't pass GBF, else passed GBF
    #do chained fit with marginalisation nterm = nexp,nexp Nmarg=Marginalisation us in p0 bits
    models = copy.deepcopy(modelsA)
    if len(modelsB[0]) !=0:
        models.append(dict(nterm=(Nexp,Nexp)))
        models.extend(modelsB)
    else:
        print('Marginalisation not applied as no parrallelised models')
    print('Models',models)
    fitter = cf.CorrFitter(models=models, maxit=maxiter, fast=False, tol=(1e-6,0.0,0.0))
    p0 = get_p0(Fit,'chained-marginalised_N{0}{0}'.format(Nexp),Marginalised,allcorrs,prior,FitCorrs)
    print(30 * '=','Chained-marginalised','Nexp =',Marginalised,'nterm = ({0},{0})'.format(Nexp),'Date',datetime.datetime.now())
    fit = fitter.chained_lsqfit(data=data, prior=prior, p0=p0, noise=noise,debug=True)
    update_p0([f.pmean for f in fit.chained_fits.values()],fit.pmean,Fit,'chained-marginalised_N{0}{0}'.format(Nexp),Marginalised,allcorrs,FitCorrs,fit.Q,True) #fittype=chained, for marg,includeN
    if GBF == None:
        print(fit)#.format(pstyle='m'))
        print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
        print_results(fit.p,prior)
        print_Z_V(fit.p,Fit,allcorrs)
        if fit.Q > 0.05 and save: #threshold for a 'good' fit
            save_fit(fit,Fit,allcorrs,'chained-marginalised_N{0}{0}'.format(Nexp),Marginalised,SvdFactor,PriorLoosener,currents,smallsave)
            #print_fit_results(fit) do this later
        return()
    elif fit.logGBF - GBF < 1 and fit.logGBF - GBF > 0:
        print('log(GBF) went up by less than 1: {0:.2f}'.format(fit.logGBF - GBF))
        return(fit.logGBF)
    elif fit.logGBF - GBF < 0:
        print('log(GBF) went down {0:.2f}'.format(fit.logGBF - GBF))
        return(fit.logGBF)
    else:
        print(fit)#.format(pstyle='m'))
        print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
        print_results(fit.p,prior)
        print_Z_V(fit.p,Fit,allcorrs)
        print('log(GBF) went up {0:.2f}'.format(fit.logGBF - GBF))
        if fit.Q > 0.05 and save: #threshold for a 'good' fit
            save_fit(fit,Fit,allcorrs,'chained-marginalised_N{0}{0}'.format(Nexp),Marginalised,SvdFactor,PriorLoosener,currents,smallsave)
            #print_fit_results(fit) do this later
        return(fit.logGBF)

######################################################################################################

def do_unchained_fit(data,prior,Nexp,models,svdcut,Fit,noise,currents,allcorrs,SvdFactor,PriorLoosener,save,smallsave,GBF):#if GBF = None doesn't pass GBF, else passed GBF
    #do chained fit with no marginalisation Nexp = NMax
    print('Models',models)
    fitter = cf.CorrFitter(models=models, maxit=maxiter, fast=False, tol=(1e-6,0.0,0.0))
    p0 = get_p0(Fit,'unchained',Nexp,allcorrs,prior,allcorrs) # FitCorrs = allcorrs
    print(30 * '=','Unchained-Unmarginalised','Nexp =',Nexp,'Date',datetime.datetime.now())
    fit = fitter.lsqfit(pdata=data, prior=prior, p0=p0, svdcut=svdcut, noise=noise,debug=True)
    update_p0(fit.pmean,fit.pmean,Fit,'unchained',Nexp,allcorrs,allcorrs,fit.Q) #fittype=chained, for marg,includeN
    if GBF == None:
        print(fit)
        print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
        print_results(fit.p,prior)#,Fit)
        print_Z_V(fit.p,Fit,allcorrs)
        if fit.Q > 0.05 and save: #threshold for a 'good' fit
            save_fit(fit,Fit,allcorrs,'unchained',Nexp,SvdFactor,PriorLoosener,currents,smallsave)
            #print_fit_results(fit) do this later
        return()
    elif fit.logGBF - GBF < 1 and fit.logGBF - GBF > 0:
        print('log(GBF) went up by less than 1: {0:.2f}'.format(fit.logGBF - GBF))
        return(fit.logGBF)
    elif fit.logGBF - GBF < 0:
        print('log(GBF) went down: {0:.2f}'.format(fit.logGBF - GBF))
        return(fit.logGBF)
    else:
        print(fit)
        print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
        print_results(fit.p,prior)#,Fit)
        print_Z_V(fit.p,Fit,allcorrs)
        print('log(GBF) went up more than 1: {0:.2f}'.format(fit.logGBF - GBF))
        if fit.Q > 0.05 and save: #threshold for a 'good' fit
            save_fit(fit,Fit,allcorrs,'unchained',Nexp,SvdFactor,PriorLoosener,currents,smallsave)
            #print_fit_results(fit) do this later
        return(fit.logGBF)

#######################################################################################################

def do_sep_mass_fit(data,prior,Nexp,models,svdcut,Fit,noise,currents,allcorrs,SvdFactor,PriorLoosener,save,smallsave,GBF):
    #if GBF = None doesn't pass GBF, else passed GBF
    #do chained fit with no marginalisation Nexp = NMax
    print('Models',models)
    #print(data)
    fitter = cf.CorrFitter(models=models, maxit=maxiter, fast=False, tol=(1e-6,0.0,0.0))
    p0 = get_p0(Fit,'sepmass',Nexp,allcorrs,prior,allcorrs) # FitCorrs = allcorrs
    print(30 * '=','Seperate Mass Fit','Nexp =',Nexp,'Date',datetime.datetime.now())
    fit = fitter.lsqfit(pdata=data, prior=prior, p0=p0, svdcut=svdcut, noise=noise,debug=True)
    update_p0(fit.pmean,fit.pmean,Fit,'sepmass',Nexp,allcorrs,allcorrs,fit.Q) #fittype=chained, for marg,includeN
    print(fit)
    print('chi^2/dof = {0:.3f} Q = {1:.3f} logGBF = {2:.0f}'.format(fit.chi2/fit.dof,fit.Q,fit.logGBF))
    print_results(fit.p,prior)#,Fit)
    return(fit)

######################################################################################################

def combine_sep_mass_fits(result,Fit,priors,allcorrs,Nexp,SvdFactor,PriorLoosener,currents,save,smallsave):
    prior = gv.BufferDict()
    combined = []
    for mass in Fit['masses']:
        smallresult = gv.BufferDict()
        fit = result[mass].p
        for key in fit:
            if key[0] == 'l':
                key2 = key.split('(')[1].split(')')[0]
                if key2.split(':')[0] =='dE':
                    smallresult[key] = [fit[key][0]]
            elif key[2] =='n' and key[3] == 'n':
                smallresult[key] = [[fit[key][0][0]]]
        combined.append(smallresult)

    prior = copy.deepcopy(priors[Fit['masses'][0]])
    for mass in Fit['masses']:
        for key in priors[mass]:
            if key not in prior:
                prior[key] = copy.deepcopy(priors[mass][key])
    #print(combined)
    final = lsqfit.wavg(combined)
    #print(gv.evalcorr([final['SVnn_m0.433_tw0.8563'][0][0],final['SVnn_m0.683_tw0.8563'][0][0]]))
    chi = 0
    Q = 0
    GBF =  0
    for mass in Fit['masses']:
        chi += (result[mass].chi2/result[mass].dof)/len(Fit['masses'])
        Q += (result[mass].Q)/len(Fit['masses'])
        GBF += result[mass].logGBF
    print('Mean chi^2/dof = {0:.3f} Q = {1:.3f}, total logGBF {2:.1f}'.format(chi,Q,GBF))
    print_results(final,prior)#,Fit)
    print_Z_V(final,Fit,allcorrs)
    if save:
        save_combined_fit(final,Fit,allcorrs,'sep_mass',Nexp,SvdFactor,PriorLoosener,currents,smallsave,chi,Q)
    return()

######################################################################################################################

def save_combined_fit(fit,Fit,allcorrs,fittype,Nexp,SvdFactor,PriorLoosener,currents,smallsave,chi,Q):
    filename = 'Fits/{0}{1}{2}{3}{4}{5}{6}_Nexp{7}_sfac{8}_pfac{9}_Q{10:.2f}_chi{11:.3f}_sm{12}'.format(Fit['conf'],Fit['filename'],strip_list(Fit['masses']),strip_list(Fit['twists']),strip_list(allcorrs),strip_list(Fit['Ts']),fittype,Nexp,SvdFactor,PriorLoosener,Q,chi,smallsave)
    for corr in allcorrs:
        if corr in currents:
            filename += '_{0}tmin{1}'.format(corr,Fit['{0}tmin'.format(corr)])
    savedict = gv.BufferDict()
    if smallsave:
        for key in fit:
            if key[0] == 'l':
                key2 = key.split('(')[1].split(')')[0]
                if key2.split(':')[0] =='dE' and key2.split(':')[1][0] != 'o':
                    savedict[key] = [fit[key][0]]
            elif key[2] =='n' and key[3] == 'n':
                savedict[key] = [[fit[key][0][0]]]
    elif smallsave == False:
        print('Error, can only do small save with sep masses' )
    #print(gv.evalcorr([savedict['SVnn_m0.433_tw0.8563'][0][0],savedict['SVnn_m0.683_tw0.8563'][0][0]]))
    print('Started gv.gdump to {1}, smallsave = {0}'.format(smallsave,'{0}.pickle'.format(filename)),datetime.datetime.now())
    gv.gdump(savedict,'{0}.pickle'.format(filename))
    print('Finished gv.gdump fit',datetime.datetime.now())
    return()
######################################################################################################

def print_p_p0(p,p0,prior):
    print('{0:<30}{1:<20}{2:<40}{3:<20}'.format('key','p','p0','prior'))
    for key in prior:
        if len(np.shape(p[key])) ==1 :
            for element in range(len(p[key])):
                if element == 0:
                    print('{0:<30}{1:<20}{2:<40}{3:<20}'.format(key,p[key][element],p0[key][element],prior[key][element]))
                else:
                    print('{0:>30}{1:<20}{2:<40}{3:<20}'.format('',p[key][element],p0[key][element],prior[key][element]))
    return()

#####################################################################################################

def print_results(p,prior):#,Fit):
    print(100*'-')
    print('{0:<30}{1:<15}{2:<15}{3:<15}{4}'.format('key','p','p error','prior','prior error'))
    print(100*'-')
    print('Ground state energies')
    print(100*'-')
    for key in prior:
        if key[0] == 'l':
            key = key.split('(')[1].split(')')[0]
        if key.split(':')[0] =='dE' and key.split(':')[1][0] != 'o':
            print('{0:<30}{1:<15}{2:<15.3%}{3:<15}{4:.2%}'.format(key,p[key][0],p[key][0].sdev/p[key][0].mean,prior[key][0],prior[key][0].sdev/prior[key][0].mean))
            #if '{0}'.format(key.split(':')[1]) == Fit['BG-Tag'].format(Fit['masses'][0]):
             #   print('split: ', p['dE:{0}'.format(Fit['BNG-Tag'].format(Fit['masses'][0]))][0]-p[key][0])
    print('')
    print('Oscillating ground state energies')
    print(100*'-')
    for key in prior:
        if key[0] == 'l':
            key = key.split('(')[1].split(')')[0]
        if key.split(':')[0] =='dE' and key.split(':')[1][0] == 'o':
            print('{0:<30}{1:<15}{2:<15.3%}{3:<15}{4:.2%}'.format(key,p[key][0],p[key][0].sdev/p[key][0].mean,prior[key][0],prior[key][0].sdev/prior[key][0].mean))
    print('')
    print('V_nn[0][0]')
    print(100*'-')
    for key in prior:
        if key[1] != '2' and key[2] =='n' and key[3] == 'n':
            print('{0:<30}{1:<15}{2:<15.3%}{3:<15}{4:.2%}'.format(key,p[key][0][0],p[key][0][0].sdev/p[key][0][0].mean,prior[key][0][0],prior[key][0][0].sdev/prior[key][0][0].mean))
    print(100*'-')
    return()
#####################################################################################################

def make_Z_V(m_h,m_s,M_parent,M_daughter,S,V):
    Z_V = (m_h-m_s)/(M_parent-M_daughter) * S/V
    return(Z_V)

#####################################################################################################
# needs generalising
#####################################################################################################
def print_Z_V(p,Fit,allcorrs):
    if 'S' in allcorrs and 'V' in allcorrs:
        print(100*'-')
        for mass in Fit['masses']:
            M_parent = p['dE:{0}'.format(Fit['{0}-Tag'.format('BG')].format(mass))][0]
            M_daughter = p['dE:{0}'.format(Fit['{0}-Tag'.format('KG')].format('0'))][0]
            S = p['SVnn_m{0}_tw0'.format(mass)][0][0]
            V = p['VVnn_m{0}_tw0'.format(mass)][0][0]
            Z_V = make_Z_V(float(mass),float(Fit['m_s']),M_parent,M_daughter,S,V)
            print("Mass = {0} Z_V = {1}".format(mass,Z_V))
        print(100*'-')
    return()

#####################################################################################################


""" Comments scraper class"""
import json
import logging
import random
from time import sleep

import numpy as np
import pandas as pd
from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait

from scripts.argparser import ArgParser
from scripts.auxiliary.scraper_aux import check_or_create_folders, save
from scripts.chrome_driver import ChromeDriver

logging.basicConfig(
    format="%(asctime)s | %(levelname)s: %(message)s", level=logging.CRITICAL
)

class CommentsScraper:
    def __init__(self):
        pass

    @staticmethod
    def iterate(args, ttw, driver):
        # ~ 24 comments loaded each iteration
        check = True
        MAX_ITER = int(args.max_iterations)
        n = 0

        while check and n <= MAX_ITER:
            sleep(int(random.choice(ttw)))

            try:
                load_more_comments_button = WebDriverWait(driver, 5).until(
                    EC.visibility_of_element_located(
                        (By.CSS_SELECTOR, "[aria-label='Load more comments']")
                    )
                )
                load_more_comments_button.click()
            except:
                check = False
                if n == MAX_ITER:
                    logging.warning(
                        "Reached the max iterations number before exhausting the post comments. \
                            You may consider to raise the max iterations number"
                    )
                else:
                    logging.info("Exhausted all the post comments")
            n = n + 1

    @classmethod
    def setup(self):
        argparser = ArgParser()
        chrome_driver = ChromeDriver()
        proxy = chrome_driver.set_up_proxy()

        args = argparser.likes_scraper_read_input()

        driver = chrome_driver.set_up_driver(proxy=proxy)
        driver = chrome_driver.make_IG_access_w_creds(
            driver=driver, ig_usr=args.username, ig_pass=args.password
        )
        return driver, proxy, args

    @classmethod
    def scrape(self, driver, proxy, args, save_raw_data=bool, **kwargs):

        if "target_post" in kwargs:
            proxy.new_har(
                kwargs.get("target_post"),
                options={"captureHeaders": True, "captureContent": True},
            )
            driver.get(kwargs.get("target_post"))
            target = kwargs.get("target_post")
        else:
            proxy.new_har(
                args.target_post,
                options={"captureHeaders": True, "captureContent": True},
            )
            driver.get(args.target_post)
            target = args.target_post

        # Random time intervals to sleep between load more comment button pushes
        ttw = []
        for i in range(0, 20):
            ttw.append(np.round(random.uniform(4, 8), 2))

        # ~ 24 comments loaded each iteration
        CommentsScraper.iterate(args=args, ttw=ttw, driver=driver)

        R = json.loads(json.dumps(proxy.har, ensure_ascii=False))

        if save_raw_data:
            save(data=R, target=target, args=args)

        return R

    @classmethod
    def parse_and_save_data(self, raw_data, args, target):
        "Parse raw scraped data and write to disk"

        RAW = {}
        for n, v in enumerate(raw_data["log"]["entries"]):
            if v["response"]["content"]["mimeType"] in [
                "application/json; charset=utf-8",
                "application/json",
            ]:
                try:
                    RAW[n] = json.loads(v["response"]["content"]["text"])["comments"]
                except:
                    pass

        comments_df = pd.DataFrame.from_dict(RAW[list(RAW.keys())[0]])

        for k in list(RAW.keys())[1:]:
            comments_df = pd.concat([comments_df, pd.DataFrame.from_dict(RAW[k])])

        comments_df = comments_df.reset_index(drop=True)

        check_or_create_folders(target=target, args=args)

        short_code = target.split("/")[-1]

        # TODO FIX get the profile name from somewhere and create the correct folder!
        comments_df.to_csv(
            f"{args.output_folder}/{short_code}_comments_clean.csv", mode="w+"
        )
        logging.info("Data correctly saved/overwrote.")
        # print(f'Dave saved in: {f"{args.output_folder}_{short_code}_comments_clean.csv"}')
        return comments_df


from itertools import combinations
from collections import Counter

def solution(orders, course):
  answer = []
  for c in course:
    temp = []
    for order in orders:
      combi = combinations(sorted(order), c)
      temp += combi
    counter = Counter(temp)
    if len(counter) != 0 and max(counter.values()) != 1:
      for cnt in counter:
        if counter[cnt] == max(counter.values()):
          answer.append("".join(cnt))
  return sorted(answer)


from django.contrib.auth.models import User
from django.db import models

# Create your models here.
from django.db.models.signals import pre_delete
from django.dispatch import receiver

class CommonInfo(models.Model):
    # 开始时间 auto_now_add=True,
    startday = models.DateField(verbose_name="下单时间", null=True)
    # 结束时间
    endday = models.DateField(verbose_name="交货时间", null=True)
    # 备注
    remark = models.TextField(default="", verbose_name="备注")

    status_choice = (
        ('pending', "未开始"),
        ('process', "进行中"),
        ('finish', "已完成")
    )
    status = models.CharField(
        max_length=10,
        choices=status_choice,
        default='pending',
        verbose_name="订单状态",
    )

    class Meta:
        abstract = True

class Productplan(CommonInfo):
    # 订单号
    orderid = models.CharField(max_length=10, verbose_name="订单号")
    # 订单分类
    category_choice = (
        ('std', "标准"),
        ('unstd', "非标")
    )
    category = models.CharField(
        max_length=5,
        choices=category_choice,
        default='std',
        verbose_name="订单类型",
    )
    # 型号
    productid = models.CharField(max_length=20, verbose_name="产品型号")
    # 序列号
    serial = models.CharField(max_length=20, verbose_name="序列号")
    # 用户
    customer = models.CharField(max_length=20, verbose_name="用户")
    # 配置单
    conffile = models.FileField(upload_to='service_files/%Y%m%d/', blank=True)
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    def __str__(self):
        return self.customer + "_SN" + self.serial

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "订单管理"

@receiver(pre_delete, sender=Productplan)
def delete(sender, instance, **kwargs):
    instance.conffile.delete(False)

class ProcessElPrepare(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='proel',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "电路板准备"

class ProcessMePrepare(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='prome',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "机械件准备"

class ProcessScPrepare(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='prosc',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "干涉仪准备"

class ProcessAssemble(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='proas',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "装配中"

class ProcessTesting(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='prots',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "测试中"

    def __str__(self):
        return self.orderid.customer

# 添加软件状态
class ProcessSoftware(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='prosw',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "软件"

    def __str__(self):
        return self.orderid.customer

# 添加付款状态
class ProcessPayment(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='propm',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "付款"

    def __str__(self):
        return self.orderid.customer

# 添加发货状态
class ProcessDeliver(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='prodi',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "发货"

    def __str__(self):
        return self.orderid.customer

# 添加开票状态
class ProcessBilling(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='probi',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "开票"

    def __str__(self):
        return self.orderid.customer

# 添加尾款状态
class ProcessDueing(CommonInfo):
    orderid = models.OneToOneField(
        Productplan,
        on_delete=models.CASCADE,
        related_name='produe',
        primary_key=True
    )
    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "尾款"

    def __str__(self):
        return self.orderid.customer

class ProductHistory(models.Model):
    # 操作者
    user = models.ForeignKey(
        User,
        null=True,
        blank=True,
        on_delete=models.SET_NULL,
        related_name='user'
    )
    # 订单
    orderid = models.ForeignKey(
        Productplan,
        verbose_name="订单号",
        on_delete=models.CASCADE,
        related_name='product_history',
        primary_key=False
    )
    # 操作项目
    proitem = models.CharField(max_length=10, verbose_name="操作项目")
    # 新内容
    newcontent = models.CharField(max_length=10, verbose_name="新内容")
    # 操作时间
    operateday = models.DateTimeField(verbose_name="操作时间", null=True)

    # 更新时间
    updated = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ('-updated',)
        verbose_name_plural = "订单历史"

    def __str__(self):
        return self.orderid.customer


import argparse

from experiment.monitor import monitor
from apps.qe import qe

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    monitor.setup_run_args(parser)
    qe.setup_run_args(parser)

    args, extra_args = parser.parse_known_args()
    app_conf = qe.QuantumEspressoAppConf(args.node_count, args.benchmark_name)
    monitor.launch(app_conf=app_conf, args=args,
                   experiment_cli_args=extra_args)


import logging
from collections import defaultdict
from typing import Dict, Optional, Tuple, Union

import numpy as np
from matplotlib import rcParams
from matplotlib.axes import SubplotBase
from matplotlib.axis import Axis
from matplotlib.colors import LogNorm
from matplotlib.ticker import AutoMinorLocator, MaxNLocator
from pymatgen.electronic_structure.bandstructure import (
    BandStructure,
    BandStructureSymmLine,
)
from pymatgen.electronic_structure.core import Spin
from pymatgen.electronic_structure.plotter import BSPlotter
from sumo.plotting import pretty_plot
from sumo.symmetry import Kpath, PymatgenKpath

from amset.constants import defaults, hbar
from amset.interpolation.bandstructure import Interpolator
from amset.interpolation.periodic import PeriodicLinearInterpolator
from amset.log import initialize_amset_logger
from amset.plot import BaseMeshPlotter, amset_base_style, styled_plot

__author__ = "Alex Ganose"
__maintainer__ = "Alex Ganose"
__email__ = "aganose@lbl.gov"

logger = logging.getLogger(__name__)

class LineshapePlotter(BaseMeshPlotter):
    def __init__(
        self,
        data,
        interpolation_factor=5,
        print_log=defaults["print_log"],
        symprec=defaults["symprec"],
    ):
        super().__init__(data)
        self.interpolation_factor = interpolation_factor

        if print_log:
            initialize_amset_logger(filename="lineshape.log")
        self.symprec = symprec

    def _get_interpolater(self, n_idx, t_idx, mode="linear"):
        props = defaultdict(dict)
        for spin in self.spins:
            # calculate total rate
            spin_rates = np.sum(self.scattering_rates[spin][:, n_idx, t_idx], axis=0)

            # easier to interpolate the log
            log_rates = np.log10(spin_rates)

            # # handle rates that close to numerical noise
            log_rates[log_rates > 18] = 15
            log_rates[np.isnan(log_rates)] = 15

            # map to full k-point mesh
            props[spin]["rates"] = log_rates

        if mode == "linear":
            return _LinearBandStructureInterpolator(
                self.kpoints,
                self.ir_to_full_kpoint_mapping,
                self.energies,
                self.structure,
                self.efermi,
                props,
            )
        elif mode == "fourier":
            bs = BandStructure(
                self.ir_kpoints,
                self.energies,
                self.structure.lattice,
                self.efermi,
                structure=self.structure,
            )

            return Interpolator(
                bs,
                self.num_electrons,
                interpolation_factor=self.interpolation_factor,
                soc=self.soc,
                other_properties=props,
            )
        raise ValueError("Unknown interpolation mode; should be 'linear' or 'fourier'.")

    @styled_plot(amset_base_style)
    def get_plot(
        self,
        n_idx,
        t_idx,
        zero_to_efermi=True,
        estep=0.01,
        line_density=100,
        height=3.2,
        width=3.2,
        emin=None,
        emax=None,
        amin=5e-5,
        amax=1e-1,
        ylabel="Energy (eV)",
        plt=None,
        aspect=None,
        kpath=None,
        cmap="viridis",
        colorbar=True,
        style=None,
        no_base_style=False,
        fonts=None,
    ):
        interpolater = self._get_interpolater(n_idx, t_idx)

        bs, prop = interpolater.get_line_mode_band_structure(
            line_density=line_density,
            return_other_properties=True,
            kpath=kpath,
            symprec=self.symprec,
        )
        bs, rates = force_branches(bs, {s: p["rates"] for s, p in prop.items()})

        fd_emin, fd_emax = self.fd_cutoffs
        if not emin:
            emin = fd_emin
            if zero_to_efermi:
                emin -= bs.efermi

        if not emax:
            emax = fd_emax
            if zero_to_efermi:
                emax -= bs.efermi

        logger.info("Plotting band structure")
        if isinstance(plt, (Axis, SubplotBase)):
            ax = plt
        else:
            plt = pretty_plot(width=width, height=height, plt=plt)
            ax = plt.gca()

        if zero_to_efermi:
            bs.bands = {s: b - bs.efermi for s, b in bs.bands.items()}
            bs.efermi = 0

        bs_plotter = BSPlotter(bs)
        plot_data = bs_plotter.bs_plot_data(zero_to_efermi=zero_to_efermi)

        energies = np.linspace(emin, emax, int((emax - emin) / estep))
        distances = np.array([d for x in plot_data["distances"] for d in x])

        # rates are currently log(rate)
        mesh_data = np.full((len(distances), len(energies)), 0.0)
        for spin in self.spins:
            for spin_energies, spin_rates in zip(bs.bands[spin], rates[spin]):
                for d_idx in range(len(distances)):
                    energy = spin_energies[d_idx]
                    linewidth = 10 ** spin_rates[d_idx] * hbar / 2
                    broadening = lorentzian(energies, energy, linewidth)
                    broadening /= 1000  # convert 1/eV to 1/meV
                    mesh_data[d_idx] += broadening

        im = ax.pcolormesh(
            distances,
            energies,
            mesh_data.T,
            rasterized=True,
            cmap=cmap,
            norm=LogNorm(vmin=amin, vmax=amax),
            shading="auto",
        )
        if colorbar:
            pos = ax.get_position()
            cax = plt.gcf().add_axes([pos.x1 + 0.035, pos.y0, 0.035, pos.height])
            cbar = plt.colorbar(im, cax=cax)
            cbar.ax.tick_params(axis="y", length=rcParams["ytick.major.size"] * 0.5)
            cbar.ax.set_ylabel(
                r"$A_\mathbf{k}$ (meV$^{-1}$)", rotation=270, va="bottom"
            )

        _maketicks(ax, bs_plotter, ylabel=ylabel)
        _makeplot(
            ax,
            plot_data,
            bs,
            zero_to_efermi=zero_to_efermi,
            width=width,
            height=height,
            ymin=emin,
            ymax=emax,
            aspect=aspect,
        )
        return plt

def _makeplot(
    ax,
    data,
    bs,
    zero_to_efermi=True,
    ymin=-3.0,
    ymax=3.0,
    height=None,
    width=None,
    aspect=None,
):
    """Tidy the band structure & add the density of states if required."""
    # draw line at Fermi level if not zeroing to e-Fermi
    if not zero_to_efermi:
        ytick_color = rcParams["ytick.color"]
        ef = bs.efermi
        ax.axhline(ef, color=ytick_color)

    # set x and y limits
    ax.set_xlim(0, data["distances"][-1][-1])
    if bs.is_metal() and not zero_to_efermi:
        ax.set_ylim(bs.efermi + ymin, bs.efermi + ymax)
    else:
        ax.set_ylim(ymin, ymax)

    # keep correct aspect ratio for axes based on canvas size
    if aspect is not False:
        x0, x1 = ax.get_xlim()
        y0, y1 = ax.get_ylim()
        if width is None:
            width = rcParams["figure.figsize"][0]
        if height is None:
            height = rcParams["figure.figsize"][1]

        if not aspect:
            aspect = height / width

        ax.set_aspect(aspect * ((x1 - x0) / (y1 - y0)))

def _maketicks(ax, bs_plotter, ylabel="Energy (eV)"):
    """Utility method to add tick marks to a band structure."""
    # set y-ticks
    ax.yaxis.set_major_locator(MaxNLocator(6))
    ax.yaxis.set_minor_locator(AutoMinorLocator(2))

    # set x-ticks; only plot the unique tick labels
    ticks = bs_plotter.get_ticks()
    unique_d = []
    unique_l = []
    if ticks["distance"]:
        temp_ticks = list(zip(ticks["distance"], ticks["label"]))
        unique_d.append(temp_ticks[0][0])
        unique_l.append(temp_ticks[0][1])
        for i in range(1, len(temp_ticks)):
            # Append label to sequence if it is not same as predecessor
            if unique_l[-1] != temp_ticks[i][1]:
                unique_d.append(temp_ticks[i][0])
                unique_l.append(temp_ticks[i][1])

    logging.info("Label positions:")
    for dist, label in list(zip(unique_d, unique_l)):
        logging.info(f"\t{dist:.4f}: {label}")

    ax.set_xticks(unique_d)
    ax.set_xticklabels(unique_l)
    ax.xaxis.grid(True)
    ax.set_ylabel(ylabel)

def lorentzian(x, x0, gamma):
    return 1 / np.pi * gamma / ((x - x0) ** 2 + gamma**2)

class _LinearBandStructureInterpolator:
    def __init__(
        self,
        full_kpoints,
        ir_to_full_idx,
        energies,
        structure,
        efermi,
        other_properties,
    ):
        self.structure = structure
        self.efermi = efermi
        self.spins = list(energies.keys())
        self.nbands = {s: len(e) for s, e in energies.items()}

        full_energies = {s: e[:, ir_to_full_idx] for s, e in energies.items()}
        self.bs_interpolator = PeriodicLinearInterpolator.from_data(
            full_kpoints, full_energies
        )

        self.property_interpolators = {}
        other_properties = _transpose_dict(other_properties)
        for prop, prop_data in other_properties.items():
            full_prop_data = {s: p[:, ir_to_full_idx] for s, p in prop_data.items()}
            self.property_interpolators[prop] = PeriodicLinearInterpolator.from_data(
                full_kpoints, full_prop_data, gaussian=0.75
            )

    def get_line_mode_band_structure(
        self,
        line_density: int = 50,
        kpath: Optional[Kpath] = None,
        symprec: Optional[float] = defaults["symprec"],
        return_other_properties: bool = False,
    ) -> Union[
        BandStructureSymmLine,
        Tuple[BandStructureSymmLine, Dict[Spin, Dict[str, np.ndarray]]],
    ]:
        """Gets the interpolated band structure along high symmetry directions.

        Args:
            line_density: The maximum number of k-points between each two
                consecutive high-symmetry k-points
            symprec: The symmetry tolerance used to determine the space group
                and high-symmetry path.
            return_other_properties: Whether to include the interpolated
                other_properties data for each k-point along the band structure path.

        Returns:
            The line mode band structure.
        """
        if not kpath:
            kpath = PymatgenKpath(self.structure, symprec=symprec)

        kpoints, labels = kpath.get_kpoints(line_density=line_density, cart_coords=True)
        labels_dict = {
            label: kpoint for kpoint, label in zip(kpoints, labels) if label != ""
        }

        rlat = self.structure.lattice.reciprocal_lattice
        frac_kpoints = rlat.get_fractional_coords(kpoints)

        energies = {}
        other_properties = defaultdict(dict)
        for spin in self.spins:
            energies[spin] = self._interpolate_spin(
                spin, frac_kpoints, self.bs_interpolator
            )

            if return_other_properties:
                for prop, property_interpolator in self.property_interpolators.items():
                    other_properties[spin][prop] = self._interpolate_spin(
                        spin, frac_kpoints, property_interpolator
                    )

        bs = BandStructureSymmLine(
            kpoints,
            energies,
            rlat,
            self.efermi,
            labels_dict,
            coords_are_cartesian=True,
            structure=self.structure,
        )

        if return_other_properties:
            return bs, other_properties
        else:
            return bs

    def _interpolate_spin(self, spin, kpoints, interpolator):
        nkpoints = len(kpoints)
        spin_nbands = self.nbands[spin]
        ibands = np.repeat(np.arange(spin_nbands), nkpoints)
        all_kpoints = np.tile(kpoints, (spin_nbands, 1))
        data = interpolator.interpolate(spin, ibands, all_kpoints)
        return data.reshape(spin_nbands, nkpoints)

def _transpose_dict(d):
    td = defaultdict(dict)
    for k1, v1 in d.items():
        for k2, v2 in v1.items():
            td[k2][k1] = v2
    return td

def force_branches(bandstructure, other_property=None):
    """Force a linemode band structure to contain branches.

    Branches give a specific portion of the path from one high-symmetry point
    to another. Branches are required for the plotting methods to function correctly.
    Unfortunately, due to the pymatgen BandStructure implementation they require
    duplicate k-points in the band structure path. To avoid this unnecessary
    computational expense, this function can reconstruct branches in band structures
    without the duplicate k-points.

    Args:
        bandstructure: A band structure object.
        other_property: Another property with the format {spin: (nbands, nkpts, ...)
            to split into branches.

    Returns:
        A band structure with brnaches.
    """
    kpoints = np.array([k.frac_coords for k in bandstructure.kpoints])
    labels_dict = {k: v.frac_coords for k, v in bandstructure.labels_dict.items()}

    # pymatgen band structure objects support branches. These are formed when
    # two kpoints with the same label are next to each other. This bit of code
    # will ensure that the band structure will contain branches, if it doesn't
    # already.
    dup_ids = []
    high_sym_kpoints = tuple(map(tuple, labels_dict.values()))
    for i, k in enumerate(kpoints):
        dup_ids.append(i)
        if (
            tuple(k) in high_sym_kpoints
            and i != 0
            and i != len(kpoints) - 1
            and (
                not np.array_equal(kpoints[i + 1], k)
                or not np.array_equal(kpoints[i - 1], k)
            )
        ):
            dup_ids.append(i)

    kpoints = kpoints[dup_ids]

    eigenvals = {}
    projections = {}
    for spin, spin_energies in bandstructure.bands.items():
        eigenvals[spin] = spin_energies[:, dup_ids]
        if len(bandstructure.projections) != 0:
            projections[spin] = bandstructure.projections[spin][:, dup_ids]

    new_property = {}
    if other_property is not None:
        for spin, spin_prop in other_property.items():
            new_property[spin] = spin_prop[:, dup_ids]

    new_bandstructure = type(bandstructure)(
        kpoints,
        eigenvals,
        bandstructure.lattice_rec,
        bandstructure.efermi,
        labels_dict,
        structure=bandstructure.structure,
        projections=projections,
    )
    return new_bandstructure, new_property


import backtrader as bt
import backtrader.analyzers as btanalyzers
import matplotlib
import matplotlib.dates as mdates
from matplotlib import pyplot as plt
from datetime import datetime
import pandas as pd
import datetime as dt

# Create a subclass of Strategy to define the indicators and logic
class SMA_CrossStrategy(bt.Strategy):

    def __init__(self):
        ma_fast = bt.ind.SMA(period = 9)
        ma_slow = bt.ind.SMA(period = 21)

        self.crossover = bt.ind.CrossOver(ma_fast, ma_slow)

    def next(self):
        if not self.position:
            if self.crossover > 0:
                self.buy()
        elif self.crossover < 0:
            self.close()

    # def next(self):
    #     if not self.position:
    #         if self.crossover > 0:
    #             self.buy()
    #         elif self.crossover < 0:
    #             self.sell()
    #     elif self.crossover < 0:
    #         self.close()

def main():
    cerebro = bt.Cerebro()

    tickers_list = ['AAPL']
    df_tic = pd.read_hdf("datasets/df_SnP_500_ohlcv.h5", "df", mode = 'r')
    df_tic = df_tic[df_tic['tic'].isin(tickers_list)]
    df_tic = df_tic.set_index('date')
    # df_tic['date'] = pd.to_datetime(df_tic['date'])
    print(df_tic.head(5))
    data = bt.feeds.PandasData(dataname = df_tic,
                                            # datetime=None,
                                            open =1,
                                            high=2,
                                            low=3,
                                            close=4,
                                            volume=6,
                                            openinterest=-1,
                                            timeframe = bt.TimeFrame.Days,
                                            fromdate=dt.datetime(2023, 1, 1),  # Specify the start date
                                            todate=dt.datetime(2023, 8, 24),   # Specify the end date
                                        )

    # data = bt.feeds.YahooFinanceData(dataname = 'AAPL', fromdate = datetime(2010, 1, 1), todate = datetime(2020, 1, 1))
    cerebro.adddata(data)

    cerebro.addstrategy(SMA_CrossStrategy)

    cerebro.broker.setcash(10000.0)

    cerebro.addsizer(bt.sizers.PercentSizer, percents = 100)

    cerebro.addanalyzer(btanalyzers.SharpeRatio, _name = "sharpe")
    cerebro.addanalyzer(btanalyzers.Transactions, _name = "trans")
    cerebro.addanalyzer(btanalyzers.TradeAnalyzer, _name = "trades")

    back = cerebro.run()

    cerebro.broker.getvalue()

    back[0].analyzers.sharpe.get_analysis()

    back[0].analyzers.trans.get_analysis()

    back[0].analyzers.trades.get_analysis()

    cerebro.plot(style='candlestick', barup='green', bardown='red')
    plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    plt.gcf().autofmt_xdate()  # Rotates the date labels for better visibility

if __name__=="__main__":
    main()


from flask import Flask, request, render_template, redirect, flash, session
from flask_debugtoolbar import DebugToolbarExtension
from surveys import surveys, satisfaction_survey, personality_quiz

app = Flask(__name__)

app.config['SECRET_KEY'] = "secret_code_here"
# app.config['DEBUG_TB_INTERCEPT_REDIRECTS'] = False
# debug = DebugToolbarExtension(app)

@app.route('/')
def choose_the_servey():
    """ start-page that allows a user to choose a survey from a survey list """
    return render_template('choose_survey.html', surveys=surveys.items())

@app.route('/start-page', methods=['POST'])
def start_survey():
    """render a page for the chosen survey that shows the user the title of the survey, the instructions, and a button to start the survey;
    setup session variables """
    current_survey=surveys[request.form.get('choice')] #current survey object
    session['title'] = current_survey.title
    session['num_of_questions'] = len(current_survey.questions)
    session['survey'] = request.form['choice'] # key of current survey in surveys list
    session['comments'] = []
    return render_template('start.html', instructions = current_survey.instructions)

@app.route('/start', methods=['POST'])
def handling_start():
    """ Set a current session responses-list to an empty list and redirect to the start of the survey """
    session['responses'] = []
    return redirect('/questions/0')

@app.route('/questions/<question_number>')
def question_page(question_number):
    """ Page shows a form asking the current question, and listing the choices as radio buttons, also comments if are allowed """
    answers = session['responses']
    if len(answers) == session['num_of_questions']:
        return render_template('thanks.html')
    if len(answers) == int(question_number):
        number = int(question_number)
    else:
        flash('Could you please answer this question, you tried to access an invalid question','error')
        number = len(answers)
    current_survey = surveys[session['survey']]
    current_question = current_survey.questions[number]
    return render_template('question.html', number = number, choices = current_question.choices, question = current_question.question, title = current_survey.title, allow_text = current_question.allow_text)

@app.route('/answer', methods=['POST'])
def handling_answer():
    """ function appends the answer to responses list, adds comments if necessary and then redirect user to the next question; if no more questions in the survey - render thanks page ."""
    comment = request.form.get('comment')
    current_answer = request.form.get('choice')
    answers = session['responses']
    if current_answer:
        answers.append(current_answer)
        session['responses'] = answers
        if comment:
            comments = session['comments']
            comments.append((len(answers),comment))
            session['comments'] = comments
    else:
        flash('We are still waiting for your response!','error')
    if len(answers) < session['num_of_questions']:
        next_question = f"/questions/{len(answers)}"
        return redirect(next_question)
    else:
        return render_template('thanks.html')


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

import datasets
import utils

FLAGS = tf.flags.FLAGS

def get_lr(global_step, base_lr, steps_per_epoch,  # pylint: disable=missing-docstring
           decay_epochs, lr_decay_factor, warmup_epochs):

  warmup_lr = 0.0
  if warmup_epochs > 0:
    warmup_lr = (tf.cast(global_step, tf.float32) *
                 (base_lr / (warmup_epochs * steps_per_epoch)))

  normal_lr = tf.train.piecewise_constant(
      global_step,
      [e * steps_per_epoch for e in decay_epochs],
      [base_lr * (lr_decay_factor ** i) for i in range(len(decay_epochs) + 1)]
  )

  lr = tf.cond(tf.less(global_step, warmup_epochs * steps_per_epoch),
               lambda: warmup_lr,
               lambda: normal_lr)

  return lr

# TODO(akolesnikov): add more logging
class Trainer(object):
  """Base trainer class."""

  def __init__(self,
               update_batchnorm_params=True):
    self.update_batchnorm_params = update_batchnorm_params

    split = FLAGS.get_flag_value('train_split', 'train')
    num_samples = datasets.get_count(split)
    steps_per_epoch = num_samples // FLAGS.batch_size

    global_step = tf.train.get_or_create_global_step()
    self.global_step_inc = tf.assign_add(global_step, 1)

    # lr_scale_batch_size defines a canonical batch size that is coupled with
    # the initial learning rate. If actual batch size is not the same as
    # canonical than learning rate is linearly scaled. This is very convinient
    # as this allows to vary batch size without recomputing learning rate.
    lr_factor = 1.0
    if FLAGS.get_flag_value('lr_scale_batch_size', 0):
      lr_factor = FLAGS.batch_size / float(FLAGS.lr_scale_batch_size)

    deps = FLAGS.get_flag_value('decay_epochs', None)
    decay_epochs = utils.str2intlist(deps) if deps else [FLAGS.epochs]

    self.lr = get_lr(
        global_step,
        base_lr=FLAGS.lr * lr_factor,
        steps_per_epoch=steps_per_epoch,
        decay_epochs=decay_epochs,
        lr_decay_factor=FLAGS.get_flag_value('lr_decay_factor', 0.1),
        warmup_epochs=FLAGS.get_flag_value('warmup_epochs', 0))

    # TODO(marvinritter): Re-enable summaries with support for TPU training.
    # tf.summary.scalar('learning_rate', self.lr)

  def get_train_op(self, loss,  # pylint: disable=missing-docstring
                   var_list=None,
                   add_reg_loss=True,
                   use_tpu=False):

    if add_reg_loss:
      l2_loss = tf.reduce_sum(tf.losses.get_regularization_losses())
      loss += l2_loss

    optimizer = FLAGS.get_flag_value('optimizer', 'sgd')
    if optimizer == 'sgd':
      optimizer = tf.train.MomentumOptimizer(learning_rate=self.lr,
                                             momentum=0.9)
    elif optimizer == 'adam':
      optimizer = tf.train.AdamOptimizer(learning_rate=self.lr)
    else:
      raise ValueError('Unknown optimizer: %s' % optimizer)

    if use_tpu:
      # Wrap optimizer in CrossShardOptimizer which takes care of
      # synchronizing the weight updates between TPU cores.
      optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

    opt_step = optimizer.minimize(loss, var_list=var_list,
                                  colocate_gradients_with_ops=True)

    if self.update_batchnorm_params:
      opt_step = tf.group([opt_step] +
                          tf.get_collection(tf.GraphKeys.UPDATE_OPS))

    opt_step = tf.group([opt_step, self.global_step_inc])

    return opt_step

def make_estimator(mode, loss=None, eval_metrics=None, predictions=None):
  """Returns an EstimatorSpec (maybe TPU) for all modes."""

  # Always use TPUEstimator, even when not using TPU, then it's (almost) no-op.
  spec_type = tf.contrib.tpu.TPUEstimatorSpec

  if mode == tf.estimator.ModeKeys.PREDICT:
    assert predictions is not None, 'Need to pass `predict` arg.'
    return spec_type(mode=mode, predictions=predictions)

  if mode == tf.estimator.ModeKeys.EVAL:
    return spec_type(mode=mode, loss=loss, eval_metrics=eval_metrics)

  if mode == tf.estimator.ModeKeys.TRAIN:
    assert loss is not None, 'Need to pass `loss` arg.'
    trainer = Trainer(update_batchnorm_params=True)
    train_op = trainer.get_train_op(loss, use_tpu=FLAGS.use_tpu)
    return spec_type(mode=mode, loss=loss, train_op=train_op)

  raise ValueError('Unsupported mode %s' % mode)


import argparse
import logging
from typing import List

import torch
import torch.nn as nn

from .probe_base import ProbeBase

logger = logging.getLogger(__name__)

class OneWordNNProbe(ProbeBase):
    """
    Computes all squared L2 norm of n words as depths after an MLP projection.
    Can be used for probing the depth of words in a tree
    """

    def __init__(
        self, model_hidden_dim: int, probe_hidden_layers: int, intermediate_dim: int
    ):
        logging.info("Constructing OneWordNNDepthProbe")
        super(OneWordNNProbe, self).__init__()
        initial_linear = nn.Linear(model_hidden_dim, intermediate_dim)
        intermediate_layers: List[nn.Module] = []
        for i in range(probe_hidden_layers):
            intermediate_layers.append(nn.Linear(intermediate_dim, intermediate_dim))
            if i != probe_hidden_layers - 1:
                intermediate_layers.append(nn.ReLU())
        self.nn_probe = nn.Sequential(initial_linear, nn.ReLU(), *intermediate_layers)

    def forward(self, batch):
        """
        Computes all squared L2 norm of n words as depths after an MLP projection
         for each sentence in a batch. predicts the depth through an MLP
        Args:
            batch: a batch of word representations of the shape
                   (batch_size, max_seq_len, representation_dim)
        Returns: A tensor of depths of shape (batch_size, max_seq_len)
        """
        batch = self.nn_probe(batch)
        batchlen, seqlen, rank = batch.size()
        norms = torch.bmm(
            batch.view(batchlen * seqlen, 1, rank),
            batch.view(batchlen * seqlen, rank, 1),
        )
        norms = norms.view(batchlen, seqlen)
        return norms

    @staticmethod
    def add_probe_specific_args(
        parent_parser: argparse.ArgumentParser,
    ) -> argparse.ArgumentParser:
        parser = parent_parser.add_argument_group("ProbeArgs")
        parser.add_argument(
            "--probe_hidden_layers",
            type=int,
            default=2,
            help="Number of laers in probe.",
        )
        parser.add_argument(
            "--intermediate_dim",
            type=int,
            default=300,
            help="Dimension the probe maps embeddings to.",
        )
        return parent_parser