#Import Necessary Packages
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats
import seaborn as sns
import ruptures as rpt
from statistics import stdev
import pandas as pd

def load_rms(path, sect, ref):
    raw_string = open('../../' + path + '/rmsd_' + sect + '_ref_' + ref + '.txt').readlines()
    #Convert data to fload
    raw = np.zeros(len(raw_string))
    for i in range(len(raw_string)):
        raw[i] = float(raw_string[i])*10
    return raw

def plot_compare(RMSD_mean, RMSD_err, Label, sect, n, ref):
    rmsd_n = RMSD_mean[:,n]
    rmsd_err_n = RMSD_err[:,n]

    section = sect[n]
#    bar_color = ['gray', 'gray', 'pink', 'blue', 'pink', 'red', 'red']

    num = np.linspace(1, len(Label)+1, num = len(Label))
    fig = plt.figure(figsize=(18,10))
    ax1 = fig.add_subplot(111)
    ax1.set_title("RMSD for " + section + ' to ' + ref)
    ax1.set_ylabel(r'RMSD($\AA$)')
    ax1.bar(num, rmsd_n)
    plt.xticks(num, Label, fontsize=14)
    plt.errorbar(num, rmsd_n, yerr= rmsd_err_n, fmt='o', color='black')
    fig.savefig('RMSD_compare_' + section + '_' + ref + '.png')
    plt.close(fig)

def plot_kernel_mut(df, sect_name, sect_file, xmin, xmax):
    ax = sns.kdeplot(data = df, fill=True, alpha=0.5, common_grid = True)
    plt.setp(ax.get_legend().get_texts(), fontsize='12') # for legend text
    plt.xlabel(r'RMSD($\AA$)', fontsize = 14)
    plt.xlim(xmin, xmax)
    plt.xticks(fontsize = 13)
    plt.yticks(fontsize = 13)
    plt.ylabel(r'Normalized Density', fontsize = 14)
    plt.title(str(sect_name) + r' RMSD Relative to WT Closed', fontsize = 15)
    plt.savefig('mutate_RMSD_' + str(sect_file) + '.png')
    plt.close()

def plot_kernel_cmpr_lig(apo_df, AD_df, BBR_df, mut, sect, n):
    df = pd.concat([apo_df, AD_df, BBR_df])

    sns.kdeplot(data = df, fill=True, alpha=0.5, common_norm = True, common_grid = True)
    plt.xlabel(r'RMSD($\AA$)', fontsize = 14)
    plt.ylabel(r'Normalized Density', fontsize = 14)
    plt.title(sect + ' RMSD Compared to WT Closed', fontsize = 15)
    plt.savefig('mutate_RMSD_' + sect + '_' + mut + '.png')
    plt.close()

def rmsd_sect(sect, file_path_close, file_path_close_AD, file_path_close_BBR, ref, n):
    rmsd_1sug = load_rms(file_path_close[0], sect, ref[n])
    rmsd_apo = load_rms(file_path_close[1], sect, ref[n])
    rmsd_L192F = load_rms(file_path_close[2], sect, ref[n])
    rmsd_E276F = load_rms(file_path_close[3], sect, ref[n])
    rmsd_F280Y = load_rms(file_path_close[4], sect, ref[n])
    rmsd_L195F = load_rms(file_path_close[5], sect, ref[n])
    rmsd_F196A = load_rms(file_path_close[6], sect, ref[n])
    rmsd_V287T = load_rms(file_path_close[7], sect, ref[n])

    rmsd_L192F_AD = load_rms(file_path_close_AD[0], sect, ref[n])
    rmsd_L192F_BBR = load_rms(file_path_close_BBR[0], sect, ref[n])
    rmsd_E276F_AD = load_rms(file_path_close_AD[1], sect, ref[n])
    rmsd_E276F_BBR = load_rms(file_path_close_BBR[1], sect, ref[n])
    rmsd_F280Y_AD = load_rms(file_path_close_AD[2], sect, ref[n])
    rmsd_F280Y_BBR = load_rms(file_path_close_BBR[2], sect, ref[n])
    rmsd_L195F_AD = load_rms(file_path_close_AD[3], sect, ref[n])
    rmsd_L195F_BBR = load_rms(file_path_close_BBR[3], sect, ref[n])
    rmsd_F196A_AD = load_rms(file_path_close_AD[4], sect, ref[n])
    rmsd_F196A_BBR = load_rms(file_path_close_BBR[4], sect, ref[n])
    rmsd_V287T_AD = load_rms(file_path_close_AD[4], sect, ref[n])
    rmsd_V287T_BBR = load_rms(file_path_close_BBR[4], sect, ref[n])

    return rmsd_1sug, rmsd_apo, rmsd_L192F, rmsd_E276F, rmsd_F280Y, rmsd_L195F, rmsd_F196A, rmsd_V287T, rmsd_L192F_AD, rmsd_E276F_AD, rmsd_F280Y_AD, rmsd_L195F_AD, rmsd_F196A_AD, rmsd_V287T_AD, rmsd_L192F_BBR, rmsd_E276F_BBR, rmsd_F280Y_BBR, rmsd_L195F_BBR, rmsd_F196A_BBR, rmsd_V287T_BBR

#File paths for all input files
file_path = ['../Apo_dis/analysis', 'L192F/Apo/analysis', 'E276F/Apo/analysis', 'F280Y/Apo/analysis', 'L195F/Apo/analysis', 'F196A/Apo/analysis', 'V287T/Apo/analysis'] #Indices to rank in order of closest activity to WT to Furthest

file_path_close = ['../Apo_1SUG/analysis/1sug', '../Apo_dis/analysis', 'L192F/Apo/analysis', 'E276F/Apo/analysis', 'F280Y/Apo/analysis', 'L195F/Apo/analysis', 'F196A/Apo/analysis', 'V287T/Apo/analysis']
file_path_close_AD = ['L192F/AD/analysis', 'E276F/AD/analysis', 'F280Y/AD/analysis', 'L195F/AD/analysis', 'F196A/AD/analysis', 'V287T/AD/analysis']
file_path_close_BBR = ['L192F/BBR/analysis', 'E276F/BBR/analysis', 'F280Y/BBR/analysis', 'L195F/BBR/analysis', 'F196A/BBR/analysis', 'V287T/BBR/analysis']

sections = ['WPD', 'WPD_a3', 'SBL', 'beg', 'P', 'CYS', 'a3', 'a3_top', 'a4', 'a5', 'a6', 'a6_bot', 'a7', 'Q']
ref = ['open', 'closed', 'self', 'F196A', 'V287T']

#open all files
RMSD_mean = np.zeros((len(file_path), len(sections))) #Mean for reference open
RMSD_err = np.zeros((len(file_path), len(sections))) #SEM for reference open
RMSD_mean_close = np.zeros((len(file_path_close), len(sections))) #Mean for reference closed
RMSD_err_close = np.zeros((len(file_path_close), len(sections))) #SEM for reference closed
RMSD_mean_close_AD = np.zeros((len(file_path_close_AD), len(sections))) #Mean for reference closed
RMSD_err_close_AD = np.zeros((len(file_path_close_AD), len(sections))) #SEM for reference closed
RMSD_mean_close_BBR = np.zeros((len(file_path_close_BBR), len(sections))) #Mean for reference closed
RMSD_err_close_BBR = np.zeros((len(file_path_close_BBR), len(sections))) #SEM for reference closed

