# -*- coding: utf-8 -*-
"""

Two part script:
### 1. Calculates the proportion of mice that had graduated to stage 2 on each day of the experiment
### 2. Plots graduation day vs expression graphs

"""

# import packages and functions
from Functions_CoreFunctions_0100 import adjust_spines, makelegend2
import numpy as np
import matplotlib.pyplot as plt
import math
from scipy import stats

# Load data containing day of graduation to stage 2 for each mouse
marray = np.loadtxt('Data_Input/Behaviour_SummaryData/Stagegraduation_0100.txt',delimiter = '\t')

# specify mouse/mice and day/s to analyse
days = ['Day' + str(int(x)) for x in np.arange(1,22.1)]
mice = ['M' + str(int(x)) for x in np.arange(1,16.1)]# choose specific day/s
tetmice = ['M' + str(int(x)) for x in [2,5,6,7,8,9,11]]
conmice = ['M' + str(int(x)) for x in [1,3,4,10]]

#plot legend
def makelegend3(fig,ax):
    handles, labels = ax.get_legend_handles_labels()
    leg = fig.legend(handles,labels, loc="right", bbox_to_anchor=(0.85, 0.58), fontsize = "small", markerscale = 1)
    for l in leg.get_lines():l.set_linewidth(2)
    frame  = leg.get_frame()
    frame.set_edgecolor('w')
    frame.set_alpha(0.2)

# calculate proportion of mice that have graduated for each training day
proportionday_store = np.zeros((len(days), 4))
n_groups = 2

for daycount,day in enumerate(marray,0):
    allmice = marray[daycount,:]
    tetmice = np.hstack((marray[daycount,1],marray[daycount,4],marray[daycount,5],marray[daycount,6],marray[daycount,7],marray[daycount,8],marray[daycount,10],marray[daycount,11],marray[daycount,12],marray[daycount,13]))
    conmice = np.hstack((marray[daycount,0],marray[daycount,2],marray[daycount,3],marray[daycount,9],marray[daycount,14],marray[daycount,15]))
    htetmice = np.hstack((marray[daycount,1],marray[daycount,5],marray[daycount,6],marray[daycount,8]))
    ltetmice = np.hstack((marray[daycount,4],marray[daycount,7],marray[daycount,10],marray[daycount,11],marray[daycount,12],marray[daycount,13]))
    p_tet = (np.nansum(tetmice))/10
    p_htet = (np.nansum(htetmice))/4
    p_ltet = (np.nansum(ltetmice))/6
    p_con = (np.nansum(conmice))/6
    proportionday_store[daycount,0] = p_tet
    proportionday_store[daycount,1] = p_con
    proportionday_store[daycount,2] = p_htet
    proportionday_store[daycount,3] = p_ltet

# plot graph of proportion of graduation vs training day

daybins = np.arange(0,21.1,1) #number of days array
bar_width = 0.4
index = np.arange(n_groups)

"""
# plot proportion for GFP and TeLC mice

fig = plt.figure(figsize = (5,3))
ax = fig.add_subplot(111)
ax.plot(daybins,proportionday_store[:,0], color = '0.3', label = 'AAV-FLEX-TeLC', linewidth = 1.5)
ax.plot(daybins,proportionday_store[:,1], color = 'PaleVioletRed', label = 'AAV-FLEX-GFP',linewidth = 1.5)
makelegend3(fig,ax)
adjust_spines(ax, ['left','bottom'])
ax.tick_params(axis='x', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.tick_params(axis='y', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.set_ylabel('Proportion', fontsize=13, labelpad = 16)
ax.set_xlabel('Training day', fontsize=14, labelpad = 16)
ax.set_ylim(-0.2,1.05)
ax.set_xlim(1,20)
ax.axvline(1,linewidth=3, color="black")
ax.axhline(-0.2,linewidth=3, color="black")
plt.yticks(index, ('0','1'))
plt.yticks(rotation=70)
ax = plt.locator_params(axis = 'x', nbins = 4)
plt.subplots_adjust(hspace = 1, wspace = 0.5,  bottom = 0.3, left = 0.15, right = 0.72, top = .85)
fig.savefig('Plots/Figure5/Proportion_graduation_0200' +'.png', dpi = 200)
plt.close()
"""

# plot proportion for GFP and lTeLC and hTeLC mice

