import numpy as np
from .correlated import correlated,correlated_ts
from .mc import mc
from .AddPaths import delta,x_hat,Sfunc
import time
import pandas as pd
from numba import njit,jit_module
from numba import prange,objmode
import sys
from datetime import datetime
import matplotlib.pyplot as plt
from scipy.stats import kurtosis

# @njit(nogil=True)
def select_levels(t0,T,M_t,eps,r,F,strategy=1,N=100,boundary=None):
    if strategy==1:
        dt_0 = np.minimum((eps)**2/r(np.array([0])),T-t0)
        levels = dt_0/M_t**np.arange(0,2)
        N_out=np.zeros(2,dtype=np.int64)
        N_diff=np.ones(2,dtype=np.int64)*N
        SS_out = np.zeros(2);C_out = np.zeros(2); E_out=np.zeros(2)
    else:
        dt_1 = np.minimum(eps**2/r(np.array([0])),T-t0)
        levels = np.append(np.array([T-t0]),dt_1/M_t**np.arange(0,2))
        N_out=np.zeros(3,dtype=np.int64)
        N_diff=np.ones(3,dtype=np.int64)*N
        SS_out = np.zeros(3);C_out = np.zeros(3); E_out=np.zeros(3)
    return levels,N_out,N_diff,E_out,SS_out,C_out

def rnd1(x, decimals, out):
    return np.round_(x, decimals, out)#.astype(np.int64)

# @njit(nogil=True,parallel=True)
def update_paths(I,E,SS,C,N,N_diff,levels,t0,T,M_t,eps,Q,M,r,F,boundary,strategy=1,rev=False,diff=False,v_ms=1,std=False,update=True):
    cores = 1 if update else 8
    n = np.maximum(2,rnd1(N_diff/cores,0,np.empty_like(N_diff/cores)).astype(np.int64))
    for j in range(len(I)):
        i=I[j]
        dt_f = levels[i]
        for k in prange(cores):
            if i!=0:
                if strategy==1 or i!=1:
                    with objmode(start1 = 'f8'):
                        start1 = time.perf_counter()
                    if rev:
                        x_f,x_c = correlated_ts(dt_f,M_t,t0,T,eps,n[i],Q,M,r,boundary=boundary,strategy=strategy,diff=diff,std=std)
                    else:
                        x_f,x_c = correlated(dt_f,M_t,t0,T,eps,n[i],Q,M,r,boundary=boundary,strategy=strategy,diff=diff)
                    with objmode(end1 = 'f8'):
                        end1 = time.perf_counter()
                else:
                    with objmode(start1 = 'f8'):
                        start1 = time.perf_counter()
                    x_f,x_c = correlated(dt_f,round(levels[0]/dt_f),t0,T,eps,n[i],Q,M,r,boundary=boundary,strategy=strategy,diff=diff,v_ms=v_ms)
                    with objmode(end1 = 'f8'):
                        end1 = time.perf_counter()
                est_f = F(x_f);est_c=F(x_c)
                C_temp = (end1-start1)/n[i]
                E_temp = np.mean(est_f-est_c)
                SS_temp = np.sum((est_f-est_c-E_temp)**2)
            else:
                with objmode(start2 = 'f8'):
                    start2 = time.perf_counter()
                x = mc(dt_f,t0,T,n[i],eps,Q,M,r,boundary=boundary,rev=rev,diff=diff,v_ms=v_ms)
                with objmode(end2 = 'f8'):
                    end2 = time.perf_counter()
                est = F(x)
                C_temp = (end2-start2)/n[i]
                E_temp = np.mean(est)
                SS_temp = np.sum((est-E_temp)**2)
            Delta = delta(E[i],E_temp,N[i],n[i])
            E[i] = x_hat(N[i],n[i],E[i],E_temp,Delta)
            if update:
                SS[i] = Sfunc(N[i],n[i],SS[i],SS_temp,Delta)
                '''Update cost like updating an average'''
                C[i] = x_hat(N[i],n[i],C[i],C_temp,delta(C[i],C_temp,N[i],n[i]))
            N[i] = N[i] + n[i]
    return E,SS,N,C

# @njit(nogil=True)
def ml(e2,Q,t0,T,M_t,eps,M,r,F,N_warm=40,boundary=None,strategy=1,alpha=None,beta=None,gamma=None,rev=False,diff=False,v_ms=1,std=False):
    '''
    e2: bound on mean square error
    Q: initial distribution
    t0: starting time
    T: end time
    M_t: dt_c/dt_f, where dt_c is coarse step size and dt_f is fine step size
    eps: mean free path, i.e. epsilon in model
    M: equilibrium distribution for velocity
    r: collision rate
    F: function used to find quantity of interest, E(F(X,V)).
    '''
    levels,N,N_diff,E,SS,C = select_levels(t0,T,M_t,eps,r,F,N=N_warm,strategy=strategy)
    '''
    levels: a list of step sizes for each level
    N: list of number of paths used
    N_diff: list of number of paths needed at every level
    E: estimation of quantity of interest at each level
    SS: sum of squares at each level
    C: Cost at each level
    '''
    L = len(levels)-1
    '''While loop to continue until RMSE fits with e2'''
    while True:
        '''Update paths based on N_diff'''
        while np.max(N_diff)>0:
            I = np.where(N_diff > 0)[0] #Index for Levels that need more paths
            N_diff = np.minimum(N_diff,np.ones(len(N_diff),dtype=np.int64)*400_000)
            # print(f'index where more paths are needed: {I}, N_diff: {N_diff}')
            first = np.sum(N)==0 or (beta is None and gamma is None)
            E,SS,N,C = update_paths(I,E,SS,C,N,N_diff,levels,t0,T,M_t,eps,Q,M,r,F,boundary,strategy,rev=rev,diff=diff,v_ms=v_ms,std=std,update=first)
            if first: V = SS/(N-1) #Update variance
            '''Determine number of paths needed with new information'''
            N_diff = np.ceil(2/e2*np.sqrt(V/C)*np.sum(np.sqrt(V*C))).astype(np.int64) - N
        '''Test bias is below e2/2 - to do so it is necessary to know exponent
        in weak error'''
        if (E.size>=4 and strategy==1) or (strategy==3 and E.size>4):
            if alpha is None:
                L1 = max(1,np.where(levels<=eps**2/r(1))[0][0])
                pa = lin_fit(np.arange(L1,L),np.log2(np.abs(E[L1:L+1]))); alpha = -pa[0]
            test = np.max(np.abs(E[-3:]))/(M_t**alpha-1) < np.sqrt(e2/2)
        else:
            test=False
        if test:
            break
        L += 1;
        if L==17:
            print('WARNING: maximum capacity reached!')
            break
        print('Level added')
        if beta is None and gamma is None:
            N_diff = np.append(N_diff,N_warm).astype(np.int64)
            N = np.append(N,0).astype(np.int64)
            E = np.append(E,0.0); V = np.append(V,0.0); C = np.append(C,0.0); SS = np.append(SS,0.0)
        else:
            #Calculate constant from Theorem
            c2 = V[-1]*2**((L-1)*beta)
            V = np.append(V,c2*2**(-L*beta))
            c3 = C[-1]*2**(-(L-1)*gamma)
            C = np.append(C,c3*2**(L*gamma))
            E = np.append(E,0.0);
            N = np.append(N,0).astype(np.int64)
            N_diff = np.ceil(2/e2*np.sqrt(V/C)*np.sum(np.sqrt(V*C))).astype(np.int64) - N
        if strategy==1:
            levels = np.append(levels,levels[0]/(M_t**(L)))
        else:
            levels = np.append(levels,levels[1]/(M_t**(L)))
    return E,V,C,N,levels

def lin_fit(x,y):
    p=1; n = x.size
    X = np.stack((np.ones(n),x),axis=1)
    normal_matrix = (X.T).dot(X)
    moment_matrix = (X.T).dot(y)
    return np.linalg.inv(normal_matrix).dot(moment_matrix)

@njit(nogil=True,parallel=True)
def convergence_tests(N,dt_list,Q,t0,T,M_t,eps,M,r,F,boundary,strategy,rev=False,diff=False,v_ms=1,std=False):
    '''Calculates values for consistency test for each level given by dt_list'''
    cores = 8 #Controls parrelisation

    L = dt_list.size
    b = np.zeros(L) #mean(F(x^f)-F(X^c))
    b2 = np.zeros(L) #mean((F(x^f)-F(X^c))^2)
    b3 = np.zeros(L) #mean((F(x^f)-F(X^c))^3)
    b4 = np.zeros(L) #mean((F(x^f)-F(X^c))^4)
    v = np.zeros(L) #mean(F(X))
    v2 = np.zeros(L) #mean(F(X)^2)
    var1 = np.zeros(L) #var(F(X))
    var2 = np.zeros(L) #var(F(X^f)-F(X^c))
    # kur1 = np.empty(L-1) #kurtosis calculated for biases
    cons = np.zeros(L) #consistency calculated for all biases
    cost1 = np.zeros(L) #cost for each level
    cost2 = np.zeros(L) #cost for each level bias

    if N%cores!=0:
        print('WARNING -  Number of samples is not divisble by 8!\n Please change N for optimal functionality')
    n = round(N/cores)
    # all_diff = np.zeros((N,L))

    for l in range(L):
        print(l)
        diff_vec = np.empty((cores,n))
        val = np.empty((cores,n))
        for j in prange(cores):
            if l<L-1:
                # print(f'l={l}')
                with objmode(start1 = 'f8'):
                    start1 = time.perf_counter()
                if rev:
                    x_f,x_c = correlated_ts(dt_list[l+1],M_t,t0,T,eps,n,Q,M,r,boundary=boundary,diff=diff,std=std)
                else:
                    x_f,x_c = correlated(dt_list[l+1],M_t,t0,T,eps,n,Q,M,r,boundary=boundary,diff=diff,v_ms=v_ms)
                with objmode(end1 = 'f8'):
                    end1 = time.perf_counter()
                cost2[l+1] += (end1-start1)
                diff_vec[j,:] = F(x_f)-F(x_c)
            with objmode(start2 = 'f8'):
                start2 = time.perf_counter()
            x = mc(dt_list[l],t0,T,n,eps,Q,M,r,boundary=boundary,rev=rev,diff=diff,v_ms=v_ms)
            with objmode(end2 = 'f8'):
                end2 = time.perf_counter()
            cost1[l] += (end2-start2)
            val[j,:] = F(x)
        cost1[l] = cost1[l]/N
        v[l] = np.mean(val)
        v2[l] = np.mean(val**2)
        var1[l]  = v2[l] - v[l]**2
        if l<L-1:
            # all_diff[:,l+1] = diff.flatten()
            b[l+1] = np.mean(diff_vec)
            b2[l+1] = np.mean(diff_vec**2)
            b3[l+1] = np.mean(diff_vec**3)
            b4[l+1] = np.mean(diff_vec**4)
            var2[l+1] = b2[l+1]-b[l+1]**2
            cost2[l+1] = cost2[l+1]/N
            # cons[l+1] = np.abs(b[l+1]+v[l]-v[l+1])/(3*(np.sqrt(var2[l+1])+ np.sqrt(var1[l])+np.sqrt(var1[l+1]))/np.sqrt(N))
    cons[1:] = np.abs(b[1:]+v[:-1]-v[1:])/(3*(np.sqrt(var2[1:])+np.sqrt(var1[:-1])+np.sqrt(var1[1:]))/np.sqrt(N))
    kur1 = np.append(0,(b4[1:]-4*b3[1:]*b[1:]+6*b2[1:]*b[1:]**2-3*b[1:]**4)/(b2[1:]-b[1:]**2)**2)
    return b,b2,b3,b4,v,v2,var1,var2,kur1,cons,cost1,cost2#,all_diff