#Save all rmsd values for a3_top, a4_top, and a6 helix
rmsd_a3_1sug, rmsd_a3_apo, rmsd_a3_L192F, rmsd_a3_E276F, rmsd_a3_F280Y, rmsd_a3_L195F, rmsd_a3_F196A, rmsd_a3_V287T, rmsd_a3_L192F_AD, rmsd_a3_E276F_AD, rmsd_a3_F280Y_AD, rmsd_a3_L195F_AD, rmsd_a3_F196A_AD, rmsd_a3_V287T_AD, rmsd_a3_L192F_BBR, rmsd_a3_E276F_BBR, rmsd_a3_F280Y_BBR, rmsd_a3_L195F_BBR, rmsd_a3_F196A_BBR, rmsd_a3_V287T_BBR = rmsd_sect('a3', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

rmsd_a3_top_1sug, rmsd_a3_top_apo, rmsd_a3_top_L192F, rmsd_a3_top_E276F, rmsd_a3_top_F280Y, rmsd_a3_top_L195F, rmsd_a3_top_F196A, rmsd_a3_top_V287T, rmsd_a3_top_L192F_AD, rmsd_a3_top_E276F_AD, rmsd_a3_top_F280Y_AD, rmsd_a3_top_L195F_AD, rmsd_a3_top_F196A_AD, rmsd_a3_top_V287T_AD, rmsd_a3_top_L192F_BBR, rmsd_a3_top_E276F_BBR, rmsd_a3_top_F280Y_BBR, rmsd_a3_top_L195F_BBR, rmsd_a3_top_F196A_BBR, rmsd_a3_top_V287T_BBR = rmsd_sect('a3_top', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

rmsd_a4_1sug, rmsd_a4_apo, rmsd_a4_L192F, rmsd_a4_E276F, rmsd_a4_F280Y, rmsd_a4_L195F, rmsd_a4_F196A, rmsd_a4_V287T, rmsd_a4_L192F_AD, rmsd_a4_E276F_AD, rmsd_a4_F280Y_AD, rmsd_a4_L195F_AD, rmsd_a4_F196A_AD, rmsd_a4_V287T_AD, rmsd_a4_L192F_BBR, rmsd_a4_E276F_BBR, rmsd_a4_F280Y_BBR, rmsd_a4_L195F_BBR, rmsd_a4_F196A_BBR, rmsd_a4_V287T_BBR = rmsd_sect('a4', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

rmsd_a6_1sug, rmsd_a6_apo, rmsd_a6_L192F, rmsd_a6_E276F, rmsd_a6_F280Y, rmsd_a6_L195F, rmsd_a6_F196A, rmsd_a6_V287T, rmsd_a6_L192F_AD, rmsd_a6_E276F_AD, rmsd_a6_F280Y_AD, rmsd_a6_L195F_AD, rmsd_a6_F196A_AD, rmsd_a6_V287T_AD, rmsd_a6_L192F_BBR, rmsd_a6_E276F_BBR, rmsd_a6_F280Y_BBR, rmsd_a6_L195F_BBR, rmsd_a6_F196A_BBR, rmsd_a6_V287T_BBR = rmsd_sect('a6', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

rmsd_a6_bot_1sug, rmsd_a6_bot_apo, rmsd_a6_bot_L192F, rmsd_a6_bot_E276F, rmsd_a6_bot_F280Y, rmsd_a6_bot_L195F, rmsd_a6_bot_F196A, rmsd_a6_bot_V287T, rmsd_a6_bot_L192F_AD, rmsd_a6_bot_E276F_AD, rmsd_a6_bot_F280Y_AD, rmsd_a6_bot_L195F_AD, rmsd_a6_bot_F196A_AD, rmsd_a6_bot_V287T_AD, rmsd_a6_bot_L192F_BBR, rmsd_a6_bot_E276F_BBR, rmsd_a6_bot_F280Y_BBR, rmsd_a6_bot_L195F_BBR, rmsd_a6_bot_F196A_BBR, rmsd_a6_bot_V287T_BBR = rmsd_sect('a6_bot', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

rmsd_CYS_1sug, rmsd_CYS_apo, rmsd_CYS_L192F, rmsd_CYS_E276F, rmsd_CYS_F280Y, rmsd_CYS_L195F, rmsd_CYS_F196A, rmsd_CYS_V287T, rmsd_CYS_L192F_AD, rmsd_CYS_E276F_AD, rmsd_CYS_F280Y_AD, rmsd_CYS_L195F_AD, rmsd_CYS_F196A_AD, rmsd_CYS_V287T_AD, rmsd_CYS_L192F_BBR, rmsd_CYS_E276F_BBR, rmsd_CYS_F280Y_BBR, rmsd_CYS_L195F_BBR, rmsd_CYS_F196A_BBR, rmsd_CYS_V287T_BBR = rmsd_sect('CYS', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

rmsd_beg_1sug, rmsd_beg_apo, rmsd_beg_L192F, rmsd_beg_E276F, rmsd_beg_F280Y, rmsd_beg_L195F, rmsd_beg_F196A, rmsd_beg_V287T, rmsd_beg_L192F_AD, rmsd_beg_E276F_AD, rmsd_beg_F280Y_AD, rmsd_beg_L195F_AD, rmsd_beg_F196A_AD, rmsd_beg_V287T_AD, rmsd_beg_L192F_BBR, rmsd_beg_E276F_BBR, rmsd_beg_F280Y_BBR, rmsd_beg_L195F_BBR, rmsd_beg_F196A_BBR, rmsd_beg_V287T_BBR = rmsd_sect('beg', file_path_close, file_path_close_AD, file_path_close_BBR, ref, 1)

for i in range(len(file_path_close)):
    for j in range(len(sections)):
        #Load Data for reference open
        rmsd_Apo = load_rms(file_path_close[i], sections[j], ref[1])

        #Mean and SEM for each trajectory
        RMSD_mean_close[i][j] = np.mean(rmsd_Apo)
        RMSD_err_close[i][j] = stats.sem(rmsd_Apo)

for i in range(len(file_path)):
    #Load Data for reference open
    rmsd = load_rms(file_path[i], sections[j], ref[0])
    #Mean and SEM for each trajectory
    RMSD_mean[i][j] = np.mean(rmsd)
    RMSD_err[i][j] = stats.sem(rmsd)

for i in range(len(file_path_close_AD)):
    #Load Data for reference open
    rmsd_AD = load_rms(file_path_close_AD[i], sections[j], ref[1])
    RMSD_mean_close_AD[i][j] = np.mean(rmsd_AD)
    RMSD_err_close_AD[i][j] = stats.sem(rmsd_AD)

    rmsd_BBR = load_rms(file_path_close_BBR[i], sections[j], ref[1])
    RMSD_mean_close_BBR[i][j] = np.mean(rmsd_BBR)
    RMSD_err_close_BBR[i][j] = stats.sem(rmsd_BBR)

#Name Labels
Label = ['WT', 'L192F', 'E276F', 'F280Y', 'L195F', 'F196A', 'V287T']
Label_close = ['WT Close', 'WT Open', 'L192F', 'E276F', 'F280Y', 'L195F', 'F196A', 'V287T']
Labels_mut = ['L192F', 'E276F', 'F280Y', 'L195F', 'F196A', 'V287T']

#Plot all compared to WT Open
for i in range(len(sections)):
    plot_compare(RMSD_mean, RMSD_err, Label, sections, i, ref[0])
    plot_compare(RMSD_mean_close, RMSD_err_close, Label_close, sections, i, ref[1])

#Determine % difference from WT
RMSD_diff = np.zeros((len(sections), len(Labels_mut)))
for i in range(1, len(Label)):
    n = i-1
    for j in range(len(sections)):
        WT = RMSD_mean[0][j]
        Mut = RMSD_mean[i][j]
        RMSD_diff[j][n] = ((Mut-WT)/((Mut+WT)/2)) * 100

#Plot table comparing residue interactions to WT
ax = plt.figure(figsize=(12, 6), frameon=False) # no visible frame
ax = sns.heatmap(RMSD_diff, annot=False, cmap = 'jet', cbar = True, cbar_kws={'label': 'Percentage Difference from WT'}, vmin = 0, vmax = 150, xticklabels = Labels_mut, yticklabels = sections)
#ax.add_artist(lines.Line2D([0, 20], [7, 7], color = 'black', linestyle= '--', linewidth = 4))
plt.title('Section RMSD Compared to WT')
plt.savefig('mutate_RMSD_Apo.png')
plt.close()

RMSD_mean_mut = np.zeros((len(Label_close), len(sections))) #Mean for reference open
RMSD_err_mut = np.zeros((len(Label_close), len(sections))) #SEM for reference open

#Plot self and two references
for i in [0, 2, 3, 4]:
    for j in range(len(sections)):
        #Load Data
        RMSD_mean_mut[0][j] = RMSD_mean_close[0][j]
        RMSD_err_mut[0][j] = RMSD_err_close[0][j]
        for k in range(1, len(Label_close)):
            rmsd = load_rms(file_path_close[k], sections[j], ref[i])
            RMSD_mean_mut[k][j] = np.mean(rmsd)
            RMSD_err_mut[k][j] = stats.sem(rmsd)
        plot_compare(RMSD_mean_mut, RMSD_err_mut, Label_close, sections, j, ref[i])

#Plot Kernel DEnsity Estimate Plot
#Compare a3_top for L192F, E276F, L195F, V287T
a3_top_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_a3_top_apo})
a3_top_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_a3_top_1sug})
a3_top_L192F_df = pd.DataFrame({'L192F': rmsd_a3_top_L192F})
a3_top_L195F_df = pd.DataFrame({'L195F': rmsd_a3_top_L195F})
a3_top_F280Y_df = pd.DataFrame({'F280Y': rmsd_a3_top_F280Y})
a3_top_E276F_df = pd.DataFrame({'E276F': rmsd_a3_top_E276F})
a3_top_F196A_df = pd.DataFrame({'F196A': rmsd_a3_top_F196A})
a3_top_V287T_df = pd.DataFrame({'V287T': rmsd_a3_top_V287T})
a3_top_L192F_AD_df = pd.DataFrame({'L192F AD': rmsd_a3_top_L192F_AD})
a3_top_L195F_AD_df = pd.DataFrame({'L195F AD': rmsd_a3_top_L195F_AD})
a3_top_F280Y_AD_df = pd.DataFrame({'F280Y AD': rmsd_a3_top_F280Y_AD})
a3_top_E276F_AD_df = pd.DataFrame({'E276F AD': rmsd_a3_top_E276F_AD})
a3_top_F196A_AD_df = pd.DataFrame({'F196A AD': rmsd_a3_top_F196A_AD})
a3_top_V287T_AD_df = pd.DataFrame({'V287T AD': rmsd_a3_top_V287T_AD})
a3_top_L192F_BBR_df = pd.DataFrame({'L192F BBR': rmsd_a3_top_L192F_BBR})
a3_top_L195F_BBR_df = pd.DataFrame({'L195F BBR': rmsd_a3_top_L195F_BBR})
a3_top_F280Y_BBR_df = pd.DataFrame({'F280Y BBR': rmsd_a3_top_F280Y_BBR})
a3_top_E276F_BBR_df = pd.DataFrame({'E276F BBR': rmsd_a3_top_E276F_BBR})
a3_top_F196A_BBR_df = pd.DataFrame({'F196A BBR': rmsd_a3_top_F196A_BBR})
a3_top_V287T_BBR_df = pd.DataFrame({'V287T BBR': rmsd_a3_top_V287T_BBR})

df = pd.concat([a3_top_Apo_open_df, a3_top_Apo_close_df, a3_top_L192F_df, a3_top_E276F_df, a3_top_V287T_df, a3_top_F196A_df, a3_top_F280Y_df, a3_top_L195F_df])
plot_kernel_mut(df, r'Top of $\alpha$3', 'a3_top_all', 0, 2)

df = pd.concat([a3_top_L192F_df, a3_top_E276F_df, a3_top_V287T_df, a3_top_F196A_df, a3_top_F280Y_df, a3_top_L195F_df])
plot_kernel_mut(df, r'Top of $\alpha$3', 'a3_top_mut_all', 0, 2)

df = pd.concat([a3_top_Apo_open_df, a3_top_Apo_close_df, a3_top_V287T_df, a3_top_F280Y_df])
plot_kernel_mut(df, r'Top of $\alpha$3', 'a3_top_extr', 0, 2)

plot_kernel_cmpr_lig(a3_top_L192F_df, a3_top_L192F_AD_df, a3_top_L192F_BBR_df, 'L192F', sections[7], 7)
plot_kernel_cmpr_lig(a3_top_L195F_df, a3_top_L195F_AD_df, a3_top_L195F_BBR_df, 'L195F', sections[7], 7)
plot_kernel_cmpr_lig(a3_top_E276F_df, a3_top_E276F_AD_df, a3_top_E276F_BBR_df, 'E276F', sections[7], 7)
plot_kernel_cmpr_lig(a3_top_V287T_df, a3_top_V287T_AD_df, a3_top_V287T_BBR_df, 'V287T', sections[7], 7)

#Compare a3_top for L192F, E276F, L195F, V287T
a3_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_a3_apo})
a3_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_a3_1sug})
a3_L192F_df = pd.DataFrame({'L192F': rmsd_a3_L192F})
a3_L195F_df = pd.DataFrame({'L195F': rmsd_a3_L195F})
a3_F280Y_df = pd.DataFrame({'F280Y': rmsd_a3_F280Y})
a3_E276F_df = pd.DataFrame({'E276F': rmsd_a3_E276F})
a3_F196A_df = pd.DataFrame({'F196A': rmsd_a3_F196A})
a3_V287T_df = pd.DataFrame({'V287T': rmsd_a3_V287T})
a3_L192F_AD_df = pd.DataFrame({'L192F AD': rmsd_a3_L192F_AD})
a3_L195F_AD_df = pd.DataFrame({'L195F AD': rmsd_a3_L195F_AD})
a3_F280Y_AD_df = pd.DataFrame({'F280Y AD': rmsd_a3_F280Y_AD})
a3_E276F_AD_df = pd.DataFrame({'E276F AD': rmsd_a3_E276F_AD})
a3_F196A_AD_df = pd.DataFrame({'F196A AD': rmsd_a3_F196A_AD})
a3_V287T_AD_df = pd.DataFrame({'V287T AD': rmsd_a3_V287T_AD})
a3_L192F_BBR_df = pd.DataFrame({'L192F BBR': rmsd_a3_L192F_BBR})
a3_L195F_BBR_df = pd.DataFrame({'L195F BBR': rmsd_a3_L195F_BBR})
a3_F280Y_BBR_df = pd.DataFrame({'F280Y BBR': rmsd_a3_F280Y_BBR})
a3_E276F_BBR_df = pd.DataFrame({'E276F BBR': rmsd_a3_E276F_BBR})
a3_F196A_BBR_df = pd.DataFrame({'F196A BBR': rmsd_a3_F196A_BBR})
a3_V287T_BBR_df = pd.DataFrame({'V287T BBR': rmsd_a3_V287T_BBR})

df = pd.concat([a3_Apo_open_df, a3_Apo_close_df, a3_L192F_df, a3_E276F_df, a3_V287T_df, a3_F196A_df, a3_F280Y_df, a3_L195F_df])
plot_kernel_mut(df, r'$\alpha$3', 'a3_all', 0, 2)

df = pd.concat([a3_L192F_df, a3_E276F_df, a3_V287T_df, a3_F196A_df, a3_F280Y_df, a3_L195F_df])
plot_kernel_mut(df, r'$\alpha$3', 'a3_mut_all', 0, 2)

df = pd.concat([a3_Apo_open_df, a3_Apo_close_df, a3_V287T_df, a3_F280Y_df])
plot_kernel_mut(df, r'$\alpha$3', 'a3_mut_extr', 0, 2)

plot_kernel_cmpr_lig(a3_L192F_df, a3_L192F_AD_df, a3_L192F_BBR_df, 'L192F', sections[6], 6)
plot_kernel_cmpr_lig(a3_L195F_df, a3_L195F_AD_df, a3_L195F_BBR_df, 'L195F', sections[6], 6)
plot_kernel_cmpr_lig(a3_E276F_df, a3_E276F_AD_df, a3_E276F_BBR_df, 'E276F', sections[6], 6)
plot_kernel_cmpr_lig(a3_V287T_df, a3_V287T_AD_df, a3_V287T_BBR_df, 'V287T', sections[6], 6)

#Compare a4 for L192F, E276F, L195F, V287T
a4_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_a4_apo})
a4_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_a4_1sug})
a4_L192F_df = pd.DataFrame({'L192F': rmsd_a4_L192F})
a4_L195F_df = pd.DataFrame({'L195F': rmsd_a4_L195F})
a4_F280Y_df = pd.DataFrame({'F280Y': rmsd_a4_F280Y})
a4_E276F_df = pd.DataFrame({'E276F': rmsd_a4_E276F})
a4_F196A_df = pd.DataFrame({'F196A': rmsd_a4_F196A})
a4_V287T_df = pd.DataFrame({'V287T': rmsd_a4_V287T})
a4_L192F_AD_df = pd.DataFrame({'L192F AD': rmsd_a4_L192F_AD})
a4_L195F_AD_df = pd.DataFrame({'L195F AD': rmsd_a4_L195F_AD})
a4_F280Y_AD_df = pd.DataFrame({'F280Y AD': rmsd_a4_F280Y_AD})
a4_E276F_AD_df = pd.DataFrame({'E276F AD': rmsd_a4_E276F_AD})
a4_F196A_AD_df = pd.DataFrame({'F196A AD': rmsd_a4_F196A_AD})
a4_V287T_AD_df = pd.DataFrame({'V287T AD': rmsd_a4_V287T_AD})
a4_L192F_BBR_df = pd.DataFrame({'L192F BBR': rmsd_a4_L192F_BBR})
a4_L195F_BBR_df = pd.DataFrame({'L195F BBR': rmsd_a4_L195F_BBR})
a4_F280Y_BBR_df = pd.DataFrame({'F280Y BBR': rmsd_a4_F280Y_BBR})
a4_E276F_BBR_df = pd.DataFrame({'E276F BBR': rmsd_a4_E276F_BBR})
a4_F196A_BBR_df = pd.DataFrame({'F196A BBR': rmsd_a4_F196A_BBR})
a4_V287T_BBR_df = pd.DataFrame({'V287T BBR': rmsd_a4_V287T_BBR})

df = pd.concat([a4_Apo_open_df, a4_Apo_close_df, a4_L192F_df, a4_E276F_df, a4_V287T_df, a4_F196A_df, a4_F280Y_df, a4_L195F_df])
plot_kernel_mut(df, r'$\alpha$4', 'a4_all', 0, 1.5)

df = pd.concat([a4_L192F_df, a4_E276F_df, a4_V287T_df, a4_F196A_df, a4_F280Y_df, a4_L195F_df])
plot_kernel_mut(df, r'$\alpha$4', 'a4_mut_all',0, 1.5)

