"""
plot_bias_timeseries_v2.py

"""

import pygrib
from dateutils import daterange, dateshift, dayofyear, splitdate
import os, sys
from datetime import datetime
import _pickle as cPickle
import numpy as np
import numpy.ma as ma
import matplotlib.pyplot as plt
from read_reanalysis_timeseries import read_reanalysis_timeseries
from read_forecast_timeseries_GEFSv12 import read_forecast_timeseries_GEFSv12

# --------------------------------------------------------------

def find_nearest(vec, value):
    idx = np.abs(vec-value).argmin()
    return idx

# --------------------------------------------------------------

def read_bias_corrections(infile):
    print ('   reading beta estimates from ', infile)
    inf = open(infile,'rb')
    beta_3d = cPickle.load(inf)
    lats = cPickle.load(inf)
    lons = cPickle.load(inf)
    inf.close()
    return beta_3d

# --------------------------------------------------------------

def initialize_date_lists(cyear, clead):

    start_date = cyear+'010100'
    end_date = cyear+'123100'
    date_list_anal = daterange(start_date, end_date, 24)
    ndates = len(date_list_anal)
    date_list_forecast = []
    for i in range(ndates):
        date_list_forecast.append(dateshift(date_list_anal[i], \
            int(clead)))
    return date_list_anal, date_list_forecast

# --------------------------------------------------------------

# ---- various initialization

clead = sys.argv[1]
clon = sys.argv[2]
clat = sys.argv[3]
rlon = float(clon)
rlat = float(clat)
iskip = int(clead)//24
cvariable = '2t'
cpath_era5 = '/Volumes/Backup Plus/ecmwf/'
cpath_forecast = '/Volumes/Backup Plus/gefsv12/t2m/'
cpath_beta = '/Volumes/Backup Plus/gefsv12/t2m/'
cpath_random = '/Volumes/Backup Plus/gefsv12/t2m/'
cpath_gain = '/Volumes/Backup Plus/gefsv12/t2m/'

now = datetime.now()
current_time = now.strftime("%H:%M:%S")
print (current_time,'   &&&&&&&&&& LEAD TIME = ',clead,'  &&&&&&&&&&')

cyear = '2019'
date_list_anal, date_list_forecast = initialize_date_lists(cyear, clead)
ndates = len(date_list_forecast)
print (date_list_forecast)

# - get lat/lon indices

infile = cpath_era5 +'2001/t2m_era5_halfdegree_2001010100.cPick'
inf = open(infile, 'rb')
analysis = cPickle.load(inf) - 273.16
analysis = np.flipud(analysis)
lats = cPickle.load(inf)
lons = cPickle.load(inf)
nlats, nlons = np.shape(lats)
lats = np.flipud(lats)
lons_1d = lons[0,:]
lats_1d = lats[:,0]
imin = find_nearest(lons_1d, rlon)
jmin = find_nearest(lats_1d, rlat)
inf.close()

print ('lons_1d = ',lons_1d)
print ('lats_1d = ',lons_1d)
imin = find_nearest(lons_1d, rlon)
jmin = find_nearest(lats_1d, rlat)
print ('jmin, imin = ', jmin, imin)

# ---- read climatology

if clead == '12' or clead == '36' or clead == '60' or clead == '84' or clead == '108':
    infile = cpath_era5 + 'ERA5_temperature_climatology_12UTC.cPick'
else:
    infile = cpath_era5 + 'ERA5_temperature_climatology_00UTC.cPick'

print ('reading from ', infile)
inf = open(infile,'rb')
climo_temps_estimated = cPickle.load(inf)
inf.close()
analysis_climatology_1d = climo_temps_estimated[:,jmin, imin]
print (analysis_climatology_1d[180:211])

#print (date_list_forecast)
analyses_3d, lats, lons = read_reanalysis_timeseries(cpath_era5, \
    date_list_forecast)
analyses_1d = analyses_3d[:,jmin, imin]
ndates, nlats, nlons = np.shape(analyses_3d)
print (analyses_1d[180:211])

# ---- read the forecast time series on the dates specified.   Pass in
#      the lead time and the initial time of the analysis.

#print (date_list_anal)
forecast_3d, latsf, lonsf = read_forecast_timeseries_GEFSv12 ( \
    cpath_forecast, date_list_anal, clead)
forecast_1d = forecast_3d[:,jmin, imin]
print (forecast_1d[180:211])

infile = cpath_forecast + '2019_qmap_lead'+clead+'.cPick'
beta_qmap = read_bias_corrections(infile)
beta_qmap_1d = beta_qmap[:,jmin,imin]
r = np.corrcoef(beta_qmap_1d[0:-1],beta_qmap_1d[1:] )
cr_qmap = r"$r_1$ = {:.2f}".format(r[0,1])
print (cr_qmap)

infile = cpath_forecast + '2019_decayavg_lead'+clead+'.cPick'
beta_decayavg = read_bias_corrections(infile)
beta_decayavg_1d = beta_decayavg[:,jmin,imin]
r = np.corrcoef(beta_decayavg_1d[0:-1],beta_decayavg_1d[1:] )
cr_decay = r"$r_1$ = {:.2f}".format(r[0,1])

infile = cpath_forecast + '2019_MOS_lead'+clead+'.cPick'
beta_uMOS = read_bias_corrections(infile)
beta_uMOS_1d = beta_uMOS[:,jmin,imin]
r = np.corrcoef(beta_uMOS_1d[0:-1],beta_uMOS_1d[1:] )
cr_uMOS = r"$r_1$ = {:.2f}".format(r[0,1])

infile = cpath_forecast + '2019_MOS_mvr_lead'+clead+'.cPick'
beta_mvMOS = read_bias_corrections(infile)
beta_mvMOS_1d = beta_mvMOS[:,jmin,imin]
r = np.corrcoef(beta_mvMOS_1d[0:-1],beta_mvMOS_1d[1:] )
cr_mvMOS = r"$r_1$ = {:.2f}".format(r[0,1])

# ---- loop thru months

cmonths = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec']

f = plt.figure(figsize=(9.,6.5))

ax = f.add_axes([.08,.57,.9,.36])
ax.set_title('(a) 00 UTC forecasts and analyses',fontsize=15)
ax.plot(range(ndates),analyses_1d,'o-',color='Red',\
    lw=1, markersize=1.5, label='ERA5 analysis', markerfacecolor='Red')
ax.plot(range(ndates),analysis_climatology_1d,'-',color='Peru',\
    lw=2, markersize=1.5, label='ERA5 climatology')
ax.plot(range(ndates),forecast_1d,'o-',color='RoyalBlue',\
    lw=1, markersize=1.5, label='Forecast', markerfacecolor='RoyalBlue')
ax.set_ylim(-20,40)
ax.grid(True,lw=0.25,color='LightGray')
ax.set_xlim(0,365)
ax.set_xticks([0,30,58,89,119,150,180, 211,242,272,303,333])
ax.set_xticklabels(['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec'])
ax.set_xlabel('Julian day number',fontsize=11)
ax.set_ylabel('Temperature (deg C)',fontsize=11)
ax.legend(loc=0)

ax = f.add_axes([.08,.07,.9,.36])
ax.set_title('(b) bias estimates',fontsize=15)
ax.plot([0,365],[0,0],lw=1,color='Black')
ax.plot(range(ndates),beta_decayavg_1d,'o-',color='Peru',\
    lw=0.5, markersize=1.5, label='DAV, '+cr_decay, markerfacecolor='Peru')
ax.plot(range(ndates),beta_qmap_1d,'o-',color='Blue',\
    lw=0.5, markersize=1.5, label='QM, '+cr_qmap, markerfacecolor='Blue')
ax.plot(range(ndates),beta_uMOS_1d,'o-',color='Green',\
    lw=0.5, markersize=1.5, label='uMOS, '+cr_uMOS, markerfacecolor='Green')
ax.plot(range(ndates),beta_mvMOS_1d,'o-',color='Gray',\
    lw=0.5, markersize=1.5, label='mvMOS, '+cr_mvMOS, markerfacecolor='Gray')
ax.set_xticks([0,30,58, 89,119,150, 180,211,242, 272,303,333])
ax.set_xticklabels(['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec'])
ax.legend(loc=0)

ax.set_xlim(0,365)
ax.grid(True,lw=0.25,color='LightGray')
ax.set_ylim(-3,2)
ax.set_xlabel('Julian day number',fontsize=11)
ax.set_ylabel('Bias estimate (deg C)',fontsize=11)

figname = 'bias_tseries_lon='+clon+'_lat='+clat+'_lead='+clead+'.png'
plt.savefig(figname,dpi=400)
print ('Plot done', figname)
plt.close()


#!/usr/bin/python3

import os
import sys
from PIL import Image
import numpy as np
from sklearn.preprocessing import StandardScaler
import sklearn.datasets, sklearn.decomposition

THRESHOLD = 0.99

#	**********************************************		CONVERTING IMAGE DATA		**************************************************************#
image_name = "opera_house.jpg"

# load the image
image = Image.open(image_name)

# convert image to numpy array
data = np.asarray(image)

#Summarize Data
print (data.shape)
print (data)

#	**********************************************		Segmenting RGB Channel		***************************************************************#

print ("Shape: %d :: %d :: %d" % (data.shape[0], data.shape[1], data.shape[2]))
row = data.shape[0]
col = data.shape[1]
channel = data.shape[2]

r_channel = data[:row, :col, :1]
r_channel = np.reshape(r_channel, (row, col))

print (r_channel.shape)
print (r_channel)

g_channel = data[:row, :col, 1:2]
g_channel = np.reshape(g_channel, (row, col))

print (g_channel.shape)
print (g_channel)

b_channel = data[:row, :col, 2:]
b_channel = np.reshape(b_channel, (row, col))

print (b_channel.shape)
print (b_channel)

#	**********************************************		Standardizing Data Set		***************************************************************#

sc = StandardScaler()
r_channel_std = sc.fit_transform(r_channel)
g_channel_std = sc.fit_transform(g_channel)
b_channel_std = sc.fit_transform(b_channel)

#	**********************************************		Covariance matrix		***************************************************************#

cov_mat_r = np.cov(r_channel_std.T)
cov_mat_g = np.cov(g_channel_std.T)
cov_mat_b = np.cov(b_channel_std.T)

#	**********************************************		Eigen properties		***************************************************************#

eigen_vals_r, eigen_vecs_r = np.linalg.eig(cov_mat_r)
eigen_vals_g, eigen_vecs_g = np.linalg.eig(cov_mat_g)
eigen_vals_b, eigen_vecs_b = np.linalg.eig(cov_mat_b)

#	**********************************************		Correctness Parameter Tuning		***************************************************************#

tot_r = sum(eigen_vals_r)
var_exp_r = [(i / tot_r) for i in sorted(eigen_vals_r, reverse=True)]
cum_var_exp_r = np.cumsum(var_exp_r)

tot_g = sum(eigen_vals_g)
var_exp_g = [(i / tot_g) for i in sorted(eigen_vals_g, reverse=True)]
cum_var_exp_g = np.cumsum(var_exp_g)

tot_b = sum(eigen_vals_b)
var_exp_b = [(i / tot_b) for i in sorted(eigen_vals_b, reverse=True)]
cum_var_exp_b = np.cumsum(var_exp_b)

count_r = count_g = count_b = 0

for num in cum_var_exp_r:
	if (num >= THRESHOLD):
		count_r += 1
		break
	count_r += 1