jit_module(nopython=True,nogil=True)

def ml_test(N,N0,dt_list,E2,Q,t0,T,M_t,eps,M,r,F,logfile,boundary=None,strategy=1,convergence=True,complexity=True,rev=False,diff=False,v_ms=1,std=False):
    ''''
    filename for logfile should always begin with 'logfile_APS'

    '''
    save_file= logfile is not None
    L = dt_list.size
    if save_file:
        if logfile.name[0:11]!='logfile_APS':
            sys.exit('ERROR: name of logfile should start with "logfile_APS"')
        now = datetime.now().strftime("%d-%B-%Y %H:%M:%S")
        logfile.write("\n")
        logfile.write("*********************************************************\n")
        logfile.write(f"***Python ml_test for APS method on {now}         ***\n")
        if rev:
            logfile.write(f"***    Reverse splitting is used!!      ***\n")
        if diff:
            logfile.write(f"***    Altered diffusive coefficient is used!!      ***\n")
        logfile.write("\n")
        logfile.write("*********************************************************\n")
        logfile.write("*** Experiemnt setup  ***\n")
        logfile.write("*** S(x,v) = 1/sqrt(2*pi)e^{-v^2/2}(1+cos(2pi*(x+1/2))) ***\n")
        logfile.write(f"*** r(x) = (-a(x-0.5)+b)(x<=0.5) + (a(x-0.5)+b)(x>0.5), with a={(r(np.array([1]))-r(np.array([0])))}, b={r(np.array([0]))} and eps = {eps}  ***\n")
        if 2!=F(2):
            logfile.write(f"*** OBS!! replace this text with specification of quantity of interest F(3) = {F(np.array([3]))} ***\n")
        else:
            logfile.write("*** Quantity of interest is F(X) = X  ***\n")
        if boundary is None:
            logfile.write("*** No boundary conditions  ***\n")
        else:
            logfile.write("*** OBS!! replace this text with specification of boundary conditions  ***\n")
        logfile.write("*********************************************************\n")
        logfile.write("Convergence tests, kurtosis, telescoping sum check \n")
        logfile.write(f"*** using {N} samples and {L} levels ***\n")
        logfile.write(" l dt^f mean(F(X^f)-F(X^c)) mean(F(X^c))  var(F(X^f)-F(X^c)) var(F(X^c))")
        logfile.write(" cost(F(X^f)-F(X^c)) cost(F(X)) kurtosis consistency \n")
        logfile.write("---------------------------------------------------------\n")

    if convergence:
        print('Convergence test')
        b,b2,b3,b4,v,v2,var1,var2,kur1,cons,cost1,cost2 = convergence_tests(N,dt_list,Q,t0,T,M_t,eps,M,r,F,boundary=boundary,strategy=strategy,rev=rev,diff=diff,v_ms=v_ms,std=std)
        print('Convergence test DONE')
        print(kur1)
        if save_file:
            for i in range(dt_list.size):
                logfile.write(f'{i} {dt_list[i]} {b[i]} {v[i]} {var2[i]} {var1[i]} {cost2[i]} {cost1[i]} {kur1[i]} {cons[i]}\n')

        # Linear regression to estimate alpha, beta and gamma. Only test for dt << eps^2
        i = 0
        old = 0
        for va in var2[1:]:
            if va < old:
                break
            i+=1
            old=va
        alpha = -np.polyfit(range(i,var2.size),np.log2(np.abs(b[i:])),1)[0]
        beta = -np.polyfit(range(i,var2.size),np.log2(np.abs(var2[i:])),1)[0]
        gamma = np.polyfit(range(i,var2.size),np.log2(np.abs(cost2[i:])),1)[0]
        print(f'alpha= {alpha}, beta = {beta}, gamma= {gamma}')
        if save_file:
            np.savetxt('resultfile'+logfile.name[7:],(dt_list,v,b,var1,var2,cost1,cost2,kur1,cons))
            logfile.write('\n*********************************************************\n')
            logfile.write('\n*** Linear regression estimates of MLMC paramters ***\n')
            logfile.write(f'\n*** regression is done for levels with dt << eps^2 = {eps**2} ***\n')
            logfile.write('*********************************************************\n')
            logfile.write(f'alpha = {alpha} (exponent for weak convergence) \n')
            logfile.write(f'beta = {beta} (exponent for variance of bias estimate) \n')
            logfile.write(f'gamma = {gamma} (exponent for cost of bias estimate) \n')

    if complexity:
        if save_file:
            logfile.write("\n*********************************************************\n")
            logfile.write("*** MLMC complexity test ***\n")
            logfile.write("*********************************************************\n")
            logfile.write(" e2 value mlmc_cost N_l dt \n")
            logfile.write("---------------------------------------------------------\n")
        data = {}
        pd.set_option('max_columns',None)
        # for e2 in E2:
        #     print(f'MSE= {e2}')
        #     E,V,C,N,levels = ml(e2,Q,t0,T,M_t,eps,M,r,F,N0,boundary=boundary,strategy=strategy,alpha=1,rev=rev,diff=diff)
        #     data[e2] = {'dt':levels,'N_l':N,'E':E,'V_l':V,'V[E]':V/N,'C_l':C,'N_l C_l':N*C}
        #     df = pd.DataFrame(data[e2])
        #     df = df.append(pd.DataFrame({'dt':' ','N_l':' ','E':[np.sum([e for e in data[e2]['E']])],'V_l':' ','V[E]':[np.sum([v for v in data[e2]['V[E]']])],'C_l':[np.sum([c for c in data[e2]['C_l']])],'N_l C_l':[np.sum(n*c for n,c in zip(data[e2]['N_l'],data[e2]['C_l']))]}))
        #     print(df)
        #     if save_file:
        #         name = f'resultfile_complexity_{e2}'+logfile.name[7:]
        #         df.to_csv(name,index=False)
        #         logfile.write(f" {e2} {np.sum(E)} {np.dot(C,N)} {N} {levels} \n")

        if convergence:
            '''Find index where they transition'''
            i = 0
            old = 0
            for va in var2[1:]:
                if va < old:
                    break
                i+=1
                old=va
            # print(f'slope of variance before:{np.polyfit(range(1,i-2),np.log2(np.abs(var2[1:i-2])),1)}')
            # print(f'slope of variance after:{np.polyfit(range(i+2,var2.size),np.log2(np.abs(var2[i+2:])),1)}')
            # print(f'slope of bias before:{np.polyfit(range(1,i-2),np.log2(np.abs(bias[1:i-2])),1)}')
            alpha = -np.polyfit(range(i,var2.size),np.log2(np.abs(b[i:])),1)[0]
            # print(f'slope of bias after:{alpha}')
            # print(f'slope of cost:{np.polyfit(range(i,cost2.size),np.log2(np.abs(cost2[i:])),1)}')
            '''Complexity test based on estimations in convergence test'''
            dfs = {}
            pd.set_option('display.float_format', '{:.2E}'.format)
            for e2 in E2:
                L=4 + i
                E = np.append(v[i],b[i:L])
                V = np.append(var1[i],var2[i:L])
                C = np.append(cost1[i],cost2[i:L])
                N = np.ceil(2/e2*np.sqrt(V/C)*np.sum(np.sqrt(V*C)))
                test = (np.max(np.abs(E[-3:]))/(2**alpha-1) < np.sqrt(e2/2))
                out=False
                while not test:
                    L+=1
                    # print(L)
                    if L>=v.size:
                        out = True
                        print(f'Warning!. Need more results to adhere to bias requirement for e2= {e2}')
                        break
                    N = np.append(N,0)
                    E = np.append(E,b[L])
                    V = np.append(V,var2[L])
                    C = np.append(C,cost2[L])
                    N += np.maximum(0,np.ceil(2/e2*np.sqrt(V/C)*np.sum(np.sqrt(V*C))).astype(np.int64))
                    test = (np.max(np.abs(E[-3:]))/(2**alpha-1) < np.sqrt(e2/2))
                if not out:
                    df = pd.DataFrame({'E':E,'V':V,'V[Y]':V/N,'C':C,'N':N,'N C':C*N,'dt':dt_list[i-1:min(L,v.size-1)]})
                    df = df.append(pd.DataFrame({'E':np.sum(E),'V':' ','V[Y]':[np.sum(V/N)],'C':[np.sum(C)],'N':' ','N C':[np.dot(C,N)],'dt':' '}),ignore_index=True)
                    # print(df)
                    dfs[e2] = df
                    if save_file:
                        name = f'resultfile_complexity_{e2}'+logfile.name[7:]
                        df.to_csv(name,index=False)
                        logfile.write(f" {e2} {np.sum(E)} {np.dot(C,N)} {N} {dt_list[i-1:min(L,v.size-1)]} \n")
                else:
                    break
            '''Plotting number of paths and computational costs as functions of MSE'''
            fig, (ax1, ax2) = plt.subplots(2, 1)
            N = []
            Cost = []
            end = len(dfs)
            start = end-min(end,6)
            for e2 in E2[start:end]:
                N += [dfs[e2]['N'][:-1]]
                Cost += [dfs[e2]['N C'][N[-1].size]]
                # print(Cost)
                ax1.plot(range(N[-1].size),N[-1],label=r'$E^2 = ${:.2E}'.format(e2))
            ax2.loglog(E2[start:end],Cost)
            c4 = np.mean(Cost*E2[start:end]/np.log(E2[start:end]))
            c4 *= 2
            ax2.plot(E2[start:end],c4*np.log(E2[start:end])/E2[start:end],label= r'$\mathcal{O}(\log{E}^2E^{-2})$')
            ax1.set_yscale('log')
            ax2.set_xlabel(r'$E^2$')
            ax2.set_ylabel(r'Total costs')
            ax1.set_xlabel(r'Levels')
            ax1.set_ylabel(r'Paths')
            ax2.legend()
            ax1.legend()
            if save_file:
                filename = f'complexity_results'+logfile.name[7:-3]
                fig.savefig(filename)

        else:
            # Projected cost based on Multilevel Theorem
            cost_e=0; l = 1
            e2 = 1/2**l
            # Wall clock time
            WC = []
            print('Compile')
            # ml(e2,Q,t0,T,M_t,eps,M,r,F,N0,boundary=boundary,strategy=strategy,alpha=1.0,beta=1.0,gamma=1.0,rev=rev,diff=diff)
            print('Compilation done')
            while cost_e<3000:
                print(f'MSE= {e2}')
                start = time.perf_counter()
                E,V,C,N,levels = ml(e2,Q,t0,T,M_t,eps,M,r,F,N0,boundary=boundary,strategy=strategy,alpha=1.0,rev=rev,diff=diff)
                WC += [time.perf_counter()-start]
                data[e2] = {'dt':levels,'N_l':N,'E':E,'V_l':V,'V[E]':V/N,'C_l':C,'N_l C_l':N*C}
                df = pd.DataFrame(data[e2])
                df = df.append(pd.DataFrame({'dt':' ','N_l':' ','E':[np.sum([e for e in data[e2]['E']])],'V_l':' ','V[E]':[np.sum([v for v in data[e2]['V[E]']])],'C_l':[np.sum([c for c in data[e2]['C_l']])],'N_l C_l':[np.sum(n*c for n,c in zip(data[e2]['N_l'],data[e2]['C_l']))]}))
                print(df)
                if save_file:
                    name = f'resultfile_complexity_{e2}'+logfile.name[7:]
                    df.to_csv(name,index=False)
                    logfile.write(f" {e2} {np.sum(E)} {np.dot(C,N)} {N} {levels} \n")
                #Actual simulation cost
                cost_s = WC[-1]#np.dot(C,N)
                #Constant from Theorem
                c4 = c4 = cost_s*e2/(np.log(np.sqrt(e2))**2)
                l+=1
                e2 = 1/2**l
                cost_e = c4*np.log(np.sqrt(e2))**2/e2
            if save_file:
                name = f'wall_clock_time_{e2}'+logfile.name[7:]
                np.savetxt(name,WC)

    if save_file:
        logfile.write('\n')
        logfile.close()

    if convergence:
        plt.figure()
        plt.plot(dt_list[1:],var2[1:],':',label='var(F(x^f)-F(X^c))')
        plt.plot(dt_list,var1,'--',color = plt.gca().lines[-1].get_color(),label='var(F(X))')
        plt.plot(dt_list[1:],np.abs(b[1:]),':',label='mean(|F(x^f)-F(X^c)|)')
        plt.plot(dt_list,v,'--',color = plt.gca().lines[-1].get_color(),label='mean(F(X))')
        # plt.title(f'Plot of variance and bias')
        plt.xscale('log')
        plt.yscale('log')
        plt.xlim(max(dt_list),min(dt_list))
        plt.legend()
        if save_file:plt.savefig(f'var_bias'+logfile.name[7:-3])
        plt.figure()
        plt.plot(range(1,kur1.size),kur1[1:],':')
        plt.title(f'Plot of kurtosis')
        if save_file:plt.savefig(f'kurtosis'+logfile.name[7:-3])
        plt.figure()
        plt.plot(range(1,kur1.size),cons[1:],':')
        plt.title(f'Plot check of consistency')
        if save_file:plt.savefig(f'consistency_check'+logfile.name[7:-3])
        if not save_file:plt.show()