df = pd.concat([a4_Apo_open_df, a4_Apo_close_df, a4_V287T_df, a4_F196A_df, a4_F280Y_df])
plot_kernel_mut(df, r'$\alpha$4', 'a4', 0, 1.5)

plot_kernel_cmpr_lig(a4_F196A_df, a4_F196A_AD_df, a4_F196A_BBR_df, 'F196A', sections[8], 8)
plot_kernel_cmpr_lig(a4_F280Y_df, a4_F280Y_AD_df, a4_F280Y_BBR_df, 'F280Y', sections[8], 8)

#a4_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_a4_apo_rapo})
#a4_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_a4_1sug_rapo})
#a4_F196A_df = pd.DataFrame({'F196A': rmsd_a4_F196A_rapo})
#a4_F196A_AD_df = pd.DataFrame({'F196A AD': rmsd_a4_F196A_AD_rapo})
#a4_F196A_BBR_df = pd.DataFrame({'F196A BBR': rmsd_a4_F196A_BBR_rapo})

#df = pd.concat([a4_Apo_open_df, a4_Apo_close_df, a4_F196A_df, a4_F196A_AD_df, a4_F196A_BBR_df])

#ax = plt.figure(figsize=(12, 6), frameon=False) # no visible frame
#sns.kdeplot(data = df, fill=True, alpha=0.5, common_norm = True, common_grid = True)
#plt.xlabel(r'RMSD($\AA$)')
#plt.ylabel(r'Normalized Density')
#plt.title(r'$\alpha$-4 RMSD Compared to Apo F196A')
#plt.savefig('mutate_RMSD_a4_ref_F196A.png')
#plt.close()

#a6 comparison
a6_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_a6_apo})
a6_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_a6_1sug})
a6_L192F_df = pd.DataFrame({'L192F': rmsd_a6_L192F})
a6_L195F_df = pd.DataFrame({'L195F': rmsd_a6_L195F})
a6_F280Y_df = pd.DataFrame({'F280Y': rmsd_a6_F280Y})
a6_E276F_df = pd.DataFrame({'E276F': rmsd_a6_E276F})
a6_F196A_df = pd.DataFrame({'F196A': rmsd_a6_F196A})
a6_V287T_df = pd.DataFrame({'V287T': rmsd_a6_V287T})
a6_L192F_AD_df = pd.DataFrame({'L192F AD': rmsd_a6_L192F_AD})
a6_L195F_AD_df = pd.DataFrame({'L195F AD': rmsd_a6_L195F_AD})
a6_F280Y_AD_df = pd.DataFrame({'F280Y AD': rmsd_a6_F280Y_AD})
a6_E276F_AD_df = pd.DataFrame({'E276F AD': rmsd_a6_E276F_AD})
a6_F196A_AD_df = pd.DataFrame({'F196A AD': rmsd_a6_F196A_AD})
a6_V287T_AD_df = pd.DataFrame({'V287T AD': rmsd_a6_V287T_AD})
a6_L192F_BBR_df = pd.DataFrame({'L192F BBR': rmsd_a6_L192F_BBR})
a6_L195F_BBR_df = pd.DataFrame({'L195F BBR': rmsd_a6_L195F_BBR})
a6_F280Y_BBR_df = pd.DataFrame({'F280Y BBR': rmsd_a6_F280Y_BBR})
a6_E276F_BBR_df = pd.DataFrame({'E276F BBR': rmsd_a6_E276F_BBR})
a6_F196A_BBR_df = pd.DataFrame({'F196A BBR': rmsd_a6_F196A_BBR})
a6_V287T_BBR_df = pd.DataFrame({'V287T BBR': rmsd_a6_V287T_BBR})

df = pd.concat([a6_Apo_open_df, a6_Apo_close_df, a6_L192F_df, a6_E276F_df, a6_V287T_df, a6_F196A_df, a6_F280Y_df, a6_L195F_df])
plot_kernel_mut(df, r'$\alpha$6', 'a6_all', 0, 2)

df = pd.concat([a6_L192F_df, a6_E276F_df, a6_V287T_df, a6_F196A_df, a6_F280Y_df, a6_L195F_df])
plot_kernel_mut(df, r'$\alpha$6', 'a6_mut_all', 0, 2)

plot_kernel_cmpr_lig(a6_L192F_df, a6_L192F_AD_df, a6_L192F_BBR_df, 'L192F', sections[11], 11)
plot_kernel_cmpr_lig(a6_L195F_df, a6_L195F_AD_df, a6_L195F_BBR_df, 'L195F', sections[11], 11)
plot_kernel_cmpr_lig(a6_E276F_df, a6_E276F_AD_df, a6_E276F_BBR_df, 'E276F', sections[11], 11)
plot_kernel_cmpr_lig(a6_V287T_df, a6_V287T_AD_df, a6_V287T_BBR_df, 'V287T', sections[11], 11)

#Just CYS215
cys_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_CYS_apo})
cys_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_CYS_1sug})
cys_L192F_df = pd.DataFrame({'L192F': rmsd_CYS_L192F})
cys_L195F_df = pd.DataFrame({'L195F': rmsd_CYS_L195F})
cys_F280Y_df = pd.DataFrame({'F280Y': rmsd_CYS_F280Y})
cys_E276F_df = pd.DataFrame({'E276F': rmsd_CYS_E276F})
cys_F196A_df = pd.DataFrame({'F196A': rmsd_CYS_F196A})
cys_V287T_df = pd.DataFrame({'V287T': rmsd_CYS_V287T})

df = pd.concat([cys_Apo_open_df, cys_Apo_close_df, cys_L192F_df, cys_E276F_df, cys_V287T_df, cys_F196A_df, cys_F280Y_df, cys_L195F_df])
plot_kernel_mut(df, 'CYS215', 'cys_all', 0, 1)

df = pd.concat([cys_L192F_df, cys_E276F_df, cys_V287T_df, cys_F196A_df, cys_F280Y_df, cys_L195F_df])
plot_kernel_mut(df, 'CYS215', 'cys_mut_all', 0, 1)

rmsd_cys = [rmsd_CYS_1sug, rmsd_CYS_apo, rmsd_CYS_F196A]
ax = plt.figure(figsize=(12, 6), frameon=False) # no visible frame
sns.kdeplot(data = rmsd_cys, fill=True, alpha=0.5)
plt.title('CYS215 RMSD Compared to WT Closed')
plt.savefig('mutate_RMSD_cys_F196A.png')
plt.close()

#BEG loop (L1)
beg_Apo_open_df = pd.DataFrame({'Apo Open':rmsd_beg_apo})
beg_Apo_close_df = pd.DataFrame({'Apo Closed': rmsd_beg_1sug})
beg_L192F_df = pd.DataFrame({'L192F': rmsd_beg_L192F})
beg_L195F_df = pd.DataFrame({'L195F': rmsd_beg_L195F})
beg_F280Y_df = pd.DataFrame({'F280Y': rmsd_beg_F280Y})
beg_E276F_df = pd.DataFrame({'E276F': rmsd_beg_E276F})
beg_F196A_df = pd.DataFrame({'F196A': rmsd_beg_F196A})
beg_V287T_df = pd.DataFrame({'V287T': rmsd_beg_V287T})

df = pd.concat([beg_Apo_open_df, beg_Apo_close_df, beg_L192F_df, beg_E276F_df, beg_V287T_df, beg_F196A_df, beg_F280Y_df, beg_L195F_df])
plot_kernel_mut(df, 'L1', 'beg_all', 0, 4)

df = pd.concat([beg_L192F_df, beg_E276F_df, beg_V287T_df, beg_F196A_df, beg_F280Y_df, beg_L195F_df])
plot_kernel_mut(df, 'L1', 'beg_mut_all', 0, 4)

#Determine p-values for each of the sections of focus
file_p = open('p_values_mut.txt', 'w')
p = np.zeros((5, 7))
st, p[0,0] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_L192F, equal_var = False) #Welch's t-test
st, p[0,1] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_E276F, equal_var = False) #Welch's t-test
st, p[0,2] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_F280Y, equal_var = False) #Welch's t-test
st, p[0,3] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_L195F, equal_var = False) #Welch's t-test
st, p[0,4] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_F196A, equal_var = False) #Welch's t-test
st, p[0,5] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_V287T, equal_var = False) #Welch's t-test
st, p[0,6] = stats.ttest_ind(rmsd_a3_top_apo, rmsd_a3_top_1sug, equal_var = False) #Welch's t-test

st, p[1,0] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_L192F, equal_var = False) #Welch's t-test
st, p[1,1] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_E276F, equal_var = False) #Welch's t-test
st, p[1,2] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_F280Y, equal_var = False) #Welch's t-test
st, p[1,3] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_L195F, equal_var = False) #Welch's t-test
st, p[1,4] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_F196A, equal_var = False) #Welch's t-test
st, p[1,5] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_V287T, equal_var = False) #Welch's t-test
st, p[1,6] = stats.ttest_ind(rmsd_a3_apo, rmsd_a3_1sug, equal_var = False) #Welch's t-test

st, p[1,0] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_L192F, equal_var = False) #Welch's t-test
st, p[1,1] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_E276F, equal_var = False) #Welch's t-test
st, p[1,2] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_F280Y, equal_var = False) #Welch's t-test
st, p[1,3] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_L195F, equal_var = False) #Welch's t-test
st, p[1,4] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_F196A, equal_var = False) #Welch's t-test
st, p[1,5] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_V287T, equal_var = False) #Welch's t-test
st, p[1,6] = stats.ttest_ind(rmsd_a4_apo, rmsd_a4_1sug, equal_var = False) #Welch's t-test

st, p[2,0] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_L192F, equal_var = False) #Welch's t-test
st, p[2,1] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_E276F, equal_var = False) #Welch's t-test
st, p[2,2] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_F280Y, equal_var = False) #Welch's t-test
st, p[2,3] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_L195F, equal_var = False) #Welch's t-test
st, p[2,4] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_F196A, equal_var = False) #Welch's t-test
st, p[2,5] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_V287T, equal_var = False) #Welch's t-test
st, p[2,6] = stats.ttest_ind(rmsd_a6_bot_apo, rmsd_a6_bot_1sug, equal_var = False) #Welch's t-test

st, p[3,0] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_L192F, equal_var = False) #Welch's t-test
st, p[3,1] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_E276F, equal_var = False) #Welch's t-test
st, p[3,2] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_F280Y, equal_var = False) #Welch's t-test
st, p[3,3] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_L195F, equal_var = False) #Welch's t-test
st, p[3,4] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_F196A, equal_var = False) #Welch's t-test
st, p[3,5] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_V287T, equal_var = False) #Welch's t-test
st, p[3,6] = stats.ttest_ind(rmsd_CYS_apo, rmsd_CYS_1sug, equal_var = False) #Welch's t-test

sections_mini = ['a3_top', 'a3', 'a4', 'a6_bot']
Labels_mut = ['L192F', 'E276F', 'F280Y', 'L195F', 'F196A', 'V287T', 'Apo Closed']
file_p.write('P values of RMSD with Apo closed reference structure Relative to Apo Open \n')
for i in range(len(sections_mini)):
    file_p.write(str(sections_mini[i]) + '\n')
    for j in range(len(Labels_mut)):
        file_p.write(Labels_mut[j] + ': ' + str(p[i,j]) + '\n')

p = np.zeros((5, 12))
st, p[0,0] = stats.ttest_ind(rmsd_a3_top_L192F, rmsd_a3_top_L192F_AD, equal_var = False) #Welch's t-test
st, p[0,1] = stats.ttest_ind(rmsd_a3_top_E276F, rmsd_a3_top_E276F_AD, equal_var = False) #Welch's t-test
st, p[0,2] = stats.ttest_ind(rmsd_a3_top_F280Y, rmsd_a3_top_F280Y_AD, equal_var = False) #Welch's t-test
st, p[0,3] = stats.ttest_ind(rmsd_a3_top_L195F, rmsd_a3_top_L195F_AD, equal_var = False) #Welch's t-test
st, p[0,4] = stats.ttest_ind(rmsd_a3_top_F196A, rmsd_a3_top_F196A_AD, equal_var = False) #Welch's t-test
st, p[0,5] = stats.ttest_ind(rmsd_a3_top_V287T, rmsd_a3_top_V287T_AD, equal_var = False) #Welch's t-test
st, p[0,6] = stats.ttest_ind(rmsd_a3_top_L192F, rmsd_a3_top_L192F_BBR, equal_var = False) #Welch's t-test
st, p[0,7] = stats.ttest_ind(rmsd_a3_top_E276F, rmsd_a3_top_E276F_BBR, equal_var = False) #Welch's t-test
st, p[0,8] = stats.ttest_ind(rmsd_a3_top_F280Y, rmsd_a3_top_F280Y_BBR, equal_var = False) #Welch's t-test
st, p[0,9] = stats.ttest_ind(rmsd_a3_top_L195F, rmsd_a3_top_L195F_BBR, equal_var = False) #Welch's t-test
st, p[0,10] = stats.ttest_ind(rmsd_a3_top_F196A, rmsd_a3_top_F196A_BBR, equal_var = False) #Welch's t-test
st, p[0,11] = stats.ttest_ind(rmsd_a3_top_V287T, rmsd_a3_top_V287T_BBR, equal_var = False) #Welch's t-test