for num in cum_var_exp_g:
	if (num >= THRESHOLD):
		count_g += 1
		break
	count_g += 1

for num in cum_var_exp_b:
	if (num >= THRESHOLD):
		count_b += 1
		break
	count_b += 1

#	**********************************************		Eigen Pairs		***************************************************************#

eigen_pairs_r = [(np.abs(eigen_vals_r[i]), eigen_vecs_r[:, i]) for i in range(len(eigen_vals_r))]
eigen_pairs_g = [(np.abs(eigen_vals_g[i]), eigen_vecs_g[:, i]) for i in range(len(eigen_vals_g))]
eigen_pairs_b = [(np.abs(eigen_vals_b[i]), eigen_vecs_b[:, i]) for i in range(len(eigen_vals_b))]

# Sort the (eigenvalue, eigenvector) tuples from high to low
eigen_pairs_r.sort(key=lambda k: k[0], reverse=True)
eigen_pairs_g.sort(key=lambda k: k[0], reverse=True)
eigen_pairs_b.sort(key=lambda k: k[0], reverse=True)

#	**********************************************		Mean properties		***************************************************************#

mu_r = np.mean(r_channel, axis = 0)
mu_g = np.mean(g_channel, axis = 0)
mu_b = np.mean(b_channel, axis = 0)

#	**********************************************		Principal Component		***************************************************************#

pca_r = sklearn.decomposition.PCA()
pca_r.fit(r_channel)
nComp = count_r
r_PCA = pca_r.transform(r_channel)[:,:nComp]

pca_g = sklearn.decomposition.PCA()
pca_g.fit(g_channel)
nComp = count_g
g_PCA = pca_g.transform(g_channel)[:,:nComp]

pca_b = sklearn.decomposition.PCA()
pca_b.fit(b_channel)
nComp = count_b
b_PCA = pca_b.transform(b_channel)[:,:nComp]

print ("R_PCA SHAPE: %s" % str(r_PCA.shape))
print (r_PCA)
print ("G_PCA SHAPE: %s" % str(g_PCA.shape))
print (g_PCA)
print ("B_PCA SHAPE: %s" % str(b_PCA.shape))
print (b_PCA)
#	**********************************************		Eigen TransPose properties		***************************************************************#

eigen_trans_r = pca_r.components_[:count_r,:]
eigen_trans_g = pca_g.components_[:count_g,:]
eigen_trans_b = pca_b.components_[:count_b,:]

#	**********************************************		Actual RGB dataset		***************************************************************#

rhat = np.dot(r_PCA, eigen_trans_r)
rhat += mu_r
ghat = np.dot(g_PCA, eigen_trans_g)
ghat += mu_g
bhat = np.dot(b_PCA, eigen_trans_b)
bhat += mu_b

print ("RHAT SHAPE: %s" % (str(rhat.shape)))
print (rhat)
print ("GHAT SHAPE: %s" % (str(ghat.shape)))
print (ghat)
print ("BHAT SHAPE: %s" % (str(bhat.shape)))
print (bhat)

#	**********************************************		Creating RGB Data Frame		***************************************************************#
if (rhat.shape == ghat.shape and ghat.shape == bhat.shape):
	rows = rhat.shape[0]
	cols = ghat.shape[1]
	data_list = []
	rhat_list = rhat.tolist()
	ghat_list = ghat.tolist()
	bhat_list = bhat.tolist()
	for r in range(row):
		rlist = [int(round(x)) for x in rhat_list[r]]
		glist = [int(round(x)) for x in ghat_list[r]]
		blist = [int(round(x)) for x in bhat_list[r]]
		row_list = []
		for c in range(cols):
			alist = []
			alist.append(rlist[c])
			alist.append(glist[c])
			alist.append(blist[c])
			row_list.append(alist)
		data_list.append(row_list)

else:
	print ("RGB Properties Mismatching while doing back propagation: Exiting!")
	sys.exit()

rgb_frame = np.array(data_list)
print ("RGB SHAPE: %s" % (str(rgb_frame.shape)))
#rgb_frame = np.reshape(rgb_frame, (rows, cols, 3))
print (rgb_frame)

#	**********************************************		Data to Image Processor		***************************************************************#

#Create Pillow Image

#image2 = Image.fromarray(rgb_frame)
image2 = Image.fromarray((rgb_frame * 255).astype(np.uint8))

print(image2.format)
print(image2.mode)
print(image2.size)

image2.show()
image_arr = image_name.split()[:-1]
new_image_name = '.'.join(image_arr) + "comp.jpg"
image2.save(new_image_name)


from tqdm import tqdm
from collections import OrderedDict, Counter
import numpy as np
from onconet.datasets.abstract_onco_dataset import Abstract_Onco_Dataset
from onconet.datasets.factory import RegisterDataset
from onconet.utils.generic import normalize_dictionary
import pdb

SUMMARY_MSG = "Contructed MGH Mammo {} year {} {} dataset with {} records, and the following class balance \n {}"

METADATA_FILENAMES = {
    'Risk': "mammo_metadata_all_years_only_breast_cancer_aug04_2018.json",
    'Detection' : "mammo_metadata_all_years_all_cancer_aug04_2018.json"
}

class Abstract_MGH_Mammo_Cancer_With_Prior_Dataset(Abstract_Onco_Dataset):
    """A risk dataset where each input consists of an mammo view, and a prior of that same view. Concretly, suppose a patient has 3 mammograms, each with a Left CC/MLO and a Right CC MLO, for a total of 12 images. This dataset would create 12 dataset samples. 4 comparing the most recent set of views to their immediate prior, 4 comparing them to the older prior, and 4 comparing the immediate and older prior."""

    def create_dataset(self, split_group, img_dir):
        """Gets the dataset from the paths and labels in the json.

        Arguments:
            split_group(str): One of ['train'|'dev'|'test'].
            img_dir(str): The path to the directory containing the images.

        Returns:
            The dataset, which is a list of dictionaries with each dictionary
            containing paths to the relevant images, the label, and any other
            additional information.
        """

        dataset = []
        class_balance = {}
        for mrn_row in tqdm(self.metadata_json):
            ssn, split, exams = mrn_row['ssn'], mrn_row['split'], mrn_row['accessions']
            if not split == split_group:
                continue

            year_to_exam = {}
            for exam in exams:
                if not self.check_label(exam):
                    continue

                # Get label
                left_label = self.get_label(exam, 'L')
                right_label = self.get_label(exam, 'R')

                for label in left_label, right_label:
                    if label not in class_balance:
                        class_balance[label] = 0

                # Determine images of left and right CCs and MLOs
                # Note: Validation of cancer side is performed in the query scripts/from_db/cancer.py in OncoQueries
                left_ccs, left_mlos, right_ccs, right_mlos = self.image_paths_by_views(exam)

                if self.args.drop_benign_side:
                    if left_label and not right_label:
                        right_ccs = []
                        right_mlos = []
                    if right_label and not left_label:
                        left_ccs = []
                        left_mlos = []

                year = exam['sdate']

                year_to_exam[year] = {
                    'L CC': left_ccs,
                    'L MLO': left_mlos,
                    'R CC': right_ccs,
                    'R MLO': right_mlos,
                    'L y': left_label,
                    'R y': right_label,
                    'exam': exam['accession']
                }

            # Go from most recent year to oldest
            all_views = ['L CC', 'L MLO', 'R CC','R MLO']
            all_years = list(reversed(sorted(year_to_exam.keys())))
            for indx, year in enumerate(all_years):

                exam = year_to_exam[year]
                prior_exams = [ (prior_year, year_to_exam[prior_year]) for prior_year in all_years[indx+1:]]
                for view in all_views:
                    side = view.split()[0]

                    if len(exam[view]) == 0:
                        continue

                    prior_with_view = [ (prior_year, prior) for prior_year, prior in prior_exams if len(prior[view]) > 0]

                    if len(prior_with_view) == 0:
                        continue

                    for prior_year, prior in prior_with_view:
                        years_and_label = "cur_year:{},prior_year:{},label:{}".format(year, prior_year, exam['{} y'.format(side)])

                        if split_group == 'train':
                            if not (year in self.args.train_years and prior_year in self.args.train_years):
                                continue
                        elif split_group == 'dev':
                            if not (year in self.args.dev_years and
                                prior_year in self.args.dev_years):
                                continue
                        else:
                            assert split_group == 'test'
                            if not (year in self.args.test_years and
                                prior_year in self.args.test_years):
                                continue

                        dataset.append({
                            'paths': [exam[view][0], prior[view][0]],
                            'y': exam['{} y'.format(side)],
                            'additionals': [],
                            'exam': exam['exam'],
                            'prior_exam': prior['exam'],
                            'dist_key': years_and_label,
                            'year': year,
                            'prior_year': prior_year,
                            'ssn':ssn
                        })

                        class_balance[exam['{} y'.format(side)]] += 1

        class_balance = normalize_dictionary(class_balance)
        print(SUMMARY_MSG.format(self.years, self.task, split_group, len(dataset), class_balance))

        return dataset

    def check_label(self, row):
        valid_pos = row['years_to_cancer'] < self.years
        valid_neg = row['years_to_last_followup'] >= self.years
        return valid_pos or valid_neg

    def get_label(self, row, side='Any'):
        if side == 'Any':
            return row['years_to_cancer'] < self.years
        elif side == 'L':
            return row['left_years_to_cancer'] < self.years
        else:
            assert(side == 'R')
            return row['right_years_to_cancer'] < self.years

    @property
    def METADATA_FILENAME(self):
        return METADATA_FILENAMES[self.task]

    @staticmethod
    def set_args(args):
        args.num_classes = 2
        args.multi_image = True
        args.num_images = 2

@RegisterDataset("mgh_mammo_5year_risk_with_prior")
class MGH_Mammo_5Year_Risk_With_Prior(Abstract_MGH_Mammo_Cancer_With_Prior_Dataset):
    def __init__(self, args, transformer, split_group):
        self.years = 5
        super(MGH_Mammo_5Year_Risk_With_Prior, self).__init__(args, transformer, split_group)

    @property
    def task(self):
        return "Risk"
@RegisterDataset("mgh_mammo_1year_risk_with_prior")
class MGH_Mammo_1Year_Risk_With_Prior(Abstract_MGH_Mammo_Cancer_With_Prior_Dataset):
    def __init__(self, args, transformer, split_group):
        self.years = 1
        super(MGH_Mammo_1Year_Risk_With_Prior, self).__init__(args, transformer, split_group)

    @property
    def task(self):
        return "Risk"

@RegisterDataset("mgh_mammo_1year_detection_with_prior")
class MGH_Mammo_1Year_Detection_With_Prior(Abstract_MGH_Mammo_Cancer_With_Prior_Dataset):
    def __init__(self, args, transformer, split_group):
        self.years = 1
        super(MGH_Mammo_1Year_Detection_With_Prior, self).__init__(args, transformer, split_group)

    @property
    def task(self):
        return "Detection"


import logging

def configure_logger():
    logger = logging.getLogger("fibery_api")
    logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler()
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger


from tkmacosx import Button
from timeit import repeat
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
import tkinter as tk
import matplotlib.pyplot as plt
import matplotlib
from matplotlib import animation
import random
import visualisations as visual
from tkinter import ttk

matplotlib.use("TkAgg")
plt.style.use("seaborn-dark")