# -*- coding: utf-8 -*-
import sys
from pathlib import Path
from typing import Any, Optional, Tuple
from at.auth.client import AuthStatus, licensed
from at.gui import *
from at.gui.worker import run_thread
from at.logger import log
from at.result import Result
from at.io.utils import load_json, write_json
from atparser.app.utils.path import paths
from atparser.app.utils.state import state
# from PyQt5.QtCore import Qt, QThreadPool, pyqtSignal
# from PyQt5.QtGui import QFont
# from PyQt5.QtWidgets import QApplication, QHBoxLayout, QVBoxLayout, QWidget

# When setting fixed width to QLineEdit ->
# -> add alignment=Qt.AlignmentFlag.AlignLeft when adding widget to layout

class SettingsTab(AtWidget):
    settingsChanged = pyqtSignal()
    serverStatusChanged = pyqtSignal(tuple)

    def __init__(self,
                 size: Tuple[Optional[int]] = (None, None),
                 parent: Optional[QWidget] = None,
                 *args,
                 **kwargs) -> None:
        super().__init__(parent=parent, *args, **kwargs)
        self.initUi(size)

    def setupUi(self, size):
        set_size(widget=self, size=size)

        layout = QVBoxLayout()
        layout.setContentsMargins(2, 4, 2, 0)
        driverLayout = QHBoxLayout()
        labelLayout = QHBoxLayout()
        buttonLayout = QHBoxLayout()

        self.app = Label(icon='app',
                         label=state.appname,
                         parent=self)
        self.version = Label(icon='hash',
                             label=state.version,
                             parent=self)
        self.driverSelect = ComboInput(label="WebDriver",
                                       labelsize=(140, 24),
                                       combosize=(100, 24),
                                       items=('Firefox', 'Chrome'),
                                       parent=self)
        self.driverExist = StatusLabel(icon='hdd-network-fill',
                                       statussize=(110, 24),
                                       parent=self)
        self.delayLaunch = IntInput(label='Launch delay',
                                    labelsize=(140, 24),
                                    editsize=(100, 24),
                                    parent=self)
        self.delayUrlChange = IntInput(label='URL change delay',
                                       labelsize=(140, 24),
                                       editsize=(100, 24),
                                       parent=self)
        self.delayPaginator = IntInput(label='Paginator delay',
                                       labelsize=(140, 24),
                                       editsize=(100, 24),
                                       parent=self)
        self.saveButton = Button(label="Αποθήκευση Αλλαγών",
                                 size=(160, 26),
                                 parent=self)
        self.status = StatusButton(parent=self)

        labelLayout.addWidget(self.app)
        labelLayout.addWidget(self.version, stretch=2, alignment=Qt.AlignmentFlag.AlignLeft)
        layout.addLayout(labelLayout)
        layout.addWidget(HLine())
        driverLayout.addWidget(self.driverSelect)
        driverLayout.addWidget(self.driverExist,
                               stretch=2,
                               alignment=Qt.AlignmentFlag.AlignLeft)
        layout.addLayout(driverLayout)
        layout.addWidget(self.delayLaunch, alignment=Qt.AlignmentFlag.AlignLeft)
        layout.addWidget(self.delayUrlChange, alignment=Qt.AlignmentFlag.AlignLeft)
        layout.addWidget(self.delayPaginator, alignment=Qt.AlignmentFlag.AlignLeft)

        buttonLayout.addWidget(self.saveButton, alignment=Qt.AlignmentFlag.AlignRight)
        layout.addLayout(buttonLayout)

        layout.addWidget(HLine())
        layout.addWidget(self.status, stretch=2, alignment=Qt.AlignmentFlag.AlignBottom)

        self.setLayout(layout)

    def initUi(self, size: tuple):
        self.setupUi(size)
        self.driverSelect.setCurrentText(state.webdriver)
        self.delayLaunch.setText(str(state.launch_delay))
        self.delayUrlChange.setText(str(state.urlchange_delay))
        self.delayPaginator.setText(str(state.paginator_delay))
        self.checkDrivers()
        self.driverSelect.subscribe(self.onDriverChange)
        self.saveButton.subscribe(self.onSave)

    def checkDrivers(self):
        if self.driverSelect.getCurrentText() == 'chrome':
            if paths.get_chromedriver(True).exists():
                self.driverExist.changeStatus("Driver exists", 'statusOk')
            else:
                self.driverExist.changeStatus("No Driver", 'statusError')
        else:
            if paths.get_firefoxdriver(True).exists():
                self.driverExist.changeStatus("Driver exists", 'statusOk')
            else:
                self.driverExist.changeStatus("No Driver", 'statusError')

    def onDriverChange(self):
        self.pickedDriver = self.driverSelect.getCurrentText()
        self.checkDrivers()

    def onSave(self):
        state.webdriver = self.driverSelect.getCurrentText()

        state.launch_delay = self.delayLaunch.getInt()
        state.urlchange_delay = self.delayUrlChange.getInt()
        state.paginator_delay = self.delayPaginator.getInt()

        self.settingsChanged.emit()

        self.popup.info("Οι ρυθμίσεις αποθηκεύτηκαν")

if __name__ == '__main__':
    cssGuide = paths.get_css(obj=True).joinpath("_style.css").read_text()
    SEGOE = QFont("Segoe UI", 9)

    app = QApplication(sys.argv)
    app.setFont(SEGOE)
    app.setStyle('Fusion')

    ui = SettingsTab(size=(600, None))
    ui.setStyleSheet(cssGuide)
    ui.show()

    sys.exit(app.exec_())


# coding:utf-8
#
# chatgpt/find_lora_module_name.py
#
# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on August 19, 2023
# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on September 4, 2023
#
# 查看模型 lora 模块名.

from transformers import AutoModel

model = AutoModel.from_pretrained("chatgpt/Atom-7B")

for name, module in model.named_modules():
    print(name)


"""
Implementaion for Beta VAE
"""
from distutils import log
from torch.utils import data
from carla_disentanglement.datasets.ground_truth_data import GroundTruthDataset
from carla_disentanglement.metrics.factor_vae_score import factor_vae_score
import os
import subprocess
from typing import Callable
from collections.abc import Sequence
from datetime import datetime
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.utils import save_image
from PIL import Image
from PIL.Image import Image as PILImage
import numpy as np
from tqdm import tqdm
from carla_disentanglement.metrics.beta_vae_score import beta_vae_score
from torch.utils.tensorboard import SummaryWriter
from disentanglement_lib.evaluation.metrics import beta_vae
from disentanglement_lib.evaluation.metrics import factor_vae
from disentanglement_lib.evaluation.metrics import dci
from disentanglement_lib.evaluation.metrics import mig
from disentanglement_lib.evaluation.metrics import sap_score

import gin.tf
gin.parse_config_file("metrics.gin")

class VAEModule(nn.Module):
    """
    pytorch Module for Variational Auto Encoders
    """

    def __init__(self, encoder: nn.Module, decoder: nn.Module, activation=nn.Sigmoid(), loss_activation: nn.Module = None):
        super().__init__()

        self.encoder = encoder
        self.decoder = decoder
        self.activation = activation
        self.use_model_regulariser = False
        self.loss_activation = loss_activation if loss_activation is not None else self.activation

    def encode(self, x: torch.Tensor):
        """Encodeds Input to latents

        Args:
            x (Tensor): input

        Returns:
            Tensor: latent encoding
        """
        return self.encoder(x)

    def decode(self, z):
        """Decondes latent z into output (with (sigmoid) activation)

        Args:
            z (Tensor): latent

        Returns:
            Tensor: output
        """
        return self.activation(self.decoder(z))

    def reparametrize(self, mu, logvar):
        """Reparamterization trick for training

        Args:
            mu (Tensor): mean of latents
            logvar (Tensor): log variance of latents

        Returns:
            Tensor: latent sample
        """
        std = logvar.mul(0.5).exp_()
        eps = std.data.new(std.size()).normal_()
        return eps.mul(std).add_(mu)

    def forward(self, x):
        """Forward pass from input to reconstruction (including reparamterization trick).
        Omits (sigmoid) activation for the use of BCEWithLogitsLoss

        Args:
            x (Tensor): input

        Returns:
            Tensor, (Tensor, Tensor): reconstruction, (mu, log_var), z
        """
        mu, logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        # NOTE: we use decodeR(!) here not decode (which includes sigmoid activation)
        return self.loss_activation(self.decoder(z)), (mu, logvar), z

    def reconstruct(self, x):
        """Forward pass from input to reconstruction, without reparamterization trick.
        Only uses mu, ignores logvar

        Args:
            x (Tensor): input

        Returns:
            Tensor, (Tensor, Tensor): reconstruction, (mu, log_var)
        """
        mu, logvar = self.encode(x)
        return self.activation(self.decoder(mu)), (mu, logvar)

    def eval_forward(self, x):
        """Forward pass from input to reconstruction, without reparamterization trick.
        Only uses mu, ignores logvar
        Omits (sigmoid) activation for the use of BCEWithLogitsLoss

        Args:
            x (Tensor): input

        Returns:
            Tensor, (Tensor, Tensor): reconstruction, (mu, log_var)
        """
        mu, logvar = self.encode(x)
        return self.loss_activation(self.decoder(mu)), (mu, logvar)

    def model_regulariser(self):
        return 0