fig = plt.figure(figsize = (5,4))
ax = fig.add_subplot(111)
ax.plot(daybins,proportionday_store[:,3], color = 'blue', label = 'lTeLC', linewidth = 1.5)
ax.plot(daybins,proportionday_store[:,2], color = 'red', label = 'hTeLC', linewidth = 1.5)
ax.plot(daybins,proportionday_store[:,1], color = 'Black', label = 'GFP',linewidth = 1.5)
makelegend3(fig,ax)
adjust_spines(ax, ['left','bottom'])
ax.tick_params(axis='x', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.tick_params(axis='y', pad = 10, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =12)
ax.set_ylabel('Proportion', fontsize=13, labelpad = 16)
ax.set_xlabel('Training day', fontsize=14, labelpad = 16)
ax.set_ylim(-0.15,1.15)
ax.set_xlim(0,21)
ax.axvline(0,linewidth=3, color="black")
ax.axhline(-0.15,linewidth=3, color="black")
plt.yticks(index, ('0','1'))
plt.yticks(rotation=70)
ax = plt.locator_params(axis = 'x', nbins = 4)
plt.subplots_adjust(hspace = 1, wspace = 0.5,  bottom = 0.3, left = 0.15, right = 0.72, top = .85)
fig.savefig('Plots/Figure5/Proportion_graduation_0200' +'.png', dpi = 200)
plt.close()

# Store data

genotype = ("TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","GFP","GFP","GFP" ,"GFP" ,"GFP" ,"GFP" )
mouse = np.arange(1,16.1,1)
days = np.arange(1,22.1,1)

lTeLC = proportionday_store[:,3]
hTeLC = proportionday_store[:,2]
GFP = proportionday_store[:,1]

data = np.vstack((days, lTeLC, hTeLC, GFP))
data =  np.transpose(data)

np.savetxt('Data_Output/Figure5/Figure5_A_right_0100.csv', data,fmt =  '%s', delimiter = ',', header = 'Day, lTeLC, hTeLC, GFP ')

## Graduation days for all mice (taken from spreadsheet for data input reasons)

m1 = 7
m2 = 20
m3 = 9
m4 = 7
m5 = 15
m6 = 20
m7 = 20
m8 = 7
m9 = 20
m10 = 7
m11 = 8
m13 = 12
m14 = 11
m15 = 10
m16 = 7
m17 = 7

# stack days for tetanus and control mice
tet = np.hstack((m2,m5,m6,m7,m8,m9,m11,m13,m14,m15))
con = np.hstack((m1,m3,m4,m10,m16,m17))

htet = np.vstack((m2,m6,m7,m9))
ltet = np.vstack((m5,m8,m11,m13,m14,m15))

# average number of mice graduated per day
avg_tet = np.nanmean(tet)
sd_tet = np.nanstd(tet)/math.sqrt(10)
avg_con = np.nanmean(con)
sd_con = np.nanstd(con)/math.sqrt(6)

tt_total = np.hstack((avg_con,avg_tet))
tt_total_sd = np.hstack((sd_con,sd_tet))

hTeLC = np.nanmean(np.hstack((m2,m6,m7,m9)))
lTeLC = np.nanmean(np.hstack((m5,m8,m11,m13,m14,m15)))

# Plot vs expression

mice = [str(int(x)) for x in np.arange(1,16.1)]
expression = np.loadtxt('Data_Input/ExpressionQuantification/T15_FinalQuantification_0300.txt')

expressionstore = np.zeros((len(mice), 3))
expressionstore_d = np.zeros((len(mice), 3))
expressionstore_v = np.zeros((len(mice), 3))
sdexpressionstore = np.zeros((len(mice), 3))
sdexpressionstore_d = np.zeros((len(mice), 3))
sdexpressionstore_v = np.zeros((len(mice), 3))

for mcount,mouse in enumerate(mice):
    marraybym = expression[expression[:,0]==mcount+1,:] # get data for each mouse
    rh = marraybym[:,1:4]
    lh = marraybym[:,4:]
    avg_RH = np.nanmean(rh[:,1])
    sd_RH = stats.sem(rh[:,1])
    avg_LH = np.nanmean(lh[:,1])
    sd_LH =stats.sem(lh[:,1])
    avg = np.nanmean(np.hstack((avg_RH,avg_LH)))
    sd = stats.sem(np.hstack((avg_RH,avg_LH)))
    expressionstore[mcount, 0] = avg
    expressionstore[mcount, 1] = avg_RH
    expressionstore[mcount, 2] = avg_LH
    sdexpressionstore[mcount, 0] = sd
    sdexpressionstore[mcount, 1] = sd_RH
    sdexpressionstore[mcount, 2] = sd_LH

    #for dorsal, mid and ventral
    rh_d = np.delete(rh, np.where(rh[:,2] != 1),0)
    rh_v = np.delete(rh, np.where(rh[:,2] != 2),0)
    lh_d = np.delete(lh, np.where(lh[:,2] != 1),0)
    lh_v = np.delete(lh, np.where(lh[:,2] != 2),0)

    avg_RH_d = np.nanmean(rh_d[:,1])
    avg_LH_d = np.nanmean(lh_d[:,1])
    avg_d = np.nanmean(np.hstack((avg_RH_d,avg_LH_d)))
    avg_RH_v = np.nanmean(rh_v[:,1])
    avg_LH_v = np.nanmean(lh_v[:,1])
    avg_v = np.nanmean(np.hstack((avg_RH_v,avg_LH_v)))

    sd_RH_d = stats.sem(rh_d[:,1])
    sd_LH_d = stats.sem(lh_d[:,1])
    sd_d = stats.sem(np.hstack((avg_RH_d,avg_LH_d)))
    sd_RH_v = stats.sem(rh_v[:,1])
    sd_LH_v = stats.sem(lh_v[:,1])
    sd_v = stats.sem(np.hstack((avg_RH_v,avg_LH_v)))

    expressionstore_d[mcount, 0] = avg_d
    expressionstore_d[mcount, 1] = avg_RH_d
    expressionstore_d[mcount, 2] = avg_LH_d
    expressionstore_v[mcount, 0] = avg_v
    expressionstore_v[mcount, 1] = avg_RH_v
    expressionstore_v[mcount, 2] = avg_LH_v

    sdexpressionstore_d[mcount, 0] = sd_d
    sdexpressionstore_d[mcount, 1] = sd_RH_d
    sdexpressionstore_d[mcount, 2] = sd_LH_d
    sdexpressionstore_v[mcount, 0] = sd_v
    sdexpressionstore_v[mcount, 1] = sd_RH_v
    sdexpressionstore_v[mcount, 2] = sd_LH_v

    mcount +=1

# Average expression

total = np.hstack((expressionstore[1,1],expressionstore[4,1],expressionstore[5,1],expressionstore[6,1],expressionstore[7,1],expressionstore[8,1],expressionstore[10,1],expressionstore[11,1],expressionstore[12,1],expressionstore[13,1]))
totalcon = np.hstack((expressionstore[0,1],expressionstore[2,1],expressionstore[3,1],expressionstore[9,1],expressionstore[14,1],expressionstore[15,1]))
total_probe = np.hstack((expressionstore[4,1],expressionstore[7,1],expressionstore[10,1],expressionstore[11,1],expressionstore[12,1],expressionstore[13,1]))
dorsal = np.hstack((expressionstore_d[1,1],expressionstore_d[4,1],expressionstore_d[5,1],expressionstore_d[6,1],expressionstore_d[7,1],expressionstore_d[8,1],expressionstore_d[10,1],expressionstore_d[11,1],expressionstore_d[12,1],expressionstore_d[13,1]))
dorsalcon = np.hstack((expressionstore_d[0,1],expressionstore_d[2,1],expressionstore_d[3,1],expressionstore_d[9,1],expressionstore_d[14,1],expressionstore_d[15,1]))
dorsal_probe = np.hstack((expressionstore_d[4,1],expressionstore_d[7,1],expressionstore_d[10,1],expressionstore_d[11,1],expressionstore_d[12,1],expressionstore_d[13,1]))
ventral = np.hstack((expressionstore_v[1,1],expressionstore_v[4,1],expressionstore_v[5,1],expressionstore_v[6,1],expressionstore_v[7,1],expressionstore_v[8,1],expressionstore_v[10,1],expressionstore_v[11,1],expressionstore_v[12,1],expressionstore_v[13,1]))
ventralcon = np.hstack((expressionstore_v[0,1],expressionstore_v[2,1],expressionstore_v[3,1],expressionstore_v[9,1],expressionstore_v[14,1],expressionstore_v[15,1]))
sdtotal = np.hstack((sdexpressionstore[1,1],sdexpressionstore[4,1],sdexpressionstore[5,1],sdexpressionstore[6,1],sdexpressionstore[7,1],sdexpressionstore[8,1],sdexpressionstore[10,1],sdexpressionstore[11,1],sdexpressionstore[12,1],sdexpressionstore[13,1]))
sdtotalcon = np.hstack((sdexpressionstore[0,1],sdexpressionstore[2,1],sdexpressionstore[3,1],sdexpressionstore[9,1],sdexpressionstore[14,1],sdexpressionstore[15,1]))
sdtotal_probe = np.hstack((sdexpressionstore[4,1],sdexpressionstore[7,1],sdexpressionstore[10,1],sdexpressionstore[11,1],sdexpressionstore[12,1],sdexpressionstore[13,1]))
sddorsal = np.hstack((sdexpressionstore_d[1,1],sdexpressionstore_d[4,1],sdexpressionstore_d[5,1],sdexpressionstore_d[6,1],sdexpressionstore_d[7,1],sdexpressionstore_d[8,1],sdexpressionstore_d[10,1],sdexpressionstore_d[11,1],sdexpressionstore_d[12,1],sdexpressionstore_d[13,1]))
sddorsalcon = np.hstack((sdexpressionstore_d[0,1],sdexpressionstore_d[2,1],sdexpressionstore_d[3,1],sdexpressionstore_d[9,1],sdexpressionstore_d[14,1],sdexpressionstore_d[15,1]))
sdventral = np.hstack((sdexpressionstore_v[1,1],sdexpressionstore_v[4,1],sdexpressionstore_v[5,1],sdexpressionstore_v[6,1],sdexpressionstore_v[7,1],sdexpressionstore_v[8,1],sdexpressionstore_v[10,1],sdexpressionstore_v[11,1],sdexpressionstore_v[12,1],sdexpressionstore_v[13,1]))
sdventralcon = np.hstack((sdexpressionstore_v[0,1],sdexpressionstore_v[2,1],sdexpressionstore_v[3,1],sdexpressionstore_v[9,1],sdexpressionstore_v[14,1],sdexpressionstore_v[15,1]))

# plot graduation day vs dorsal expression

slope,intercept,r_value, p_value, std_err = stats.linregress(dorsal,tet)
ablinevalues = []
for i in dorsal:
    ablinevalues.append(slope*i+intercept)

slope,intercept,r_value, p_value, std_err = stats.linregress(dorsalcon,con)
ablinevalues1 = []
for i in dorsalcon:
    ablinevalues1.append(slope*i+intercept)
print(dorsal)
fig = plt.figure(figsize = (5,4))
ax = fig.add_subplot(111)
ax.plot(dorsal, tet, 'o',color = 'red',markersize = 5,label = '   Low')
ax.errorbar(dorsal, tet, xerr=sddorsal, fmt='o', markersize = 5,color = 'red',ecolor = 'red')
ax.plot(dorsal,ablinevalues, '-',color = 'red', linewidth = 1)
ax.plot(dorsalcon, con, 'o',color = 'Black',markersize = 5,label = '   Low')
ax.errorbar(dorsalcon, con, xerr=sddorsalcon, fmt='o', markersize = 5,color = 'Black',ecolor = 'Black')
ax.plot(dorsalcon,ablinevalues1, '-',color = 'Black', linewidth = 1)
ax.set_xlim(-50,2000)
ax.set_ylim(-1,21)
adjust_spines(ax, ['left','bottom'])
ax.tick_params(axis='x', pad = 7, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =14)
ax.tick_params(axis='y', pad = 7, which = 'both', top='off', right = 'off', direction = 'out', length = 7, width = 2, labelsize =14)
plt.xticks(rotation=70)
plt.locator_params(axis = 'y', nbins  = 5)
plt.locator_params(axis = 'x', nbins  = 5)#plt.yticks(n_groups,('Ventral','','','Dorsal'))
ax.axhline(-1,linewidth=3, color="black")
ax.axvline(-50,linewidth=3, color="black")
ax.set_xlabel('Flourescence', fontsize=16, labelpad = 20)
ax.set_ylabel('G day', fontsize=16, labelpad = 20)
plt.subplots_adjust(hspace = 1, wspace = 0.5,  bottom = 0.3, left = 0.15, right = 0.72, top = .85)
fig.savefig('Plots/Figure5/GraduationVDorsalExpression' +'.png', dpi = 200)
plt.close()

# SAVE DATA FOR R

genotype = ("TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","TeLC","GFP","GFP","GFP" ,"GFP" ,"GFP" ,"GFP" )
mouse = np.arange(1,16.1,1)
fluorescence = np.hstack((dorsal,dorsalcon))
gday = np.hstack((tet,con))
data = np.vstack((mouse,genotype,fluorescence,gday))
data = np.transpose(data)
np.savetxt('Data_Output/Graduation_dorsal.csv', data,delimiter = '\t', header = 'Mouse\tGenotype\tFlourescence\tGraduation', comments='',fmt='%s')

dfluorescence = np.hstack((dorsal,dorsalcon))
vfluorescence = np.hstack((ventral,ventralcon))

# SAVE DATA FOR MANUSCIPT

data = np.vstack((genotype, gday, dfluorescence, vfluorescence)); data = np.transpose(data)
np.savetxt('Data_Output/Figure5/Figure5_A_left_0100.csv', data,fmt =  '%s', delimiter = ',', header = 'Genotype, Graduation day, Dorsal Fluorescence, Ventral Fluorescence')


import matplotlib.pyplot as plt
import numpy as np
import matplotlib.cm as cm
from matplotlib import colors
import random as rand

class IsingModel:
    '''a class to investigate the properties of the ising model'''
    def __init__(self, grid_type = 'square', dim = 10, J=[1], pbcs = True, Jrule = 'nn', temp=1, no_vacancies = 0, h = 0):

        self.grid_type = grid_type
        self.dim = dim
        self.pbcs = pbcs
        self.Jrule = Jrule #List of J values
        self.h = h #External Field

        self.neighbours = [] #List of neighbours

        #Below takes into account what type of grid it is and makes a list of nearest neighbours
        if Jrule == 'nn':
            if self.grid_type == 'BCC': #If you want a square lattice type square!
                self.BCC()
                self.num = 3

            if self.grid_type == 'square': #If you want a square lattice type square!
                self.create_square_grid()
                self.num = 2

            if self.grid_type == '1D': #If you want a 1D lattice
                self.create_1D_grid()
                self.num =1

            if self.grid_type == 'triangular': #If you want a triangular lattice
                self.create_triangular_grid()

            if self.grid_type == '3D': #If you want to be fancy and go 3d
                self.num = 3
                self.create_three_grid()

            #I created this self.num parameter as there was a wee bit of trouble in the code later on when working with 1d and 3d
            #It made life simpler to just define this function
            if self.grid_type != '3D' and self.grid_type != '1D' and self.grid_type != 'BCC':
                self.num = 2

        if Jrule == 'nnn' and self.grid_type == 'square': #next nearest neighbours
            self.extra_neighbours()
            self.num = 2

        if Jrule == 'nnn' and self.grid_type == 'BCC': #next nearest neighbours
            self.BCC()
            self.num = 3

        #Random Initial Moments are calculated
        self.init_moments = np.random.choice([-1, 1], self.X.shape) #Initial moments are calculated. They are randomised
        self.J = self.create_Jlist(J) #Creates list of J
        self.current_moments = np.copy(self.init_moments)
        self.magnetisation = self.current_moments.flatten().sum() #Calculates the magnetisation of the whole lattice
        self.maglist = [] #List for average magnetisation per site (will be used over range of temperatures)
        self.maglist.append(abs(self.magnetisation)/self.num_sites) #Calculates average magnetisation per site
        self.energy_list = []
        print(self.current_moments)
        self.current_energy =  self.grid_energy(self.current_moments) #Calculates the energy of the lattice
        self.energy_list.append(self.current_energy/self.num_sites) #Calculates the average energy per lattice site

        self.temp = temp #Temperature at which the Ising Model is investigated

    def reset(self):
        '''Resets the lattice, randomly giving each point a spin'''
        self.current_moments = np.random.choice([-1, 1], self.Y.shape)

    def extra_neighbours(self):
        '''Gives each point an additional neareat neighbours four neighbours'''
        x, y = np.linspace(0, self.dim-1, self.dim), np.linspace(0, self.dim-1, self.dim)
        X, Y = np.meshgrid(y, x)
        self.X, self.Y = X.flatten(), Y.flatten()
        self.num_sites = self.dim*self.dim

        for i, (x1, x2) in enumerate(zip(self.X, self.Y)):

            if self.pbcs == True:
                    templist = [[int((x1 - 1) % self.dim), x2] ,
                                [int((x1 + 1) % self.dim), x2] ,
                                [x1, int((x2 - 1) % self.dim) ],
                                [x1, int((x2 + 1) % self.dim)],
                                [int((x1 - 1) % self.dim), int((x2 - 1) % self.dim)] ,
                                [int((x1 + 1) % self.dim), int((x2 + 1) % self.dim)] ,
                                [int((x1 + 1) % self.dim),int((x2 - 1) % self.dim) ],
                                [int((x1 - 1) % self.dim), int((x2 + 1)%(self.dim))] ]
                    #print(templist)
                    reallist = [int(a[1]*self.dim + a[0]) for a in templist] #indices of nearest neighbour
                    self.neighbours.append(reallist)

    def create_1D_grid(self):
        '''Creates a square grid, which will then be used for the lattice investigated'''
        self.X = np.linspace(0, self.dim-1, self.dim)
        self.Y = np.zeros_like(self.X)
        self.num_sites = self.dim

        for i, (x1) in enumerate((self.X)):
            if self.pbcs == True:
                templist = [int((x1 - 1) % self.dim), int((x1 + 1) % self.dim) ]
                self.neighbours.append(templist)

    def BCC(self):
        '''creates a three dimensional lattice square lattice, also known as the simple cubic lattice'''
        x, y, z = np.linspace(0, (self.dim)-1, self.dim), np.linspace(0, (self.dim)-1, self.dim), np.linspace(0, (self.dim)-1, self.dim)
        X, Y, Z = np.meshgrid(y, x, z)
        self.X, self.Y, self.Z = X.flatten(), Y.flatten(), Z.flatten()
        self.num_sites = self.dim*self.dim*self.dim

        if self.Jrule == 'nn':

            for i, (x1, x2, x3) in enumerate(zip(self.X, self.Y, self.Z)):
                if self.pbcs == True:
                    templist = [[x1, x2, int((x3 - 1) % self.dim)] ,  #say atom on x1 plane
                                [x1, x2, int((x3 + 1) % self.dim)] ,
                                [x1, int((x2 + 1) % self.dim), int((x3 - 1) % self.dim)],
                                [x1, int((x2 + 1) % self.dim), int((x3 + 1) % self.dim)],
                                [int((x1 + 1) % self.dim), x2, int((x3 - 1) % self.dim)] ,  #say atom on x1 plane
                                [int((x1 + 1) % self.dim), x2, int((x3 + 1) % self.dim)] ,
                                [int((x1 + 1) % self.dim), int((x2 + 1) % self.dim), int((x3 - 1) % self.dim)],
                                [int((x1 + 1) % self.dim), int((x2 + 1) % self.dim), int((x3 + 1) % self.dim)]]

                    reallist = [int(a[0]*self.dim + a[1]*self.dim**2 + a[2]) for a in templist] #indices of nearest neighbour

                    self.neighbours.append(reallist)

        if self.Jrule == 'nnn':

            for i, (x1, x2, x3) in enumerate(zip(self.X, self.Y, self.Z)):
                if self.pbcs == True:
                    templist = [[x1, x2, int((x3 - 1) % self.dim)] ,  #say atom on x1 plane
                                [x1, x2, int((x3 + 1) % self.dim)] ,
                                [x1, int((x2 + 1) % self.dim), int((x3 - 1) % self.dim)],
                                [x1, int((x2 + 1) % self.dim), int((x3 + 1) % self.dim)],
                                [int((x1 + 1) % self.dim), x2, int((x3 - 1) % self.dim)] ,  #say atom on x1 plane
                                [int((x1 + 1) % self.dim), x2, int((x3 + 1) % self.dim)] ,
                                [int((x1 + 1) % self.dim), int((x2 + 1) % self.dim), int((x3 - 1) % self.dim)],
                                [int((x1 + 1) % self.dim), int((x2 + 1) % self.dim), int((x3 + 1) % self.dim)],
                                [x1, int((x2 + 1) % self.dim), x3], [x1, int((x2 - 1) % self.dim), x3],
                                [int((x1 + 1) % self.dim), x2, x3], [int((x1 - 1) % self.dim), x2, x3]]
                    #last four is next nearest neighbours

                    reallist = [int(a[0]*self.dim + a[1]*self.dim**2 + a[2]) for a in templist] #indices of nearest neighbour

                    self.neighbours.append(reallist)

    def create_square_grid(self):
        '''Creates a square grid, which will then be used for the lattice investigated. Non periodic boundary
        conditions can also be applied to the square grid lattice in this model'''
        x, y = np.linspace(0, self.dim-1, self.dim), np.linspace(0, self.dim-1, self.dim)
        X, Y = np.meshgrid(x, y)
        self.X, self.Y = X.flatten(), Y.flatten()
        self.num_sites = self.dim*self.dim

        for i, (x1, x2) in enumerate(zip(self.X, self.Y)):

            if self.pbcs == True:
                templist = [[int((x1 - 1) % self.dim), x2] , [int((x1 + 1) % self.dim), x2] , [x1,int((x2 - 1) % self.dim) ],
                            [x1, int((x2 + 1) % self.dim)] ]
                reallist = [int(a[1]*self.dim + a[0]) for a in templist]
                self.neighbours.append(reallist)

            elif self.pbcs == False: #This is when non periodic conditions are applied. Note that the number of neighbours is not equal
                templist = [[x1 - 1, x2] , [x1+1, x2] , [x1,x2-1],
                            [x1, x2+1] ] #list of the nearest neighbours
                ar = []
                for a in templist:

                    if a[0] >= 0 and a[1] >= 0 and a[0] < self.dim and a[1] < self.dim:
                        ar.append([a[0], a[1]])
                reallist = [int(c[1]*self.dim + c[0]) for c in ar]

                self.neighbours.append(reallist)

            elif self.pbcs == 'Cylindrical':
                templist = [[int((x1 - 1) % self.dim), x2] , [int((x1 + 1) % self.dim), x2] , [x1,x2-1],
                            [x1, x2+1] ] #list of the nearest neighbours
                ar = []
                for a in templist:
                    if a[1] >= 0 and a[1] < self.dim:
                        ar.append([a[0], a[1]])
                reallist = [int(c[1]*self.dim + c[0]) for c in ar]

                self.neighbours.append(reallist)

    def create_triangular_grid(self):
        '''creates a triangular lattice, where each point has six nearest neighbours in the lattice'''
        x, y = np.linspace(0, self.dim-1, self.dim), np.linspace(0, self.dim-1, self.dim)
        X, Y = np.meshgrid(y, x)
        self.X, self.Y = X.flatten(), Y.flatten()
        self.num_sites = self.dim*self.dim

        for i, (x1, x2) in enumerate(zip(self.X, self.Y)):
            if self.pbcs == True:
                templist = [[int((x1 - 1) % self.dim), x2] , [int((x1 + 1) % self.dim), x2] , [x1,int((x2 - 1) % self.dim) ],
                            [x1, int((x2 + 1) % self.dim)],[int((x1 + 1) % self.dim),int((x2 - 1) % self.dim) ],
                            [int((x1 - 1) % self.dim), int((x2 + 1) % self.dim)] ]
                reallist = [int(a[1]*self.dim + a[0]) for a in templist] #indices of nearest neighbour

                self.neighbours.append(reallist)

    def create_three_grid(self):
        '''creates a three dimensional lattice square lattice, also known as the simple cubic lattice'''
        x, y, z = np.linspace(0, (self.dim)-1, self.dim), np.linspace(0, (self.dim)-1, self.dim), np.linspace(0, (self.dim)-1, self.dim)
        X, Y, Z = np.meshgrid(y, x, z)
        self.X, self.Y, self.Z = X.flatten(), Y.flatten(), Z.flatten()
        self.num_sites = self.dim*self.dim*self.dim

        for i, (x1, x2, x3) in enumerate(zip(self.X, self.Y, self.Z)):
            if self.pbcs == True:
                templist = [[x1, x2, int((x3 - 1) % self.dim)] ,
                            [x1, x2, int((x3 + 1) % self.dim)] ,
                            [int((x1 + 1) % self.dim),x2 , x3],
                            [int((x1 - 1) % self.dim),x2 , x3],
                           [x1, int((x2 - 1) % self.dim),  x3] ,
                            [x1, int((x2 + 1) % self.dim), x3
                            ] ]

                reallist = [int(a[0]*self.dim + a[1]*self.dim**2 + a[2]) for a in templist] #indices of nearest neighbour

                self.neighbours.append(reallist)

    def grid_energy(self, moments):
        energy = 0
        '''Returns scalar for total energy for a given lattice configuration'''
        momentlist = self.current_moments
        for i in range((self.dim**self.num)): #Notice here how I am using self.num. When summing over the nieghbours for npbc it was helpful
            #print(self.J)
            energy += 0.5*momentlist[i]*np.sum(self.J[i]*momentlist[self.neighbours[i]])

            +self.h*np.sum(momentlist) #Note here that an external magnetic field can be applied
        return energy

    def create_Jlist(self, J):
        '''This function creates a list of the J value to be used when calculating the energy at each point. Notes that
        if there is more than one J value, this can only be done when the dimension of the lattice is even'''

        if len(J) ==1 and self.Jrule == 'nn': #if there is only one interaction energy
            return np.array([J[0] for i in range((self.dim**self.num))])

        if self.Jrule == 'nnn':
            if self.grid_type == 'square':
                '''First four will be 1, next four will be sqrt2 apart'''
                jlist = []
                #First create place for each point
                for i in range(self.dim*self.dim*self.dim): #each point
                    for j in range(4):
                        jlist.append(J[0]/np.sqrt(2))
                    for j in range(4):
                        jlist.append(J[0]/1)
                return jlist
            if self.grid_type == 'BCC':
                jlist = []
                #First create place for each point
                for i in range(self.dim*self.dim*self.dim): #each point
                    for j in range(8):
                        jlist.append(J[0]*np.sqrt(3)/2)
                    for j in range(4):
                        jlist.append(J[0]/1)
                return jlist

        #Please note: although this is not necessary, I still created this as I would have liked to look into how different interaction energies between points
        #can effect the curie temperature. Unfortunately I did not have enough time, but I would like to carry this out in the furture and this
        #J list would then be important as it will tell me the interaction energy between each neighbour

    def update_single(self, temp):
        '''Update by going through every point in the lattice'''
        self.current_energy =  self.grid_energy(self.current_moments) #Gets current energy
        newE = self.current_energy #This newE will be updated throughout the loop
        if self.Jrule =='nn' or self.Jrule =='nnn': #Checks what the Jrule is
            for q in range(self.dim**self.num):
                neigh = self.current_moments[self.neighbours[q]]
                neighsum = np.sum(self.J[q]*self.current_moments[q]*neigh)
                deltaE = 2*neighsum + 2*self.h*self.current_moments[q] #energy for a given lattice site
                deltaM = 0 #change in magnetisation
                #print(deltaE)
                if (deltaE <= 0):
                    newE += deltaE #update energy
                    self.current_moments[q] *= -1 #Flip
                    deltaM = 2*self.current_moments[q]
                else:
                    newE +=0
                    self.current_moments[q] = self.current_moments[q]
                    deltaM = 0

                if(deltaE>0): #Now it is defined according to boltzmann probability
                    random_num = float(rand.uniform(0, 1))
                    if (random_num < float(np.exp(-1*float(deltaE)/float(self.temp)))):
                        newE += deltaE #Update energy
                        self.current_moments[q] *= -1 #flip
                        deltaM = 2*self.current_moments[q] #change in magnetisation
                    else:
                        newE +=0
                        self.current_moments[q] = self.current_moments[q]
                        deltaM = 0

                self.magnetisation += deltaM #Update Magnetisation

        self.current_energy = newE #update energy

        self.energy_list.append(self.current_energy/self.num_sites) #average energy per site
        self.maglist.append(abs(self.magnetisation)/self.num_sites) #Absolute value used due to symmetry of hamiltonian

    def run(self, num_its=100, temp=-1):
        '''Performs multiple updates for a given temperature'''
        self.temp = temp
        for it in range(num_its):
            self.update_single(temp = self.temp)

    def grid_figure(self):
        '''This prints the lattice for a given configuration'''
        fig = plt.figure(figsize=(6,5))
        left, bottom, width, height = 0.1, 0.1, 0.8, 0.8
        ax = fig.add_axes([left, bottom, width, height])

        start, stop, n_values = 0, self.dim, self.dim

        x_vals = np.linspace(start, stop, n_values)
        y_vals = np.linspace(start, stop, n_values)
        X, Y = np.meshgrid(x_vals, y_vals)

        lattice = np.zeros((self.dim, self.dim))
        n = 0
        for i in range(self.dim):
            for j in range(self.dim):
                lattice[i, j] = self.current_moments[n]
                n +=1

        cp = plt.pcolormesh(X, Y, lattice)

        ax.set_title('100 x 100 Grid')
        plt.show()

    def print_equil(self, num_its= 300, temp = -1, to_print = 'all'):
        '''This runs the Ising Model and puts it into equilibrium and prints the average energy and magnetisation list'''
        self.run(temp = temp, num_its = num_its)

        if to_print == 'all' or to_print == 'energy':
            plt.plot(self.energy_list)
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Number of Iterations')
            plt.xlabel('Average Energy')
            plt.show()

        if to_print == 'all' or to_print == 'magnetisation':
            plt.plot(self.maglist)
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Number of Iterations')
            plt.xlabel('Average Energy')
            plt.show()

    def props_with_std(self, tempmin = 0.1, tempmax = 4, to_print = 'all'):
        '''This investigates the various properties of the Ising Model over a wide range of temperatures and prints the properties
        over a wide range of temperatures'''
        testruns, countruns = 100, 100
        self.run(temp = self.temp, num_its =3000)
        templist = np.linspace(tempmin,tempmax, 20) #Input range of temperatures here
        aveng, avmag, stdeng, stdmag = np.zeros(len(templist)), np.zeros(len(templist)), np.zeros(len(templist)), np.zeros(len(templist))
        heat_capacity, mag_suscept = np.zeros(20), np.zeros(20)
        list1, list2, list3, list4 = np.zeros((20, 20)), np.zeros((20, 20)), np.zeros((20, 20)), np.zeros((20, 20))
        for i, temp0 in enumerate(templist):
            for j in range(len(list1)):
                print(j, i)
                self.run(num_its = testruns + countruns, temp=temp0)
                aveng[i] = (np.mean(self.energy_list[-countruns:]))
                avmag[i] = (np.mean(self.maglist[-countruns:]))
                stdeng[i] =  (np.std(self.energy_list[-countruns:]))
                stdmag[i] = (np.std(self.maglist[-countruns:]))

                heat_capacity[i] = (stdeng[i]/(temp0**2))
                mag_suscept[i] = (stdmag[i]/(temp0))
                list1[i][j] = np.mean(self.energy_list[-countruns:]) #this repeats it many times and gets the list
                list2[i][j] = np.mean(self.maglist[-countruns:])
                list3[i][j] = heat_capacity[i] #Needed for magnetic susceptibility and heat capacity
                list4[i][j] = mag_suscept[i]

        stdE, stdM, stdH, stdS = [], [], [], []
        print(list1[0])
        for i in range(len(list1)):
            for j in range(len(list1)):

                if abs(list1[i][j] - list1[i][int((j - 1) % 20)]) > 0.2:
                    list1[i][j] = 0

        #now we have to get the standard deviation. Go through array and get each average for given ttemp
        #Now we have to create a list of each of the average of the averages
        for j in range(len(list1)):
            stdE.append(np.std(list1[j]))
            stdM.append(np.std(list2[j]))
            stdH.append(np.std(list3[j]))
            stdS.append(np.std(list4[j]))

        if to_print == 'all' or to_print =='magnetisation':
            plt.errorbar(templist[0:], avmag[0:], yerr = stdM,  marker='o', mfc='red', mec='green')
            plt.grid(linestyle= '-', linewidth=-.2)

            plt.xlabel('Temperature J/k')
            plt.show()

        if to_print == 'all' or to_print =='energy':
            #plt.plot(templist[0:], aveng[0:], 'ro')
            plt.errorbar(templist[0:], aveng[0:], yerr = stdE,  marker='o', mfc='red', mec='green')
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Average Energy per Spin')
            plt.xlabel('Temperature J/k')
            plt.show()

        if to_print == 'all' or to_print =='heat capacity':
            plt.errorbar(templist[0:], heat_capacity[0:], yerr = stdH,  marker='o', mfc='red', mec='green')
            #plt.plot(templist[0:], heat_capacity[0:], 'ro')
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Heat Capactity')
            plt.xlabel('Temperature J/k')
            plt.show()

        if to_print == 'all' or to_print =='magnetic susceptibility':
            #plt.plot(templist[0:], mag_suscept[0:], 'ro')
            plt.errorbar(templist[0:], mag_suscept[0:], yerr = stdS,  marker='o', mfc='red', mec='green', capthick = 2)
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Magnetic Susceptibility')
            plt.xlabel('Temperature J/k')

            plt.show()

        return aveng, avmag,heat_capacity, mag_suscept

    def props(self, tempmin = 0.1, tempmax = 4, to_print = 'all'):
        '''This investigates the various properties of the Ising Model over a wide range of temperatures'''
        testruns, countruns = 100, 200
        templist = np.linspace(tempmin,tempmax, 20) #Input range of temperatures here
        aveng, avmag, stdeng, stdmag = np.zeros(len(templist)), np.zeros(len(templist)), np.zeros(len(templist)), np.zeros(len(templist))
        self.run(temp = self.temp, num_its =100)
        for i, temp0 in enumerate(templist):
            #self.reset()
            print(i)
            self.run(num_its = testruns + countruns, temp=temp0)
            aveng[i] = (np.mean(self.energy_list[-countruns:]))
            avmag[i] = (np.mean(self.maglist[-countruns:]))
            stdeng[i] =  (np.std(self.energy_list[-countruns:]))
            stdmag[i] = (np.std(self.maglist[-countruns:]))

        heat_capacity = stdeng/(templist**2)
        mag_suscept = stdmag/(templist)

        if to_print == 'all' or to_print =='energy':
            plt.plot(templist[0:], avmag[0:], 'ro')
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Average Magnetisation per Spin')
            plt.xlabel('Temperature J/k')
            plt.show()

        if to_print == 'all' or to_print =='magnetisation':
            plt.plot(templist[0:], aveng[0:], 'ro')
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Average Energy per Spin')
            plt.xlabel('Temperature J/k')
            plt.show()

        if to_print == 'all' or to_print =='heat capacity':
            plt.plot(templist[0:], heat_capacity[0:], 'ro')
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Heat Capactity')
            plt.xlabel('Temperature J/k')
            plt.show()

        if to_print == 'all' or to_print =='magnetic susceptibility':
            plt.plot(templist[0:], mag_suscept[0:], 'ro')
            plt.grid(linestyle= '-', linewidth=-.2)
            plt.ylabel('Magnetic Susceptibility')
            plt.xlabel('Temperature J/k')
            plt.show()

        return aveng, avmag,heat_capacity, mag_suscept

#Equilibrium over various different temperatures
'''
ising1 = IsingModel(dim = 25, pbcs = True, temp = 1, J = [1], h = 0)
ising1.run(temp = 1, num_its = 100)

ising2 = IsingModel(dim = 25, pbcs = True, temp = 2, J = [1], h = 0)
ising2.run(temp = 2, num_its = 100)

ising3 = IsingModel(dim = 25, pbcs = True, temp = 10, J = [1], h = 0)
ising3.run( temp = 10, num_its = 100)

plt.plot(ising1.energy_list, label = 'T = 1')
plt.plot(ising2.energy_list, label = 'T = 2')
plt.plot(ising3.energy_list, label = 'T = 10')
plt.legend()
plt.title('Energy of lattice vs Number of Iterations')
plt.xlabel('Number of Iterations')
plt.ylabel('Average Energy per Site')
plt.show()
'''

#Investigate various properties of the Ising Model
'''
ising1 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0)
ising1.print_equil(temp = 1, num_its = 100)

ising1.props_with_std()
'''

#Ising Model for different dimensions
'''
ising1 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0, grid_type = '1D')
ising1.print_equil(temp = 1, num_its = 50)

ising1.props_with_std(tempmax = 6)

ising2 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0, grid_type = 'square')
ising2.print_equil(temp = 1, num_its = 50)

ising2.props_with_std(tempmax = 6)

ising3 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0, grid_type = '3D')
ising3.print_equil(temp = 1, num_its = 50)

ising3.props_with_std(tempmax = 6)
'''

#Include if you want to see how the size effects the Curie Temperature

ising_3 = IsingModel(dim = 3, pbcs = True, temp = 1, J = [1], h = 0)
m13 = ising_3.props(tempmax = 10)[1]

m3 = ising_3.maglist[0:]

ising_10 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0)
m110 = ising_10.props(tempmax = 10)[1]