def make_root() -> tk.Tk:
    """Creates a root window

    Parameters
    ----------
    None

    Returns
    -------
    root : tk.Tk
    a Tkinter Window

    """

    root = tk.Tk()
    root.title("Home")
    screen_width = root.winfo_screenwidth()
    screen_height = root.winfo_screenheight()
    root.geometry(f"{screen_width}x{screen_height}")
    root.configure(background="black")
    return root

def visualize(array: list, name: str, titles: dict):
    """Visualizes the comparison of algorithms for a given array

    Parameters
    ----------
    array : list
        The array to be sorted
    name : str
        The name of the algorithm
    titles : dict
        The titles for the visualisations

    Returns
    -------
    None
    """

    # creates a generator object containing all
    # the states of the array while performing
    # sorting algorithm
    if name == "bubble_sort":
        generator = visual.bubble_sort(array)
    elif name == "insertion_sort":
        generator = visual.insertion_sort(array)
    elif name == "quick_sort":
        generator = visual.quick_sort(array, 0, len(array) - 1)
    elif name == "merge_sort":
        generator = visual.merge_sort(array, 0, len(array) - 1)
    elif name == "count_sort":
        generator = visual.count_sort(array)

    # creates a figure and subsequent subplots
    fig, ax = plt.subplots()
    fig.set_facecolor("black")
    fig.canvas.manager.set_window_title(titles[name])
    ax.set_title(f"{titles[name]} {complexity[name]}", color="white")

    ax.set_xlabel("Array Index", color="white")
    ax.set_ylabel("Array Value", color="white")
    ax.set_facecolor("black")
    ax.tick_params(axis="y", colors="white")
    ax.tick_params(axis="x", colors="white")

    bar_sub = ax.bar(range(len(array)), array, align="edge", color="#6669eb")

    # sets the maximum limit for the x-axis
    ax.set_xlim(0, len(array))
    iteration = [0]

    # helper function to update each frame in plot
    def update(A, rects, iteration):
        for rect, val in zip(rects, A):
            rect.set_height(val)
        iteration[0] += 1

    # creating animation object for rendering the iteration
    anim = animation.FuncAnimation(
        fig,
        func=update,
        fargs=(bar_sub, iteration),
        frames=generator,
        repeat=True,
        blit=False,
        interval=0,
        save_count=90000,
    )

    # for showing the animation on screen
    plt.show()
    plt.close()

def init_window(title: str = "Comparison when unsorted ") -> None:
    """Creates a window

    Parameters
    ----------
    title : str,optional
    The title of the window

    Returns
    -------
    None
    """
    root = make_root()

    def run_algorithm(
        algorithm: str,
        array: list = [random.randint(0, 10000) for i in range(5000)],
        n: int = 5000,
        to_show: bool = True,
    ) -> float:
        """Runs an algorithm out of the list of algorithms

        Parameters
        ----------
        algorithm : str
        The name of the algorithm to run
        array : list, optional
        The array to carry out the operation on
        n : int, optional
        The number of elements in the array
        to_show: bool, optional
        Whether to show the graph with visual depiction

        Returns
        -------
        time : float
        Minimum time taken for the algorithm to run on the array
        """
        if array is None:
            array = [random.randint(0, 2 * n) for i in range(n)]
        if button_frame:
            button_frame.destroy()
            heading.destroy()
        else:
            pass
        titles = {
            "bubble_sort": "Bubble Sort",
            "insertion_sort": "Insertion Sort",
            "merge_sort": "Merge Sort",
            "quick_sort": "Quick Sort",
            "count_sort": "Count Sort",
            "sorted": "Built-in Sort(Tim Sort)",
        }

        if algorithm == "quick_sort" or algorithm == "merge_sort":
            # Because Sorting is done in place and we need to keep the original array
            copy_of_array = array
            stmt = f"{algorithm}{copy_of_array,0,len(array)-1}"
        else:
            stmt = f"{algorithm}({array})"

        # Set up the context and prepare the call to the specified
        # algorithm using the supplied array. Only import the
        # algorithm function if it's not the built-in `sorted()`.
        setup_code = (
            f"from algorithms import {algorithm}" if algorithm != "sorted" else ""
        )

        # Execute the code ten different times and return the time
        # in seconds that each execution took
        times = repeat(setup=setup_code, stmt=stmt, repeat=3, number=10)
        # Finally, display the name of the algorithm and the
        # minimum time it took to run
        root.title(titles[algorithm])

        algo_name = tk.Label(
            root,
            text=titles[algorithm],
            bg="black",
            fg="#C0C0C0",
            font=("Helvetica", 36),
        )
        algo_name.place(relx=0.5, rely=0.05, anchor=tk.N)
        min_time = tk.Label(
            root,
            text=f"""
            Minimum time to sort {n} elements: \
{ (min(times))/10} seconds""",
            bg="black",
            fg="green",
            font=("Helvetica", 26),
        )
        min_time.place(relx=0.5, rely=0.55, anchor=tk.CENTER)
        time_complexity = tk.Label(
            root,
            text=f"Time Complexity: {complexity[algorithm]}",
            bg="black",
            fg="#057DCD",
            font=("Helvetica", 26),
        )
        time_complexity.place(relx=0.5, rely=0.5, anchor=tk.CENTER)
        description = tk.Label(
            root,
            text=descriptions[algorithm],
            bg="black",
            fg="#C0C0C0",
            font=("Helvetica", 24),
            justify=tk.LEFT,
        )
        description.place(relx=0.5, rely=0.2, anchor=tk.N)

        def back():
            """Returns to the main menu

            Parameters
            ----------
            None

            Returns
            -------
            None
            """
            root.destroy()
            init_window()

        back_btn = Button(
            root,
            text="Back",
            command=back,
            bg="#C0C0C0",
            borderless=True,
            fg="black",
            activebackground="black",
            activeforeground="#C0C0C0",
        )
        back_btn.place(x=10, y=10)
        if to_show and n < 101:
            visualize(array, algorithm, titles)
        return min(times) / 10

    heading = tk.Label(
        root,
        text="Analysis of Sorting Algorithms",
        bg="black",
        fg="white",
        font=("Helvetica", 36),
    )
    heading.place(relx=0.5, rely=0.05, anchor=tk.N)
    button_frame = tk.Frame(root, bg="black")
    button_frame.grid(row=0, column=0)
    root.grid_rowconfigure(0, weight=1)
    root.grid_columnconfigure(0, weight=1)

    def get_value(win: tk.Tk, name: str, function: any) -> None:
        """Gets the value of the input box

        Parameters
        ----------
        win : tk.Tk
        The previous window to pass

        name : str
        The name for creating title for the next window

        function : any
        Used to differentiate between the function to run based on origin of the call

        Returns
        -------
        None
        """

        n = 0
        is_make = id(function) == id(make_graph)

        def submit():
            """Runs the algorithm based on the input

            Parameters
            ----------
            None

            Returns
            -------
            None
            """

            def update_time(remaining: int) -> None:
                """Updates the timer on screen

                Parameters
                ----------
                remaining : int
                The number of seconds remaining

                Returns
                -------
                None
                """
                if remaining > -1:
                    close_label.config(
                        text=f"Automatically closing in {remaining} seconds"
                    )
                    hint_window.after(1000, update_time, remaining - 1)
                else:
                    after_hint()

            def after_hint():
                """
                Displays a hint about the algorithm recommended to the user

                Parameters
                ----------
                None

                Returns
                -------
                None

                """
                hint_window.destroy()
                root.destroy()
                array = [random.randint(0, 2 * n) for i in range(n)]
                run_algorithm(function, array, n)

            nonlocal n
            n = int(no_of_elements.get())
            if is_make:
                make_graph(win, name, n)
            else:
                hint_window = tk.Tk()
                hint_window.title("Hint")
                hint_window.geometry("600x300")
                hint_window.configure(bg="black ")
                hint_label = tk.Label(
                    hint_window,
                    text="We recommend you use the built-in sorting algorithm in Python",
                    bg="black",
                    fg="cyan",
                    font=("Helvetica", 16),
                )
                hint_label.place(relx=0.5, rely=0.5, anchor=tk.CENTER)
                close_label = tk.Label(
                    hint_window,
                    text="Automatically closing in 5 seconds",
                    bg="black",
                    fg="red",
                    font=("Helvetica", 16),
                )
                close_label.place(relx=0.5, rely=0.7, anchor=tk.CENTER)
                update_time(5)
                hint_window.mainloop()

        root = make_root()
        tk.Label(
            root,
            text="Enter the number of elements you wish to sort",
            bg=root["bg"],
            font=("Helvetica", 24),
        ).place(
            relx=0.5,
            rely=0.1,
            anchor=tk.N,
        )
        if not is_make:
            tk.Label(
                root,
                text="To see a live graph, enter a value less than 100",
                bg=root["bg"],
                fg="#FAD02C",
                font=("Helvetica", 20),
            ).place(
                relx=0.5,
                rely=0.7,
                anchor=tk.N,
            )
        no_of_elements = tk.Entry(
            root,
            insertbackground="black",
            bg="black",
            fg="#C0C0C0",
            font=("Helvetica", 24),
        )
        no_of_elements.place(relx=0.5, rely=0.2, anchor=tk.N)
        submit_btn = Button(
            root,
            text="Submit",
            command=submit,
            bg="#C0C0C0",
            borderless=True,
            fg="black",
            activebackground="black",
            activeforeground="#C0C0C0",
            padx=10,
            pady=10,
            font=("Helvetica", 24),
        )
        submit_btn.place(relx=0.5, rely=0.35, anchor=tk.CENTER)
        if name == "Compare between 2 algorithms":
                algos = ["Bubble Sort",
                    "Insertion Sort",
                    "Merge Sort",
                    "Quick Sort",
                    "Count Sort",
                    "Built-in Sort(Tim Sort)",
                    ]
                drop_down1 = ttk.Combobox(root, value = algos)
                drop_down1.current(0)
                drop_down1.pack(pady=20)
                drop_down2 = ttk.Combobox(root, value = algos)
                drop_down2.current(0)
                drop_down2.pack(pady=20)
                print("the above one")
                def what(e):
                    global alg1
                    alg1 = (drop_down1.get())
                def what1(e):
                    global alg2
                    alg2 =(drop_down2.get())
                drop_down1.bind("<<ComboboxSelected>>",what)
                drop_down2.bind("<<ComboboxSelected>>",what1)

        root.mainloop()

    def make_graph(win: tk.Tk, name: str, n: int) -> None:
        """Creates a graph of the time taken to sort an array

        Parameters
        ----------
        win : tk.Tk
        The previous window

        name : str
        The title of the window

        Returns
        -------
        None
        """

        array = [random.randint(0, 2 * n) for i in range(n)]
        sorted_array = sorted(array)
        if name == "Comparison when fully sorted":
            data = {
                "Bubble Sort": run_algorithm("bubble_sort", sorted_array, n, False),
                "Insertion Sort": run_algorithm(
                    "insertion_sort", sorted_array, n, False
                ),
                "Merge Sort": run_algorithm("merge_sort", sorted_array, n, False),
                "Quick Sort": run_algorithm("quick_sort", sorted_array, n, False),
                "Count Sort": run_algorithm("count_sort", sorted_array, n, False),
                "Built-in Sort(Tim Sort)": run_algorithm(
                    "sorted", sorted_array, n, False
                ),
            }
        elif name == "Comparison when partially sorted":
            partially_sorted_array = sorted_array[: len(array) // 2] + [
                random.randint(0, n) for i in range(len(array) // 2)
            ]
            data = {
                "Bubble Sort": run_algorithm(
                    "bubble_sort", partially_sorted_array, n, False
                ),
                "Insertion Sort": run_algorithm(
                    "insertion_sort", partially_sorted_array, n, False
                ),
                "Merge Sort": run_algorithm(
                    "merge_sort", partially_sorted_array, n, False
                ),
                "Quick Sort": run_algorithm(
                    "quick_sort", partially_sorted_array, n, False
                ),
                "Count Sort": run_algorithm(
                    "count_sort", partially_sorted_array, n, False
                ),
                "Built-in Sort(Tim Sort)": run_algorithm(
                    "sorted", partially_sorted_array, n, False
                ),
            }
        elif name == "Comparison when unsorted ":
            data = {
                "Bubble Sort": run_algorithm("bubble_sort", array, n, False),
                "Insertion Sort": run_algorithm("insertion_sort", array, n, False),
                "Merge Sort": run_algorithm("merge_sort", array, n, False),
                "Quick Sort": run_algorithm("quick_sort", array, n, False),
                "Count Sort": run_algorithm("count_sort", array, n, False),
                "Built-in Sort(Tim Sort)": run_algorithm("sorted", array, n, False),
            }

        elif name == "Compare between 2 algorithms":
            data1 = {
                "Bubble Sort": run_algorithm("bubble_sort", array, n, False),
                "Insertion Sort": run_algorithm("insertion_sort", array, n, False),
                "Merge Sort": run_algorithm("merge_sort", array, n, False),
                "Quick Sort": run_algorithm("quick_sort", array, n, False),
                "Count Sort": run_algorithm("count_sort", array, n, False),
                "Built-in Sort(Tim Sort)": run_algorithm("sorted", array, n, False),
            }
            data = {
                alg1:data1[alg1],
                alg2:data1[alg2]
            }

        show_graph(root, name, data, n)

    unsorted_btn = Button(
        root,
        text=title,
        font=("Helvetica", 20),
        command=lambda: get_value(root, title, make_graph),
        pady=15,
        borderless=True,
        bg="yellow",
        fg="#252930",
        activebackground="white",
        activeforeground="black",
    )
    unsorted_btn.place(anchor=tk.N, relx=0.8, rely=0.7)

    sorted_btn = Button(
        root,
        text="Comparison when fully sorted",
        font=("Helvetica", 20),
        borderless=True,
        command=lambda: get_value(
            root, "Comparison when fully sorted", make_graph),
        pady=15,
        bg="#32CBF1",
        fg="#252930",
        activebackground="white",
        activeforeground="black",
    )
    sorted_btn.place(anchor=tk.N, relx=0.5, rely=0.7)

    partially_sorted_btn = Button(
        root,
        text="Comparison when partially sorted",
        font=("Helvetica", 20),
        borderless=True,
        command=lambda: get_value(
            root, "Comparison when partially sorted", make_graph),
        pady=15,
        bg="#6ECB5A",
        fg="#252930",
        activebackground="white",
        activeforeground="black",
    )
    partially_sorted_btn.place(anchor=tk.N, relx=0.2, rely=0.7)

    Compare_between_2 = Button(
        root,
        text="Compare between 2 algorithms",
        font=("Helvetica", 20),
        borderless=True,
        command=lambda: get_value(
           root, "Compare between 2 algorithms", make_graph),
        pady=15,
        bg="#6ECB5A",
        fg="#252930",
        activebackground="white",
        activeforeground="black",
    )
    Compare_between_2.place(anchor=tk.N, relx=0.4, rely=0.8)

    bubble_sort_btn = Button(
        button_frame,
        text="Bubble Sort",
        font=("Helvetica", 20),
        command=lambda: get_value(root, "Bubble Sort", "bubble_sort"),
        borderless=True,
        pady=25,
        padx=25,
        bg="#1b1b1b",
        fg="#C0C0C0",
        activebackground="white",
        activeforeground="black",
    )
    bubble_sort_btn.grid(row=0, column=1, sticky="eW")

    insertion_sort_btn = Button(
        button_frame,
        text="Insertion Sort",
        font=("Helvetica", 20),
        command=lambda: get_value(root, "Insertion Sort", "insertion_sort"),
        borderless=True,
        pady=25,
        padx=25,
        bg="#1b1b1b",
        fg="#C0C0C0",
        activebackground="white",
        activeforeground="black",
    )
    insertion_sort_btn.grid(row=0, column=2, sticky="eW")

    merge_sort_btn = Button(
        button_frame,
        text="Merge Sort",
        font=("Helvetica", 20),
        command=lambda: get_value(root, "Merge Sort", "merge_sort"),
        borderless=True,
        pady=25,
        bg="#1b1b1b",
        fg="#C0C0C0",
        activebackground="white",
        activeforeground="black",
    )
    merge_sort_btn.grid(row=1, column=1, sticky="eW")

    quick_sort_btn = Button(
        button_frame,
        text="Quick Sort",
        font=("Helvetica", 20),
        command=lambda: get_value(root, "Quick Sort", "quick_sort"),
        borderless=True,
        pady=25,
        bg="#1b1b1b",
        fg="#C0C0C0",
        activebackground="white",
        activeforeground="black",
    )
    quick_sort_btn.grid(row=1, column=2, sticky="eW")

    count_sort_btn = Button(
        button_frame,
        text="Count Sort",
        font=("Helvetica", 20),
        command=lambda: get_value(root, "Count Sort", "count_sort"),
        borderless=True,
        pady=25,
        bg="#1b1b1b",
        fg="#C0C0C0",
        activebackground="white",
        activeforeground="black",
    )
    count_sort_btn.grid(row=3, column=1, sticky="eW")

    button_frame.grid_columnconfigure(0, weight=1)
    button_frame.grid_rowconfigure(0, weight=1)
    root.configure(bg="black")
    root.mainloop()

def show_graph(win: tk.Tk, name: str, data: dict, n=5000) -> None:
    """Shows the graph of the comparison amongst algorithms for a given array

    Parameters
    ----------
    win : tk.Tk
        Previous window that is to be destroyed
    name : str
        The title for the new window
    data : dict
        The data to be plotted

    Returns
    -------
    None
    """
    win.destroy()
    root = make_root()
    root.title(name)
    languages = data.keys()
    time = data.values()

    # create a figure
    figure = Figure(figsize=(8, 8), dpi=100,
                    facecolor="black", edgecolor="black")
    # create FigureCanvasTkAgg object
    figure_canvas = FigureCanvasTkAgg(figure, root)
    # create the toolbar
    toolbar = NavigationToolbar2Tk(
        figure_canvas,
        root,
    )
    toolbar.configure(background="black")

    def back():
        """Returns to the main menu

        Parameters
        ----------
        None

        Returns
        -------
        None
        """
        root.destroy()
        init_window()

    topLeftFrame = tk.Frame(root, bg="black")
    topLeftFrame.pack(side=tk.TOP, fill=tk.BOTH, expand=1)
    back_btn = Button(
        topLeftFrame,
        text="Back",
        command=back,
        borderless=True,
        activebackground=root["bg"],
        activeforeground="#C0C0C0",
    )
    back_btn.place(x=10, y=10)
    # create axes
    axes = figure.add_subplot()
    # create the barchart
    axes.bar(languages, time, color="#6669eb")
    axes.set_title(name, fontsize=36, pad=25, color="#C0C0C0")
    axes.set_ylabel(
        f"Time taken to sort {n} elements (seconds)",
        labelpad=15,
        fontsize=26,
        color="#C0C0C0",
    )
    axes.set_xlabel("Sorting Algorithms", labelpad=15,
                    fontsize=26, color="#C0C0C0")
    axes.set_facecolor("black")
    axes.tick_params(axis="y", colors="#C0C0C0")
    axes.tick_params(axis="x", colors="#C0C0C0")

    figure_canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)

    root.mainloop()