class BetaVAE(object):
    """
    Beta VAE model
    https://openreview.net/forum?id=Sy2fzU9gl
    """

    def __init__(self, model: VAEModule, beta=1.0, reconstruction='bce', tag=None, model_regulariser_weight=0.01, tag_suffix: str = None):

        self.model = model
        if reconstruction == 'mse_no_activation':
            self.reconstruction_loss = nn.MSELoss(reduction='sum')
            self.model.loss_activation = nn.Identity()
        elif reconstruction == 'mse':
            self.reconstruction_loss = nn.MSELoss(reduction='sum')
        else:
            self.reconstruction_loss = nn.BCEWithLogitsLoss(reduction='sum')
            self.model.loss_activation = nn.Identity()

        self.num_channels = model.encoder._num_channels
        self.z_dim = model.encoder._latent_dim
        self.image_size = model.encoder._image_size

        self.beta = beta
        self.model_regulariser_weight = model_regulariser_weight

        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(self.device)
        self.model.to(self.device)

        tag = str(tag)+'_' if tag else ''
        self.full_tag = tag+str(self.__class__.__name__)+self.getParamString()
        if tag_suffix is not None:
            self.full_tag += "_" + tag_suffix
        print(self.full_tag, tag_suffix)
        self.num_log_values = 4

        self.initOptimizer()

    def getParamString(self):
        return "_beta"+str(self.beta)

    def gauss_vae_regulariser(self, mu: torch.Tensor, log_var: torch.Tensor):
        """
        Compute the regularisation factor for training a VAE
        with a Gaussian latent space.

        Parameters
        ----------
        mu: Tensor
            Mean predictions from variational encoder.
        log_var: Tensor
            log_var predictions from variational encoder.

        Returns
        -------
        kl_div : Tensor
            The Kullback-Leibler divergence between a Gaussian distribution
            with the given parameters and the standard normal distribution.
        log info : None
        """
        return ((torch.exp(log_var) + mu*mu - log_var - 1)/2).sum(1).mean(0)

    def _visualize_reconstruction(self, ds: torch.utils.data.Dataset, count=10, indices=None):
        """Create visualatin of VAE reconstruction

        Args:
            ds (torch.utils.data.Dataset): dataset
            count (int, optional): Number of sampels to be visualized. Defaults to 10.

        Returns:
            torch.Tensor: image
        """

        if indices != None:
            ds = torch.utils.data.Subset(ds, indices)
            count = len(indices)

        data_loader = DataLoader(ds, batch_size=count)
        xs = next(iter(data_loader))[0]

        with torch.no_grad():
            # get reconstructions
            preds = self.model.reconstruct(xs.to(self.device))[0].cpu()

            # concatenate all data
            x_cat = torch.cat(
                [torch.cat([x, torch.zeros(self.num_channels, self.image_size, 1)], dim=-1) for x in xs], dim=-1)
            pred_cat = torch.cat(
                [torch.cat([x, torch.zeros(self.num_channels, self.image_size, 1)], dim=-1) for x in preds], dim=-1)

            return torch.cat([x_cat, pred_cat], dim=1)

    def visualize_reconstruction(self, ds: torch.utils.data.Dataset, count=10, indices=None):
        """Create visualatin of VAE reconstruction
        How to use: display(image, metadata={'width': '100%'})

        Args:
            ds (torch.utils.data.Dataset): dataset
            count (int, optional): Number of sampels to be visualized. Defaults to 10.

        Returns:
            PIL.Image: image
        """

        image = self._visualize_reconstruction(ds, count, indices)

        to_image = transforms.Compose([
            # inverts normalization of image # implement unnormalize when we use normalization
            # transforms.Normalize(-means / stds, 1. / stds),
            transforms.Lambda(lambda x: torch.clamp(x, 0, 1)),
            transforms.ToPILImage()
        ])

        return to_image(image)

    def visualize_latents_static(self, ds: torch.utils.data.Dataset, index=0, limit=3):
        image = self._visualize_latents_static(ds, index, limit=limit)

        to_image = transforms.Compose([
            # inverts normalization of image # implement unnormalize when we use normalization
            # transforms.Normalize(-means / stds, 1. / stds),
            transforms.Lambda(lambda x: torch.clamp(x, 0, 1)),
            transforms.ToPILImage()
        ])

        return to_image(image)

    def _visualize_latents_static(self, ds: torch.utils.data.Dataset, index=0, limit=3):
        inter = 2/3
        interpolation = torch.arange(-limit, limit+0.1, inter)

        img = ds.__getitem__(index)[0].cuda().unsqueeze(0)
        z_org = self.model.encode(img)[0]
        z_dim = z_org.shape[-1]
        if len(z_org.shape) < 2:
            z_org = z_org.unsqueeze(0)

        samples = []
        for z_index in range(z_dim):
            z = z_org.clone()
            inters = []
            for val in interpolation:
                z[:, z_index] = val
                sample = self.model.decode(z).data
                inters.append(sample)
                inters.append(torch.full((1, sample.shape[1], sample.shape[2], 1), .5).cuda())

            samples.append(torch.cat(inters, -1))
            samples.append(torch.full((1, sample.shape[1], 1, samples[-1].shape[-1]), .5).cuda())

        samples = torch.cat(samples, -2).cpu()
        return samples.squeeze()

    def visualize_latents(self, ds: torch.utils.data.Dataset, index=0, output_dir='./viz/', delay=100):
        limit = 3
        inter = 2/3
        interpolation = torch.arange(-limit, limit+0.1, inter)

        img = ds.__getitem__(index)[0].cuda().unsqueeze(0)
        z_org = self.model.encode(img)[0]
        z_dim = z_org.shape[-1]
        if len(z_org.shape) < 2:
            z_org = z_org.unsqueeze(0)

        samples = []
        for z_index in range(z_dim):
            z = z_org.clone()
            for val in interpolation:
                z[:, z_index] = val
                sample = self.model.decode(z).data
                samples.append(sample)

        samples = torch.cat(samples).cpu()
        gif = samples.view(z_dim, len(interpolation), self.num_channels, self.image_size, self.image_size).transpose(0, 1)

        for j in range(gif.shape[0]):
            save_image(tensor=gif[j].cpu(),
                       fp=os.path.join(output_dir, '{}.png'.format(j)),
                       nrow=z_dim, pad_value=1)

        image_str = os.path.join(output_dir, '*.png')
        output_gif = os.path.join(output_dir, self.full_tag+'_'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")+'.gif')
        str1 = 'magick convert -delay '+str(delay)+' -loop 0 ' + image_str + ' ' + output_gif
        subprocess.call(str1, shell=True)

        return output_gif

    def initOptimizer(self, lr=1e-4, betas=None):
        """Initialize adam optimizer

        Args:
            lr (float, optional): Learning rate. Defaults to 1e-4.
            betas ((float,float), optional): betas param for Adam.
        """
        if betas is not None:
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr, betas=betas)
        else:
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)

    def beforeTrain(self):
        # keep track of training
        self.writer = SummaryWriter(
            "runs/"+self.full_tag+"_"+datetime.now().strftime('%b%d_%H-%M-%S'))

    def afterTrainStep(self, epoch_data, epoch_step):
        pass

    def train(self, ds: torch.utils.data.Dataset,
              batch_size=256,
              epochs=100,
              eval_every=10,
              score_every=10,
              reconstruction_every=10,
              reconstruct_indices: Sequence = None,
              save_every=20,
              save_path="./checkpoints/",
              num_workers=8,
              score_ds: GroundTruthDataset = None,
              max_iter=None,
              log_deadNeurons=False,
              log_stepdata=False):
        """Start training

        Args:
            ds (torch.utils.data.Dataset): Dataset
            batch_size (int, optional): batch size. Defaults to 256.
            epochs (int, optional): number of epoches. Defaults to 100.
            score (int, optional): after how many epochs score metrics will be calculated. Defaults to 10.
            eval_every (int, optional): after how many epochs evaluation will be perforemd. Defaults to 10.
            reconstruction_every (int, optional): after how many epochs reconstruction_callback will be called. Defaults to 10.
            save_every (int, optional): after how many epochs model will be saved. Defaults to 20.
            save_path (str, optional): path for model state dict. Defaults to "./checkpoints/".
        """

        self.beforeTrain()

        train_size = int(0.75 * len(ds))
        test_size = len(ds) - train_size
        train_dataset, test_dataset = torch.utils.data.random_split(ds, [train_size, test_size], generator=torch.Generator().manual_seed(0))

        data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers)
        self.global_iter = 0
        self.model.train()

        last_save = None

        if log_deadNeurons:
            dn_weights = []
            dn_update_matix = []
            for name, param in self.model.named_parameters():
                if 'weight' in name:
                    dn_weights.append(name+" "+str(param.shape)[11:-1])
                    # two update maps: num updates, last update

                    dn_update_matix.append([torch.zeros(param.shape), torch.zeros(param.shape), torch.zeros([epochs] + list(param.shape),dtype=torch.bool)])

        for epoch in range(epochs):
            self.epoch_start = self.global_iter
            self.train_iter = epoch
            epoch_data = torch.empty((len(data_loader), self.num_log_values), dtype=torch.float32, device=self.device)

            if log_deadNeurons:
                dn_epoch_data = torch.empty((len(data_loader), len(dn_weights)), dtype=torch.float32, device=self.device)

            with tqdm(data_loader) as bar:
                for i, (x_true, _) in enumerate(bar):
                    self.global_iter += 1

                    x_true = x_true.to(self.device)

                    loss = self.calc_loss(x_true, batch_size, epoch_data, i)

                    self.optimizer.zero_grad(set_to_none=True)
                    loss.backward(retain_graph=False)
                    if log_deadNeurons:
                        dn_i=0
                        for name, param in self.model.named_parameters():
                            if 'weight' in name:
                                dn_filter = (param.grad == 0)
                                dn_epoch_data[i, dn_i] = dn_filter.sum().div(param.grad.numel())
                                if log_stepdata:
                                    self.writer.add_scalar('Dead Neurons_/weight.'+str(dn_i)+'.'+name+" "+str(param.shape)[11:-1], dn_epoch_data[i, dn_i].cpu().numpy(), self.global_iter)

                                dn_filter_inv = ~dn_filter
                                update_map, iteration_map, deadneuron_map_epochs = dn_update_matix[dn_i]
                                update_map[dn_filter_inv] +=1 #update_map
                                iteration_map[dn_filter_inv] = self.global_iter #iteration_map
                                deadneuron_map_epochs[epoch, dn_filter] = True

                                dn_i += 1

                    self.optimizer.step()

                    self.afterTrainStep(epoch_data, i)

                avg_loss = self.log_data(epoch, epoch_data, 'train')
                if log_deadNeurons:
                    dn_epoch_data = dn_epoch_data.detach().mean(axis=0).cpu().numpy()
                    for dn_i, name in enumerate(dn_weights):
                        self.writer.add_scalar('Dead Neurons/weight.'+str(dn_i)+'.'+name, dn_epoch_data[dn_i], epoch)
                        temp = ((dn_update_matix[dn_i][1] < self.epoch_start)*1.).mean().detach().cpu().numpy()
                        self.writer.add_scalar('Dead Neurons/weight.'+str(dn_i)+'.'+name+"_", temp, epoch)

                        self.writer.add_scalar('Dead_Neurons/weight.'+str(dn_i)+'.'+name, (dn_update_matix[dn_i][2][epoch]*1.0).mean(), epoch)

                if epoch % save_every == 0 and epoch > 0:
                    last_save = self.full_tag + "_" + str(epoch) + '_' + \
                        datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_"+str(avg_loss)+".save"
                    torch.save(self.model.state_dict(), save_path+last_save)

                print({"epoch": epoch, "step": self.global_iter, "loss": avg_loss, "last_save": last_save})
                # bar.set_postfix(epoch_stat=epoch, global_iter=self.global_iter, loss=avg_loss, last_save=last_save)

            do_eval = epoch % eval_every == 0
            do_recon = epoch % reconstruction_every == 0
            do_score = epoch % score_every == 0
            if do_eval or do_recon or do_score:
                self.model.eval()

            if do_eval:
                eval_loss = self.eval(test_dataset, batch_size=batch_size)
                self.log_data(epoch, eval_loss, 'eval')

            if do_recon:
                img = self._visualize_reconstruction(ds, indices=reconstruct_indices)
                self.writer.add_image('Reconstruction', torch.clamp(img, 0, 1), epoch)

            if do_score:
                scores = self.score(score_ds if score_ds is not None else ds)
                print(scores)
                for score_name, score_values in scores.items():
                    # self.writer.add_scalars(score_name, score_values, global_step=epoch)
                    for score_value_name, value in score_values.items():
                        self.writer.add_scalar(score_name+"/"+score_value_name, value, global_step=epoch)

            if do_eval or do_recon or do_score:
                self.model.train()

            if max_iter and self.global_iter >= max_iter:
                break

        # if we didn't calc scores in last iteration do it now
        if not do_score:
            scores = self.score(score_ds if score_ds is not None else ds)
            print(scores)
            for score_name, score_values in scores.items():
                # self.writer.add_scalars(score_name, score_values, global_step=epoch)
                for score_value_name, value in score_values.items():
                    self.writer.add_scalar(score_name+"/"+score_value_name, value, global_step=epoch)

        last_save = self.full_tag + "_" + \
            datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + "_"+str(avg_loss)+".save"
        torch.save(self.model.state_dict(), save_path+last_save)
        print("saved "+last_save)

        return_dict = {}
        if log_deadNeurons:
            return_dict['deadNeurondata'] = dn_update_matix
            return_dict['weights'] = dn_weights

        return return_dict

    def score(self, ds):
        random_state = np.random.RandomState(0)
        beta_vae_score = beta_vae.compute_beta_vae_sklearn(
            ds, make_representor(self),
            random_state, batch_size=64, num_train=10000, num_eval=5000)

        random_state = np.random.RandomState(0)
        factor_vae_score = factor_vae.compute_factor_vae(
            ds, make_representor(self),
            random_state, num_variance_estimate=10000)

        random_state = np.random.RandomState(0)
        mig_score = mig.compute_mig(ds, make_representor(self), random_state)

        # random_state = np.random.RandomState(0)
        # dci_score = dci.compute_dci(ds, make_representor(self), random_state)

        random_state = np.random.RandomState(0)
        sap = sap_score.compute_sap(ds, make_representor(self), random_state)

        return {
            "beta_vae_score": beta_vae_score,
            "factor_vae_score": factor_vae_score,
            "mig_score": mig_score,
            # "dci_score": dci_score,
            "sap_score": sap,
        }

    def calc_loss(self, x_true, batch_size, epoch_data, epoch_step):
        y, (mu, log_var), _ = self.model.forward(x_true)
        kld = self.gauss_vae_regulariser(mu, log_var)
        reg_loss = self.beta * kld
        recon_loss = self.reconstruction_loss(y, x_true).div(batch_size)

        loss = recon_loss + reg_loss

        if self.model.use_model_regulariser:
            loss += self.model_regulariser_weight * self.model.model_regulariser()

        epoch_data[epoch_step, 0] = recon_loss.detach()
        epoch_data[epoch_step, 1] = reg_loss.detach()
        epoch_data[epoch_step, 2] = kld.detach()

        with torch.no_grad():
            mse = nn.functional.mse_loss(
                self.model.activation(y) if isinstance(self.model.loss_activation, nn.Identity) else y,
                x_true,
                reduction='sum').div(batch_size)
            epoch_data[epoch_step, 3] = mse.detach()

        return loss

    def log_data(self, epoch, epoch_data, type='train'):
        epoch_data = epoch_data.detach().mean(axis=0).cpu().numpy()
        l1, l2 = epoch_data[0], epoch_data[1]

        self.writer.add_scalar('Loss reconstruction/'+type, l1, epoch)
        self.writer.add_scalar('Loss regularisation/'+type, l2, epoch)
        self.writer.add_scalar('KL divergence/'+type, epoch_data[2], epoch)
        self.writer.add_scalar('MSE reconstruction/'+type, epoch_data[3], epoch)

        return l1 + l2

    def eval(self, ds, batch_size=256):
        self.model.eval()
        data_loader = DataLoader(ds, batch_size=batch_size, shuffle=True)

        with torch.no_grad():
            epoch_loss = torch.empty((len(data_loader), self.num_log_values), dtype=torch.float32, device=self.device)
            for i, (x_true, _) in enumerate(tqdm(data_loader)):
                x_true = x_true.to(self.device)

                self.calc_loss(x_true, batch_size, epoch_loss, i)

            return epoch_loss

    def loadModel(self, checkpoint, folder="./checkpoints/"):
        """Loads model checkpint

        Args:
            checkpoint (str): Filename of checkpoint
            folder (str, optional): dir where checkpoint is located. Defaults to "./checkpoints/".
        """
        self.model.load_state_dict(torch.load(folder + checkpoint))
        self.model.eval()

    def encode(self, x):
        return self.model.encode(x.to(self.device))[0]