m10 = ising_10.maglist[0:]

ising_25 = IsingModel(dim = 25, pbcs = True, temp = 1, J = [1], h = 0)
m125 = ising_25.props(tempmax = 10)[1]

m25 = ising_25.maglist[0:]

templist = np.linspace(0.1, 10, 20)

plt.plot(templist[0:], m13[0:], 'ro', label = 'N = 3')
plt.plot(templist[0:], m110[0:], 'bo', label = 'N = 10')
plt.plot(templist[0:], m125[0:], 'go', label = 'N = 25')
plt.legend()
plt.xlabel('Temperature')
plt.ylabel('Magnetisation ')
plt.show()

#Include if you want to see how boundary conditions effect the Curie Temp
'''
ising_3 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [1], h = 0)
m13 = ising_3.props(tempmax = 4)[1]

m3 = ising_3.maglist[0:]

ising_10 = IsingModel(dim = 10, pbcs = False, temp = 0.1, J = [1], h = 0)
m110 = ising_10.props(tempmax = 4)[1]

m10 = ising_10.maglist[0:]

ising_25 = IsingModel(dim = 10, pbcs = 'Cylindrical', temp = 0.1, J = [1], h = 0)
m125 = ising_25.props(tempmax = 4)[1]

m25 = ising_25.maglist[0:]

templist = np.linspace(0.1, 4, 20)

plt.plot(templist[0:], m13[0:], 'ro', label = 'Perioidic')
plt.plot(templist[0:], m110[0:], 'bo', label = 'Non-Periodic')
plt.plot(templist[0:], m125[0:], 'go', label = 'Cylindrical')
plt.legend()
plt.xlabel('Temperature')
plt.ylabel('Magnetisation ')
plt.show()
'''
#Include this is you want to determine the size of the lattice for pbc and nbc
'''
ising_3 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0)
m13 = ising_3.props(tempmax = 4)[1]

m3 = ising_3.maglist[0:]

ising_10 = IsingModel(dim = 10, pbcs = False, temp = 1, J = [1], h = 0)
m110 = ising_10.props(tempmax = 4)[1]

m10 = ising_10.maglist[0:]

ising_25 = IsingModel(dim = 50, pbcs = False, temp = 1, J = [1], h = 0)
m125 = ising_25.props(tempmax = 4)[1]

m25 = ising_25.maglist[0:]

templist = np.linspace(0.1, 10, 20)

plt.plot(templist[0:], m13[0:], 'ro', label = 'PBC; N = 10')
plt.plot(templist[0:], m110[0:], 'bo', label = 'NPBC; N = 10')
plt.plot(templist[0:], m125[0:], 'go', label = 'NPBC; N = 50')
plt.legend()
plt.xlabel('Temperature J/k')
plt.ylabel('Magnetisation')
plt.show()
'''
'''
#Now we introduce how distsnce dependancy effects the Curie Temperature
#Expect fro larger distances the CT will decrease

ising_3 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [1], h = 0)
m13 = ising_3.props(tempmax = 20)[1]

m3 = ising_3.maglist[0:]

ising_10 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [2.5], h = 0)
m110 = ising_10.props(tempmax = 20)[1]

m10 = ising_10.maglist[0:]

ising_25 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [5], h = 0)
m125 = ising_25.props(tempmax = 20)[1]

m25 = ising_25.maglist[0:]

templist = np.linspace(0.1, 10, 20)

plt.plot(templist[0:], m13[0:], 'ro', label = 'd = 5')
plt.plot(templist[0:], m110[0:], 'bo', label = 'd = 2')
plt.plot(templist[0:], m125[0:], 'go', label = 'd = 5' )
plt.legend()
plt.xlabel('Temperature')
plt.ylabel('Magnetisation')
plt.show()
'''