descriptions = {
    "bubble_sort": """Bubble Sort, sometimes referred to as sinking sort,
is a simple sorting algorithm that repeatedly steps through the list,
compares adjacent elements and swaps them if they are in the wrong order.""",
    "insertion_sort": """Insertion Sort is a simple sorting algorithm,
that builds the final sorted array (or list) one item at a time.""",
    "merge_sort": """Merge Sort comparison-based sorting algorithm. Most implementations
produce a stable sort, which means that the order of equal elements is the
same in the input and output. Merge sort is a divide and conquer algorithm""",
    "quick_sort": """Quicksort is a divide-and-conquer algorithm.
It works by selecting a 'pivot' element from the array and partitioning the other elements
into two sub-arrays, according to whether they are less than or greater than the pivot.
For this reason, it is sometimes called partition-exchange sort.
The sub-arrays are then sorted recursively.""",
    "count_sort": """Counting sort is an algorithm for sorting a collection of objects based on
    keys that are small positive integers, it is not a comparison sort; it uses key values as
indexes into an array, the Ω(n log n) lower bound for comparison sorting does not apply.""",
    "sorted": """The built in sort function in Python is Tim Sort.
It implements the idea that the real-world data sets almost always contain already ordered subsequences,
so the sorting strategy is to identify them and sort them further using both merge and insert methods""",
}

complexity = {
    "bubble_sort": "O(n^2)",
    "insertion_sort": "O(n^2)",
    "merge_sort": "O(n log n)",
    "quick_sort": "O(n log n)",
    "count_sort": "O(n+k) ∈ k is range of values",
    "sorted": "O(n log n) ",
}
init_window()


from django.db import models
from jsonfield import JSONField
from django.conf import settings
import google_auth_oauthlib.flow
from google.oauth2.credentials import Credentials
import google.auth.transport.requests
import json
import google.auth.exceptions

from poll.models.poll import Poll
from accounts.models import User
from poll.models.surveypassing import SurveyPassing

class GoogleSheetIntegration(models.Model):
    id = models.OneToOneField(Poll, primary_key=True, on_delete=models.CASCADE,
                              verbose_name='poll id')
    user = models.ForeignKey(User, on_delete=models.CASCADE, verbose_name='User')
    is_active = models.BooleanField(default=False,
                                    verbose_name='active')
    spreadsheet_id = models.CharField(max_length=255, default='', blank=True, null=True,
                                      verbose_name='spreadsheetId')
    spreadsheet_url = models.URLField(blank=True, null=True,
                                      verbose_name='spreadsheet url')
    row_count = models.PositiveIntegerField(default=1, verbose_name='row count')
    survey_id = models.ManyToManyField(SurveyPassing, blank=True, verbose_name='survey_id')

    REQUIRED_FIELDS = []

    class Meta:
        verbose_name = 'GoogleSheetIntegration'
        verbose_name_plural = 'GoogleSheetIntegrations'

    def __str__(self):
        return repr(self.id)