def make_representor(model: BetaVAE):
    """
    Encloses the pytorch ScriptModule in a callable that can be used by `disentanglement_lib`.
    Parameters
    ----------
    model : torch.nn.Module or torch.jit.ScriptModule
        The Pytorch model.
    cuda : bool
        Whether to use CUDA for inference. Defaults to the return value of the `use_cuda`
        function defined above.
    Returns
    -------
    callable
        A callable function (`representation_function` in dlib code)
    """
    # Deepcopy doesn't work on ScriptModule objects yet:
    # https://github.com/pytorch/pytorch/issues/18106
    # model = deepcopy(model)
    # cuda = use_cuda() if cuda is None else cuda
    # model = model.cuda() if cuda else model.cpu()

    # Define the representation function
    def _represent(x):
        # assert isinstance(x, np.ndarray), \
        #     "Input to the representation function must be a ndarray."
        # assert x.ndim == 4, \
        #     "Input to the representation function must be a four dimensional NHWC tensor."
        # # Convert from NHWC to NCHW
        # x = np.moveaxis(x, 3, 1)
        # # Convert to torch tensor and evaluate
        # x = torch.from_numpy(x).float().to('cuda' if cuda else 'cpu')
        with torch.no_grad():
            y = model.encode(x)
        y = y.cpu().numpy()
        assert y.ndim == 2, \
            "The returned output from the representor must be two dimensional (NC)."
        return y

    return _represent


import numpy as np
import torch

joints = [[0,1], [1, 2], [2, 3], [3, 4], [4, 5], [6, 7], [7, 8], [8, 9],
            [9, 10], [10, 11], [12, 13], [2, 12]]
# sample_path = '3d_sample.npy'

# # should be 16 * 3 dim xyz
# prediction_3d = np.load(sample_path)
# limbs = [0] * len(mpii_edges)
# # apex for computing the angle
# relative_angles = []
# absolute_angles = np.zeros([16, 3])
# offset = prediction_3d[0]

def joint_angles(prediction_3d):
    jangles = torch.zeros([12, 3])
    for i, joint in enumerate(joints):
        e1, e2 = joint
        i1, i2 = mpii_edges[e1]
        i3, i4 = mpii_edges[e2]
        v1 = prediction_3d[i1] - prediction_3d[i2]
        v2 = prediction_3d[i4] - prediction_3d[i3]

        jangles[i] = torch.acos(torch.dot(v1, v2) / torch.norm(v1) / torch.norm(v2))
    return jangles

def absolute_angles(prediction_3d):
    absolute_angles = np.zeros([14, 3])
    offset = prediction_3d[8] ## offset !!!!!!!
    limbs = np.zeros([14, 1])
    for i in range(len(mpii_edges)):
        i1, i2 = mpii_edges[i]
        e1 = prediction_3d[i1] - prediction_3d[i2]
        l = np.linalg.norm(e1)
        limbs[i] = l
        absolute_angles[i] = np.arccos(e1/l)
    return absolute_angles, limbs, offset

def anglelimbtoxyz(offset, absolute_angles, limbs):
    limbs = limbs.float()
    res_3d = torch.zeros([4, 16, 3]).cuda()

    norm_direction = torch.cos(absolute_angles)
    limbs = limbs.repeat(1,1,3)
    direction = limbs * norm_direction

    res_3d[:, 0] = offset
    res_3d[:, 1] = res_3d[:, 0] - direction[:, 0]
    res_3d[:, 2] = res_3d[:, 1] - direction[:, 1]
    res_3d[:, 6] = res_3d[:, 2] - direction[:, 2]
    res_3d[:, 3] = res_3d[:, 6] - direction[:, 3]
    res_3d[:, 4] = res_3d[:, 3] - direction[:, 4]
    res_3d[:, 5] = res_3d[:, 4] - direction[:, 5]
    res_3d[:, 8] = res_3d[:, 6] - direction[:, 12]
    res_3d[:, 9] = res_3d[:, 8] - direction[:, 13]
    res_3d[:, 13] = res_3d[:, 8] - direction[:, 9]
    res_3d[:, 14] = res_3d[:, 13] - direction[:, 10]
    res_3d[:, 15] = res_3d[:, 14] - direction[:, 11]
    res_3d[:, 12] = res_3d[:, 8] + direction[:, 8]
    res_3d[:, 11] = res_3d[:, 12] + direction[:, 7]
    res_3d[:, 10] = res_3d[:, 11] + direction[:, 6]

    return res_3d

mpii_edges = [[0, 1], [1, 2], [2, 6], [6, 3], [3, 4], [4, 5],
                [10, 11], [11, 12], [12, 8], [8, 13], [13, 14], [14, 15],
              [6, 8], [8, 9]]

def anglelimbtoxyz2(offset, absolute_angles, limbs):
    limbs = limbs.float()
    b = limbs.shape[0]
    res_3d = torch.zeros([b, 16, 3]).cuda()

    norm_direction = torch.cos(absolute_angles)
    limbs = limbs.repeat(1,1,3)
    direction = limbs * norm_direction

    res_3d[:, 8] = offset
    res_3d[:, 13] = res_3d[:, 8] - direction[:, 9]
    res_3d[:, 14] = res_3d[:, 13] - direction[:, 10]
    res_3d[:, 15] = res_3d[:, 14] - direction[:, 11]
    res_3d[:, 12] = res_3d[:, 8] + direction[:, 8]
    res_3d[:, 11] = res_3d[:, 12] + direction[:, 7]
    res_3d[:, 10] = res_3d[:, 11] + direction[:, 6]
    res_3d[:, 9] = res_3d[:, 8] - direction[:, 13]
    res_3d[:, 6] = res_3d[:, 8] + direction[:, 12]
    res_3d[:, 3] = res_3d[:, 6] - direction[:, 3]
    res_3d[:, 4] = res_3d[:, 3] - direction[:, 4]
    res_3d[:, 5] = res_3d[:, 4] - direction[:, 5]
    res_3d[:, 2] = res_3d[:, 6] + direction[:, 2]
    res_3d[:, 1] = res_3d[:, 2] + direction[:, 1]
    res_3d[:, 0] = res_3d[:, 1] + direction[:, 0]

    return res_3d