#Next we do the the dependence between neighest neighbours.
'''
ising_3 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [1], h = 0)
m13 = ising_3.props(tempmax = 20)[1]

m3 = ising_3.maglist[0:]

ising_10 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [1], h = 0, Jrule = 'nnn')
m110 = ising_10.props(tempmax = 20)[1]

plt.plot(templist[0:], m13[0:], 'ro', label = 'First NN')
plt.plot(templist[0:], m110[0:], 'bo', label = 'First and Second NN')
plt.xlabel('Temperature J/k')
plt.ylabel('Magnetisation')
plt.legend()
plt.show()
'''

#Now we investigate BCC vs Square
'''
ising_bcc = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0, Jrule = 'nnn', grid_type = 'BCC')
m13 = ising_bcc.props(tempmax = 20)[1]

ising_bcc2 = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], h = 0, Jrule = 'nn', grid_type = 'BCC')
m132 = ising_bcc2.props(tempmax = 20)[1]

ising_s = IsingModel(dim = 10, pbcs = True, temp = 1, J = [1], grid_type = '3D', Jrule = 'nn')
m110 = ising_s.props(tempmax = 20)[1]

plt.plot(templist[0:], m13[0:], 'ro', label = 'BCC with NNN')
plt.plot(templist[0:], m132[0:], 'go', label = 'BCC')
plt.plot(templist[0:], m110[0:], 'bo', label = 'Cubic')
plt.xlabel('Temperature J/k')
plt.ylabel('Magnetisation')
plt.legend()
plt.show()
'''

#Triangular 2d vs square 2d
'''
ising_3 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [1], h = 0, grid_type = 'triangular')
m13 = ising_3.props(tempmax = 20)[1]

m3 = ising_3.maglist[0:]

ising_10 = IsingModel(dim = 10, pbcs = True, temp = 0.1, J = [1], h = 0)
m110 = ising_10.props(tempmax = 20)[1]

plt.plot(templist[0:], m13[0:], 'ro', label = 'Triangular 2D grid')
plt.plot(templist[0:], m110[0:], 'bo', label = 'Square 2D grid')
plt.xlabel('Temperature J/k')
plt.ylabel('Magnetisation')
plt.legend()
plt.show()
'''


from collections import defaultdict

class Solution:
    def first_unique_character(self, arr: str) -> int:
        hash = defaultdict(int)
        for i in range(len(arr)):
            if hash[arr[i]] == 0:
                hash[arr[i]] = i + 1
            elif hash[arr[i]] < len(arr) + 1:
                hash[arr[i]] = len(arr) + 1
        return min(hash.values()) if min(hash.values()) < len(arr) + 1 else -1

so = Solution()
arr = "abcdccacb"
print(so.first_unique_character(arr))


import requests
import re

CONTEST_NAMES_PATTERN = r"(\d+)	([A-Za-z\d\/., -]+)  	([A-Za-z\d. -]+)  	"

def getContestNames() -> dict:

    # Fetch contestcalendar
    resp = requests.get("https://www.contestcalendar.com/cabnames.php")

    if int(resp.status_code / 100) != 2:
        return None

    # Use regex to get all data
    contests = re.findall(CONTEST_NAMES_PATTERN, resp.text, re.M)

    return [
        {
            "id": int(contest[0]),
            "name":contest[1],
            "cabrillo_name":contest[2],
        }
        for contest in contests
    ]