class GoogleSheetCredentials(models.Model):
    user = models.OneToOneField(User, on_delete=models.CASCADE)
    google_sheet_credentials = JSONField(default='', blank=True, null=True)

    REQUIRED_FIELDS = []

    class Meta:
        verbose_name = 'GSCredentials'
        verbose_name_plural = 'GSCredentials'

    def __str__(self):
        return repr(self.id)

    @staticmethod
    def get_google_auth_url(credentials_fname):
        flow = google_auth_oauthlib.flow.Flow.from_client_secrets_file(
            credentials_fname,
            scopes=['https://www.googleapis.com/auth/drive.file',
                    'https://www.googleapis.com/auth/spreadsheets',
                    'https://www.googleapis.com/auth/userinfo.email',
                    'https://www.googleapis.com/auth/userinfo.profile',
                    'openid'])

        flow.redirect_uri = f'{settings.DOMAIN}/api/v1/googleauth/creategooglesheet/'

        return flow.authorization_url(access_type='offline',
                                      prompt='consent',
                                      include_granted_scopes='true')

    def create_connection(self, request, credentials_fname):
        flow = google_auth_oauthlib.flow.Flow.from_client_secrets_file(
            credentials_fname,
            scopes=['https://www.googleapis.com/auth/drive.file',
                    'https://www.googleapis.com/auth/spreadsheets',
                    'https://www.googleapis.com/auth/userinfo.email',
                    'https://www.googleapis.com/auth/userinfo.profile',
                    'openid'], )

        flow.redirect_uri = f'{settings.DOMAIN}/api/v1/googleauth/creategooglesheet/'

        authorization_response = request.build_absolute_uri('/') + request.get_full_path()

        if authorization_response[:5] != 'https':
            authorization_response = 'https' + authorization_response[4:]

        flow.fetch_token(authorization_response=authorization_response)
        credentials = flow.credentials
        self.google_sheet_credentials = credentials.to_json()

        self.save()
        return self.google_sheet_credentials

    def refresh_token(self, request):

        credentials = Credentials.from_authorized_user_info(
            json.loads(self.google_sheet_credentials))
        request_auth = google.auth.transport.requests.Request()

        try:
            credentials.refresh(request_auth)
            self.google_sheet_credentials = credentials.to_json()
        except google.auth.exceptions.RefreshError:
            self.google_sheet_credentials = ''

        self.save()
        return self.google_sheet_credentials


import numpy as np
import os
from pathlib import Path
import scipy
import torch
from torch.nn import functional as F
from typing import *

from src.enums import *