corres = [10,9,8,11,12,13,17,17,1,17,4,3,2,5,6,7]

def cords_to_map(cords, img_size, sigma=6):
    result = torch.zeros([cords.size(0), 18, 256, 176])
    for i, points in enumerate(cords):
        for j in range(16):
            co = corres[j]
            point = points[j]
            xx, yy = torch.meshgrid(torch.arange(img_size[0], dtype=torch.int32).cuda(), torch.arange(img_size[1],dtype=torch.int32).cuda())
            xx = xx.float()
            yy = yy.float()
            res = torch.exp(-((yy - point[0]) ** 2 + (xx - point[1]) ** 2) / (2 * sigma ** 2))
            result[i, co] = res

    return result


import pandas as pd
import seaborn as sns
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.dates as mdates
from datetime import date, timedelta
from simtools.Analysis.BaseAnalyzers import BaseAnalyzer

mpl.rcParams['pdf.fonttype'] = 42

class InsetAnalyzer(BaseAnalyzer):
    def __init__(self, exp_name, working_dir='.', sweep_variables=None):  # sweep_variables=None,

        super(InsetAnalyzer, self).__init__(working_dir=working_dir,
                                            filenames=["output/InsetChart.json"]
                                            )
        self.sweep_variables = sweep_variables or ["scenario", "ITN_coverage"]
        self.inset_channels = ['Infected', 'New Infections', 'Adult Vectors', 'Daily EIR']
        #self.sweep_variables = sweep_variables or ["start_days", 'Run_Number', 'krate_scale_factor_s']
        self.exp_name = exp_name
        # self.inset_channels = ['Infected', 'New Infections', 'Adult Vectors', 'Daily EIR']

    # def filter(self, simulation):
    #    return simulation.tags["Run_Number"] == 0

    def select_simulation_data(self, data, simulation):
        """
        Extract data from output data and accumulate in same bins as reference.
        """

        # Load data from simulation
        simdata = pd.DataFrame( { x : data[self.filenames[0]]['Channels'][x]['Data'] for x in self.inset_channels})
        simdata['day'] = simdata.index

        for sweep_var in self.sweep_variables :
            simdata[sweep_var] = simulation.tags[sweep_var]

        return simdata

    def finalize(self, all_data):

        selected = [data for sim, data in all_data.items()]
        if len(selected) == 0:
            print("No data have been returned... Exiting...")
            return

        df = pd.concat(selected, sort=False).reset_index(drop=True)

       # df['krate'] = df['krate_scale_factor_s'].apply(lambda x : 'old' if x == 1 else 'new') # comment out

        sns.set_style('whitegrid', {'axes.linewidth': 0.5})

        fig = plt.figure(figsize=(12, 8))
        formatter = mdates.DateFormatter("%m-%Y")
        axes = [fig.add_subplot(2,2, x+1) for x in range(len(self.inset_channels))]
        fig.autofmt_xdate()
        palette = sns.color_palette('Set1')

        df['date'] = df['day'].apply(lambda x : date(2013, 1, 1) + timedelta(days=x))

        for ia, (a, adf) in enumerate(df.groupby('scenario')): # comment out block

            for ax, channel in zip(axes, self.inset_channels) :
                gdf = adf.groupby('date')[channel].agg([np.min, np.max, np.mean]).reset_index()
                ax.plot(gdf['date'], gdf['mean'], '-', label=a, color=palette[ia])
                #ax.fill_between(gdf['date'], gdf['amin'], gdf['amax'], linewidth=0, color=palette[ia], alpha=0.3)
                ax.set_ylabel(channel)
                ax.legend()
                ax.xaxis.set_major_formatter(formatter)
                ax.xaxis.set_major_locator(mdates.MonthLocator(interval=3))
        plt.show()

if __name__ == "__main__":

    from simtools.Analysis.AnalyzeManager import AnalyzeManager
    from simtools.SetupParser import SetupParser

    SetupParser.default_block = 'HPC'
    SetupParser.init()

    analyzer = InsetAnalyzer(exp_name='Spatial IR Rotation ITNs')

    am = AnalyzeManager('eef6b34b-f4d2-ec11-a9f8-b88303911bc1', analyzers=analyzer)
    am.analyze()


import json

person_string = '{"name" : "Savaş" , "languages":["Python","Java"] }'

#JSON stirng to Dict

#result = json.loads(person_string)

#print(type(result))
#print(result)

#with open("jsonString",encoding="utf-8") as f :
#    data = json.load(f)
#    print(data)

person_dict = {

    "name": "Ali",
   "languages":"[Python,Java]"
}

#Dict to String
#result = json.dumps(person_dict)
#print(type(person_dict))

with open("jsonString","w") as file :
    json.dump(person_dict,file)


from __future__ import division

import s2s
import argparse
import torch
import torch.nn as nn
from torch import cuda
from torch.autograd import Variable
import math
import time
import logging
import random
import numpy as np

try:
    import ipdb
except ImportError:
    pass
from nltk.translate import bleu_score
from s2s.xinit import xavier_normal, xavier_uniform
import os
import xargs
torch.backends.cudnn.enabled=False

parser = argparse.ArgumentParser(description='train.py')
xargs.add_data_options(parser)
xargs.add_model_options(parser)
xargs.add_train_options(parser)

opt = parser.parse_args()

logging.basicConfig(format='%(asctime)s [%(levelname)s:%(name)s]: %(message)s', level=logging.INFO)
log_file_name = time.strftime("%Y%m%d-%H%M%S") + '.log.txt'
if opt.log_home:
    log_file_name = os.path.join(opt.log_home, log_file_name)
file_handler = logging.FileHandler(log_file_name, encoding='utf-8')
file_handler.setFormatter(logging.Formatter('%(asctime)s [%(levelname)-5.5s:%(name)s] %(message)s'))
logging.root.addHandler(file_handler)
logger = logging.getLogger(__name__)

logger.info('My PID is {0}'.format(os.getpid()))
logger.info('PyTorch version: {0}'.format(str(torch.__version__)))
logger.info(opt)

if torch.cuda.is_available() and not opt.gpus:
    logger.info("WARNING: You have a CUDA device, so you should probably run with -gpus 0")

if opt.seed > 0:
    torch.manual_seed(opt.seed)

if opt.gpus:
    if opt.cuda_seed > 0:
        torch.cuda.manual_seed(opt.cuda_seed)
    cuda.set_device(opt.gpus[0])

logger.info('My seed is {0}'.format(torch.initial_seed()))
logger.info('My cuda seed is {0}'.format(torch.cuda.initial_seed()))

def NMTCriterion(vocabSize):
    weight = torch.ones(vocabSize)
    weight[s2s.Constants.PAD] = 0

    crit = nn.NLLLoss(weight, size_average=False)
    if opt.gpus:
        crit.cuda()
    return crit

def reverse_NMTCriterion(vocabSize):
    weight = torch.zeros(vocabSize)
    weight[s2s.Constants.PAD] = 0

    crit = nn.NLLLoss(weight, size_average=False)
    if opt.gpus:
        crit.cuda()
    return crit

def loss_function(g_outputs, g_targets, generator, crit, eval=False):
    batch_size = g_outputs.size(1)

    g_out_t = g_outputs.view(-1, g_outputs.size(2))
    g_prob_t = generator(g_out_t)

    g_loss = crit(g_prob_t, g_targets.view(-1))
    total_loss = g_loss
    report_loss = total_loss.item()
    return total_loss, report_loss, 0

def generate_copy_loss_function(g_outputs, c_outputs, g_targets,
                                c_switch, c_targets, c_gate_values,
                                generator, nlu_generator, classifier, crit, copyCrit,
                                reverse_crit, guideCrit, nluCrit, guides, g_hiddens, nlu_outputs,type_emb):

    # g_out is output [(TargetLength, BatchSize, Hidden)]
    batch_size = g_outputs.size(1)
    targetLength = g_outputs.size(0)

    guides = guides.transpose(1,0)

    # transform typed target with the special <ent> symbol to mention level target
    m_targets = g_targets.clone().detach()
    m_targets = m_targets.transpose(1,0)
    for i in range(batch_size):
        count = 0
        for j in range(targetLength):
            if m_targets[i][j] == s2s.Constants.SS:
                m_targets[i][j] = guides[i][count]
                count += 1
            else:
                m_targets[i][j] = s2s.Constants.PAD

    m_targets = m_targets.transpose(1,0)

    # calculate nlu prediction distribution
    nlu_out_t = nlu_outputs.view(-1, nlu_outputs.size(2))
    nlu_prob_t = nlu_generator(nlu_out_t)
    nlu_prob_t = nlu_prob_t.view(-1, batch_size, nlu_prob_t.size(1))

    # calculate generation prediction distribution
    g_out_t = g_outputs.view(-1, g_outputs.size(2))
    g_prob_t = generator(g_out_t)
    g_prob_t = g_prob_t.view(-1, batch_size, g_prob_t.size(1))
    g_hiddens = g_hiddens.squeeze(1).transpose(1,0)

    # calculate the mention predictor loss
    cl_loss = 0
    type_emb = type_emb.transpose(1,0) # injected type embedding from source
    g_prob_t_trans = g_prob_t.transpose(1,0)

    g_targets = g_targets.transpose(1,0)
    guide_probs = []
    for i in range(batch_size):
        guide_prob = []
        guide = []
        count = 0
        for j in range(targetLength):
            if g_targets[i][j] == s2s.Constants.SS:
                injected = torch.cat([g_hiddens[i][j],type_emb[i][count]],dim=0)

                guide_prob.append(torch.log(classifier(injected.unsqueeze(0))))
                guide.append(guides[i][count])
                count += 1
        guide_prob = torch.stack(guide_prob).squeeze(1)
        guide = torch.stack(guide).view(-1)
        if i == 0:
            cl_loss = guideCrit(guide_prob,guide)
        else:
            cl_loss += guideCrit(guide_prob,guide)
    g_targets = g_targets.transpose(1,0)

    g_output_prob = g_prob_t + 1e-8

    g_output_prob_log = torch.log(g_output_prob)
    nlu_prob_log = torch.log(nlu_prob_t)
    g_output_prob_log = g_output_prob_log * ((1 - c_switch).unsqueeze(2).expand_as(g_output_prob_log))

    g_output_prob_log = g_output_prob_log.view(-1, g_output_prob_log.size(2))
    nlu_prob_log = nlu_prob_log.view(-1, nlu_prob_log.size(2))

    # g_output_prob_log [(TargetLength * BatchSize, vocab_size)]
    g_loss = crit(g_output_prob_log, g_targets.view(-1))
    nlu_loss = nluCrit(nlu_prob_log, m_targets.view(-1))

    total_loss = g_loss + nlu_loss * 1.5 + cl_loss * 2.5
    report_loss = total_loss.item()
    return total_loss, report_loss, 0, g_loss, cl_loss, nlu_loss