import tensorflow as tf
from utils import Loader, Option_data
import pandas as pd
import numpy as np
import datetime
import QuantLib as ql
import matplotlib.pyplot as plt
from tensorflow.keras.layers import TimeDistributed
from tensorflow.python.client import device_lib

class DHGAN(Loader, Option_data):
    def __init__(self, symbol_list, settings, dh_settings, option_data, option_settings):
        Loader.__init__(self, symbol_list = symbol_list)
        Option_data.__init__(self, option_data ,option_settings)
        self.settings = settings
        self.dh_settings = dh_settings
        self.input_data = np.asarray(self.data)
        self.module = self.settings['module']
        self.hidden_dim = self.settings['hidden_dim']
        self.latent_dim = self.settings['latent_dim']
        self.num_layers = self.settings['num_layers']
        self.iterations = self.settings['iterations']
        self.batch_size = self.settings['batch_size']
        self.learning_rate = self.settings['learning_rate']
        self.no, self.seq_len, self.dim = self.input_data.shape
        # self.time, self.max_seq_len = extract_time(self.input_data)
        self.dh_sims = self.dh_settings['hedge_simulations']
        self.dh_features = self.dh_settings['hedge_features']
        self.dh_epochs = self.dh_settings['hedge_epochs']
        self.dh_batch_size = self.dh_settings['hedge_batch_size']
        self.dh_neurons = self.dh_settings['hedge_neurons']
        self.dh_learning_rate = self.dh_settings['hedge_learning_rate']
        self.dh_validation_proportion = self.dh_settings['hedge_validation_proportion']
        self.dh_test_proportion = self.dh_settings['hedge_test_proportion']
        self.dh_nb_test = int(self.dh_test_proportion * self.dh_sims)
        self.dh_nb_valid = int(self.dh_test_proportion * self.dh_sims)
        self.dh_nb_train = int(self.dh_sims - (self.dh_nb_test+ self.dh_nb_valid))
        self.alpha = self.dh_settings['alpha']
        self.epsilon = self.dh_settings['epsilon']
        self.K = self.dh_settings['strike']
        self.T = self.seq_len
        self.W_dim = self.dim
        # models
        self.Embedder = self.EmbedderNet()
        self.Embedder.summary()
        self.Embedder_optimizer = tf.keras.optimizers.Adam(self.learning_rate)
        self.Recovery = self.RecoveryNet()
        self.Recovery.summary()
        self.Recovery_optimizer = tf.keras.optimizers.Adam(self.learning_rate)
        self.Supervisor = self.SupervisorNet()
        self.Supervisor.summary()
        self.Supervisor_optimizer = tf.keras.optimizers.Adam(self.learning_rate)
        self.Generator = self.GeneratorNet()
        self.Generator.summary()
        self.Generator_optimizer = tf.keras.optimizers.Adam(self.learning_rate)
        self.Discriminator = self.DiscriminatorNet()
        self.Discriminator.summary()
        self.Discriminator_optimizer = tf.keras.optimizers.Adam(self.learning_rate)
        self.dhmodel = self.DeepHedgeNet()
        self.dhmodel_optimizer = tf.keras.optimizers.Adam(self.dh_learning_rate)
        self.dhmodel.summary()
        # Training datasets

        #renorm
        # self.dta = self.norm_data * self.max_val
        # self.dta = self.dta + self.min_val
        # self.dta = np.exp(self.dta)
        # print(self.dta[:, :,:] * np.asarray(self.starting_values))
        # print(self.starting_values)

    def BrownianMotion(self,nb_sims, timesteps, W_dim, dt = 1/365):
        W = np.zeros(shape = (nb_sims, timesteps, W_dim))
        for t in range(1, timesteps):
            W[:,t,:] = W[:,t-1,:] + np.random.randn(nb_sims, W_dim)* np.sqrt(dt)
        return W
    def EmbedderNet(self):
        '''
        Latent space map from function space X (bounded random processes F-adapted)
        to target latent space representation
        '''
        model = tf.keras.Sequential(name = 'EmbedderNet')
        if self.module == 'lstm':
            model.add(tf.keras.layers.LSTM(self.dim, input_shape = (self.T, self.dim), activation = 'tanh',return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.LSTM(self.hidden_dim, activation = 'tanh', return_sequences = True))
        elif self.module == 'gru':
            model.add(tf.keras.layers.GRU(self.dim, input_shape = (self.T, self.dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.GRU(self.hidden_dim, activation = 'tanh', return_sequences = True))
        else:
            raise AssertionError('Please Choose gru or lstm as module')
        model.add(TimeDistributed(tf.keras.layers.Dense(self.latent_dim, activation = 'sigmoid')))
        return model
    def RecoveryNet(self):
        '''
        Inverse map from latent (HxT) space representation to state space of bounded F-adapted random processes
        '''
        model = tf.keras.Sequential(name = 'RecoveryNet')
        if self.module == 'lstm':
            model.add(tf.keras.layers.LSTM(self.hidden_dim, input_shape = (self.T, self.latent_dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.LSTM(self.hidden_dim, activation = 'tanh', return_sequences = True))
        elif self.module == 'gru':
            model.add(tf.keras.layers.GRU(self.hidden_dim, input_shape = (self.T, self.latent_dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.GRU(self.hidden_dim, activation = 'tanh', return_sequences = True))
        model.add(TimeDistributed(tf.keras.layers.Dense(self.dim, activation = 'sigmoid')))
        return model
    def SupervisorNet(self):
        '''
        map from latent space (HxT)to latent space: Generate next sequence based on previous
        '''
        model = tf.keras.Sequential(name = 'SupervisorNet')
        if self.module == 'lstm':
            model.add(tf.keras.layers.LSTM(self.hidden_dim, input_shape = (self.T, self.latent_dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.LSTM(self.hidden_dim, activation = 'tanh', return_sequences = True))
        elif self.module == 'gru':
            model.add(tf.keras.layers.GRU(self.hidden_dim, input_shape = (self.T, self.latent_dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.GRU(self.hidden_dim, activation = 'tanh', return_sequences = True))
        model.add(TimeDistributed(tf.keras.layers.Dense(self.latent_dim, activation = 'sigmoid')))
        return model
    def GeneratorNet(self):
        '''
        Map from d dimensional Brownian motion (W(omega)xT) to latent space
        '''
        model = tf.keras.Sequential(name = 'GeneratorNet')
        if self.module == 'lstm':
            model.add(tf.keras.layers.LSTM(self.hidden_dim, input_shape = (self.T, self.dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.LSTM(self.hidden_dim, activation = 'tanh', return_sequences = True))
        elif self.module == 'gru':
            model.add(tf.keras.layers.GRU(self.hidden_dim, input_shape = (self.T, self.dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.GRU(self.hidden_dim, activation = 'tanh', return_sequences = True))
        model.add(TimeDistributed(tf.keras.layers.Dense(self.latent_dim, activation = 'sigmoid')))
        return model
    def DiscriminatorNet(self):
        '''
        Map from latent space (HxT) representation to [0,1] classification

        '''
        model = tf.keras.Sequential(name = 'DiscriminatorNet')
        if self.module == 'lstm':
            model.add(tf.keras.layers.LSTM(self.hidden_dim, input_shape = (self.T, self.latent_dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.LSTM(self.hidden_dim, activation = 'tanh', return_sequences = True))
        elif self.module == 'gru':
            model.add(tf.keras.layers.GRU(self.hidden_dim, input_shape = (self.T, self.latent_dim), activation = 'tanh', return_sequences = True))
            for layer in range(self.num_layers):
                model.add(tf.keras.layers.GRU(self.hidden_dim, activation = 'tanh', return_sequences = True))
        model.add(TimeDistributed(tf.keras.layers.Dense(self.latent_dim, activation = 'linear')))
        return model
    def DeepHedgeNet(self):
        inp = tf.keras.layers.Input(shape = (self.seq_len, self.dim))
        if self.module == 'lstm':
            X = tf.keras.layers.LSTM(self.dh_neurons, return_sequences = False)(inp)
            X = tf.keras.layers.BatchNormalization()(X)
            T1 = tf.keras.layers.RepeatVector(self.maturities[0]-1)(X)
            T2 = tf.keras.layers.RepeatVector(self.maturities[1]-1)(X)
            T3 = tf.keras.layers.RepeatVector(self.maturities[2]-1)(X)
            T1 = tf.keras.layers.LSTM(self.dh_neurons, activation = 'elu', return_sequences = True)(T1)
            T2 = tf.keras.layers.LSTM(self.dh_neurons, activation = 'elu', return_sequences = True)(T2)
            T3 = tf.keras.layers.LSTM(self.dh_neurons, activation = 'elu', return_sequences = True)(T3)
        elif self.module == 'gru':
            X = tf.keras.layers.GRU(self.dh_neurons, return_sequences = False)(inp)
            X = tf.keras.layers.BatchNormalization()(X)
            T1 = tf.keras.layers.RepeatVector(self.maturities[0]-1)(X)
            T2 = tf.keras.layers.RepeatVector(self.maturities[1]-1)(X)
            T3 = tf.keras.layers.RepeatVector(self.maturities[2]-1)(X)
            T1 = tf.keras.layers.GRU(self.dh_neurons, activation = 'elu', return_sequences = True)(T1)
            T2 = tf.keras.layers.GRU(self.dh_neurons, activation = 'elu', return_sequences = True)(T2)
            T3 = tf.keras.layers.GRU(self.dh_neurons, activation = 'elu', return_sequences = True)(T3)
        T1 = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(self.dh_features * len(self.strikes), activation = 'linear'))(T1)
        T2 = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(self.dh_features * len(self.strikes), activation = 'linear'))(T2)
        T3 = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(self.dh_features * len(self.strikes), activation = 'linear'))(T3)
        model = tf.keras.Model(inputs = inp, outputs = [T1, T2, T3], name = 'DeepHedgeNet')
        return model
    @tf.function
    def EmbedderNetLosst0(self, X, X_tilde):
        E_loss_t0 = tf.compat.v1.losses.mean_squared_error(X, X_tilde)
        return E_loss_t0
    @tf.function
    def EmbedderNetLoss(self, X, X_tilde, G_loss_S):
        E_loss_t0 = self.EmbedderNetLosst0(X, X_tilde)
        E_loss_0 = 10*tf.sqrt(E_loss_t0)
        E_loss = E_loss0 + 0.1 * G_loss_S
        return E_loss_t0, E_loss_0, E_loss
    def RecoveryNetLoss(self):
        pass
    def SupervisorNetLoss(self):
        pass
    @tf.function
    def GeneratorNetLoss(self, y_fake, y_fake_e, H, H_hat_supervise, X_hat, X, gamma = 1):
        G_loss_U = tf.compat.v1.losses.sigmoid_cross_entropy(tf.ones_like(y_fake), y_fake)
        G_loss_U_e = tf.compat.v1.losses.sigmoid_cross_entropy(tf.ones_like(y_fake_e), y_fake_e)
        G_loss_S = self.GeneratorNet_SupervisedLoss(H, H_hat_supervise)
        G_loss_V1 = tf.reduce_mean(tf.abs(tf.sqrt(tf.nn.moments(X_hat,[0])[1] + 1e-6) - tf.sqrt(tf.nn.moments(X,[0])[1] + 1e-6)))
        G_loss_V2 = tf.reduce_mean(tf.abs((tf.nn.moments(X_hat,[0])[0]) - (tf.nn.moments(X,[0])[0])))
        G_loss_V = G_loss_V1 + G_loss_V2
        G_loss = G_loss_U + gamma * G_loss_U_e + 100 * tf.sqrt(G_loss_S) + 100 * G_loss_V
        return G_loss_U, G_loss_U_e, G_loss_S, G_loss_V1, G_loss_V2, G_loss_V, G_loss
    @tf.function
    def GeneratorNet_SupervisedLoss(self, H, H_hat_supervise):
        G_loss_S = tf.compat.v1.losses.mean_squared_error(H[:, 1:, :], H_hat_supervise[:, :-1, :])
        return G_loss_S
    @tf.function
    def DiscriminatorNetLoss(self, y_real,y_fake, y_fake_e, gamma = 1):
        D_loss_real = tf.compat.v1.losses.sigmoid_cross_entropy(tf.ones_like(y_real), y_real)
        D_loss_fake = tf.compat.v1.losses.sigmoid_cross_entropy(tf.zeros_like(y_fake), y_real)
        D_loss_fake_e = tf.compat.v1.losses.sigmoid_cross_entropy(tf.zeros_like(y_fake_e), y_fake_e)
        D_loss = D_loss_real + D_loss_fake + gamma * D_loss_fake_e
        return D_loss_real, D_loss_fake, D_loss_fake_e, D_loss
    @tf.function
    def ReplicationPortfolio(self, X, H):
        dX = X[:, 1:, :self.dh_features] - X[:, :-1, :self.dh_features]
        discrete_integral = tf.reduce_sum(tf.multiply(H, dX), axis = 1)
        trading_strategy = tf.reduce_sum(discrete_integral, axis = 1)
        return trading_strategy
    @tf.function
    def TransactionCosts(self, epsilon, X, H):
        X = X[:, :, :self.dh_features]
        zeros = tf.zeros(shape = (H.shape[0], 1, H.shape[2]))
        H = tf.concat([zeros, H, zeros], axis = 1)
        dH = H[:,1:,:]- H[:,:-1, :]
        dH = tf.abs(dH)
        transaction_costs = tf.multiply(tf.reduce_sum(tf.reduce_sum(tf.multiply(X, dH), axis = 1), axis = 1), epsilon)
        return transaction_costs
    @tf.function
    def DeepHedgeNetLoss1(self, y, trading_strategy, transaction_costs):
        loss,_ =tf.math.top_k(-(-tf.squeeze(y) + trading_strategy - transaction_costs), tf.cast(self.alpha * self.batch_size, tf.int32))
        loss = tf.reduce_mean(tf.abs(loss))
        return loss
    @tf.function
    def DeepHedgeNetLoss2(self, y, trading_strategy, transaction_costs):
        loss,_ =tf.math.top_k(-(-tf.squeeze(y) + trading_strategy - transaction_costs), tf.cast(self.alpha * self.batch_size, tf.int32))
        loss = tf.reduce_mean(loss)
        return loss
    @tf.function
    def DHTermStructureLoss(self, H, X, loss_type = 'type2'):
        L1, L2, L3 = [], [], []
        P01, P02, P03 = [], [], []
        H1 = H[0]
        H2 = H[1]
        H3 = H[2]
        strikes = self.price_array[:,0]
        HlistT1 = tf.split(H[0], len(self.strikes), axis = 2)
        HlistT2 = tf.split(H[1], len(self.strikes), axis = 2)
        HlistT3 = tf.split(H[2], len(self.strikes), axis = 2)
        HKT1 = zip(HlistT1, strikes)
        HKT2 = zip(HlistT2, strikes)
        HKT3 = zip(HlistT3, strikes)
        X1 = X[:,:self.maturities[0], :]
        X2 = X[:,:self.maturities[1], :]
        X3 = X[:,:self.maturities[2], :]
        for (H1, K1),(H2, K2),(H3, K3) in zip(HKT1, HKT2, HKT3):
            strategy1 = self.ReplicationPortfolio(X1, H1)
            cost1 = self.TransactionCosts(self.epsilon, X1, H1)
            Y1 = tf.math.maximum(X1[:,-1:, 0] - K1, 0)
            loss1 = self.DeepHedgeNetLoss1(Y1, strategy1, cost1)
            p01 = strategy1 - cost1
            P01.append(p01[0, 0])
            L1.append(loss1)

            strategy2 = self.ReplicationPortfolio(X2, H2)
            cost2 = self.TransactionCosts(self.epsilon, X2, H2)
            Y2 = tf.math.maximum(X2[:,-1:, 0] - K2, 0)
            loss2 = self.DeepHedgeNetLoss1(Y2, strategy2, cost2)
            p02 = strategy2 - cost2
            P02.append(p02[0, 0])
            L2.append(loss2)

            strategy3 = self.ReplicationPortfolio(X3, H3)
            cost3 = self.TransactionCosts(self.epsilon, X3, H3)
            Y3 = tf.math.maximum(X3[:,-1:, 0] - K3, 0)
            loss3 = self.DeepHedgeNetLoss1(Y3, strategy3, cost3)
            p03 = strategy3 - cost3
            P03.append(p03[0, 0])
            L3.append(loss3)
        P01 = tf.convert_to_tensor(P01, dtype = tf.float32)
        P02 = tf.convert_to_tensor(P02, dtype = tf.float32)
        P03 = tf.convert_to_tensor(P03, dtype = tf.float32)
        L1 = tf.convert_to_tensor(L1, dtype = tf.float32)
        L2 = tf.convert_to_tensor(L2, dtype = tf.float32)
        L3 = tf.convert_to_tensor(L3, dtype = tf.float32)
        market_prices = tf.cast(self.price_array[:, 2], dtype = tf.float32)
        if loss_type == 'type1':
            replication_prices = tf.concat([P01, P02, P03], axis = 0)
            replication_loss = tf.keras.losses.mse(replication_prices, market_prices)
            price_error = None
        elif loss_type == 'type2':
            replication_prices = tf.concat([L1, L2, L3], axis = 0)
            replication_loss = tf.keras.losses.mse(replication_prices, market_prices)
            price_error = None
        elif loss_type == 'type3':
            losses = tf.concat([L1, L2, L3], axis = 0)
            replication_loss = tf.reduce_mean(tf.square(losses))
            replication_prices = tf.concat([P01, P02, P03], axis = 0)
            price_error = tf.keras.losses.mse(replication_prices, market_prices)
        else:
            raise NotImplementedError
        return replication_loss, replication_prices, price_error
    @tf.function
    def Embedder_train_step(self, X):
        with tf.GradientTape() as tape:
            H = self.Embedder(X, training = True)
            X_tilde = self.Recovery(H, training = True)
            loss = self.EmbedderNetLosst0(X, X_tilde)
        grads = tape.gradient(loss, self.Embedder.trainable_variables + self.Recovery.trainable_variables)
        self.Embedder_optimizer.apply_gradients(zip(grads, self.Embedder.trainable_variables + self.Recovery.trainable_variables))
        return loss
    @tf.function
    def Supervised_Generator_train_step(self, X, W):
        with tf.GradientTape() as tape:
            E_hat = self.Generator(W)
            H = self.Embedder(X)
            H_hat_supervise = self.Supervisor(H)
            loss = self.GeneratorNet_SupervisedLoss(H, H_hat_supervise)
        grads = tape.gradient(loss, self.Supervisor.trainable_variables + self.Generator.trainable_variables)
        self.Supervisor_optimizer.apply_gradients(zip(grads, self.Supervisor.trainable_variables + self.Generator.trainable_variables))
        return loss
    @tf.function
    def Unsupervised_Generator_train_step(self, X1, W1):
        with tf.GradientTape(persistent = True) as tape:
            H = self.Embedder(X1)
            E_hat = self.Generator(W1, training = True)
            H_hat_supervise = self.Supervisor(H)
            H_hat = self.Supervisor(E_hat)
            X_hat = self.Recovery(H_hat)
            y_fake = self.Discriminator(H_hat)
            y_fake_e = self.Discriminator(E_hat)
            G_loss_U, _, G_loss_S, _, _, G_loss_V,G_loss = self.GeneratorNetLoss(y_fake, y_fake_e, H, H_hat_supervise, X_hat, X1)
            # Embedder training
            X_tilde = self.Recovery(H, training = True)
            embedder_loss = self.EmbedderNetLosst0(X1, X_tilde)
        generator_grads = tape.gradient(G_loss, self.Generator.trainable_variables + self.Supervisor.trainable_variables)
        self.Generator_optimizer.apply_gradients(zip(generator_grads, self.Generator.trainable_variables + self.Supervisor.trainable_variables))
        embedder_grads = tape.gradient(embedder_loss, self.Embedder.trainable_variables + self.Recovery.trainable_variables)
        self.Embedder_optimizer.apply_gradients(zip(embedder_grads, self.Embedder.trainable_variables + self.Recovery.trainable_variables))
        return G_loss_U, G_loss_S, G_loss_V, G_loss, embedder_loss
    @tf.function
    def Discriminator_train_step(self, X2, W2):
        with tf.GradientTape() as tape:
            H2 = self.Embedder(X2)
            E_hat2 = self.Generator(W2)
            H_hat2 = self.Supervisor(E_hat2)
            y_fake2 = self.Discriminator(H_hat2)
            y_real2 = self.Discriminator(H2)
            y_fake_e2 = self.Discriminator(E_hat2)
            _,_,_, discriminator_loss = self.DiscriminatorNetLoss(y_fake2, y_real2, y_fake_e2)
        if (discriminator_loss > 0.15):
            discriminator_grads = tape.gradient(discriminator_loss, self.Discriminator.trainable_variables)
            self.Discriminator_optimizer.apply_gradients(zip(discriminator_grads, self.Discriminator.trainable_variables))
        return discriminator_loss
    @tf.function
    def DH_train_step(self, W):
        with tf.GradientTape()  as tape:
            E_hat = self.Generator(W, training = True)
            H_hat = self.Supervisor(E_hat, training = True)
            X_hat = self.Recovery(H_hat, training = True)
            X_renorm = self.Renormalize_tf(X_hat)
            hedge_weights = self.dhmodel(X_renorm, training = True)
            termstructure_replication_loss,replication_prices, price_error  = self.DHTermStructureLoss(hedge_weights, X_renorm, loss_type = 'type3')
        dh_grads = tape.gradient(termstructure_replication_loss, self.dhmodel.trainable_variables)
        self.dhmodel_optimizer.apply_gradients(zip(dh_grads, self.dhmodel.trainable_variables))
        return termstructure_replication_loss, replication_prices, price_error
    @tf.function
    def DHGAN_train_step(self,X, W):
        with tf.GradientTape(persistent = True)  as tape:
            H = self.Embedder(X, training = True)
            E_hat = self.Generator(W, training = True)
            H_hat_supervise = self.Supervisor(H, training = True)
            H_hat = self.Supervisor(E_hat, training = True)
            X_hat = self.Recovery(H_hat, training = True)
            y_fake = self.Discriminator(H_hat, training = True)
            y_fake_e = self.Discriminator(E_hat, training = True)
            G_loss_U, _, G_loss_S, _, _, G_loss_V,G_loss = self.GeneratorNetLoss(y_fake, y_fake_e, H, H_hat_supervise, X_hat, X)
            X_tilde = self.Recovery(H, training = True)
            embedder_loss = self.EmbedderNetLosst0(X, X_tilde)
            X_renorm = self.Renormalize_tf(X_hat)
            # X_renorm = X_renorm[:,:,:self.dh_features]
            hedge_weights = self.dhmodel(X_renorm, training = True)
            termstructure_replication_loss,replication_prices ,price_error  = self.DHTermStructureLoss(hedge_weights, X_renorm, loss_type = 'type3')
        generator_grads = tape.gradient(G_loss, self.Generator.trainable_variables + self.Supervisor.trainable_variables)
        self.Generator_optimizer.apply_gradients(zip(generator_grads, self.Generator.trainable_variables + self.Supervisor.trainable_variables))
        embedder_grads = tape.gradient(embedder_loss, self.Embedder.trainable_variables + self.Recovery.trainable_variables)
        self.Embedder_optimizer.apply_gradients(zip(embedder_grads, self.Embedder.trainable_variables + self.Recovery.trainable_variables))
        dhvars = self.Generator.trainable_variables + self.Supervisor.trainable_variables + self.Recovery.trainable_variables + self.dhmodel.trainable_variables
        dh_grads = tape.gradient(termstructure_replication_loss,dhvars)
        self.dhmodel_optimizer.apply_gradients(zip(dh_grads, dhvars))
        return termstructure_replication_loss, replication_prices,price_error ,G_loss, embedder_loss

    def train(self, model_names = None, save = False):
        save_dir = 'model_saves/DHGAN'
        print('Initiating EmbedderNet Pre-training')
        for iteration in range(self.iterations):
            X_batch = self.sample_batch(self.norm_data)
            _embedder_loss = self.Embedder_train_step(X_batch)
            if iteration % 1000 == 0:
                print('Batch training (embedder) loss at iteration %d: %.4f'% (iteration, float(np.sqrt(_embedder_loss))))
                print('---------------------------------------------------------------')
        print('EmbedderNet Pre-training Complete')
        print('Initiating Generator Pre-training (Supervised Loss Only)')
        for iteration in range(self.iterations):
            X_batch = self.sample_batch(self.norm_data)
            W_batch = self.BrownianMotion(X_batch.shape[0],X_batch.shape[1], X_batch.shape[2])
            _generator_loss = self.Supervised_Generator_train_step(X_batch, W_batch)
            if iteration % 1000 == 0:
                print('Batch training supervised (generator) loss at iteration %d: %.4f'% (iteration, float(np.sqrt(_generator_loss))))
                print('---------------------------------------------------------------')
        print('GeneratorNet Pre-training Complete')
        print('Initiating Joint GAN Pre-Training')
        for iteration in range(self.iterations):
            for kk in range(2):
                X1 = self.sample_batch(self.norm_data)
                W1 = self.BrownianMotion(X1.shape[0], X1.shape[1], X1.shape[2])
                G_loss_U, G_loss_S, G_loss_V, generator_loss, embedder_loss = self.Unsupervised_Generator_train_step(X1, W1)
            X2 = self.sample_batch(self.norm_data)
            W2 = self.BrownianMotion(X2.shape[0], X2.shape[1], X2.shape[2])
            discriminator_loss = self.Discriminator_train_step(X2, W2)
            if iteration % 1000 == 0:
                print('---------------------------------------------------------------')
                print('Batch training discriminator loss at iteration %d: %.4f'% (iteration, float(discriminator_loss)))
                print('Batch training generator loss at iteration %d: %.4f'% (iteration, float(generator_loss)))
                print('Batch training unsupervised loss at iteration %d: %.4f'% (iteration, float(G_loss_U)))
                print('Batch training supervised loss at iteration %d: %.4f'% (iteration, float(np.sqrt(G_loss_S))))
                print('Batch training embedder loss at iteration %d: %.4f'% (iteration, float(np.sqrt(embedder_loss))))
                print('---------------------------------------------------------------')
        if save:
            model_dirs = [save_dir + '/{}'.format(model_name) for model_name in model_names]
            # print('Saving model to directory: {}'.format(model_name) for model name in model_names)
            self.Generator.save(model_dirs[0])
            self.Supervisor.save(model_dirs[1])
            self.Recovery.save(model_dirs[2])
        print('Joint GAN Pre-Training Complete')
        print('Initiating Deep Hedge Pre-Training')
        for iteration in range(self.iterations):
            W = self.BrownianMotion(self.batch_size, self.seq_len, self.W_dim)
            termstructure_replication_loss, replication_prices, price_error = self.DH_train_step(W)
            if iteration % 1000 == 0:
                print('---------------------------------------------------------------')
                print('Batch training termstructure replication loss at iteration %d: %.4f'% (iteration, float(termstructure_replication_loss)))
                print('Batch training pricing error at iteration %d: %.4f'% (iteration, float(price_error)))
                print('Batch training termstructure replication prices: {}'.format(replication_prices))
                print('---------------------------------------------------------------')
        if save:
            model_dirs = [save_dir + '/{}'.format(model_name) for model_name in model_names]
            self.dhmodel.save(model_dirs[3])
        print('Deep Hedge Pre-Training Complete')
        print('Initiating Joint DHGAN training')
        for iteration in range(self.iterations):
            for kk in range(2):
                X1 = self.sample_batch(self.norm_data)
                W1 = self.BrownianMotion(X1.shape[0], X1.shape[1], X1.shape[2])
                termstructure_replication_loss, replication_prices, price_error, generator_loss, embedder_loss = self.DHGAN_train_step(X1, W1)
            X2 = self.sample_batch(self.norm_data)
            W2 = self.BrownianMotion(X2.shape[0], X2.shape[1], X2.shape[2])
            discriminator_loss = self.Discriminator_train_step(X2, W2)
            if iteration % 1000 == 0:
                print('---------------------------------------------------------------')
                print('Batch training discriminator loss at iteration %d: %.4f'% (iteration, float(discriminator_loss)))
                print('Batch training generator loss at iteration %d: %.4f'% (iteration, float(generator_loss)))
                print('Batch training embedder loss at iteration %d: %.4f'% (iteration, float(np.sqrt(embedder_loss))))
                print('Batch training termstructure replication loss at iteration %d: %.4f'% (iteration, float(termstructure_replication_loss)))
                print('Batch training pricing error at iteration %d: %.4f'% (iteration, float(price_error)))
                print('Batch training termstructure replication prices: {}'.format(replication_prices))
                market_prices = tf.cast(self.price_array[:, 2], dtype = tf.float32)
                print('Actual Prices: {}'.format(market_prices))
                print('---------------------------------------------------------------')
        if save:
            model_dirs = [save_dir + '/{}'.format(model_name) for model_name in model_names]
            # print('Saving model to directory: {}'.format(model_name) for model name in model_names)
            self.Generator.save(model_dirs[0])
            self.Supervisor.save(model_dirs[1])
            self.Recovery.save(model_dirs[2])
            self.dhmodel.save(model_dirs[3])

    def sample_batch(self, data):
        idx = np.random.randint(low = 0, high = data.shape[0], size = self.batch_size)
        dta = data[idx, :, :]
        dta = tf.cast(dta, tf.float32)
        return dta
    def generate_data(self, W, generator, supervisor, recovery):
        E = generator(W, training = False)
        H = supervisor(E, training = False)
        X = recovery(H, training = False)
        return X
    def Renormalize(self, X):
        X = X * (self.max_val + 1e-7)
        X = X + self.min_val
        X = np.exp(X)
        X = X * np.asarray(self.starting_values)
        return X
    @tf.function
    def Renormalize_tf(self, X):
        max_val = tf.cast(self.max_val + 1e-7, tf.float32)
        min_val = tf.cast(self.min_val, tf.float32)
        X = X * max_val
        X = X + min_val
        X = tf.math.exp(X)
        X = X * tf.convert_to_tensor(np.asarray(self.starting_values), dtype = tf.float32)
        return X
    def Simulate_Hedge(self, W, generator, supervisor, recovery, dhmodel):
        X = self.generate_data(W, generator, supervisor, recovery)
        X = self.Renormalize(X)
        H = dhmodel(X)
        return X, H
    def Hedge_on_data(self, data):
        pass
    def Simulate_Data(self, W, generator, supervisor, recovery):
        X = self.generate_data(W, generator, supervisor, recovery)
        X = self.Renormalize(X)
        return X
    def restore_model(self, model_name):
        save_dir = 'model_saves/DHGAN'
        model = tf.keras.models.load_model(save_dir + '/{}'.format(model_name))
        return model
    def RestoreGAN(self, model_names):
        save_dir = 'model_saves/DHGAN'
        models = [tf.keras.models.load_model(save_dir + '/{}'.format(model_name)) for model_name in model_names]
        return models[0], models[1], models[2]
    def RestoreDHGAN(self, model_names):
        save_dir = 'model_saves/DHGAN'
        models = [tf.keras.models.load_model(save_dir + '/{}'.format(model_name)) for model_name in model_names]
        return models[0], models[1], models[2], models[3]
if __name__ == '__main__':
    print(device_lib.list_local_devices())
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
    symbol_list = ['^GSPC', '^VIX', '^NDX', '^RUT', '^DJI',
                '^FVX','^TNX', '^TYX', 'EURUSD=X','JPY=X']
    settings = {'module':'lstm',
                'hidden_dim':24,
                'latent_dim': 5,
                'num_layers': 3,
                'iterations': 30000,
                'batch_size': 100,
                'learning_rate':0.001}
    DH_settings = {'hedge_simulations': 1000000,
                    'hedge_features': 2,
                    'hedge_epochs': 2,
                    'hedge_batch_size': 1000,
                    'hedge_neurons': 70,
                    'hedge_learning_rate': 0.00001,
                    'hedge_validation_proportion':0.1,
                    'hedge_test_proportion':0.001,
                    'alpha':0.99,
                    'epsilon':0.0,
                    'strike':3200}
    datafolder_path = 'datafolder'
    filename = 'quotedata_SPX.csv'
    filepath = datafolder_path + '/' + filename
    option_data = pd.read_csv(filepath, sep = ';')
    calendar = ql.UnitedStates()
    Ks = [3000, 3100, 3200, 3300, 3400]
    ttms = [10, 21, 31]
    option_settings = {'start': datetime.date(2020,9,11),
                'day_count': ql.Actual365Fixed(),
                'calendar': ql.UnitedStates(),
                'strikes': Ks,
                'maturities':ttms}
    Model = DHGAN(symbol_list,settings, DH_settings, option_data ,option_settings)

    model_names = ['Generator', 'Supervisor', 'Recovery', 'dhmodel']
    Model.train(model_names, save = True)
    gen, sup, rec, dh = Model.RestoreDHGAN(model_names)
    W = Model.BrownianMotion(100, 31, 10)
    X, H = Model.Simulate_Hedge(W, gen, sup, rec, dh)
    print(X)
    print('aldkfjalskjfd')
    print(H)
    for i in range(15):
        plt.plot(X[i, :,0])
    plt.show()
    # Model.test(Model.Generator_data)


import subprocess
import os
from typing import Any
import yaml
import time
import random

def exec_command(command, env=None) -> str:
    if env is None:
        env = os.environ.copy()

    process = subprocess.run(command, capture_output=True, env=env)

    if process.returncode == 0:
        return process.stdout.decode()

    raise Exception('command {} failed\nstdout:\n{}\nstderr:\n{}'.format(
        command,
        process.stdout.decode(),
        process.stderr.decode(),
    ))

def generate_random_config() -> Any:
    ml_config = {
        'adam_lr': 0.01,
        'assets': ['SANHOK', 'MIRAMAR'],
        'batch_size': 1024,
        'hidden_dim': random.randint(20, 30),
        'n_batches_per_epoch': 50,
        'n_epochs': 100,
        'num_lstm_layers': 4,
        'droupout_percent': random.randint(20, 40),
        'sample_top_features': {
            'n_top_features': random.randint(60, 80),
            'feature_importance_path': "feature_importance.csv",
        },
        'scheduler': {'type': 'linear'},
    }
    return {
        'ml': {
            'config': ml_config,
            'path': 'model/ml_regressor',
        }
    }

def main():
    begin = time.time()
    with open('config.yml', 'w') as file:
        yaml.safe_dump(generate_random_config(), file)
    exec_command(['python3', 'train_scripts/train_lstm_regressor.py'])
    print(f'finished in {int(time.time() - begin)}s')

if __name__ == '__main__':
    main()


import typing

import copy
import itertools

import numpy as np
import pandas as pd
import scipy.stats
import sklearn.ensemble

class AbstractSelector:
    def fit(
        self,
        X: pd.DataFrame,
        y: pd.DataFrame,
        minima: typing.Dict[int, typing.Dict[str, float]],
        maxima: typing.Dict[int, typing.Dict[str, float]],
    ) -> None:
        raise NotImplementedError()

    def predict(
        self, X: pd.DataFrame, y: typing.Optional[pd.DataFrame] = None
    ) -> pd.DataFrame:
        prediction = self._predict(X, y)
        for col, series in prediction.iteritems():
            assert series.dtype == float, (col, series)
        np.testing.assert_array_almost_equal(
            prediction.sum(axis="columns").to_numpy(),
            np.ones(X.shape[0]),
            err_msg=prediction.to_csv(),
        )
        return prediction

    def _predict(
        self, X: pd.DataFrame, y: typing.Optional[pd.DataFrame]
    ) -> pd.DataFrame:
        raise NotImplementedError()

class OneVSOneSelector(AbstractSelector):
    def __init__(self, configuration, random_state, tie_break_order):
        self.configuration = configuration
        self.single_strategy_idx = None
        self.rng = random_state
        self.tie_break_order = tie_break_order
        self.models = None
        self.target_indices = None
        self.selectors_ = None
        self.weights_ = {}
        self.strategies_ = None

    def fit(
        self,
        X: pd.DataFrame,
        y: pd.DataFrame,
        minima: typing.Dict[int, typing.Dict[str, float]],
        maxima: typing.Dict[int, typing.Dict[str, float]],
    ) -> None:

        self.maxima = copy.deepcopy(maxima)
        target_indices = np.array(list(range(y.shape[1])))
        models = dict()
        weights = dict()

        self.strategies_ = copy.deepcopy(y.columns.to_list())
        y_pd = y.copy()
        y = y.to_numpy()

        for i in range(len(target_indices)):
            models[i] = dict()
            weights[i] = dict()
            for j in range(i + 1, len(target_indices)):

                if self.configuration["normalization"] in (
                    "all",
                    "binary",
                    "y",
                    "all1",
                    "binary1",
                ):
                    minimum2 = np.ones(len(X)) * np.inf
                    maximum2 = np.zeros(len(X))

                    if self.configuration["normalization"] in ("all", "all1"):
                        for idx, task_id in enumerate(X.index):
                            for method_id in range(len(target_indices)):
                                minimum2[idx] = np.nanmin(
                                    (
                                        minimum2[idx],
                                        minima[task_id][self.strategies_[method_id]],
                                    )
                                )
                                maximum2[idx] = np.nanmax(
                                    (
                                        maximum2[idx],
                                        maxima[task_id][self.strategies_[method_id]],
                                    )
                                )
                        if self.configuration["normalization"] == "all1":
                            maximum2 = np.ones_like(maximum2)
                    elif self.configuration["normalization"] in ("binary", "binary1"):
                        for idx, task_id in enumerate(X.index):
                            for method_id in (i, j):
                                minimum2[idx] = np.nanmin(
                                    (
                                        minimum2[idx],
                                        minima[task_id][self.strategies_[method_id]],
                                    )
                                )
                                maximum2[idx] = np.nanmax(
                                    (
                                        maximum2[idx],
                                        maxima[task_id][self.strategies_[method_id]],
                                    )
                                )
                        if self.configuration["normalization"] == "binary1":
                            maximum2 = np.ones_like(maximum2)
                    elif self.configuration["normalization"] == "y":
                        for idx, task_id in enumerate(X.index):
                            minimum2[idx] = np.nanmin(
                                (minimum2[idx], y_pd.loc[task_id].min())
                            )
                            maximum2[idx] = np.nanmax(
                                (maximum2[idx], y_pd.loc[task_id].max())
                            )
                    else:
                        raise ValueError(self.configuration["normalization"])

                    y_i_j = y[:, i] < y[:, j]
                    mask = np.isfinite(y[:, i]) & np.isfinite(y[:, j])
                    X_ = X.to_numpy()[mask]
                    y_i_j = y_i_j[mask]

                    minimum = minimum2[mask]
                    maximum = maximum2[mask]

                    diff = maximum - minimum
                    diff[diff == 0] = 1
                    normalized_y_i = (y[:, i][mask].copy() - minimum) / diff
                    normalized_y_j = (y[:, j][mask].copy() - minimum) / diff

                    weights_i_j = np.abs(normalized_y_i - normalized_y_j)

                elif self.configuration["normalization"] == "rank":
                    y_i_j = y[:, i] < y[:, j]
                    mask = np.isfinite(y[:, i]) & np.isfinite(y[:, j])
                    X_ = X.to_numpy()[mask]
                    y_i_j = y_i_j[mask]
                    ranks = scipy.stats.rankdata(y[mask], axis=1)
                    weights_i_j = np.abs(ranks[:, i] - ranks[:, j])

                elif self.configuration["normalization"] == "None":
                    y_i_j = y[:, i] < y[:, j]
                    mask = np.isfinite(y[:, i]) & np.isfinite(y[:, j])
                    X_ = X.to_numpy()[mask]
                    y_i_j = y_i_j[mask]
                    weights_i_j = np.ones_like(y_i_j).astype(int)

                else:
                    raise ValueError(self.configuration["normalization"])

                if len(y_i_j) == 0:
                    models[i][j] = None
                    weights[i][j] = None
                    continue

                if np.all([target == y_i_j[0] for target in y_i_j]):
                    n_zeros = int(np.ceil(len(y_i_j) / 2))
                    n_ones = int(np.floor(len(y_i_j) / 2))
                    import sklearn.dummy

                    base_model = sklearn.dummy.DummyClassifier(
                        strategy="constant", constant=y_i_j[0]
                    )
                    base_model.fit(
                        X_,
                        np.array(([[0]] * n_zeros) + ([[1]] * n_ones)).flatten(),
                        sample_weight=weights_i_j,
                    )
                else:
                    if self.configuration.get("max_depth") == 0:
                        import sklearn.dummy

                        loss_i = np.sum((y_i_j == 0) * weights_i_j)
                        loss_j = np.sum((y_i_j == 1) * weights_i_j)

                        base_model = sklearn.dummy.DummyClassifier(
                            strategy="constant",
                            constant=1 if loss_i < loss_j else 0,
                        )
                        base_model.fit(
                            X_,
                            np.ones_like(y_i_j) * int(loss_i < loss_j),
                            sample_weight=weights_i_j,
                        )
                    else:
                        base_model = self.fit_pairwise_model(
                            X_,
                            y_i_j,
                            weights_i_j,
                            self.rng,
                            self.configuration,
                        )
                models[i][j] = base_model
                weights[i][j] = weights_i_j
        self.models = models
        self.weights_ = weights
        self.target_indices = target_indices

    def _predict(
        self, X: pd.DataFrame, y: typing.Optional[pd.DataFrame]
    ) -> pd.DataFrame:

        if y is not None:
            raise ValueError("y must not be provided")

        raw_predictions = dict()
        raw_probas = dict()
        for i in range(len(self.target_indices)):
            for j in range(i + 1, len(self.target_indices)):
                if self.models[i][j] is not None:
                    raw_predictions[(i, j)] = self.models[i][j].predict(X)
                    raw_probas[(i, j)] = self.models[i][j].predict_proba(X)

        if len(raw_predictions) == 0:
            predictions = pd.DataFrame(
                0, index=X.index, columns=self.strategies_
            ).astype(float)
            predictions.iloc[:, self.single_strategy_idx] = 1.0
            return predictions

        predictions = {}
        for x_idx in range(X.shape[0]):
            wins = np.zeros(self.target_indices.shape)
            for i in range(len(self.target_indices)):
                for j in range(i + 1, len(self.target_indices)):
                    if (i, j) in raw_predictions:
                        if self.configuration["prediction"] == "soft":
                            if raw_probas[(i, j)].shape[1] == 1:
                                proba = raw_probas[(i, j)][x_idx][0]
                            else:
                                proba = raw_probas[(i, j)][x_idx][1]
                            wins[i] += proba
                            wins[j] += 1 - proba
                        elif self.configuration["prediction"] == "hard":
                            prediction = raw_predictions[(i, j)][x_idx]
                            if prediction == 1:
                                wins[i] += 1
                            else:
                                wins[j] += 1
                        else:
                            raise ValueError(self.configuration["prediction"])

            n_prev = np.inf
            # Tie breaking
            while True:
                most_wins = np.max(wins)
                most_wins_mask = most_wins == wins
                if n_prev == np.sum(most_wins_mask):
                    n_prev = np.sum(most_wins_mask)
                    hit = False
                    for method in self.tie_break_order:
                        if method not in self.strategies_:
                            continue
                        method_idx = self.strategies_.index(method)
                        if most_wins_mask[method_idx]:
                            wins[method_idx] += 1
                            hit = True
                            break
                    if not hit:
                        wins[
                            int(self.rng.choice(np.argwhere(most_wins_mask).flatten()))
                        ] += 1
                elif np.sum(most_wins_mask) > 1:
                    n_prev = np.sum(most_wins_mask)
                    where = np.argwhere(most_wins_mask).flatten()
                    for i, j in itertools.combinations(where, 2):
                        if (i, j) in raw_predictions:
                            prediction = raw_predictions[(i, j)][x_idx]
                            if prediction == 1:
                                wins[i] += 1
                            else:
                                wins[j] += 1
                        else:
                            method_i = self.strategies_[i]
                            method_j = self.strategies_[j]
                            if self.tie_break_order.index(
                                method_i
                            ) < self.tie_break_order.index(method_j):
                                wins[i] += 1
                            else:
                                wins[j] += 1
                else:
                    break

            wins = wins / np.sum(wins)
            predictions[X.index[x_idx]] = wins

        return_value = {
            task_id: {
                strategy: predictions[task_id][strategy_idx]
                for strategy_idx, strategy in enumerate(self.strategies_)
            }
            for task_id in X.index
        }
        return_value = pd.DataFrame(return_value).transpose().astype(float)
        return_value = return_value[self.strategies_]
        return_value = return_value.fillna(0.0)
        return return_value

    def fit_pairwise_model(self, X, y, weights, rng, configuration):
        raise NotImplementedError()

class OVORF(OneVSOneSelector):
    def __init__(self, n_estimators, **kwargs):
        super().__init__(**kwargs)
        self.n_estimators = n_estimators

    def fit_pairwise_model(self, X, y, weights, rng, configuration):
        base_model = sklearn.ensemble.RandomForestClassifier(
            random_state=rng,
            n_estimators=self.n_estimators,
            bootstrap=True if configuration["bootstrap"] == "True" else False,
            min_samples_split=configuration["min_samples_split"],
            min_samples_leaf=configuration["min_samples_leaf"],
            max_features=int(configuration["max_features"]),
            max_depth=configuration["max_depth"],
        )
        base_model.fit(X, y, sample_weight=weights)
        return base_model

class FallbackWrapper(AbstractSelector):
    def __init__(self, selector, default_strategies: typing.List[str]):
        self.selector = selector
        self.default_strategies = default_strategies

    def fit(
        self,
        X: pd.DataFrame,
        y: pd.DataFrame,
        minima: typing.Dict[int, typing.Dict[str, float]],
        maxima: typing.Dict[int, typing.Dict[str, float]],
    ) -> None:
        self.X_ = X
        self.strategies_ = y.columns
        self.return_value_ = np.array(
            [
                (len(self.strategies_) - self.default_strategies.index(strategy) - 1)
                / (len(self.strategies_) - 1)
                for strategy in self.strategies_
            ]
        )
        self.return_value_ = self.return_value_ / np.sum(self.return_value_)
        self.selector.fit(X, y, minima, maxima)

    def _predict(
        self, X: pd.DataFrame, y: typing.Optional[pd.DataFrame]
    ) -> pd.DataFrame:

        if y is not None:
            prediction = self.selector.predict(X, y)
        else:
            prediction = self.selector.predict(X)
        for task_id, x in X.iterrows():
            counter = 0
            te = x.copy()
            assert len(te) == self.X_.shape[1]
            for _, tr in self.X_.iterrows():
                tr = tr.to_numpy()
                if all([tr[i] >= te[i] for i in range(self.X_.shape[1])]):
                    counter += 1

            if counter == 0:
                prediction.loc[task_id] = pd.Series(
                    {
                        strategy: value
                        for strategy, value in zip(self.strategies_, self.return_value_)
                    }
                )

        return prediction


"""This module defines class UserLogout."""
from rest_framework_simplejwt.tokens import RefreshToken
from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework.permissions import IsAuthenticated
from rest_framework.views import APIView
from django.http import JsonResponse
from user_activity.models import UserActivity
from car_app.models import User
from car_app.views.views_helper_functions import decode_token, decode_refresh_token

class UserLogout(APIView):
    """
    This class defines a method that logout a user from
    the application.
    """
    #pylint: disable=no-member

    authentication_classes = [JWTAuthentication]
    permission_classes = [IsAuthenticated]

    def post(self, request):
        """This method blacklists refresh token when a user logs out."""
        if request.content_type != 'application/json':
            return JsonResponse({'error': 'The Content-Type must be json.'}, status=415)

        refresh_token = request.data.get('refresh')
        if not refresh_token:
            return JsonResponse({'error': 'Refresh token is required.'}, status=400)

        result = decode_token(request)
        if isinstance(result, tuple):
            user_id, is_superuser, _ = result

            result = decode_refresh_token(refresh_token)
            if isinstance(result, JsonResponse):
                return result

            refresh_token_user_id = result
            try:
                user_logging_out = User.objects.get(id=refresh_token_user_id)
            except User.DoesNotExist:
                return JsonResponse({'error': 'Provided token has no valid user.'}, status=401)

            manager = user_logging_out.team_manager if hasattr(
                user_logging_out, 'team_manager') else None

            if manager:
                user_manager = True if manager.id == user_id else False
            else:
                user_manager = False

            if refresh_token_user_id != user_id and is_superuser is False and user_manager is False:
                error_message = 'You do not have the permission to perform this action.'
                return JsonResponse({'error': error_message}, status=403)

            try:
                token = RefreshToken(refresh_token)
                token.blacklist()

                UserActivity.objects.create(
                    user = user_logging_out,
                    activity_type = 'Logout',
                )

                return JsonResponse({'message': 'Logged out successfully.'}, status=200)
            except BaseException as error: # pylint: disable=broad-exception-caught
                return JsonResponse({'error': str(error)}, status=400)

        error_message = result
        return error_message