class Transformer:
    def __init__(self, purpose: str, transforms: list):
        self.purpose = purpose

        self.transforms = transforms

        if self.purpose == "test":
            self.deterministic = True
        else:
            self.deterministic = False

        self.rng = np.random.RandomState(seed=1)

    def __call__(self, data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        transformed_data = data

        for transform_config in self.transforms:
            transform_type = TransformEnum(transform_config["type"])

            transform_fct = getattr(self, transform_type.value)
            transformed_data = transform_fct(transform_config, transformed_data)

        return transformed_data

    def white_noise(self, transform_config: dict,
                    data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        stdev = transform_config["stdev"]

        noise = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if noise is None:
                    noise_value = self.rng.normal(loc=0, scale=stdev, size=tuple(value.size()))
                    noise = value.new_tensor(noise_value, dtype=value.dtype)

                transformed_value = value + noise

                data[channel] = transformed_value

        return data

    def range_adjusted_white_noise(self, transform_config: dict,
                                   data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        res_grid = data[ChannelEnum.RES_GRID]
        robot_position = data[ChannelEnum.REL_POSITION]

        sample_grid = data[ChannelEnum.OCC_DEM]

        # distance of every pixel from the robot
        lin_x = np.arange(start=-sample_grid.shape[0] / 2, stop=sample_grid.shape[0] / 2, step=1) * res_grid[0].numpy()
        lin_y = np.arange(start=-sample_grid.shape[1] / 2, stop=sample_grid.shape[1] / 2, step=1) * res_grid[1].numpy()
        off_y, off_x = np.meshgrid(lin_x, lin_y)

        dist_x = off_x - robot_position[0].numpy()
        dist_y = off_y - robot_position[1].numpy()
        dist_p2_norm = np.sqrt(np.square(dist_x) + np.square(dist_y))

        stdev = transform_config["stdev"]
        range = transform_config["range"]

        scale = stdev * np.square(1 / range * dist_p2_norm)

        noise = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if noise is None:
                    noise_value = rng.normal(loc=0, scale=scale, size=tuple(value.size()))
                    noise = value.new_tensor(noise_value, dtype=value.dtype)

                transformed_value = value + noise

                data[channel] = transformed_value

        return data

    def gaussian_filtered_white_noise(self, transform_config: dict,
                                      data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        probability = transform_config["probability"]
        horizontal_stdev = transform_config["horizontal_stdev"]
        vertical_stdev = transform_config["vertical_stdev"]

        filtered_noise = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if filtered_noise is None:
                    peaks = rng.choice([0, 1], p=[1 - probability, probability], size=tuple(value.size()))
                    white_noise = rng.normal(loc=0, scale=vertical_stdev, size=tuple(value.size()))
                    noise = peaks * white_noise
                    filtered_noise = scipy.ndimage.gaussian_filter(input=noise, sigma=horizontal_stdev)

                    filtered_noise = value.new_tensor(filtered_noise, dtype=value.dtype)

                transformed_value = value + filtered_noise

                data[channel] = transformed_value

        return data

    def random_vertical_scale(self, transform_config: dict,
                              data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        min, max = transform_config["min"], transform_config["max"]

        scale = rng.uniform(low=min, high=max)

        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if channel is ChannelEnum.REL_POSITION:
                    transformed_value = value.clone()
                    transformed_value[2] = scale * transformed_value[2]
                else:
                    transformed_value = scale * value

                data[channel] = transformed_value

        return data

    def random_vertical_offset(self, transform_config: dict,
                               data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        min, max = transform_config["min"], transform_config["max"]

        offset = rng.uniform(low=min, high=max)

        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if channel is ChannelEnum.REL_POSITION:
                    transformed_value = value.clone()
                    transformed_value[2] = offset + transformed_value[2]
                else:
                    transformed_value = offset + value

                data[channel] = transformed_value

        return data

    def random_occlusion(self, transform_config: dict,
                         data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        probability = transform_config["probability"]

        occlusion = None
        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                if occlusion is None:
                    occlusion = rng.choice([0, 1], p=[1 - probability, probability], size=tuple(value.size()))
                    occlusion = torch.tensor(occlusion)

                if channel == ChannelEnum.OCC_MASK:
                    # assert value.dtype == torch.bool
                    #
                    # transformed_value = value.to(dtype=torch.int) + occlusion
                    #
                    # transformed_value[transformed_value > 1] = 1
                    #
                    # data[channel] = transformed_value.to(dtype=torch.bool)

                    value[occlusion == 1] = True
                elif channel in [ChannelEnum.OCC_DEM, ChannelEnum.OCC_DATA_UM]:
                    value[occlusion == 1] = np.nan
                else:
                    raise NotImplementedError

                data[channel] = value

        return data

    def random_occlusion_dilation(self, transform_config: dict,
                                  data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        mask = data[ChannelEnum.OCC_MASK]

        if mask.dtype == torch.bool:
            mask = mask.to(dtype=torch.float)

        if mask.dim() == 2:
            # we need a channel composition with minibatches for 2d convolutions
            mask = mask.unsqueeze(dim=0).unsqueeze(dim=0)
        else:
            raise NotImplementedError

        kernel_size = int(transform_config["kernel_size"])
        # probability to set kernel cell to one
        probability = transform_config["probability"]
        kernel_choice = rng.choice([0, 1], p=[1 - probability, probability], size=(kernel_size, kernel_size))
        # we always set the center of the kernel to 1
        kernel_choice[int(kernel_size // 2), int(kernel_size // 2)] = 1
        kernel = torch.tensor(kernel_choice, dtype=torch.float).unsqueeze(dim=0).unsqueeze(dim=0)

        with torch.no_grad():
            dilated_mask = F.conv2d(mask, weight=kernel, padding=1) > 0

        dilated_occlusion = dilated_mask.squeeze(dim=0).squeeze(dim=0)

        data[ChannelEnum.OCC_MASK] = dilated_occlusion

        if ChannelEnum.OCC_DEM in data:
            data[ChannelEnum.OCC_DEM][dilated_occlusion == 1] = np.nan
        if ChannelEnum.OCC_DATA_UM in data:
            data[ChannelEnum.OCC_DATA_UM][dilated_occlusion == 1] = np.nan

        return data

    def random_occlusion_raycasting(self, transform_config: dict,
                               data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        rng = self.rng if self.deterministic else np.random

        import grid_map_raycasting

        occ_mask = data[ChannelEnum.OCC_MASK]
        gt_dem = data[ChannelEnum.GT_DEM]
        res_grid = data[ChannelEnum.RES_GRID]

        vantage_point_u = 0
        vantage_point_v = 0
        vantage_point_elevation = np.NaN
        # we only accept coordinates for which we have elevation information
        while np.isnan(vantage_point_elevation):
            u_choices = np.arange(start=0, stop=gt_dem.size(0))
            v_choices = np.arange(start=0, stop=gt_dem.size(1))

            vantage_point_u = rng.choice(u_choices)
            vantage_point_v = rng.choice(v_choices)

            vantage_point_elevation = gt_dem[vantage_point_u, vantage_point_v]

        if transform_config.get("occluded_area_constraint") is not None:
            constraint_config = transform_config["occluded_area_constraint"]
            max_iterations = constraint_config["max_iterations"]
            min_fraction = constraint_config["min"]
            max_fraction = constraint_config["max"]
            min_step = constraint_config.get("min_step", 0.05)
        else:
            max_iterations = 1
            min_fraction = 0
            max_fraction = 1
            min_step = 0

        low = transform_config["height_viewpoint"]["min"]
        high = transform_config["height_viewpoint"]["max"]
        found_suitable_occlusion = False
        i = 0
        while found_suitable_occlusion is False and i < max_iterations:
            i += 1

            height_viewpoint = rng.uniform(low=low, high=high)

            vantage_point_x = (-gt_dem.size(0) / 2 + vantage_point_u) * res_grid[0]
            vantage_point_y = (-gt_dem.size(1) / 2 + vantage_point_v) * res_grid[1]
            vantage_point_z = vantage_point_elevation + height_viewpoint

            vantage_point = np.array([vantage_point_x, vantage_point_y, vantage_point_z], dtype=np.double)

            np_gt_dem = gt_dem.detach().numpy().astype(np.double)

            np_raycasted_occ_mask = grid_map_raycasting.rayCastGridMap(vantage_point, np_gt_dem, res_grid)

            occ_mask_shape = np_raycasted_occ_mask.shape
            occlusion_fraction = np_raycasted_occ_mask.sum() / (occ_mask_shape[0] * occ_mask_shape[1])

            if transform_config.get("occluded_area_constraint") is not None:
                if min_fraction <= occlusion_fraction <= max_fraction:
                    found_suitable_occlusion = True
                else:
                    if occlusion_fraction < min_fraction:
                        high = height_viewpoint

                        if np.abs(high - low) < min_step:
                            low = max(0, low - min_step)

                    elif max_fraction < occlusion_fraction:
                        low = height_viewpoint

                        if np.abs(high - low) < min_step:
                            high += min_step
                    else:
                        raise ValueError
            else:
                found_suitable_occlusion = True

        data[ChannelEnum.OCC_MASK] = torch.logical_or(occ_mask, torch.tensor(np_raycasted_occ_mask, dtype=torch.bool))

        if ChannelEnum.OCC_DEM in data:
            data[ChannelEnum.OCC_DEM][data[ChannelEnum.OCC_MASK] == 1] = np.nan
        if ChannelEnum.OCC_DATA_UM in data:
            data[ChannelEnum.OCC_DATA_UM][data[ChannelEnum.OCC_MASK] == 1] = np.nan

        return data

    def range_data_uncertainty(self, transform_config: dict,
                               data: Dict[ChannelEnum, torch.Tensor]) -> Dict[ChannelEnum, torch.Tensor]:
        res_grid = data[ChannelEnum.RES_GRID]
        robot_position = data[ChannelEnum.REL_POSITION]

        sample_grid = data[ChannelEnum.OCC_DATA_UM]

        # distance of every pixel from the robot
        lin_x = np.arange(start=-sample_grid.shape[0] / 2, stop=sample_grid.shape[0] / 2, step=1) * res_grid[0].numpy()
        lin_y = np.arange(start=-sample_grid.shape[1] / 2, stop=sample_grid.shape[1] / 2, step=1) * res_grid[1].numpy()
        off_y, off_x = np.meshgrid(lin_x, lin_y)

        dist_x = off_x - robot_position[0].numpy()
        dist_y = off_y - robot_position[1].numpy()
        dist_p2_norm = np.sqrt(np.square(dist_x) + np.square(dist_y))

        stdev = transform_config["stdev"]
        range = transform_config["range"]

        range_data_variance = stdev**2 * np.square(1 / range * dist_p2_norm)

        for channel, value in data.items():
            if channel.value in transform_config["apply_to"]:
                transformed_value = value + value.new_tensor(range_data_variance, dtype=value.dtype)

                data[channel] = transformed_value

        return data


import os, torch, argparse, sys
sys.path.append("..")
import torch.optim as optim
import torch.utils.data as utils
import pandas as pd
from model import *
from tqdm import tqdm
from sklearn.model_selection import train_test_split, StratifiedKFold
from torch.utils import data
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm
from generator import *
from sklearn.metrics import precision_score, recall_score

parser = argparse.ArgumentParser()

parser.add_argument("--root", type=str,required=True)
parser.add_argument("--k_folds", type=int, default=1)
parser.add_argument("--weight_llrgb", type=bool, default=False)
parser.add_argument("--llrgb_factor", type=float, default=0.5)
parser.add_argument("--holdout_frac", type=float, default=0.15)
parser.add_argument("--obs_len", type=float, default=82)
parser.add_argument("--init_lr", type=float, default=0.001)
parser.add_argument("--num_epochs", type=int, default=250)

config = parser.parse_args()

def train(model, model_optimizer, input_image, input_label, input_dnu, batch_idx, train_dataloader, config):

    model_optimizer.zero_grad()

    # Combined forward pass
    outputs, kl = model.probforward(input_image, input_dnu)  # Bayesian

    # Calculate loss and backpropagate
    loss = elbo(outputs, input_label, kl, get_beta(math.ceil(len(train_dataloader) / 32), beta_type="Blundell", batch_idx=batch_idx))  # 32 is batch size

    if config.weight_llrgb: # upweight for Kepler (5), but downweight for K2 (0.25)?
        if loss[input_dnu.squeeze(1) > 9].size()[0] < 1:
            loss = torch.mean(loss[input_dnu.squeeze(1) <= 9])
        elif loss[input_dnu.squeeze(1) <= 9].size()[0] < 1:
            loss = config.llrgb_factor*torch.mean(loss[input_dnu.squeeze(1) > 9])
        else:
            loss = torch.mean(loss[input_dnu.squeeze(1) <= 9]) + config.llrgb_factor*torch.mean(loss[input_dnu.squeeze(1) > 9])
    else:
        loss = torch.mean(loss)
    loss.backward()

    # Update parameters
    model_optimizer.step()
    pred = torch.max(outputs, dim=1)[1]
    correct = torch.sum(pred.eq(input_label)).item()
    total = input_label.numel()
    acc = 100. * correct / total
    return loss.item(), acc, pred

def validate(model, val_dataloader, config):
    model.eval()  # set to evaluate mode
    val_loss = 0
    val_batch_acc = 0
    val_batches = 0
    val_pred = []
    val_truth = []
    val_prob = []
    val_kic_array = []
    val_dnuz = []
    val_logits = []
    for batch_idy, val_data in enumerate(val_dataloader, 0):  # indices,scaled_indices, numax, teff, fe_h, age, tams_age

        val_image = val_data[0].cuda().float()
        val_label = val_data[1].cuda().long().squeeze(1)
        val_dnu_var = val_data[2].cuda().float()
        val_flag = val_data[3].cuda().float()
        val_kic = val_data[4].cuda().float()

        val_image = val_image[val_flag != 1]
        val_label = val_label[val_flag != 1]
        val_dnu_var = val_dnu_var[val_flag != 1]
        val_kic = val_kic[val_flag != 1]

        if len(val_label) < 1:
            continue

        with torch.no_grad():

            outputs, kl = model.probforward(val_image, val_dnu_var) # Bayesian
            val_batch_loss = elbo(outputs, val_label, kl,
                                  get_beta(math.ceil(len(val_dataloader)) / 32, beta_type="Blundell", batch_idx=batch_idy))

            if config.weight_llrgb: # upweight for Kepler (5), but downweight for K2 (0.25)?
                if val_batch_loss[val_dnu_var.squeeze(1) > 9].size()[0] < 1:
                    val_batch_loss = torch.mean(loss[val_dnu_var.squeeze(1) <= 9])
                elif val_batch_loss[val_dnu_var.squeeze(1) <= 9].size()[0] < 1:
                    val_batch_loss = config.llrgb_factor*torch.mean(loss[val_dnu_var.squeeze(1) > 9])
                else:
                    val_batch_loss = torch.mean(val_batch_loss[val_dnu_var.squeeze(1) <= 9]) + config.llrgb_factor*torch.mean(loss[val_dnu_var.squeeze(1) > 9])
            else:
                val_batch_loss = torch.mean(val_batch_loss)

            val_logits.append(outputs.data.cpu().numpy())

        pred = torch.max(outputs, dim=1)[1]
        correct = torch.sum(pred.eq(val_label)).item()
        total = val_label.numel()
        val_loss += val_batch_loss.mean().item()
        val_batch_acc += 100. * correct / total
        val_batches += 1
        val_pred.append(pred.data.cpu().numpy())
        val_dnuz.append(val_dnu_var.data.cpu().numpy())
        val_truth.append(val_label.data.cpu().numpy())
        val_prob.append(F.softmax(outputs, dim=1).data.cpu().numpy())
        val_kic_array.append(val_kic.data.cpu().numpy())
    val_kic_array = np.concatenate(val_kic_array, 0)
    val_dnuz = np.concatenate(val_dnuz, 0)
    val_logits = np.concatenate(val_logits, 0)

    return (val_loss / val_batches), (val_batch_acc / val_batches), np.concatenate(val_pred, axis=0), np.concatenate(
        val_truth, axis=0), np.concatenate(val_prob, axis=0), val_kic_array, val_logits

def train_model(config):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print('Using device:', device)

    model = Bayes_Classifier()
    model.to(device)

    folder_kic = []
    root_folder = config.root
    for dirpath, dirnames, filenames in os.walk(root_folder):
        for filex in filenames:
            folder_kic.append(int(filex.split('-')[0]))

    folder_kic = np.unique(folder_kic)

    train_kics, val_kics = train_test_split(folder_kic, test_size=config.holdout_frac, random_state=137)
    print('Number of Train Stars: ', len(train_kics))
    print('Number of Validation Stars: ', len(val_kics))
    train_gen = npz_generator(root=root_folder, select_kic=train_kics, obs_len=config.obs_len)
    train_dataloader = utils.DataLoader(train_gen, shuffle=True, batch_size=32)

    val_gen = npz_generator(root=root_folder, select_kic=val_kics, obs_len=config.obs_len)
    val_dataloader = utils.DataLoader(val_gen, shuffle=False, batch_size=32)

    learning_rate = config.init_lr
    model_optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = ReduceLROnPlateau(model_optimizer, mode='min', factor=0.5, patience=20, verbose=True, min_lr=1E-6)

    n_epochs = config.num_epochs
    best_loss = 1.e9
    model_checkpoint = True
    print(model_checkpoint)
    for epoch in range(1, n_epochs + 1):
        print('---------------------')
        print('Epoch: ', epoch)
        train_loss = 0
        train_batches = 0
        acc_cum = 0
        train_kic_array = []
        dnu_array = []
        pred_array = []
        model.train()  # set to training mode

        for batch_idx, data in tqdm(enumerate(train_dataloader, 0), total=len(train_dataloader), unit='batches'):

            image = data[0].cuda().float()
            label = data[1].cuda().long().squeeze(1)
            dnu = data[2].cuda().float()
            flag = data[3].cuda().float()
            train_kic = data[4].cuda().float()

            image = image[flag != 1] # flag 1 is bad stuff
            label = label[flag != 1]
            dnu = dnu[flag != 1]
            train_kic = train_kic[flag != 1]

            if len(label) < 2:
                print('Insufficient Batch!')
                continue

            loss, acc, predz = train(model, model_optimizer, image, label, dnu, batch_idx, train_dataloader, config)
            train_loss += loss  # Summing losses across all batches, so if you want the mean for EACH sample, divide by number of batches
            train_batches += 1
            acc_cum += acc
            train_kic_array.append(train_kic.data.cpu().numpy())
            pred_array.append(predz.data.cpu().numpy())
            dnu_array.append(dnu.data.cpu().numpy())
        train_loss = train_loss / train_batches
        train_acc = acc_cum / train_batches
        train_kic_array = np.concatenate(train_kic_array, 0)
        pred_array = np.concatenate(pred_array, axis = 0)
        dnu_array = np.concatenate(dnu_array, axis = 0)

        val_loss, val_acc, val_pred, val_truth, val_prob, val_kic, val_logits = validate(model, val_dataloader, config)
        scheduler.step(train_loss)  # reduce LR on loss plateau

        print('\n\nTrain Loss: ', train_loss)
        print('Train Acc: ', train_acc)

        val_loss = val_loss
        val_acc = val_acc
        print('Val Loss: ', val_loss)
        print('Val Acc: ', val_acc)

        for param_group in model_optimizer.param_groups:
            print('Current Learning Rate: ', param_group['lr'])

        if model_checkpoint:
            is_best = val_loss < best_loss
            best_loss = min(val_loss, best_loss)

            if is_best:
                filepath = 'checkpoint/Epoch%d-ACC:%.2f-Loss:%.3f' % (
                    epoch,val_acc, val_loss)
                print('Model saved to %s' % filepath)
                torch.save(model.state_dict(), filepath)
            else:
                print('No improvement over the best of %.4f' % best_loss)

def train_model_k_fold(config):

    folder_kic = []
    root_folder = config.root
    for dirpath, dirnames, filenames in os.walk(root_folder):
        for filex in filenames:
            folder_kic.append(int(filex.split('-')[0]))

    train_kics = np.copy(folder_kic)
    folder_kic = np.unique(folder_kic)

    label_data = pd.read_csv('Elsworth_Jie_2019_ID_Label.dat', header=0, delim_whitespace=True)
    label_kic = label_data['KIC'].values
    label_truth = label_data['Label'].values
    train_label = np.array([label_truth[np.where(label_kic == kicz)][0] for kicz in train_kics])

    cv = StratifiedKFold(n_splits=config.k_folds, shuffle=True, random_state=137)
    prob_over_folds = []
    truth_over_folds = []
    acc_over_folds = []
    rgb_precision_over_folds = []
    heb_precision_over_folds = []
    rgb_recall_over_folds = []
    heb_recall_over_folds = []
    fold_count = 0

    for train_idx, valid_idx in cv.split(train_kics, train_label):
        fold_count += 1
        print('Validation Fold: ', fold_count)

        model = Bayes_Classifier()
        model.cuda()
        torch.backends.cudnn.benchmark = True

        train_gen = npz_generator(root=root_folder, select_kic=train_kics[train_idx], obs_len=config.obs_len)
        train_dataloader = utils.DataLoader(train_gen, shuffle=True, batch_size=32)

        val_gen = npz_generator(root=root_folder, select_kic=train_kics[valid_idx], obs_len=config.obs_len)
        val_dataloader = utils.DataLoader(val_gen, shuffle=False, batch_size=32)

        learning_rate = config.init_lr
        model_optimizer = optim.Adam(model.parameters(), lr=learning_rate)
        scheduler = ReduceLROnPlateau(model_optimizer, mode='min', factor=0.5, patience=3, verbose=True, min_lr=1E-6)

        n_epochs = config.num_epochs
        for epoch in range(1, n_epochs + 1):
            print('---------------------')
            print('Epoch: ', epoch)
            print('Validation Fold: ', fold_count)
            train_loss = 0
            train_batches = 0
            acc_cum = 0

            model.train()  # set to training mode

            for batch_idx, data in tqdm(enumerate(train_dataloader, 0), total=len(train_dataloader), unit='batches'):

                image = data[0].cuda().float()
                label = data[1].cuda().long().squeeze(1)
                dnu = data[2].cuda().float()
                flag = data[3].cuda().float()

                image = image[flag != 1]
                label = label[flag != 1]
                dnu = dnu[flag != 1]

                if len(label) < 2:
                    print('Insufficient Batch!')
                    continue

                loss, acc, predz = train(model, model_optimizer, image, label, dnu, batch_idx, train_dataloader, config)
                train_loss += loss  # Summing losses across all batches, so if you want the mean for EACH sample, divide by number of batches
                train_batches += 1
                acc_cum += acc

            train_loss = train_loss / train_batches
            train_acc = acc_cum / train_batches

            val_loss, val_acc, val_pred, val_truth, val_prob, val_kic, val_logits = validate(model, val_dataloader, config)
            scheduler.step(train_loss)  # reduce LR on loss plateau

            print('\n\nTrain Loss: ', train_loss)
            print('Train Acc: ', train_acc)

            print('Val Loss: ', val_loss)
            print('Val Acc: ', val_acc)

            for param_group in model_optimizer.param_groups:
                print('Current Learning Rate: ', param_group['lr'])

        # After training, one last validation pass after training
        _, final_acc, final_pred, final_truth, final_prob, final_kic, final_logits = validate(model, val_dataloader, config)
        fold_precision_heb = precision_score(y_true=final_truth, y_pred=final_pred, pos_label=1)
        fold_precision_rgb = precision_score(y_true=final_truth, y_pred=final_pred, pos_label=0)
        fold_recall_heb = recall_score(y_true=final_truth, y_pred=final_pred, pos_label=1)
        fold_recall_rgb = recall_score(y_true=final_truth, y_pred=final_pred, pos_label=0)

        print('Accuracy for this fold: ', final_acc)
        print('HeB Precision for this fold: ', fold_precision_heb)
        print('HeB Recall for this fold: ', fold_recall_heb)
        print('RGB Precision for this fold: ', fold_precision_rgb)
        print('RGB Recall for this fold: ', fold_recall_rgb)

        prob_over_folds.append(final_prob)
        truth_over_folds.append(final_truth)
        acc_over_folds.append(final_acc)
        print('Accumulated Accuracy Thus Far: ', acc_over_folds)
        rgb_precision_over_folds.append(fold_precision_rgb)
        rgb_recall_over_folds.append(fold_recall_rgb)
        heb_precision_over_folds.append(fold_precision_heb)
        heb_recall_over_folds.append(fold_recall_heb)

    prob_over_folds = np.concatenate(prob_over_folds, axis=0)
    truth_over_folds = np.concatenate(truth_over_folds, axis=0)

    print(prob_over_folds.shape)
    print(truth_over_folds.shape)

    print('Validation Accuracy: %.3f +/- %.3f' % (np.mean(acc_over_folds), np.std(acc_over_folds)))
    print('Validation HeB Precision: %.3f +/- %.3f' % (
    np.mean(heb_precision_over_folds), np.std(heb_precision_over_folds)))
    print('Validation HeB Recall: %.3f +/- %.3f' % (np.mean(heb_recall_over_folds), np.std(heb_recall_over_folds)))
    print('Validation RGB Precision: %.3f +/- %.3f' % (
    np.mean(rgb_precision_over_folds), np.std(rgb_precision_over_folds)))
    print('Validation RGB Recall: %.3f +/- %.3f' % (np.mean(rgb_recall_over_folds), np.std(rgb_recall_over_folds)))

if __name__ == '__main__':
    if config.k_folds == 1:
        train_model(config)
    else:
        train_model_k_fold(config)


from flask import Flask, render_template
from predict import summarize
from time import time
from flask import jsonify, request
import os

app = Flask(__name__)

# @app.route('/')
# def home():
#     return render_template('home.html')

@app.route('/', methods=['GET', 'POST'])
def api_exec():
    if request.method == 'POST':
        text = request.values.get('text')
    else:
        text = request.args.get('text')
    if text == None:
        print('Summarizing Sample Description')
        text = "This app is developed by Nam and ZZ for NUS HackNRoll 2020. We are trying to build the most awesome app."

    response = {'status': 'ok'}
    try:
        start_time = time()
        summary = summarize(text)
        infer_time = time() - start_time
    except Exception as e:
        print("Error when predicting", e)
        response = {"status": "fail", "message": "Error when summarizing: {}".format(str(e))}

    if response["status"] == "ok":
        response.update({
            "text": text,
            "summary": summary,
            "infer_time": infer_time
        })

    return jsonify(response)

if __name__ == '__main__':
    app.run(debug=True)


import numpy as np
from sklearn import model_selection, svm
import pandas as pd
import pickle

bearings = ['Bearing1_1', 'Bearing1_2', 'Bearing1_6', 'Bearing1_7']
train_bearings = bearings[:-1]
test_bearings = bearings[-1]

df_train = pd.DataFrame()
for bearing in train_bearings:
    temp = pd.read_csv(bearing + '_final_features.csv')
    df_train = df_train.append(temp)

df_test = pd.DataFrame()
temp = pd.read_csv(test_bearings + '_final_features.csv')
df_test = df_test.append(temp)

X_train = np.array(df_train.drop(['RUL'], 1))
y_train = np.array(df_train['RUL'])
X_test = np.array(df_test.drop(['RUL'], 1))
y_test = np.array(df_test['RUL'])

clf = svm.SVC()
clf.fit(X_train, y_train)

with open('svm.pickle', 'wb') as f:
    pickle.dump(clf, f)

pickle_in = open('svm.pickle', 'rb')
clf = pickle.load(pickle_in)

accuracy = clf.score(X_test, y_test)
print('Accuracy: {}'.format(accuracy))


# Author: Nikol Chantzi

from setuptools import setup, find_packages

with open("requirements.txt", "r") as f:
    requirements = [file.strip() for file in f.readlines()]

setup(
        version="1.0",
        author="Nikol Chantzi",
        packages=find_packages(),
        install_requires=requirements,
        scripts=['bin/sassign']
)


from flask import Flask
from db_connect import db
from flask_migrate import Migrate
from key.db_secret import flask_pw, secret_key
from views import main_view, auth_view

migrate = Migrate()

def create_app():
    app = Flask(__name__)
    app.config['SQLALCHEMY_DATABASE_URI'] = f"mysql+pymysql://root:{flask_pw}@127.0.0.1:3306/elice_library"
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
    app.secret_key = secret_key
    # app.jinja_env.add_extension('jinja2.ext.loopcontrols')    # jinja2 for 문에서 break 가능

    db.init_app(app)
    migrate.init_app(app, db)  # Flask-Migrate 객체를 만들기 위해 플라스크 객체와 DB 객체 전달

    app.register_blueprint(main_view.board)
    app.register_blueprint(auth_view.join)

    return app

if __name__ == '__main__':
    create_app().run(debug=True)


import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import mean_absolute_error, accuracy_score
from sklearn.ensemble import RandomForestClassifier
import pickle
from tqdm import tqdm
from lime.lime_tabular import LimeTabularExplainer

use_names = ['rf_vcrim', 'rf_othsw', 'rf_attir', 'cs_objcs', 'cs_descr', 'cs_casng',
             'cs_lkout', 'rf_vcact', 'cs_cloth', 'cs_drgtr', 'ac_evasv', 'ac_assoc',
             'cs_furtv', 'rf_rfcmp', 'ac_cgdir', 'rf_verbl', 'cs_vcrim', 'cs_bulge',
             'cs_other', 'ac_incid', 'ac_time',  'rf_knowl', 'ac_stsnd', 'ac_other',
             'rf_furt',  'rf_bulg',  'sex',      'race']

cat_features = ['rf_vcrim', 'rf_othsw', 'rf_attir', 'cs_objcs', 'cs_descr', 'cs_casng',
                'cs_lkout', 'rf_vcact', 'cs_cloth', 'cs_drgtr', 'ac_evasv', 'ac_assoc',
                'cs_furtv', 'rf_rfcmp', 'ac_cgdir', 'rf_verbl', 'cs_vcrim', 'cs_bulge',
                'cs_other', 'ac_incid', 'ac_time',  'rf_knowl', 'ac_stsnd', 'ac_other',
                'rf_furt',  'rf_bulg',  'sex',      'race']

res_names2 = ['frisked',  'searched']
res_names = [ 'contrabn',  'pistol',   'riflshot', 'asltweap', 'knifcuti', 'machgun',  'othrweap']

#res_names = ['frisked', 'searched']

def load_data(columns, filename='data/stop_and_frisk/2012.csv'):
    sqf_2012 = pd.read_csv(filename)
    sqf_2012 = sqf_2012[columns]
    sqf_2012 = sqf_2012.dropna(how='any')
    return sqf_2012

def process_categorical_feature(df_sqf, cat_feat_name):
    cat_names = {}
    #store feature mapping
    for feature in cat_feat_name:
        le = LabelEncoder()
        le.fit(df_sqf.loc[:, feature])
        df_sqf.loc[:, feature] = le.transform(df_sqf.loc[:, feature])
        cat_names[feature] = le.classes_
    return cat_names
'''
# encode category features
cat_features = df_sqf[cat_feat_name]
oh_enc = OneHotEncoder()
oh_enc.fit(cat_features)
cat_feats_array = oh_enc.transform(cat_features).toarray()
cat_feats = pd.DataFrame(cat_feats_array, columns=oh_enc.get_feature_names())
'''

def calc_y(label):
    '''
    known model:
    S = 3 * 1_ps + 1 * 1_as + 1 * 1_bulge
    ps: cs_objcs
    as: ac_stsnd
    bulge: cs_bulge | rf_bulge
    '''
    y = []
    for i in label.index:
        curr_label = 0
        for res_name in res_names:
            val = 1 if label.loc[i, res_name] == 'Y' else 0
            curr_label |= val
        y.append(curr_label)
    return np.array(y)

def fit_model(X_train, y_train):
    lr = LogisticRegression()
    lr.fit(X_train, y_train)
    return lr

def data_preprocessing():
    try:
        X = pickle.load(open('data/stop_and_frisk/sqf_X_2012.pickle', 'rb'))
        y = pickle.load(open('data/stop_and_frisk/sqf_y_2012.pickle', 'rb'))
        categorical_names = pickle.load(open('data/stop_and_frisk/sqf_X_categorical_names.pickle', 'rb'))
    except:
        X = load_data(columns=use_names, filename='data/stop_and_frisk/2012.csv')
        y_raw = load_data(columns=res_names, filename='data/stop_and_frisk/2012.csv')
        X.replace(to_replace=['**'], value=np.nan, inplace=True)
        categorical_names = process_categorical_feature(X, cat_features)
        X.dropna(inplace=True)
        y = calc_y(y_raw)

        pickle.dump(X, open('data/stop_and_frisk/sqf_X_2012.pickle', 'wb'))
        pickle.dump(y, open('data/stop_and_frisk/sqf_y_2012.pickle', 'wb'))
        pickle.dump(categorical_names, open('data/stop_and_frisk/sqf_X_categorical_names.pickle', 'wb'))
    return X, y, categorical_names

def run_model(X, y, model='lr'):
    X = X.astype(float)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

    onehot_enc = OneHotEncoder(categorical_features=range(len(use_names)))
    onehot_enc.fit(X)
    encoded_X_train = onehot_enc.transform(X_train)

    if model == 'lr':
        clf = LogisticRegression()
    elif model == 'rf':
        clf = RandomForestClassifier(n_estimators=500)

    clf = fit_model(encoded_X_train, y_train)

    train_accuracy = clf.score(encoded_X_train, y_train)
    test_accuracy = accuracy_score(y_test, clf.predict(onehot_enc.transform(X_test)))
    print('train accuracy:', train_accuracy)
    print('test_accuracy:', test_accuracy)

    return clf, X_train.values, X_test.values, y_train, y_test, onehot_enc

#if __name__ == '__main__':

X, y, categorical_names = data_preprocessing()
blackbox_model, X_train, X_test, y_train, y_test, oh_enc = run_model(X, y, 'rf')
predict_fn = lambda x: blackbox_model.predict_proba(oh_enc.transform(x))
np.random.seed(1)
explainer = LimeTabularExplainer(X_train, class_names=['no_weapon', 'weapon'], feature_names=use_names,
                                 categorical_features=range(len(use_names)), categorical_names=categorical_names,
                                 mode='classification')
y_pred = predict_fn(X_test)
has_weapon_idx = []
for i in range(y_pred.shape[0]):
    if y_pred[i, 0] < y_pred[i, 1]:
        has_weapon_idx.append(i)

idx = has_weapon_idx[100]
idx = has_weapon_idx[350]
fi = dict()
for i in range(100):
    exp = explainer.explain_instance(X_test[idx], predict_fn, num_features=8, labels=[0, 1])
    #print('test sample {} prediction: {}, explanation for its predicted class:'.format(i, y_pred[idx]),
    #       exp.as_list(y_pred[idx, 0] < y_pred[idx, 1]))
    for item in exp.as_list(1):
        if item[0] not in fi:
            fi[item[0]] = []
        fi[item[0]].append(item[1])

for feature in fi.keys():
    temp = np.array(fi[feature])
    print('*' * 20)
    print('feature name:{}'.format(feature))
    print('count:{}'.format(temp.shape[0]))
    print('std:{}'.format(temp.std()))
    print('max:{}'.format(temp.max()))
    print('min:{}'.format(temp.min()))
    print('ave:{}'.format(temp.mean()))

for idx in range(0, 100, 2):
    exp = explainer.explain_instance(X_test[idx], predict_fn, num_features=8, labels=[0, 1])
    print('test sample {} prediction: {}, explanation for its predicted class:'.format(i, y_pred[has_weapon_idx[idx]]),
           exp.as_list(1))


from collections import deque
q = deque()
check = [0] * 100001
x, y = list(map(int, input().split(' ')))

def dfs(nx):
    q.append(nx)

    while q:
        x = q.popleft()
        if x == y:
            return check[x]
        for j in (x+1, x-1, 2*x):
            if 0 <= j < 100001 and check[j] == 0:
                check[j] = check[x] + 1
                q.append(j)

print(dfs(x))


import urllib
import json
class tokens(object):
   name = ""
   adress = ""
   decimal = ""

   # The class "constructor" - It's actually an initializer
   def __init__(self, name, adress, decimal):
       self.name = name
       self.adress = adress
       self.decimal = decimal

def getTokens():
    edtokens = "https://raw.githubusercontent.com/forkdelta/forkdelta.github.io/master/config/main.json"
    edTokenResponse = urllib.urlopen(edtokens)
    edTokendata = json.load(edTokenResponse)
    edTokendata = edTokendata['tokens']

    allTokens={}
    allTokens['0x0000000000000000000000000000000000000000'] = tokens('ETH', '0x0000000000000000000000000000000000000000', 18)
    for i in range(len(edTokendata)):
      allTokens[edTokendata[i]['addr']] = tokens(edTokendata[i]['name'], edTokendata[i]['addr'], edTokendata[i]['decimals'])
    return allTokens


import os
import clip
import pdb
import torch

classes = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
            'cat', 'chair', 'cow']
device = 'cuda' if torch.cuda.is_available() else 'cpu'

model, preprocess = clip.load('ViT-B/32', device)

text_inputs = torch.cat([clip.tokenize(f'a photo of a {c}') for c in classes]).to(device)
with torch.no_grad():
    text_features = model.encode_text(text_inputs)
    print(type(text_features))
    torch.save(text_features.cpu(), './Classes_features_10.pt')
    print(f'text features:{text_features.shape}')
    # pdb.set_trace()


# coding: utf-8
"""

@File    :size_only.py

@Author  : Yi Li

@E-mail  : li_yi@hust.edu.cn

Created on 2019/12/31

"""

from torchvision.models import resnet50, resnet18
import os
import torch.nn as nn
import logging

import torch
import torch.nn as nn
import torch.nn.functional as F

BN_MOMENTUM = 0.1
logger = logging.getLogger(__name__)

def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)

def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class GroupConv(nn.Module):

    def __init__(self, inplanes, planes, width=64, norm_layer=None):
        super(GroupConv, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes)
        self.bn3 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        out = self.relu(out)

        return out

class Fuse_module(nn.Module):
    def __init__(self, cfg, split='size', **kwargs):
        super(Fuse_module, self).__init__()

        output_channel = cfg.FUSE_MODULE.CHANNEL
        if split == 'size':
            self.group_conv = GroupConv(3, 64)
        elif split == 'wing':
            self.group_conv = GroupConv(2, 64)
        elif split == 'tail':
            self.group_conv = GroupConv(2, 64)

        # stem net
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1,
                               bias=False)
        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
        self.conv3 = nn.Conv2d(64, output_channel, kernel_size=3, stride=1, padding=1,
                               bias=False)
        self.bn3 = nn.BatchNorm2d(output_channel, momentum=BN_MOMENTUM)
        self.relu = nn.ReLU(inplace=True)

        logger.info('=> init fuse module weights from normal distribution')
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                nn.init.normal_(m.weight, std=0.001)
                for name, _ in m.named_parameters():
                    if name in ['bias']:
                        nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.ConvTranspose2d):
                nn.init.normal_(m.weight, std=0.001)
                for name, _ in m.named_parameters():
                    if name in ['bias']:
                        nn.init.constant_(m.bias, 0)

    def forward(self, class_output, landmark_input):

        class_output = torch.softmax(class_output, dim=1)

        class_output = torch.unsqueeze(class_output, dim=2)
        class_output = torch.unsqueeze(class_output, dim=3)
        # broadcasted_class_output = torch.zeros_like(landmark_input)
        broadcasted_class_output = torch.zeros((class_output.shape[0], class_output.shape[1],
                                                landmark_input.shape[2], landmark_input.shape[3])).cuda()
        broadcasted_class_output += class_output

        x = self.group_conv(broadcasted_class_output)
        x = F.adaptive_avg_pool2d(x, (1, 1))
        x = self.relu(x)
        global_attention = torch.sigmoid(x)

        y = self.conv1(landmark_input)
        y = self.bn1(y)
        y = self.conv2(y)
        y = self.bn2(y)

        out = y * global_attention.expand_as(y)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)
        out = self.relu(out)

        # fuse_out = torch.cat((out, landmark_input), dim=1)

        return out

def get_net(cfg, is_train, split='size', **kwargs):
    model = Fuse_module(cfg=cfg, split=split)
    return model


import numpy as np
import scipy.io as sc
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
import sys
from sklearn import preprocessing
from sklearn.neighbors import NearestNeighbors
# from sklearn.metrics import average_precision_score
# from sklearn.metrics import precision_recall_curve
from scipy.stats import wasserstein_distance
from scipy.spatial import distance
from sklearn.decomposition import PCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
#from metric_learn import LMNN, RCA
from munkres import Munkres
from sklearn.metrics.cluster import contingency_matrix
from numpy.linalg import matrix_rank

def describe_matrix(matrix,name):
	rank = matrix_rank(matrix)
	print(name,"has shape",matrix.shape,"and rank",rank,":\n",matrix,"\n")

def getAccuracy(ytrain, labels):
	m = Munkres()									# Hungarian algorithm
	cmat = contingency_matrix(ytrain, labels+1)
	print(cmat)
	print(cmat.shape)
	mapping = m.compute(cmat.max() - cmat)			# Maps old label to new labels

	print("Hungarian Algorithm Mapping:")
	print(mapping,'\n')

	di = {}
	for a, b in mapping: 							# Creates a dictionary according to the mapping
		di.setdefault(b, []).append(a)

	print("Hungarian Algorithm Mapping Dictionary:")
	print(di)

	new_labels = []
	for i in labels:
		if i in di:
			new_labels.append(di[i][0])
		else:
			new_labels.append(35)
	new_labels = np.asarray(new_labels).ravel() + 1

	print(len(new_labels))
	sys.exit()

	accuracy = 0

	for x in range(len(new_labels)):
		if new_labels[x] == ytrain[x]:
			accuracy +=1

	return (accuracy/320)

data = sc.loadmat('face(1).mat')

# N: number of images
# D: number of pixels
X = data['X']  # shape: [D x N]
y = data['l']  # shape: [1 x N]

# Number of images
D, N = X.shape

# Fix the random seed
np.random.seed(13)

# We have 52 classes and we want to achieve a split of 320 (classes 1-32) and 200 (classes 33-52)
Xtrain = []
Xtest = []
ytrain = []
ytest = []

for x in range(N):
    if x < 320:
        Xtrain.append(X[:, x])
        ytrain.append(y[:, x])
    else:
        Xtest.append(X[:, x])
        ytest.append(y[:, x])

Xtrain_norm = preprocessing.normalize(Xtrain, norm='l2', axis=1)
Xtest_norm = preprocessing.normalize(Xtest, norm='l2', axis=1)

from sklearn.cluster import KMeans
#from scipy.optimize import linear_sum_assignment

# Number of clusters
kmeans = KMeans(n_clusters=33, n_init=10, init='random')

# Fitting the input data
kmeans = kmeans.fit(Xtrain)
# Getting the cluster label
labels = kmeans.labels_
print(labels,'\n')

# Centroid values
centroids = kmeans.cluster_centers_

acc = getAccuracy(ytrain, labels)
print(acc)


import pytest
import os
from db_connect_test import FlexEnv

# The following will check to see if this is running in the Databricks environment.
# If it is, then we will use the Databricks Spark environment.
# If it is not, then we will use a DatabricksSession through Databricks-connect.
if "DATABRICKS_RUNTIME_VERSION" not in os.environ:
    pytest_plugins = ["spark.db_connect"]
else:
    pytest_plugins = ["spark.db_spark"]

@pytest.fixture(scope="function")
def flex_env(spark_env):
    yield FlexEnv(spark_env)