st, p[1,0] = stats.ttest_ind(rmsd_a3_L192F, rmsd_a3_L192F_AD, equal_var = False) #Welch's t-test
st, p[1,1] = stats.ttest_ind(rmsd_a3_E276F, rmsd_a3_E276F_AD, equal_var = False) #Welch's t-test
st, p[1,2] = stats.ttest_ind(rmsd_a3_F280Y, rmsd_a3_F280Y_AD, equal_var = False) #Welch's t-test
st, p[1,3] = stats.ttest_ind(rmsd_a3_L195F, rmsd_a3_L195F_AD, equal_var = False) #Welch's t-test
st, p[1,4] = stats.ttest_ind(rmsd_a3_F196A, rmsd_a3_F196A_AD, equal_var = False) #Welch's t-test
st, p[1,5] = stats.ttest_ind(rmsd_a3_V287T, rmsd_a3_V287T_AD, equal_var = False) #Welch's t-test
st, p[1,6] = stats.ttest_ind(rmsd_a3_L192F, rmsd_a3_L192F_BBR, equal_var = False) #Welch's t-test
st, p[1,7] = stats.ttest_ind(rmsd_a3_E276F, rmsd_a3_E276F_BBR, equal_var = False) #Welch's t-test
st, p[1,8] = stats.ttest_ind(rmsd_a3_F280Y, rmsd_a3_F280Y_BBR, equal_var = False) #Welch's t-test
st, p[1,9] = stats.ttest_ind(rmsd_a3_L195F, rmsd_a3_L195F_BBR, equal_var = False) #Welch's t-test
st, p[1,10] = stats.ttest_ind(rmsd_a3_F196A, rmsd_a3_F196A_BBR, equal_var = False) #Welch's t-test
st, p[1,11] = stats.ttest_ind(rmsd_a3_V287T, rmsd_a3_V287T_BBR, equal_var = False) #Welch's t-test

st, p[2,0] = stats.ttest_ind(rmsd_a4_L192F, rmsd_a4_L192F_AD, equal_var = False) #Welch's t-test
st, p[2,1] = stats.ttest_ind(rmsd_a4_E276F, rmsd_a4_E276F_AD, equal_var = False) #Welch's t-test
st, p[2,2] = stats.ttest_ind(rmsd_a4_F280Y, rmsd_a4_F280Y_AD, equal_var = False) #Welch's t-test
st, p[2,3] = stats.ttest_ind(rmsd_a4_L195F, rmsd_a4_L195F_AD, equal_var = False) #Welch's t-test
st, p[2,4] = stats.ttest_ind(rmsd_a4_F196A, rmsd_a4_F196A_AD, equal_var = False) #Welch's t-test
st, p[2,5] = stats.ttest_ind(rmsd_a4_V287T, rmsd_a4_V287T_AD, equal_var = False) #Welch's t-test
st, p[2,6] = stats.ttest_ind(rmsd_a4_L192F, rmsd_a4_L192F_BBR, equal_var = False) #Welch's t-test
st, p[2,7] = stats.ttest_ind(rmsd_a4_E276F, rmsd_a4_E276F_BBR, equal_var = False) #Welch's t-test
st, p[2,8] = stats.ttest_ind(rmsd_a4_F280Y, rmsd_a4_F280Y_BBR, equal_var = False) #Welch's t-test
st, p[2,9] = stats.ttest_ind(rmsd_a4_L195F, rmsd_a4_L195F_BBR, equal_var = False) #Welch's t-test
st, p[2,10] = stats.ttest_ind(rmsd_a4_F196A, rmsd_a4_F196A_BBR, equal_var = False) #Welch's t-test
st, p[2,11] = stats.ttest_ind(rmsd_a4_V287T, rmsd_a4_V287T_BBR, equal_var = False) #Welch's t-test

sections_mini = ['a3_top', 'a3', 'a4']
Labels_mut = ['L192F', 'E276F', 'F280Y', 'L195F', 'F196A', 'V287T']
file_p.write('P values of RMSD with Apo closed reference structure Relative to Apo Mut \n')
for i in range(len(sections_mini)):
    file_p.write(str(sections_mini[i]) + '\n')
    for j in range(len(Labels_mut)):
        n = j+6
        file_p.write(Labels_mut[j] + ' AD: ' + str(p[i,j]) + '\n')
        file_p.write(Labels_mut[j] + ' BBR: ' + str(p[i,n]) + '\n')


import logging
import shutil
import sys

import click

from cekit.cache.artifact import ArtifactCache
from cekit.config import Config
from cekit.crypto import SUPPORTED_HASH_ALGORITHMS
from cekit.descriptor.resource import create_resource
from cekit.log import setup_logging
from cekit.tools import Map
from cekit.version import __version__

setup_logging()
LOGGER = logging.getLogger("cekit")
CONFIG = Config()

@click.group(context_settings=dict(max_content_width=100))
@click.option("-v", "--verbose", help="Enable verbose output.", is_flag=True)
@click.option(
    "--config",
    metavar="PATH",
    help="Path to configuration file.",
    default="~/.cekit/config",
    show_default=True,
)
@click.option(
    "--work-dir",
    metavar="PATH",
    help="Location of the working directory.",
    default="~/.cekit",
    show_default=True,
)
@click.version_option(message="%(version)s", version=__version__)
def cli(config, verbose, work_dir):
    pass

@cli.command(name="ls", short_help="List cached artifacts")
def ls():
    CacheCli.prepare().ls()

@cli.command(name="add", short_help="Add artifact to cache")
@click.argument("location", metavar="LOCATION")
@click.option("--md5", metavar="CHECKSUM", help="The md5 checksum of the artifact.")
@click.option("--sha1", metavar="CHECKSUM", help="The sha1 checksum of the artifact.")
@click.option(
    "--sha256", metavar="CHECKSUM", help="The sha256 checksum of the artifact."
)
@click.option(
    "--sha512", metavar="CHECKSUM", help="The sha512 checksum of the artifact."
)
def add(location, md5, sha1, sha256, sha512):
    if not (md5 or sha1 or sha256 or sha512):
        raise click.UsageError("At least one checksum must be provided")

    CacheCli.prepare().add(location, md5, sha1, sha256, sha512)

@cli.command(name="rm", short_help="Remove artifact from cache")
@click.argument("uuid", metavar="UUID")
def rm(uuid):
    CacheCli.prepare().rm(uuid)

@cli.command(name="clear", short_help="Remove all artifacts from the cache")
def clear():
    CacheCli.prepare().clear()

class CacheCli:
    @staticmethod
    def prepare():
        """Returns an initialized object of CacheCli"""
        return CacheCli(Map(click.get_current_context().parent.params))

    def __init__(self, args):
        # TODO: logging is used only when adding the artifact, we need to find out if it would be possible to do it better
        if args.verbose:
            LOGGER.setLevel(logging.DEBUG)
        else:
            LOGGER.setLevel(logging.INFO)

        CONFIG.configure(args.config, {"work_dir": args.work_dir})

    def add(self, location, md5, sha1, sha256, sha512):
        artifact_cache = ArtifactCache()

        resource = {}
        resource["url"] = location

        if md5:
            resource["md5"] = md5

        if sha1:
            resource["sha1"] = sha1

        if sha256:
            resource["sha256"] = sha256

        if sha512:
            resource["sha512"] = sha512

        artifact = create_resource(resource)

        cached = artifact_cache.cached(artifact)

        if cached:
            click.echo("Artifact {} is already cached!".format(location))
            sys.exit(0)

        try:
            artifact_id = artifact_cache.add(artifact)
            click.echo(
                "Artifact {} cached with UUID '{}'".format(location, artifact_id)
            )
        except Exception as ex:
            click.secho(
                "Cannot cache artifact {}: {}".format(location, str(ex)), fg="red"
            )
            sys.exit(1)

    def ls(self):
        artifact_cache = ArtifactCache()
        artifacts = artifact_cache.list()
        if artifacts:
            for artifact_filename, artifact in artifacts.items():
                click.echo(
                    "\n{}:".format(
                        click.style(
                            artifact_filename.split(".")[0], fg="green", bold=True
                        )
                    )
                )
                for alg in SUPPORTED_HASH_ALGORITHMS:
                    if alg in artifact and artifact[alg]:
                        click.echo(
                            "  {}: {}".format(
                                click.style(alg, bold=True), artifact[alg]
                            )
                        )

                if artifact["names"]:
                    click.echo("  {}:".format(click.style("names", bold=True)))
                    for name in artifact["names"]:
                        click.echo("    - %s" % name)
        else:
            click.echo("No artifacts cached!")

    def rm(self, uuid: str):
        artifact_cache = ArtifactCache()

        try:
            artifact_cache.delete(uuid)
            click.echo("Artifact with UUID '{}' removed".format(uuid))
        except Exception:
            click.secho(
                "Artifact with UUID '{}' doesn't exists in the cache".format(uuid),
                fg="yellow",
            )
            sys.exit(1)

    def clear(self):
        """
        Removes the artifact cache directory with all artifacts.

        Use with caution!
        """
        artifact_cache = ArtifactCache()

        if not click.confirm(
            "Are you sure to remove all artifacts from cache?", show_default=True
        ):
            return

        try:
            shutil.rmtree(artifact_cache.cache_dir)
            click.echo("Artifact cache cleared!")
        except Exception:
            click.secho(
                "An error occurred while removing the artifact cache directory '{}'".format(
                    artifact_cache.cache_dir
                ),
                fg="red",
            )
            sys.exit(1)

if __name__ == "__main__":
    cli()


import logging
import os
import sys

def configLogger():
    root = logging.getLogger()
    root.setLevel(logging.DEBUG)

    file_handler = logging.FileHandler(_get_logfile_name())
    file_handler.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    root.addHandler(file_handler)

def _get_logfile_name():
    executable_name = os.path.basename(sys.argv[0]).split('.')[0]
    return '/tmp/logs/%s.log' % executable_name


import sys
import datetime
import csv
from os import path

""""
This section describes all functions called when the program is started.
"""
def startup():
    create_date_file()
    create_bought_file()
    create_sold_file()

def create_date_file():
    """Check if there is already a file present containing the date set as current date. If not: create one."""
    if path.exists('date.txt') == False:
        date = str(datetime.date.today())
        file = open('date.txt', 'w')
        file.write(date)
        file.close()

def create_bought_file():
    """Check if there is already a bought.csv file present. If not: create one"""
    if path.exists('bought.csv') == False:
        with open('bought.csv', 'w', newline='') as csvfile:
            bought_creation = csv.writer(csvfile)
            bought_creation.writerow(['id', 'product_name', 'buy_date', 'buy_price', 'expiration_date'])

def create_sold_file():
    """Check if there is already a sold.csv file present. If not: create one"""
    if path.exists('sold.csv') == False:
        with open('sold.csv', 'w', newline='') as csvfile:
            sold_creation = csv.writer(csvfile)
            sold_creation.writerow(['id', 'bought_id', 'sell_date', 'sell_price'])


import requests
import re
import warnings
import json