def addPair(f1, f2, f3):
    for x, y1, z in zip(f1, f2, f3):
        yield (x, y1, z)
    yield (None, None, None)

def load_dev_data(translator, src_file, bio_file, feat_files, tgt_file, guide_file):
    dataset, raw = [], []
    srcF = open(src_file, encoding='utf-8')
    guideF = open(guide_file, encoding='utf-8')
    tgtF = open(tgt_file, encoding='utf-8')
    bioF = open(bio_file, encoding='utf-8')
    featFs = [open(x, encoding='utf-8') for x in feat_files]

    src_batch, tgt_batch = [], []
    bio_batch, feats_batch = [], []
    guide_batch = []
    for line, tgt, guide in addPair(srcF, tgtF, guideF):
        if (line is not None) and (tgt is not None) and (guide is not None):
            src_tokens = line.strip().split(' ')
            src_batch += [src_tokens]
            guide_tokens = guide.strip().split(' ')
            guide_batch += [guide_tokens]
            tgt_tokens = tgt.strip().split(' ')
            tgt_batch += [tgt_tokens]
            bio_tokens = bioF.readline().strip().split(' ')
            bio_batch += [bio_tokens]
            feats_tokens = [reader.readline().strip().split((' ')) for reader in featFs]
            feats_batch += [feats_tokens]

            if len(src_batch) < opt.batch_size:
                continue
        else:
            # at the end of file, check last batch
            if len(src_batch) == 0:
                break
        data = translator.buildData(src_batch, bio_batch, feats_batch, tgt_batch, guide_batch)
        dataset.append(data)
        raw.append((src_batch, tgt_batch))
        src_batch, tgt_batch = [], []
        bio_batch, feats_batch = [], []
        guide_batch = []
    srcF.close()
    bioF.close()
    for f in featFs:
        f.close()
    tgtF.close()
    guideF.close()
    return (dataset, raw)

evalModelCount = 0
totalBatchCount = 0

def evalModel(model, translator, evalData):
    global evalModelCount
    evalModelCount += 1
    ofn_prefix = opt.ofn_prefix
    if evalModelCount % 2 == 0:
        ofn_prefix += '_test'
    else:
        ofn_prefix += '_dev'
    ofn = ofn_prefix + '.{0}'.format(evalModelCount)
    if opt.save_path:
        ofn = os.path.join(opt.save_path, ofn)

    predict, gold = [], []
    processed_data, raw_data = evalData
    for batch, raw_batch in zip(processed_data, raw_data):
        #  (1) read src and tgt (bio and feats are part of the legacy code and are never used in this project)
        src, bio, feats, tgt, guide, indices = batch[0]
        src_batch, tgt_batch = raw_batch

        #  (2) translate
        pred, predScore, predIsCopy, predCopyPosition, attn, _ = translator.translateBatch(src, bio, feats, tgt, guide)
        pred, predScore, predIsCopy, predCopyPosition, attn = list(zip(
            *sorted(zip(pred, predScore, predIsCopy, predCopyPosition, attn, indices),
                    key=lambda x: x[-1])))[:-1]

        #  (3) convert indexes to words
        predBatch = []
        for b in range(src[0].size(1)):
            n = 0
            predBatch.append(
                translator.buildTargetTokens(pred[b][n], src_batch[b],
                                             predIsCopy[b][n], predCopyPosition[b][n], attn[b][n])
            )
        # nltk BLEU evaluator needs tokenized sentences
        gold += [[r] for r in tgt_batch]
        predict += predBatch
        for i, j in zip([[r] for r in tgt_batch], predBatch):
            if random.random() > 1:
                print('{}\n{}\n{}\n'.format(i, j, '*'*50))

    no_copy_mark_predict = [[word.replace('[[', '').replace(']]', '') for word in sent] for sent in predict]
    bleu = bleu_score.corpus_bleu(gold, no_copy_mark_predict)
    report_metric = bleu

    # (4) save generation result
    with open(ofn, 'w', encoding='utf-8') as of:
        for p in predict:
            of.write(' '.join(p) + '\n')
    print(ofn)
    return report_metric

def trainModel(model, translator, trainData, validData, testData, dataset, optim):
    logger.info(model)
    model.train()
    torch.cuda.empty_cache()
    for name,para in model.named_parameters():
        if para.requires_grad:
            print(name,para.data.shape)
    logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'eval'))

    # define criterion of each GPU
    criterion = NMTCriterion(dataset['dicts']['tgt'].size())
    guideCriterion = NMTCriterion(dataset['dicts']['guide_src'].size())
    nluCriterion = NMTCriterion(dataset['dicts']['guide_src'].size())
    reverse_criterion = reverse_NMTCriterion(dataset['dicts']['tgt'].size())
    copyLossF = nn.NLLLoss(size_average=False)

    start_time = time.time()

    def saveModel(metric=None):
        model_state_dict = model.module.state_dict() if len(opt.gpus) > 1 else model.state_dict()
        model_state_dict = {k: v for k, v in model_state_dict.items() if 'generator' not in k}
        generator_state_dict = model.generator.module.state_dict() if len(
            opt.gpus) > 1 else model.generator.state_dict()
        #  (5) drop a checkpoint
        checkpoint = {
            'model': model_state_dict,
            'generator': generator_state_dict,
            'dicts': dataset['dicts'],
            'opt': opt,
            'epoch': epoch,
            'optim': optim
        }
        save_model_path = 'model'
        if opt.save_path:
            if not os.path.exists(opt.save_path):
                os.makedirs(opt.save_path)
            save_model_path = opt.save_path + os.path.sep + save_model_path
        if metric is not None:
            torch.save(checkpoint, '{0}_dev_metric_{1}_e{2}.pt'.format(save_model_path, round(metric, 4), epoch))
        else:
            torch.save(checkpoint, '{0}_e{1}.pt'.format(save_model_path, epoch))

    def trainEpoch(epoch):
        if opt.extra_shuffle and epoch > opt.curriculum:
            logger.info('Shuffling...')
            trainData.shuffle()

        # shuffle mini batch order
        batchOrder = torch.randperm(len(trainData))

        total_loss, total_words, total_num_correct = 0, 0, 0
        report_loss, report_tgt_words, report_src_words, report_num_correct = 0, 0, 0, 0
        start = time.time()
        epoch_loss = 0
        epoch_cl_loss = 0
        epoch_nlu_loss = 0
        for i in range(len(trainData)):
            global totalBatchCount
            totalBatchCount += 1
            """
            [0] ([0][0] wrap(srcBatch), [0][1] lengths), [1] ([1][0] wrap(bioBatch), [1][1] lengths), [2] ([2][0] (wrap(x) for x in featBatches), [2][1] lengths),
            [3] ([3][0] wrap(tgtBatch), [3][1] wrap(copySwitchBatch), [3][2] wrap(copyTgtBatch)),
            [4] [4] indices (not used) # original indices will be excluded !
            """
            batchIdx = batchOrder[i] if epoch > opt.curriculum else i
            batch = trainData[batchIdx][:-1]  # exclude original indices

            model.zero_grad()
            '''
            # g_out is output [(TargetLength, BatchSize, Hidden)]
            # c_out is attn_out [(TargetLength, BatchSize, InputLength)]
            # c_gate is copy_prob [(TargetLength, BatchSize, 1)]
            '''
            g_outputs, c_outputs, c_gate_values, g_hiddens, nlu_outputs, type_emb = model(batch, False)
            targets = batch[3][0][1:]  # exclude <s> from targets
            copy_switch = batch[3][1][1:]
            c_targets = batch[3][2][1:]
            g_guides = batch[4][0]

            loss, res_loss, num_correct, g_loss, cl_loss, nlu_loss = generate_copy_loss_function(
                g_outputs, c_outputs, targets, copy_switch, c_targets, c_gate_values, model.generator, model.nlu_generator, model.classifier,
                criterion, copyLossF, reverse_criterion, guideCriterion, nluCriterion, g_guides, g_hiddens, nlu_outputs, type_emb)
            torch.cuda.empty_cache()
            epoch_loss += g_loss
            epoch_cl_loss += cl_loss
            epoch_nlu_loss += nlu_loss

            if math.isnan(res_loss) or res_loss > 1e20:
                logger.info('catch NaN')
                ipdb.set_trace()

            # update the parameters
            loss.backward()
            optim.step()

            num_words = targets.data.ne(s2s.Constants.PAD).sum().item()
            report_loss += res_loss
            report_num_correct += num_correct
            report_tgt_words += num_words
            report_src_words += batch[0][-1].data.sum()
            total_loss += res_loss
            total_num_correct += num_correct
            total_words += num_words
            if i % opt.log_interval == 0:
                logger.info(
                    "Epoch %2d, %6d/%5d/%5d; acc: %6.2f; loss: %6.2f; words: %5d; ppl: %6.2f; %3.0f src tok/s; %3.0f tgt tok/s; %6.0f s elapsed" %
                    (epoch, totalBatchCount, i + 1, len(trainData),
                     report_num_correct / report_tgt_words * 100,
                     report_loss,
                     report_tgt_words,
                     math.exp(min((report_loss / report_tgt_words), 16)),
                     report_src_words / max((time.time() - start), 1.0),
                     report_tgt_words / max((time.time() - start), 1.0),
                     time.time() - start))

                report_loss = report_tgt_words = report_src_words = report_num_correct = 0
                start = time.time()

            if validData is not None and totalBatchCount % opt.eval_per_batch == 0 and totalBatchCount >= opt.start_eval_batch:
                model.eval()
                logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'eval'))
                valid_bleu = evalModel(model, translator, validData)
                model.train()
                logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'eval'))
                model.decoder.attn.mask = None
                logger.info('Validation Score: %g' % (valid_bleu * 100))
                if valid_bleu >= optim.best_metric:
                    saveModel(valid_bleu)
                optim.updateLearningRate(valid_bleu, epoch)

            if testData is not None and totalBatchCount % opt.eval_per_batch == 0 and totalBatchCount >= opt.start_eval_batch:
                model.eval()
                logger.warning("Set model to {0} mode".format('train' if model.decoder.dropout.training else 'test'))
                test_bleu = evalModel(model, translator, testData)
                logger.info('Test Score: %g' % (test_bleu * 100))
        print('epoch g loss: ')
        print(epoch_loss)
        print('epoch cl loss: ')
        print(epoch_cl_loss)
        print('epoch nlu loss: ')
        print(epoch_nlu_loss)

        return total_loss / total_words, total_num_correct / total_words

    for epoch in range(opt.start_epoch, opt.epochs + 1):
        logger.info('')
        # train for one epoch on the training set
        train_loss, train_acc = trainEpoch(epoch)
        # calculate metric for each epoch training
        train_ppl = math.exp(min(train_loss, 100))
        logger.info('Train perplexity: %g' % train_ppl)
        logger.info('Train accuracy: %g' % (train_acc * 100))
        logger.info('Saving checkpoint for epoch {0}...'.format(epoch))
        saveModel()