class NHGIS:
    '''API wrapper for the IPUMS NHGIS API.

    API Documentation: https://developer.ipums.org/docs/get-started/

    Arguments:
        api_key: Authorization key required for use of the IPUMS API. *Required*
                 API keys can be obtained here: https://account.ipums.org/api_keys
    '''
    def __init__(self, api_key):
        NHGIS.header = {"Authorization": api_key}
        self.documentation_link = 'https://developer.ipums.org/docs/get-started/'
        self.geographic_extent_options = {"010": "Nation",
                                    "020": "Region",
                                    "030": "Division2",
                                    "040": "State",
                                    "050": "State-County",
                                    "140": "State-County-Census Tract",
                                    "155": "State-Place-County",
                                    "160": "State-Place",
                                    "250": "American Indian Area/Alaska Native Area Hawaiian Home Land",
                                    "310": "Metropolitan Statistical Area/Micropolitan Statistical Area",
                                    "500": "State-Congressional District"}

        self.dataset = Dataset
        self.time_series = TimeSeries

    def dataset_metadata(self, dataset=None, data_table=None):
        '''
        Returns metadata for IPUMS datasets.

        Arguments:

        dataset: Default: None. If None, metadata for all available datasets is returned. If the name of a
                                a dataset is given, metadata for the specified dataset is returned.
        data_table: Default: None. Data table can only be specified when `dataset` != None.
                                    If specified, metadata for the specified table is returned.
        '''
        if not dataset and not data_table:
            url = "https://api.ipums.org/metadata/nhgis/datasets?version=v1"
        elif dataset and not data_table:
            url = f"https://api.ipums.org/metadata/nhgis/datasets/{dataset}?version=v1"
        elif dataset and data_table:
            url = f"https://api.ipums.org/metadata/nhgis/datasets/{dataset}/data_tables/{data_table}?version=v1"
        elif not dataset and data_table:
            raise ValueError('A dataset must be provided when data_table != None')

        response = requests.get(url, headers=self.header)
        return self._validate_response(response)

    def time_series_metadata(self, data_table=None):
        '''
        Returns metadata for available time-series tables.

        Arguments:

        data_table: Default: None. If None, metadata for all data tables is returned.
                                    If specified, metadata for the specified table is returned.
        '''

        if not data_table:
            url = "https://api.ipums.org/metadata/nhgis/time_series_tables?version=v1"
        else:
            url = f"https://api.ipums.org/metadata/nhgis/time_series_tables/{data_table}?version=v1"

        response = requests.get(url, self.header)
        return self._validate_response(response)

    def shapefile_metadata(self):
        '''
        Returns metadata for available shapefiles.

        Arguments:

        No arguments are available for this method.
        '''

        url = "https://api.ipums.org/metadata/nhgis/shapefiles?version=v1"
        response = requests.get(url, self.header)
        return self._validate_response(response)

    def _validate_response(self, response):
        if response.status_code == 200:
            return response.json()
        else:
            raise ValueError(f'''A {response.status_code} error code was returned.\n

            The following reason was given: {response.reason}

            API Keys can be obtained here: https://account.ipums.org/api_keys

            API documentation is here: {self.documentation_link}

            If need be, api error code documentation can be found here
            https://developer.mozilla.org/en-US/docs/Web/HTTP/Status''')

    def _validate_shapefiles(self, shapefiles):
        supported_files = [file['name'] for file in self.shapefile_metadata()]
        for file in shapefiles:
            if file not in supported_files:
                raise ValueError(f'Shapefile: {file} could not be found.')

    def _create_payload(self, datasets=None, time_series_tables=None, shapefiles=None,
                      data_format='csv_no_header', breakdown_and_data_type_layout='separate_files',
                      time_series_table_layout=None, geographic_extents=None, description=None):
        '''
        Receives Dataset class objects and returns a json payload formatted according to the specifications
        of the IPUMS API.
        '''

        payload = {}
        if datasets:
            if isinstance(datasets, list):
                payload['datasets'] = {}
                for dataset in datasets:
                    if isinstance(dataset, Dataset):
                        if dataset.extent_required and not geographic_extents:
                            warnings.warn(f'''
                                            Geographic extent is required for
                                            Dataset {dataset.dataset} with the provided
                                            geographic levels. It is recommended that this
                                            extent be provided. By default, all geo extents are requested.

                                            Geograpghic extent options can be accessed with the
                                            NHGIS.geographic_extent_options attribute.''')
                            geographic_extents = ['*']
                        payload['datasets'].update(dataset.to_payload())
                    elif isinstance(dataset, dict):
                        payload['datasets'].update(dataset)
                        warnings.warn('''Validation is not provided for
                                        dictionary inputs.

                                        Use of NHGIS.dataset is recommended.''')
                    else:
                        raise ValueError(f'datasets cannot be datatype: {type(dataset)}')

        if time_series_tables:
            payload['time_series_tables'] = {}
            if isinstance(time_series_tables, list):
                for table in time_series_tables:
                    if isinstance(table, TimeSeries):
                        payload['time_series_tables'].update(table.to_payload())
                    elif isinstance(table, dict):
                        payload['time_series_tables'].update(table)
                        warnings.warn('''Validation is not provided for
                                        dictionary inputs.

                                        Use of NHGIS.time_series is recommended.''')
                    else:
                        raise ValueError(f'time_series_tables cannot be datatype: {type(table)}')
            elif isinstance(time_series_tables, TimeSeries):
                payload['time_series_tables'].update(time_series_tables.to_payload())
            else:
                raise ValueError('time_series_tables must be a list or a TimeSeries instance.')

        if shapefiles:
            payload['shapefiles'] = shapefiles

        if time_series_tables:
            payload['time_series_table_layout'] = time_series_table_layout

        if geographic_extents:
            payload['geographic_extents'] = geographic_extents

        payload['data_format'] = data_format

        if description:
            payload['description'] = description
        else:
            payload['description'] = 'ipumspy extract'

        if breakdown_and_data_type_layout:
            payload['breakdown_and_data_type_layout'] = breakdown_and_data_type_layout
        payload_json = json.dumps(payload)
        payload_json = json.loads(payload_json)

        return payload_json

    def create_extract(self, datasets=None, time_series_tables=None, shapefiles=None,
                      data_format='csv_no_header', breakdown_and_data_type_layout='separate_files',
                      time_series_table_layout=None, geographic_extents=None, description=None):
        '''
        Submit a data extract request to the IPUMS NHGIS API.
        Currently, the IPUMS API does not support downloading directly from the API.
        An email notification will be received confirming your extract request.
        '''
        url = "https://api.ipums.org/extracts/?product=nhgis&version=v1"
        if shapefiles:
            self._validate_shapefiles(shapefiles)

        payload = self._create_payload(datasets=datasets, time_series_tables=time_series_tables, shapefiles=shapefiles,
                      data_format=data_format, breakdown_and_data_type_layout=breakdown_and_data_type_layout,
                      time_series_table_layout=time_series_table_layout,
                      geographic_extents=geographic_extents, description=description)

        result = requests.post(url, headers=self.header, json=payload).json()
        if 'number' in result:
            self.extract_number = result['number']
        return result

    def extract_status(self, status_only = True):
        '''
        Returns that status of the most recent data extract request.
        '''

        url = f"https://api.ipums.org/extracts/{self.extract_number}?product=nhgis&version=v1"
        response = requests.get(url, headers=self.header)
        if status_only:
            return response.json()['status']
        else:
            return response.json()

class Dataset(NHGIS):
    '''A wrapper for creating validating requests to the IPUMS NHGIS API.
        This class is used to format the json data structure for the NHGIS class.'''

    def __init__(self, dataset, data_tables, geog_levels, years=None, breakdowns=[]):
        self._validate(dataset, data_tables, geog_levels, years=years, breakdowns=breakdowns)
        self.dataset = dataset
        self.data_tables = data_tables
        self.geog_levels = geog_levels
        self.breakdowns = breakdowns

    def _validate(self, dataset, data_tables, geog_levels, years=None, breakdowns=[]):
        self.validate_types(dataset, data_tables, geog_levels, years, breakdowns)
        metadata = self.dataset_metadata(dataset=dataset)
        self.validate_data_tables(metadata, data_tables)
        self.validate_geog_levels(metadata, geog_levels)
        self.validate_years(metadata, years)
        self.validate_breakdowns(metadata, breakdowns)

    def validate_data_tables(self, metadata, data_tables):
        supported_tables = [value['name'] for value in metadata['data_tables']]
        for table in data_tables:
            if table not in supported_tables:
                raise ValueError(f'''Data Table: {table} is not supported for dataset: {metadata["name"]}''')

    def validate_geog_levels(self, metadata, geog_levels):
        supported_levels = [value['name'] for value in metadata['geog_levels']]
        for level in geog_levels:
            if level not in supported_levels:
                raise ValueError(f'''Geo level: {level} is not supported for dataset: {metadata["name"]}''')

        self.extent_required = False
        for level in metadata['geog_levels']:
            if level['name'] in geog_levels:
                if level['has_geog_extent_selection']:
                    warnings.warn(f"""

                    Geographic level: '{level['name']}' for Dataset: '{metadata['name']}'
                    requires geog_extent selection when extraction is made.
                    Available geographic extent options can be accessed with the
                    `NHGIS.geographic_extent_options` attribute.
                    The `NHGIS.create_extract` method has a default geog_extent of ['*']

                    """)
                    self.extent_required = True

    def validate_breakdowns(self, metadata, breakdowns):
        if breakdowns:
            if not 'breakdowns' in metadata:
                raise ValueError(f'Breakdowns are not supported for {metadata["name"]}')
            supported_breakdowns_list = metadata['breakdowns']
            supported_breakdowns = []
            for entry in supported_breakdowns_list:
                supported_breakdowns += [value['name'] for value in entry['breakdown_values']]
            for breakdown in breakdowns:
                if breakdown not in supported_breakdowns:
                    raise ValueError(f'''Breakdown: {breakdown} is not supported for dataset: {metadata["name"]}''')

    def is_multiyear(self, metadata):
        year_count = re.findall('(\d{4})', metadata['name'])
        if year_count:
            count = len(set(year_count))
        else:
            count = 1
        if count > 1:
            return True

    def year_range(self, metadata):
        years = re.findall('(\d{4})', metadata['name'])
        if years:
            years = [int(year) for year in years]
            return [year for year in range(years[0], years[1] + 1)]

    def validate_years(self, metadata, years):
        multiyear = self.is_multiyear(metadata)
        if multiyear and not years:
            raise ValueError(f'{metadata["name"]} has multiple years. `years` cannot be set to `None`.')
        if multiyear and years:
            supported_years = self.year_range(metadata)
            for year in years:
                if year not in supported_years:
                    raise ValueError(f'''Year: {year} is not support for dataset: {metadata["name"]}.''')
        if not multiyear and years:
            if len(years) > 1:
                raise ValueError(f'Dataset: {metadata["name"]} is not a multi year dataset, but multiple years were given.')
            supported_year = int(re.findall('(\d{4})', metadata["name"])[0])
            if not supported_year == years[0]:
                raise ValueError(f'Dataset: {metadata["name"]} supports the year {supported_year}, but {years[0]} was given.')
        if not multiyear and not years:
            self.years = []

        self.years = years

    def validate_types(self, dataset, data_tables, geog_levels, years, breakdowns):
        if type(dataset) != str:
            raise ValueError('dataset variable must be string.')
        if not type(data_tables) == list:
            raise ValueError('data_tables variable must be a list.')
        if not all(isinstance(item, str) for item in data_tables):
            raise ValueError('data_tables variable must be a list of strings.')
        if not type(geog_levels) == list:
            raise ValueError('geog_levels variable must be a list.')
        if not all(isinstance(item, str) for item in geog_levels):
            raise ValueError('geog_levels variable must be a list of strings.')
        if years:
            if type(years) != list:
                raise ValueError('year variable must be a list for multi year datasets.')
            if not all(isinstance(item, int) for item in years):
                raise ValueError('year variable must be a list of integers for multi year datasets.')
        if breakdowns:
            if type(breakdowns) != list:
                raise ValueError('If breakdowns != None, the variable must be a list.')
            if not all(isinstance(item, str) for item in breakdowns):
                raise ValueError('If breakdowns != None, the variable must be a list of strings.')

    def __repr__(self):
        return f'''Dataset(dataset: {self.dataset},
        Number of tables: {len(self.data_tables)},
        Number of geographies: {len(self.geog_levels)},
        Number of breakdowns: {len(self.breakdowns) if self.breakdowns else self.breakdowns},
        years: {self.years})'''

    def to_payload(self):
        payload = {self.dataset: {
            "years": [str(year) for year in self.years] if self.years else [],
            "breakdown_values": self.breakdowns,
            "data_tables": self.data_tables,
            "geog_levels": self.geog_levels
        }}
        return payload

class TimeSeries(NHGIS):
    '''A wrapper for creating validating requests to the IPUMS NHGIS API.
    This class is used to format the json data structure for the NHGIS class.'''
    def __init__(self, data_table, geog_levels='macro'):
        self.validate(data_table, geog_levels)

    def validate(self, data_table, geog_levels):
        self.validate_types(data_table, geog_levels)
        metadata = self.time_series_metadata(data_table=data_table)
        self.data_table = data_table
        self.validate_geog_levels(metadata, geog_levels)

    def validate_types(self, data_table, geog_levels):
        if type(data_table) != str:
            raise ValueError('`data_table` variable must be a string.')
        if geog_levels != 'macro':
            if type(geog_levels) != list:
                raise ValueError('If `geog_levels` != "macro" `geog_levels` must be a list.')
            if not all(isinstance(item, str) for item in geog_levels):
                raise ValueError('If `geog_levels` != "macro" `geog_levels` must be a list of strings.')

    def validate_geog_levels(self, metadata, geog_levels):
        if geog_levels == 'macro':
            self.geog_levels = [metadata['geog_levels'][0]['name']]
        else:
            supported_levels = [value['name'] for value in metadata['geog_levels']]
            for level in geog_levels:
                if level not in supported_levels:
                    raise ValueError(f'Time Series {metadata["name"]} does not support geo level: "{level}"')
            self.geog_levels = geog_levels

    def to_payload(self):
        payload = {self.data_table: {
            'geog_levels': self.geog_levels}}

        return payload


import torch
from torch import nn, optim
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import random
import os
import glob
import sys
import wandb

import gru_models
import build_vocab

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# TODO: Get rid of these, just pass them in to dataset.
BATCH_SIZE=16
TRAIN_TOKEN_LEN=256

#from vocab_building impor  t load_tokenized_file, load_vocab, decode_vocab, nlp, get_vocab_indx_vector

def find_files(path): return glob.glob(path)

class RNN_Dataset_multiple_sources(torch.utils.data.Dataset):
    def __init__(
        self,
        sequence_length,
        type
    ):
        folder = "../vocabs_and_tokens/" + type + "/"
        data_folder = "../data/" + type + "/"
        vocab_file = folder + "*.pt"
        token_files = folder + "*.pkl"
        self.sequence_length = sequence_length

        self.all_categories, self.n_categories = self.setup_categories(data_folder)
        self.load_words(vocab_file, token_files)

        self.uniq_words = len(self.vocab)

    # data_folder needs to be like '../data/reviews/'
    def setup_categories(self, data_folder):
        all_categories = []
        for filename in find_files(data_folder + '*.txt'):
            category = os.path.splitext(os.path.basename(filename))[0]
            all_categories.append(category)

        n_categories = len(all_categories)

        if n_categories == 0:
            raise RuntimeError('Data not found.')

        print('# categories:', n_categories, all_categories)
        #all_categories.remove('garden')
        #all_categories.remove('music')
        #all_categories.remove('small_combined')
        #n_categories_languages = len(all_categories)
        #print('# categories:', n_categories_languages, all_categories)

        return all_categories, n_categories

    def load_words(self, vocab_file, token_files):
        # We want the vocab to be constructed from all sources, but we need the raw token sets for each seperately.
        # The category vector can just be a simple index vector.
        self.vocab = build_vocab.load_vocab(find_files(vocab_file)[0])

        token_files = find_files(token_files)
        # This is only setup to handle two different categories right now
        self.raw_tokens_1 = build_vocab.load_tokenized_file(token_files[0])
        self.raw_tokens_2 = build_vocab.load_tokenized_file(token_files[1])

        self.num_samples_1 = len(self.raw_tokens_1)
        self.num_samples_2 = len(self.raw_tokens_2)

        # This is iffy, because we aren't actually going through all of the "samples"
        self.num_samples = max(1, ((self.num_samples_1 + self.num_samples_2) // TRAIN_TOKEN_LEN)) # Split raw tokens into groups of TRAIN_TOKEN_LEN
        self.num_batches = max(1, self.num_samples // BATCH_SIZE)

        print('Number of raw_tokens: ', len(self.raw_tokens_1 + self.raw_tokens_2))
        print('Number of samples in a batch: ', self.num_samples)
        print('Number of batches: ', self.num_batches)

        return 1

    def random_choice(self, l):
        return l[random.randint(0, len(l)-1)]

    def category_tensor(self, category):
        li = self.all_categories.index(category)
        if li == 0:
            tensor = torch.zeros(self.sequence_length).to(device).long()
        else:
            tensor = torch.ones(self.sequence_length).to(device).long()
        return tensor, li

    def __len__(self):
        return self.num_samples

    def __getitem__(self, index):
        # This should pick a random source, grab it's category, and then grab a sequence associated with it.
        # Pick random category
        string_category= self.random_choice(self.all_categories)
        category, category_index = self.category_tensor(string_category)

        # Pick the right token samples based on the category
        if category_index == 0:
            current_sample = self.raw_tokens_1
        else:
            current_sample = self.raw_tokens_2

        # We cut off the potential of it being too long
        random_index = random.randint(0, len(current_sample) - (self.sequence_length + 1))
        end_index = random_index + self.sequence_length
        return ( # might break if it gets the very end?
            torch.tensor(current_sample[random_index:end_index]).to(device), # x
            torch.tensor(current_sample[random_index+1:end_index+1]).to(device), # y
            category
        )

def train(dataset, model, max_epochs, batch_size, cat = False):

    train_losses = []

    model.train()

    dataloader = DataLoader(dataset, batch_size=batch_size, drop_last=True)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(max_epochs):

        total_loss = 0

        for batch, (x, y, category) in enumerate(dataloader):
            hidden_states = model.init_hidden(batch_size)

            #print('x size: ', x.size()) # 16, 256
            #print('category size: ', category.size()) # 16, 256

            optimizer.zero_grad()

            if cat:
                y_pred, hidden_states = model(x, hidden_states, batch_size, category)
            else:
                y_pred, hidden_states = model(x, hidden_states, batch_size)

            #print('y_pred size: ', y_pred.size()) # [16, 4822] for cells, [16, 256, 4822] normal
            #print('y_pred transposed size: ', y_pred.transpose(1, 2).size()) # [16, 4822, 256]

            loss = criterion(y_pred.transpose(1, 2), y)
            total_loss += loss.item()

            loss.backward()
            optimizer.step()

            print({ 'epoch': epoch, 'batch': batch, 'loss': loss.item() })
            wandb.log({"loss":loss.item()})

        batch_loss = total_loss/batch_size
        train_losses.append(batch_loss)
        wandb.log({"batch_loss":batch_loss})

    return train_losses

def predict_with_category(dataset, model, text, category, next_words=100):
    model.eval()

    prediction = build_vocab.get_vocab_indx_vector(dataset.vocab, build_vocab.load_spacy, text)
    tokens = torch.tensor(prediction).to(device)

    # Get category tensor
    li = dataset.all_categories.index(category)
    if li == 0:
        category = torch.zeros(len(prediction)).to(device).long()
    else:
        category = torch.ones(len(prediction)).to(device).long()

    print('cat size: ', category.size())
    print('prediction size: ', tokens.size())

    state_h = model.init_hidden(1) # num_layers, batch_size, lstm_size

    # Prime generation by feeding in initial input:
    for p in range(len(tokens)-1):
        _, state_h = model(tokens[p].view(1,-1), state_h, 1, category[p].view(1,-1))
        #print('state_h size: ', state_h.size())

    last_token = tokens[-1]
    for i in range(0, next_words):
        y_pred, state_h = model(last_token.view(1,-1), state_h, 1, category[0].view(1,-1))
        #print('y_pred size: ', y_pred.size()) # [16, 256, 12923], should be [1, 1, 12923]
        #print('y_pred[0][-1] size: ', y_pred[0][-1].size())

        last_word_logits = y_pred[0][-1]

        # These are the probabilities
        p = torch.nn.functional.softmax(last_word_logits, dim=0)
        word_index = torch.multinomial(p, 1)[0]
        top_values = torch.topk(p, 5)
        #top_words = top_values.indices
        #top_probs = top_values.values

        #print('word index: ', word_index)
        #print('top_words: ', top_words.tolist())
        #top_word_pred = decode_vocab(dataset.vocab, [word_index])
        #top_words_pred = decode_vocab(dataset.vocab, top_words.tolist())

        #print('The top word predicted was: ', top_word_pred)
        #print('The top five predictions were: ', top_words_pred)
        #print('Their probabilites are: ', top_probs)

        prediction.append(word_index)

        last_token = torch.tensor([word_index]).to(device)

    final_prediction = build_vocab.decode_vocab(dataset.vocab, prediction)
    return final_prediction

def train_wrapper(type, hidden_size, num_epochs):
    # Create dataset
    dataset = RNN_Dataset_multiple_sources(TRAIN_TOKEN_LEN, type)
    input_size = dataset.uniq_words # Should be size of vocab?
    n_layers = 3

    print('----------------------')
    print('Original GRU')

    run = wandb.init(name='Original GRU',
                     project='controllableRNN',
                     config={
                        'dataset':type,
                        'epochs':num_epochs,
                        'hidden_size':hidden_size
                     },
                    reinit=True
                     )

    # Model with normal pytorch GRU
    category_model = gru_models.GRU_category(input_size, hidden_size, input_size, n_layers).to(device)

    file_path = f"gru_trained_cat_" + type + ".pt"

    losses_cat = train(dataset, category_model, num_epochs, BATCH_SIZE, cat=True)

    torch.save(category_model.state_dict(), file_path)

    run.finish()

    """print('----------------------')
    print('Original GRU with cells')

    # Model with GRU Cells
    cells_category_model = gru_models.GRU_with_cells_category(input_size, hidden_size, input_size, n_layers).to(device)

    file_path = f"gru_trained_cat_cells_" + type + ".pt"

    losses_cat_cells = train(dataset, cells_category_model, num_epochs, BATCH_SIZE, True)

    torch.save(cells_category_model.state_dict(), file_path)"""

    print('----------------------')
    print('Edited GRU')

    run = wandb.init(name='Edited GRU',
                     project='controllableRNN',
                     config={
                        'dataset':type,
                        'epochs':num_epochs,
                        'hidden_size':hidden_size
                     },
                reinit=True
            )

    # Model with edited GRU Cells
    cells_category_edited_model = gru_models.GRU_with_cells_category_edited(input_size, hidden_size, input_size, n_layers).to(device)

    file_path = f"gru_trained_cat_cells_edited_" + type + ".pt"

    losses_cat_cells_edited = train(dataset, cells_category_edited_model, num_epochs, BATCH_SIZE, True)

    torch.save(cells_category_edited_model.state_dict(), file_path)
    run.finish()

    # Create loss graph and save
    """fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(range(len(losses_cat)), losses_cat, label="original")
    #ax.plot(range(len(losses_cat_cells)), losses_cat_cells, label="original with cells")
    ax.plot(range(len(losses_cat_cells_edited)), losses_cat_cells_edited, label="edited")
    plt.title("Loss over time")
    plt.xlabel("Time")
    plt.ylabel("Loss")
    plt.legend()
    plt.savefig('loss_' + str(type) + "_" + str(num_epochs) + "_" + str(hidden_size) + '.png')"""

def main():
    wandb.login() # login to wandb
    # Uncomment these to use arguments
    #arguments = sys.argv[1:]
    #type, num_epochs, hidden_size = arguments
    #num_epochs = int(num_epochs)
    #hidden_size = int(hidden_size)

    print('TRAINING LANGUAGES- HIDDEN_SIZE-256 NUM_EPOCHS-300')
    train_wrapper(type='books', hidden_size=256, num_epochs=50)
    print('TRAINING LANGUAGES- HIDDEN_SIZE-512 NUM_EPOCHS-300')
    train_wrapper(type='books', hidden_size=512, num_epochs=50)
    print('TRAINING LANGUAGES- HIDDEN_SIZE-1024 NUM_EPOCHS-300')
    train_wrapper(type='books', hidden_size=1024, num_epochs=50)

if __name__ == "__main__":
    main()


"""Holds environmental variables, sets up custom logger."""

import logging
import os

log = logging.getLogger(name="log")

# declare environment constants
COSMOSDB_CONNECTION_STRING: str = os.environ["COSMOSDB_CONNECTION_STRING"]
COSMOSDB_DATABASE_ID: str = os.environ["COSMOSDB_DATABASE_ID"]
COSMOSDB_CONTAINER_ID: str = os.environ["COSMOSDB_CONTAINER_ID"]
BLOB_SERVICE_CONNECTION_STRING: str = os.environ["BLOB_CONNECTION_STRING"]
BLOB_CONTAINER_NAME: str = os.environ["BLOB_CONTAINER_NAME"]

def logger(
    logging_format: str = "%(levelname)s, %(name)s.%(funcName)s: %(message)s",
    level: int = logging.INFO,
) -> None:
    """
    Sets up custom logger.

    Parameters:
        format (str, optional): Logging format. Defaults to "%(name)s%(funcName)s: %(message)s".
        level (int, optional): Logging level. Defaults to logging.INFO.

    Returns:
        None
    """
    log.debug(msg="Setting up custom logger.")

    log.setLevel(level=level)

    handler = logging.StreamHandler(stream=None)

    formatter = logging.Formatter(fmt=logging_format)
    handler.setFormatter(fmt=formatter)

    if log.hasHandlers():
        log.handlers.clear()

    log.addHandler(handler)


import json
import os
import pygame
from pprint import pformat

from pyggui.core import TRANSPARENT

pygame.font.init()

_DEFAULT_THEME_PATH = 'assets/themes/default_theme.json'
_DEFAULT_THEME_PATH = os.path.join(os.path.dirname(__file__), _DEFAULT_THEME_PATH)

NULL_THEME_DEFAULTS = {
    "col": TRANSPARENT,
    "width": 0,
    "padding": [0, 0],
    "radius": 0,
    "font": {
      "name": "calibri",
      "size": 32
    }
}

class Theme:
    def __init__(self, file=None):
        self.file = file

        if self.file is None:
            self.file = _DEFAULT_THEME_PATH

        self.changed = False
        self._all_styles = self._load_theme_json(self.file)
        self._top_level_theme = self._all_styles.get("*")
        self._styles = self._top_level_theme.copy()

    @staticmethod
    def _load_theme_json(file):
        with open(file, mode='r') as theme_file:
            themes = json.load(theme_file)

        return themes

    def get_widget_theme(self, widget=None):
        widget_theme = self._top_level_theme.copy()

        if widget is None:
            return widget_theme

        widget_type = widget.__class__.__name__.lower()

        # Update the widget theme with the widget type level attributes
        widget_type_theme = self._all_styles.get(widget_type, {})
        widget_theme |= widget_type_theme

        new_theme = Theme()
        new_theme._styles = widget_theme.copy()

        return new_theme

    def __getitem__(self, item):
        value = self._styles.get(item, None)
        if value is None and any(key in item for key in NULL_THEME_DEFAULTS.keys()):
            value = [v for k, v in NULL_THEME_DEFAULTS.items() if k == item or k in item][0]

        return value

    def __setitem__(self, item, value):
        old_styles = self._styles.copy()
        if item in self._styles:
            self._styles["bg-col"] = value

        self.changed = self.changed or old_styles != self._styles

    def __repr__(self):
        return pformat(self._styles)

    @property
    def font(self):
        font_dict = self['font']
        name = font_dict.get('name')
        size = font_dict.get('size')

        try:
            font_ = pygame.font.Font(name, size)
        except FileNotFoundError:
            font_ = pygame.font.SysFont(name, size)

        font_.set_bold(font_dict.get('bold', False))
        font_.set_italic(font_dict.get('italic', False))
        font_.set_underline(font_dict.get('underline', False))

        return font_

    def copy(self):
        copy = Theme()
        copy._styles = self._styles.copy()

        return copy

if __name__ == '__main__':
    theme = Theme()
    print(theme)


import boto3
import sys
import time
# input value 'ansible-controller' while running the instance
#import json
ec2_client = boto3.client('ec2', region_name = "us-east-1")
instances = ec2_client.describe_instances()
for reservation in instances['Reservations']:
    for instance in reservation["Instances"]:
        if instance["Tags"][0]["Value"] == sys.argv[1]:
            response = ec2_client.start_instances(InstanceIds=[instance["InstanceId"]])
            istance_id = instance["InstanceId"]

while True:
    print('checking the status........')
    time.sleep(3)
    status_value = ec2_client.describe_instance_status(InstanceIds=[istance_id])
    if len(status_value['InstanceStatuses']) != 0:
        if status_value['InstanceStatuses'][0]['InstanceState']['Name'] == 'running':
            print(f"{sys.argv[1]} EC2 Instance got:- {status_value['InstanceStatuses'][0]['InstanceState']['Name']}")
            break

print ("Press Enter to continue ..." )
input()


import logging
import json

from discord import Interaction, app_commands, Role
from discord.app_commands import Choice
from discord.ext.commands import Bot, Cog

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

with open("config.json") as cfg_json:
    cfg = json.loads(cfg_json.read())

owner_id = cfg["owner_id"]

async def log_reply(ctx: Interaction, response: str, ephemeral=True):
    log = f"From {ctx.user}: {response}"
    if ctx.guild:
        log = f"From {ctx.user} in {ctx.guild.name}: {response}"
    logger.info(log)
    await ctx.response.send_message(response, ephemeral=ephemeral)

async def is_owner(ctx: Interaction) -> bool:
    if ctx.user.id != owner_id:
        await log_reply(ctx, f"**Error:** Only my owner can use this command")
        return False
    return True

class CommandError(Exception):
    def __init__(self, msg: str):
        self.msg = msg
        super().__init__(msg)

    @classmethod
    async def send_err(cls, ctx: Interaction, msg: str):
        self = cls(msg)
        await log_reply(ctx, self.msg)
        return self

class OwnerCog(Cog):
    """Commands that can only be used by the bot's owner"""

    def __init__(self, bot: Bot):
        self.bot = bot

    @app_commands.command(name="sync")
    @app_commands.check(is_owner)
    async def sync_commands(self, ctx: Interaction):
        """(Owner only) Syncs app command info with Discord"""
        await self.bot.tree.sync()
        await log_reply(ctx, "Synced command tree with Discord")

    emote = app_commands.Group(
        name="emote",
        description=("(Owner only) Modify DGG emote translations"),
    )

    @emote.command(name="add")
    @app_commands.describe(
        dgg_version="The emote as it's used in DGG",
        disc_version="The emote as it's used in Discord",
    )
    @app_commands.check(is_owner)
    async def add_emote(self, ctx: Interaction, dgg_version: str, disc_version: str):
        """(Owner only) Add or modify a DGG emote translation"""
        self.bot.emotes[dgg_version] = disc_version
        self.bot.save_cfg()
        await log_reply(ctx, f"Translating {dgg_version} to {str(disc_version)}")

    @emote.command(name="remove")
    @app_commands.describe(dgg_version="The emote to remove (in DGG format)")
    @app_commands.check(is_owner)
    async def remove_emote(self, ctx: Interaction, dgg_version: str):
        """(Owner only) Remove a DGG emote translation"""
        if dgg_version in self.bot.emotes.keys():
            removed_emote = self.bot.emotes.pop(dgg_version)
            self.bot.save_cfg()
            await log_reply(ctx, f"Removed {removed_emote} from emotes")
        else:
            await log_reply(ctx, f"Couldn't find emote {dgg_version}")

    config = app_commands.Group(
        name="config",
        description=("(Owner only) Modify the bot's config file"),
    )

    @config.command(name="remove")
    @app_commands.choices(
        mode=[
            Choice(name="phrase", value="phrase"),
            Choice(name="relay", value="relay"),
        ]
    )
    @app_commands.check(is_owner)
    async def config_remove(self, ctx: Interaction, mode: str, value: str):
        """Remove a relay or phrase from the config file"""
        if mode == "phrase" and value in self.bot.phrases:
            del self.bot.phrases[value]
            self.bot.save_cfg()
            await log_reply(ctx, f"Removed '{value}' from phrases", ephemeral=False)
        elif mode == "relay" and value in self.bot.relays:
            del self.bot.relays[value]
            self.bot.save_cfg()
            await log_reply(ctx, f"Removed '{value}' from relays", ephemeral=False)
        else:
            await log_reply(ctx, f"Couldn't find '{value}' in {mode}s")

class PublicCog(Cog):
    """Commands that can be used by anybody"""

    def __init__(self, bot: Bot):
        self.bot = bot

    async def get_relay_channel(self, ctx: Interaction) -> int:
        if not ctx.guild:
            err = "**Error:** This command is only usable in servers"
            raise await CommandError(err).send_err(ctx, err)
        if "dgg-relay-mod" not in (role.name for role in ctx.user.roles):
            err = "**Error:** This command requires the 'dgg-relay-mod' role"
            raise await CommandError(err).send_err(ctx, err)
        relay_channel = None
        for channel in ctx.guild.channels:
            if channel.name == "dgg-relay":
                relay_channel = channel.id
                break
        if not relay_channel:
            err = f"**Error:** No '#dgg-relay' channel found in '{ctx.guild.name}'"
            raise await CommandError(err).send_err(ctx, err)
        return relay_channel

    relay = app_commands.Group(
        name="relay",
        description="Relays DGG messages to servers",
    )

    @relay.command(name="add")
    @app_commands.describe(dgg_username="The DGG user you want to relay messages from")
    async def relay_add(self, ctx: Interaction, dgg_username: str):
        """Add a DGG user whose messages get forwarded to this server (case sensitive!)"""
        relay_channel = await self.get_relay_channel(ctx)
        if dgg_username not in self.bot.relays:
            self.bot.relays[dgg_username] = []
            logger.info(f"Added new relay list '{dgg_username}'")
        if relay_channel not in self.bot.relays[dgg_username]:
            self.bot.relays[dgg_username].append(relay_channel)
            response = (
                f"Messages from '{dgg_username}' will be relayed to '{ctx.guild.name}'"
            )
        else:
            response = f"**Error:** '{dgg_username}' is already being relayed to '{ctx.guild.name}'"
        self.bot.save_cfg()
        await log_reply(ctx, response, ephemeral=False)

    @relay.command(name="remove")
    @app_commands.describe(dgg_username="The DGG user you want to stop relaying")
    async def relay_remove(self, ctx: Interaction, dgg_username: str):
        """Remove a DGG user's relay from this server"""
        relay_channel = await self.get_relay_channel(ctx)
        response = None
        if dgg_username in self.bot.relays.keys():
            if relay_channel in self.bot.relays[dgg_username]:
                self.bot.relays[dgg_username].remove(relay_channel)
                response = f"Removed '{dgg_username}' relay from '{ctx.guild.name}'"
                if not self.bot.relays[dgg_username]:
                    self.bot.relays.pop(dgg_username)
                    logger.info(f"Removed empty relay list for '{dgg_username}'")
                self.bot.save_cfg()
        if not response:
            response = (
                f"**Error:** '{dgg_username}' isn't being relayed to '{ctx.guild.name}'"
                " (try the '/relay list' command)"
            )

        await log_reply(ctx, response, ephemeral=False)

    @relay.command(name="list")
    async def relay_list(self, ctx: Interaction):
        """Lists DGG users currently being relayed to this server."""
        relay_channel = await self.get_relay_channel(ctx)
        relays = []
        for nickname in self.bot.relays:
            for channel in self.bot.relays[nickname]:
                if channel == relay_channel:
                    relays.append(nickname)
        relays = "', '".join(relays)
        response = f"This server gets messages from: '{relays}'"
        if not relays:
            response = "No relays are active for this server."
        await log_reply(ctx, response, ephemeral=False)

    live_notifications = app_commands.Group(
        name="live-notifications",
        description="Configure live notifications for Destiny",
    )

    @live_notifications.command(name="on")
    async def live_notifications_on(self, ctx: Interaction):
        """Enable live notifications for this server"""
        relay_channel = await self.get_relay_channel(ctx)
        if relay_channel not in self.bot.live["channels"].keys():
            self.bot.live["channels"][relay_channel] = {"role": None}
        self.bot.live["channels"][relay_channel]["enabled"] = True
        self.bot.save_cfg()
        response = f"Live notifications enabled for {ctx.guild.name}"
        await log_reply(ctx, response, ephemeral=False)

    @live_notifications.command(name="off")
    async def live_notifications_off(self, ctx: Interaction):
        """Disable live notifications for this server"""
        relay_channel = await self.get_relay_channel(ctx)
        if relay_channel not in self.bot.live["channels"].keys():
            self.bot.live["channels"][relay_channel] = {"role": None}
        self.bot.live["channels"][relay_channel]["enabled"] = False
        self.bot.save_cfg()
        response = f"Live notifications disabled for {ctx.guild.name}"
        await log_reply(ctx, response, ephemeral=False)

    @live_notifications.command(name="role")
    @app_commands.describe(role="The role that will be pinged")
    async def live_notifications_role(self, ctx: Interaction, role: Role):
        """Set a role that gets pinged for live notifications"""
        relay_channel = await self.get_relay_channel(ctx)
        if relay_channel not in self.bot.live["channels"].keys():
            self.bot.live["channels"][relay_channel] = {"enabled": True}
        self.bot.live["channels"][relay_channel]["role"] = role.id
        self.bot.save_cfg()
        response = (
            f'"<@&{role.id}>" will be pinged for live notifications in {ctx.guild.name}'
        )
        await log_reply(ctx, response, ephemeral=False)

    def check_prefs(self, disc_user):
        if disc_user not in self.bot.user_prefs.keys():
            self.bot.user_prefs[disc_user] = {"detect_presence": False, "ignores": []}
            logger.info(f"Added new user '{disc_user}' to preferences list")

    phrase = app_commands.Group(
        name="phrase",
        description="Relays DGG messages to users",
    )

    @phrase.command(name="add")
    @app_commands.describe(
        phrase="The phrase you want forwarded to you (most likely your DGG username)"
    )
    async def phrase_add(self, ctx: Interaction, phrase: str):
        """Add a phrase (usually a username) that will be forwarded
        to you when it's used in DGG (case insensitive)"""
        self.check_prefs(ctx.user.id)
        if phrase not in self.bot.phrases:
            self.bot.phrases[phrase] = []
            logger.info(f"Added new phrase list for '{phrase}'")
        if ctx.user.id not in self.bot.phrases[phrase]:
            self.bot.phrases[phrase].append(ctx.user.id)
            response = f"Forwarding '{phrase}' to {ctx.user}"
        else:
            response = f"**Error:** '{phrase}' is already being forwarded to {ctx.user}"
        self.bot.save_cfg()
        await log_reply(ctx, response)

    @phrase.command(name="remove")
    @app_commands.describe(phrase="The phrase you want to stop being forwarded")
    async def phrase_remove(self, ctx: Interaction, phrase: str):
        """Stop a phrase from being forwarded to you"""
        self.check_prefs(ctx.user.id)
        response = None
        if phrase in self.bot.phrases:
            if ctx.user.id in self.bot.phrases[phrase]:
                self.bot.phrases[phrase].remove(ctx.user.id)
                response = f"No longer forwarding '{phrase}' to {ctx.user}"
                if not self.bot.phrases[phrase]:
                    self.bot.phrases.pop(phrase)
                    logger.info(f"Removed empty phrase list '{phrase}'")
                self.bot.save_cfg()
        if not response:
            response = (
                f"**Error:** '{phrase}' isn't being forwarded to {ctx.user}"
                " (try the '/phrase list' command)"
            )
        await log_reply(ctx, response)

    @phrase.command(name="list")
    async def phrase_list(self, ctx: Interaction):
        """List the phrases currently being forwarded to you"""
        disc_user = ctx.user.id
        user_phrases = []
        for phrase in self.bot.phrases:
            for user_id in self.bot.phrases[phrase]:
                if user_id == disc_user:
                    user_phrases.append(phrase)
        user_phrases = "', '".join(user_phrases)
        response = f"Your phrases: '{user_phrases}'"
        if not user_phrases:
            response = "No phrases are being forwarded to you."
        await log_reply(ctx, response)

    @phrase.command(name="detect-dgg-presence")
    @app_commands.describe(mode="Set to True to detect DGG presence. Default is False.")
    async def detect_dgg_presence(self, ctx: Interaction, mode: bool):
        """Change behavior of the /phrase command by controlling when the bot messages you."""
        self.check_prefs(ctx.user.id)
        self.bot.user_prefs[ctx.user.id]["detect_presence"] = mode
        self.bot.save_cfg()
        word = "enabled" if mode else "disabled"
        response = f"Presence detection {word} for {ctx.user.name}"
        await log_reply(ctx, response)

    ignore = app_commands.Group(
        name="ignore",
        description="Configure your DGG Relay ignore list",
    )

    @ignore.command(name="add")
    @app_commands.describe(dgg_username="The user in DGG you want to ignore")
    async def add_ignore(self, ctx: Interaction, dgg_username: str):
        """Ignore messages from a DGG user"""
        self.check_prefs(ctx.user.id)
        ignores = self.bot.user_prefs[ctx.user.id]["ignores"]
        ignores.append(dgg_username)
        self.bot.user_prefs[ctx.user.id]["ignores"] = list(set(ignores))
        self.bot.save_cfg()
        response = f"'{dgg_username}' added to your ignore list"
        await log_reply(ctx, response)

    @ignore.command(name="remove")
    @app_commands.describe(dgg_username="The user in DGG you want to unignore")
    async def add_ignore(self, ctx: Interaction, dgg_username: str):
        """Remove someone from your ignore list"""
        self.check_prefs(ctx.user.id)
        ignores = self.bot.user_prefs[ctx.user.id]["ignores"]
        if dgg_username not in ignores:
            await log_reply(ctx, f"'{dgg_username}' is not in your ignore list")
            return
        self.bot.user_prefs[ctx.user.id]["ignores"].remove(dgg_username)
        self.bot.save_cfg()
        response = f"'{dgg_username}' removed from your ignore list"
        await log_reply(ctx, response)


import numpy as np
from sklearn.datasets import make_classification
import pytest
from pygbm.binning import BinMapper
from pygbm.grower import TreeGrower
from pygbm import GradientBoostingRegressor
from pygbm import GradientBoostingClassifier

X, y = make_classification(n_samples=150, n_classes=2, n_features=5,
                           n_informative=3, n_redundant=0,
                           random_state=0)

def test_plot_grower(tmpdir):
    pytest.importorskip('graphviz')
    from pygbm.plotting import plot_tree

    X_binned = BinMapper().fit_transform(X)
    gradients = np.asarray(y, dtype=np.float32).copy()
    hessians = np.ones(1, dtype=np.float32)
    grower = TreeGrower(X_binned, gradients, hessians, max_leaf_nodes=5)
    grower.grow()
    filename = tmpdir.join('plot_grower.pdf')
    plot_tree(grower, view=False, filename=filename)
    assert filename.exists()

def test_plot_estimator(tmpdir):
    pytest.importorskip('graphviz')
    from pygbm.plotting import plot_tree

    n_trees = 3
    est = GradientBoostingRegressor(max_iter=n_trees)
    est.fit(X, y)
    for i in range(n_trees):
        filename = tmpdir.join('plot_predictor.pdf')
        plot_tree(est, tree_index=i, view=False, filename=filename)
        assert filename.exists()

def test_plot_estimator_and_lightgbm(tmpdir):
    pytest.importorskip('graphviz')
    lightgbm = pytest.importorskip('lightgbm')
    from pygbm.plotting import plot_tree

    n_classes = 3
    X, y = make_classification(n_samples=150, n_classes=n_classes,
                               n_features=5, n_informative=3, n_redundant=0,
                               random_state=0)

    n_trees = 3
    est_pygbm = GradientBoostingClassifier(max_iter=n_trees,
                                           n_iter_no_change=None)
    est_pygbm.fit(X, y)
    est_lightgbm = lightgbm.LGBMClassifier(n_estimators=n_trees)
    est_lightgbm.fit(X, y)

    n_total_trees = n_trees * n_classes
    for i in range(n_total_trees):
        filename = tmpdir.join('plot_mixed_predictors.pdf')
        plot_tree(est_pygbm, est_lightgbm=est_lightgbm, tree_index=i,
                  view=False, filename=filename)
        assert filename.exists()


from unittest import mock

from django.test import TestCase
from django.urls import resolve, reverse
from nltk import word_tokenize

from .models import ScrapeHistory
from .views import WordCountView

class ScrapeTest(TestCase):

    def _mock_response(
            self,
            status=200,
            content="CONTENT",
            json_data=None,
            raise_for_status=None):

        mock_resp = mock.Mock()

        # mock raise_for_status call w/optional error
        mock_resp.raise_for_status = mock.Mock()

        if raise_for_status:
            mock_resp.raise_for_status.side_effect = raise_for_status
        # set status code and content
        mock_resp.status_code = status
        mock_resp.content = content

        # add json data if provided
        if json_data:
            mock_resp.json = mock.Mock(
                return_value=json_data
            )
        return mock_resp

    def setUp(self):
        self.url = reverse("wordcount")
        self.post_params = {
            "word": "fit",
            "url": "https://www.virtusize.jp/"
        }

    def test_wordcount_url_resolves(self):
        self.assertEqual(resolve(self.url).func.view_class, WordCountView)

    def test_missing_param(self):
        params = {
            "word": "ABC"
        }

        response = self.client.post(
            self.url, params, content_type='application/json')

        self.assertEqual(response.status_code, 400)

    def test_invalid_link_value(self):
        params = {
            "word": "WORd",
            "url": "https://virtuse"
        }

        response = self.client.post(
            self.url, params, content_type='application/json')

        self.assertEqual(response.status_code, 500)

    @mock.patch("scrape.views.requests.get")
    def test_wordcount_request_to_url(self, m):
        mock_resp = self._mock_response(
            content="this is fit")

        m.return_value = mock_resp

        response = WordCountView.scrap_url_and_word_count(self,
                                                          "fit", "https://www.virtusize.jp/")

        self.assertEqual(response, 1)

    @mock.patch("scrape.views.WordCountView.scrap_url_and_word_count", return_value=1)
    def test_wordcount_success_result(self, m):
        response = self.client.post(
            self.url, self.post_params, content_type='application/json')

        result = response.json()

        test_content = {
            "status": "ok",
            "count": 1
        }

        self.assertEqual(result, test_content)

    @mock.patch("scrape.views.WordCountView.scrap_url_and_word_count", return_value=1)
    def test_save_to_db(self, m):
        _ = ScrapeHistory.objects.create(
            url="https://www.virtusize.jp/",
            word="fit",
            word_count=1
        )

        _ = self.client.post(
            self.url, self.post_params, content_type='application/json')

        self.assertEqual(ScrapeHistory.objects.last().word_count, 1)

        self.assertNotEqual(ScrapeHistory.objects.last().word_count, 13)

    def test_word_tokenize(self):
        txt = "Virtusize works for innovative idea. <'new idea'> idea-ly Although there are lot of new ideas but it focuses e-commerce"

        words = word_tokenize(txt)

        self.assertEqual(words.count("idea"), 2)
        self.assertNotEqual(words.count("idea"), 1)


import numpy as np
import netCDF4
from datetime import datetime, timedelta
from glob import glob
import os, sys
"""
    This program is used to read input data.
"""
#******************************************
# Edit here (input file directories)
#------------------------------------------
slpbasedir = "/mnt/nas02/data/CMIP6"
tabasedir  = "/mnt/nas02/data/CMIP6"
uabasedir  = "/mnt/nas02/data/CMIP6"
vabasedir  = "/mnt/nas02/data/CMIP6"
tsbasedir  = "/mnt/nas02/data/CMIP6"
topobasedir= "/mnt/nas02/data/CMIP6"
landbasedir= "/mnt/nas02/data/CMIP6"
prbasedir  = "/mnt/nas02/data/CMIP6"   # not used for detection

#******************************************

dbasedir = {
    "slp":slpbasedir,
    "ta" :tabasedir,
    "ua" :uabasedir,
    "va" :vabasedir,
    "sst":tsbasedir,
    "topo":topobasedir,
    "land":landbasedir,
    "pr" :prbasedir,     # not used for detection
    }

dvar = {
    "slp":"psl",
    "ta" :"ta",
    "ua" :"ua",
    "va" :"va",
    "sst":"ts",
    "topo":"orog",
    "land":"sftlf",
    "pr" :"pr",    # not used for detection
    }

def ret_lats(model):
    return np.load(slpbasedir + "/%s/lat.npy"%(model))
    # MIROC6: -88.92773535 ~ 88.92773535, d=~1.4007664

def ret_lons(model):
    return np.load(slpbasedir + "/%s/lon.npy"%(model))
    # MIROC6: 0 ~ 358.59375, d=1.40625

def ret_ny(model):
    return len(ret_lats(model))

def ret_nx(model):
    return len(ret_lons(model))
    # MIROC6:   (128, 256)
def ret_miss(model):
    modelname = model.split(".")[0]
    if   modelname=="MIROC6":       miss_in= 9.969209968386869e+36
    elif modelname=="MRI-ESM2-0":   miss_in= 9.969209968386869e+36
    elif modelname=="MPI-ESM1-2-HR":miss_in= 9.969209968386869e+36

    return miss_in

def Load_6hrPlev(model, var, DTime, plev):
    modelname, expr, ens = model.split(".")
    vname = dvar[var]
    iplev = [850, 500, 250].index(plev)
    # Search file
    srcdir  = dbasedir[var] + "/%s"%(model)
    ssearch = srcdir + "/%s_6hrPlev*.nc"%(vname)
    lsrcpath = glob(ssearch)
    for srcpath in lsrcpath:
        stime = os.path.basename(srcpath).split("_")[6].split(".")[0]
        stime0, stime1 = stime.split("-")

        dtime0 = datetime.strptime(stime0, "%Y%m%d%H%M")
        dtime1 = datetime.strptime(stime1, "%Y%m%d%H%M")
        if (dtime0<=DTime)&(DTime<=dtime1):
            break
    nc = netCDF4.Dataset(srcpath)

    #print(nc.variables)
    #print(srcpath)
    # Find time index
    basetime = {
        ("MIROC6","piControl"):         datetime(3200,1,1),

        ("MRI-ESM2-0","piControl"):     datetime(1850,1,1),
        ("MRI-ESM2-0","historical"):    datetime(1850,1,1),

        ("MPI-ESM1-2-HR","piControl"):  datetime(1850,1,1),
        }[modelname,expr]

    dtime0 = basetime + timedelta(days=float(nc.variables["time"][0]))
    idxtime = int((DTime - dtime0).total_seconds()/21600)    # 6-hour = 21600 sec
    #print(DTime, dtime0)
    #print(idxtime)
    return nc.variables[vname][idxtime, iplev]

def Load_6hrSfc(model, var, DTime):
    modelname, expr, ens = model.split(".")
    vname = dvar[var]

    # Search file
    srcdir  = dbasedir[var] + "/%s"%(model)
    ssearch = srcdir + "/%s_6hrPlev*.nc"%(vname)
    lsrcpath = np.sort(glob(ssearch))
    for srcpath in lsrcpath:
        stime = os.path.basename(srcpath).split("_")[6].split(".")[0]
        stime0, stime1 = stime.split("-")

        dtime0 = datetime.strptime(stime0, "%Y%m%d%H%M")
        dtime1 = datetime.strptime(stime1, "%Y%m%d%H%M")
        if (dtime0<=DTime)&(DTime<=dtime1):
            break
    nc = netCDF4.Dataset(srcpath)
    #print(nc.variables)
    #print(srcpath)

    # Find time index
    basetime = {
        ("MIROC6","piControl"):         datetime(3200,1,1),
        ("MRI-ESM2-0","piControl"):     datetime(1850,1,1),
        ("MRI-ESM2-0","historical"):    datetime(1850,1,1),
        }[modelname,expr]

    dtime0 = basetime + timedelta(days=float(nc.variables["time"][0]))
    idxtime = int((DTime - dtime0).total_seconds()/21600)    # 6-hour = 21600 sec

    return nc.variables[vname][idxtime]
    #return nc.variables[vname].shape

def Load_monSfc(model, var, Year, Mon):
    modelname, expr, ens = model.split(".")
    vname = dvar[var]
    DTime = datetime(Year,Mon,1)
    # Search file
    srcdir  = dbasedir[var] + "/%s"%(model)
    ssearch = srcdir + "/%s_Amon*.nc"%(vname)
    lsrcpath = np.sort(glob(ssearch))
    for srcpath in lsrcpath:
        stime = os.path.basename(srcpath).split("_")[6].split(".")[0]
        stime0, stime1 = stime.split("-")

        dtime0 = datetime.strptime(stime0, "%Y%m")
        dtime1 = datetime.strptime(stime1, "%Y%m")
        if (dtime0<=DTime)&(DTime<=dtime1):
            break
    nc = netCDF4.Dataset(srcpath)
    #print(nc.variables)
    #print(srcpath)

    #print(nc.variables["time"][:])
    #print(len(nc.variables["time"][:]))

    # Find time index
    Year0,Mon0 = dtime0.timetuple()[:2]
    Year1,Mon1 = dtime1.timetuple()[:2]
    idxtime = int(Year-Year0)*12 -Mon0 + Mon
    #print(idxtime)
    return nc.variables[vname][idxtime]

def Load_const(model, var):
    vname = dvar[var]
    srcdir  = dbasedir[var] + "/%s"%(model)
    ssearch = srcdir + "/%s_*.nc"%(vname)
    lsrcpath= glob(ssearch)
    srcpath = lsrcpath[0]
    nc = netCDF4.Dataset(srcpath)
    #print(nc.variables)
    return nc.variables[vname][:]