def main():

    import onlinePreprocess
    onlinePreprocess.lower = opt.lower_input
    onlinePreprocess.seq_length = opt.max_sent_length
    onlinePreprocess.shuffle = 1 if opt.process_shuffle else 0
    from onlinePreprocess import prepare_data_online

    # opt.train_src (source file of sequence) 'it is a replica of the grotto at lourdes , france where the virgin mary reputedly appeared to saint bernadette soubirous in 1858 .'
    # opt.src_vocab (source file of vocab) 'the(word) 4(index) 256272(frequency) 0.06749202214022335'
    # opt.train_bio (answer position embedding) 'O O O O O O O O O O O O O O O O O O B I I O O O'
    # opt.bio_vocab (source file of answer position vocab) 'O(bio) 4(index) 2525015(frequency) 0.8958601572376024'
    # opt.train_feats (source file of postag/ner/case) 'PERSON/UPCASE/NN ...' (3 different embeddings)
    # opt.feat_vocab (source file of answer feat vocab)
    # opt.train_tgt (source file of question) 'to whom did the virgin mary allegedly appear in 1858 in lourdes france ?'
    # opt.tgt_vocab (source file of vocab) same file with opt.src_vocab !!
    dataset = prepare_data_online(opt.train_src, opt.src_vocab, opt.train_bio, opt.bio_vocab, opt.train_feats,
                                  opt.feat_vocab, opt.train_tgt, opt.tgt_vocab, opt.train_guide_src, opt.guide_src_vocab)

    trainData = s2s.Dataset(dataset['train']['src'], dataset['train']['bio'], dataset['train']['feats'],
                            dataset['train']['tgt'], dataset['train']['switch'], dataset['train']['c_tgt'], opt.batch_size, opt.gpus, dataset['train']['guide_src'])

    dicts = dataset['dicts']
    logger.info(' * vocabulary size. source = %d; target = %d' % (dicts['src'].size(), dicts['tgt'].size()))
    logger.info(' * number of training sentences. %d' % len(dataset['train']['src']))
    logger.info(' * maximum batch size. %d' % opt.batch_size)

    logger.info('Building Model ...')
    encoder = s2s.Models.Encoder(opt, dicts['src'], dicts['guide_src'])
    decoder = s2s.Models.Decoder(opt, dicts['tgt'])
    decIniter = s2s.Models.DecInit(opt)

    ''' generator map output embedding to vocab size vector then softmax'''
    generator = nn.Sequential(
        nn.Linear(opt.dec_rnn_size // opt.maxout_pool_size, dicts['tgt'].size()),
        nn.Softmax(dim=1)
    )
    classifier = nn.Sequential(
        nn.Linear(opt.dec_rnn_size + 300, dicts['guide_src'].size()),
        nn.Softmax(dim=1)
    )
    nlu_generator = nn.Sequential(
        nn.Linear(opt.dec_rnn_size * 2, dicts['guide_src'].size()),
        nn.Softmax(dim=1)
    )

    model = s2s.Models.NMTModel(encoder, decoder, decIniter)
    model.generator = generator
    model.classifier = classifier
    model.nlu_generator = nlu_generator
    translator = s2s.Translator(opt, model, dataset)

    if len(opt.gpus) >= 1:
        model.cuda()
        generator.cuda()
        classifier.cuda()
        nlu_generator.cuda()
    else:
        model.cpu()
        generator.cpu()
        classifier.cpu()
        nlu_generator.cpu()

    # if len(opt.gpus) > 1:
    #     model = nn.DataParallel(model, device_ids=opt.gpus, dim=1)
    #     generator = nn.DataParallel(generator, device_ids=opt.gpus, dim=0)

    for pr_name, p in model.named_parameters():
        logger.info(pr_name)
        # p.data.uniform_(-opt.param_init, opt.param_init)
        if p.dim() == 1:
            # p.data.zero_()
            p.data.normal_(0, math.sqrt(6 / (1 + p.size(0))))
        else:
            nn.init.xavier_normal_(p, math.sqrt(3))

    encoder.load_pretrained_vectors(opt)
    decoder.load_pretrained_vectors(opt)

    optim = s2s.Optim(
        opt.optim, opt.learning_rate,
        max_grad_norm=opt.max_grad_norm,
        max_weight_value=opt.max_weight_value,
        lr_decay=opt.learning_rate_decay,
        start_decay_at=opt.start_decay_at,
        decay_bad_count=opt.halve_lr_bad_count
    )
    optim.set_parameters(model.parameters())

    validData = None
    if opt.dev_input_src and opt.dev_ref:
        validData = load_dev_data(translator, opt.dev_input_src, opt.dev_bio, opt.dev_feats, opt.dev_ref, opt.dev_guide_src)

    testData = None
    if opt.test_input_src and opt.test_ref:
        testData = load_dev_data(translator, opt.test_input_src, opt.test_bio, opt.test_feats, opt.test_ref, opt.test_guide_src)

    trainModel(model, translator, trainData, validData, testData, dataset, optim)

if __name__ == "__main__":
    main()


# Waitrose online shopping webservice

__author__ = 'robdobsn'

from flask import Flask
from ShoppingManager import ShoppingManager
import logging
import time

app = Flask(__name__)

@app.route('/')
def index():
    return appName

@app.route('/waitrose/api/v1.0/<accountCode>/trolleyinfo')
def getTrolleyInfo(accountCode):
    return shoppingManager.getTrolleyInfo(accountCode)

if __name__ == '__main__':

    # Set logging level
    logging.basicConfig(level=logging.INFO)

    # App name
    appName = "Waitrose Service 0.1"
    logging.info(appName + " Starting")

    # Start headless display if required
    try:
        from pyvirtualdisplay import Display
        display = Display(visible=0, size=(800, 600))
        display.start()
        logging.info(appName + " is running headless")
    except ImportError:
        logging.info(appName + " is running on a GUI")

    # Start WaitroseManager
    shoppingManager = ShoppingManager()

    # Start server
    # use_reloader=False avoids initialising objects twice
    # note that server is single-threaded - and it doesn't seem to help to turn theading on - maybe due to GIL
    # so every request has to return quickly and not block
    # host=0.0.0.0 results in the machine's IP address being used
    # as opposed to localhost for the web server
    app.run(debug=True, use_reloader=False, host="0.0.0.0")


# numCourses = 2
# prerequisites = [[1,0]]

numCourses = 2
prerequisites = [[1,0],[0,1]]

from collections import deque

# 위상 정렬 => 위상 정렬하려면 사이클이 없어야 함, 사이클 존재하면 false
# 모든 노드를 방문하기 전에 큐가 비어 있어 있을 때 사이클 존재 => 사이클이 존재하여 더 이상 진행 불가
# 강의가 노드, 전제 조건이 간선 정보가 된다.
'''
위상 정렬 구현 로직
1. 진입 차수가 0인 노드를 큐에 삽입
2. 큐에서 원소를 꺼내며 해당 원소와 연결된 노드의 진입 차수 - 1
3. 새롭게 진입 차수가 0이 된 노드를 큐에 삽입
사이클이 존재할 경우 모든 노드 방문 전에 큐가 비게 됨(진입 차수가 모두 0이 되기 전에 큐가 비게 됨)
'''

q = deque()
indegree = [0] * numCourses  # 진입 차수
graph = [[] for _ in range((numCourses))]  # 그래프

for i in range(len(prerequisites)):
    graph[prerequisites[i][1]].append(prerequisites[i][0])  # prerequisites[i][1]에서 prerequisites[i][0] 이동
    indegree[prerequisites[i][0]] += 1  # 진입 차수 + 1

# 위상 정렬 구현
# 진입 차수가 0인 노드 큐에 삽입
for i in range(numCourses):
    if indegree[i] == 0:
        q.append(i)

if len(q) == 0 and sum(indegree) != 0:
    print(False)
    exit(0)

while q:
    # 큐에서 원소 빼고 해당 노드와 연결된 노드의 진입 차수 - 1
    currNode = q.popleft()
    for i in graph[currNode]:
        indegree[i] -= 1

        # 진입 차수가 0이 된 노드 큐에 삽입
        if indegree[i] == 0:
            q.append(i)

    if len(q) == 0 and sum(indegree) != 0:
        print(False)
        exit(0)

print(True)

# 참고
# https://freedeveloper.tistory.com/390
# https://spongeb0b.tistory.com/215


import numpy as np
import matplotlib
matplotlib.use('Agg')
matplotlib.rc('font', size=10)
import matplotlib.pyplot as plt

fn = '../expres/results_std.npy'
sep = 1
total = 10000
lw = 1.5
smf = 10
sep *= smf
results = np.load(fn, allow_pickle=True).item()
results = {
    'x': results['Proposed_64_r'],
    '0.5x': results['Proposed_64_r'] * 0.5,
    'k=2, F(x)': results['Proposed_64_Fr'],
    'k=2, F(G(z))': results['Proposed_64_Fg'],
    'k=1, F(x)': results['Proposed_64_k1_Fr'],
    'k=1, F(G(z))': results['Proposed_64_k1_Fg'],
}
t = np.arange(total // sep) * sep + sep
for key, value in results.items():
    value = np.reshape(value, [-1, total // smf, smf])
    value = np.mean(value, axis=2)
    v, e = np.mean(value, axis=0), np.std(value, axis=0)
    if len(v) != len(t):
        continue
#    plt.errorbar(t, v, yerr=e, label=key, linewidth=lw, fmt='.--', markersize='2', elinewidth=2)
    plt.plot(t, v, '--', label=key, linewidth=lw, markersize='5')
    plt.fill_between(t, v-e, v+e, alpha=0.4, antialiased=True)

plt.xlabel('Iterations')
plt.ylabel('Standard Deviation')
plt.xlim([sep, 2500])
plt.ylim([0, 0.8])
plt.legend(loc='upper right')

plt.savefig('../ggg.png')


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import redis
from flask import Flask
app = Flask(__name__)

"""
思路：
使用推技术，每个用户或仓库有变动的时候把消息推给粉丝的timelines里面

生产消息：
    1. 使用redis incr给每个消息生成唯一的msgID
    2. 变动的事件包括关注，fork，创建仓库，以及user_id, 变动对象的信息都写入到一个字典里面，然后序列化后存到redis的msgId里面
    3. 获取当前用户的粉丝列表，数据结构是set，用smember获取粉丝id
    4. 每个粉丝都维护一个timeline队列，遍历粉丝列表，给他们的timeline队列添加msgID

获取消息：
    1. 每个粉丝上来就获取自己的timeline队列

"""

def msg_produce():
    redis_ins = redis()
    # 使用Redis incr生成唯一的msgID
    msg_id = redis_ins.incr("global:msgID")

    # 消息内容
    user_id = 1
    msg_data = {
        "user_id": user_id,
        "timestamp": "1587601324",
        "action": "关注",
        "target_info": {
            "url": "xx.xx",
            "name": "xx",
            "user": ""
        },
        "from_info": {
            "url": "xx.xx",
            "name": "xx"
        }
    }

    # redis存储消息，msg_id->msg_data
    msg_data = json.dumps(msg_data)
    redis_ins.set(msg_id, msg_data)

    # 获取自己的粉丝列表, 并把消息写入他们的时间线
    follower_ids = redis_ins.smember("{user_id}:followers".format(user_id))
    for follower_id in follower_ids:
        redis_ins.lpush("{follower_id}:timeline".format(follower_id), msg_id)

@app.route('/')
def get_timelines(page, limit):
    start = (page-1) * limit
    redis_ins = redis()
    user_id = 2
    # 根据页数和limit逐页获取
    msg_ids = redis_ins.lrange("{user_id}:timeline".format(user_id), start, end+limit)
    result = []
    for msg_id in msg_ids:
        data = json.loads(redis.get(msg_id))
        result.append(data)

    return result