from kivy.graphics import svg
from selenium import webdriver
from time import sleep, perf_counter
import pandas as pd

driverPath = 'app/webdriver/chromedriver'

def createFile():
    tokenFile = f"app/assets/tokenData_test03.csv"
    dict = {
        'S.No.': [],
        'Token': [],
        'Total Tokens Locked': [],
        'TVL': [],
        'Volume 24h': [],
        '24h Fees': [],
        'equation1': [],
        'equation2': [],
    }
    df = pd.DataFrame(dict)
    df.to_csv(tokenFile, index=False)
    return tokenFile

def scrapeData(url):
    csvfile = createFile()
    driver = webdriver.Chrome(driverPath)
    driver.get(url)
    sleep(3)
    try:
        getPages = driver.find_element_by_xpath(
            '//*[@id="root"]/div/div[2]/div[3]/div/div[4]/div/div[13]')
        pagination = getPages.find_elements_by_tag_name('div')
        pages = len(pagination)
        # print(pagination[3].text)
        #print(pages)
        print('Page ready for extracting!!')

        alltokens = []
        for _ in range(pages):

            getAllpools = driver.find_element_by_xpath(
                '//*[@id="root"]/div/div[2]/div[3]/div/div[4]/div')

            anchors = getAllpools.find_elements_by_tag_name('a')
        #    print(anchors[0].text)
            for anchor in anchors:
                sNo = anchor.text.split('\n')[0]
                href = anchor.get_attribute('href')
                alltokens.append(f'{sNo},{href}')
        #        print(href)
        #         alltokens.append(href)
        #        print(alltokens)
            pagination[-1].click()

        # print(alltokens)
        print(len(alltokens))

        df = pd.read_csv(csvfile)
        for i in alltokens:
            splitdata = i.split(',')
            sNo = splitdata[0]
            url = splitdata[1]

            print(f'{sNo} : {url}')
            driver.get(url)
            sleep(2)

            upBlock = driver.find_element_by_xpath(
                '//*[@id="root"]/div/div[2]/div[3]/div[2]/div[2]/div[1]/div[2]'
            )
            pairs = upBlock.find_elements_by_tag_name('a')
            # print(pairs[1].text)

            leftBlock = driver.find_element_by_xpath(
                '//*[@id="root"]/div/div[2]/div[3]/div[2]/div[3]/div[1]/div')
            # print(leftBlock.text)
            # print(len(allDivs))
            lockedTokens = leftBlock.find_elements_by_class_name('sc-brqgnP')
            totalToken = lockedTokens[0].text.split('\n')
            first = lockedTokens[1].text.split('\n')
            second = lockedTokens[2].text.split('\n')
            third = lockedTokens[3].text.split('\n')
            print('Writing Data in CSV...')
            df.loc[sNo, 'S.No.'] = sNo,
            df.loc[sNo, 'Token'] = f'{totalToken[1]}-{totalToken[3]}',
            df.loc[sNo,
                   'Total Tokens Locked'] = f'{totalToken[2]}/{totalToken[4]}',
            df.loc[sNo, 'TVL'] = f'{first[1]}',
            df.loc[sNo, 'Volume 24h'] = f'{second[1]}',
            df.loc[sNo, '24h Fees'] = f'{third[1]}',
            df.loc[sNo, 'equation1']= f'{pairs[0].text}',
            df.loc[sNo, 'equation2'] = f'{pairs[1].text}',
            df.to_csv(csvfile, index=False)
            print(f'Data successfully written in CSV file!')
            print(
                f'{totalToken[0]}\n{totalToken[1]}:{totalToken[2]},{totalToken[3]}:{totalToken[4]}')
            print(f'First: {first}')
            print(f'Second: {second}')
            print(f'Third: {third}')

    except Exception as e:
        print(e)
    finally:
        driver.quit()


import pandas as pd
import folium
from folium.plugins import MarkerCluster

data = pd.read_csv("./arabica_data_new add variable.csv")
data.dataframeName = "arabica_data.csv"

world_map= folium.Map(tiles="cartodbpositron")
marker_cluster = MarkerCluster().add_to(world_map)

for i in range(len(data)):
        lat = data.iloc[i]['latitude']
        long = data.iloc[i]['longitude']
        radius=5
        popup_text = """Country : {}

                    Total Cup Points : {}
"""
        popup_text = popup_text.format(data.iloc[i]['Country.of.Origin'],
                                   data.iloc[i]['Total.Cup.Points']
                                   )
        folium.CircleMarker(location = [lat, long], radius=radius, popup= popup_text, fill =True).add_to(marker_cluster)

world_map.save("coffee_map.html")


import sys
from pathlib import Path
import pandas as pd
import pickle

directory = sys.argv[1]
pathlist = Path(directory).glob("*.tsv")

acc_data = {}
for pre_path in pathlist:
    path = str(pre_path)
    pre_name = path.split('.tsv')[0]
    name = pre_name.split('/')[-1]

    raw_data = []
    with open(path, 'r') as file:
        for cnt, line in enumerate(file):
            if cnt > 1 and line != '\n':
                line = line.replace('\n','')
                raw = line.split('\t')
                raw_data.append(raw)

    acc = [[dt[0],dt[1]] for dt in raw_data if int(dt[4]) == 1]
    acc_data[name] = acc

with open('./data/chip/acc_data.pickle', 'wb') as pickle_out:
    pickle.dump(acc_data, pickle_out)


# -*- coding: utf-8 -*-
"""
Created on Mon Jun  5 10:34:27 2023

@author: sooji
"""

from typing import Sequence, Tuple, Union

import numpy as np

import torch
import torch.nn as nn

from monai.utils import ensure_tuple_rep
from monai.networks.blocks.dynunet_block import UnetResBlock, get_conv_layer
from monai.networks.blocks.convolutions import Convolution

from networks.crossattention import CABlock
from networks.mobilevit import fold_proj, unfold_proj
from networks.crosstransformer import CrossTransformerBlock

#cross attention upsampling block
class CAUpBlock(nn.Module):
    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[Sequence[int], int],
        upsample_kernel_size: Union[Sequence[int], int],
        norm_name: Union[Tuple, str],
        patch_size: Union[Sequence[int], int],
        img_size: Union[Sequence[int], int],
        #transformer params
        transformer_dim: int,
        hidden_dim: int,
        num_heads: int,
        dropout_rate: float,
        act_name: Union[Tuple, str] = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
    ) -> None:
        """
        Args:
            spatial_dims: number of spatial dimensions.
            in_channels: number of input channels.
            out_channels: number of output channels.
            kernel_size: convolution kernel size.
            upsample_kernel_size: convolution kernel size for transposed convolution layers.
            norm_name: feature normalization type and arguments.
            res_block: bool argument to determine if residual block is used.

        """

        super().__init__()
        self.dimensions = 3
        self.img_size = ensure_tuple_rep(img_size, self.dimensions)
        self.patch_size = ensure_tuple_rep(patch_size, self.dimensions)
        self.transformer_dim = transformer_dim

        upsample_stride = upsample_kernel_size
        self.transp_conv = get_conv_layer(
            spatial_dims,
            in_channels,
            out_channels,
            kernel_size=upsample_kernel_size,
            stride=upsample_stride,
            conv_only=True,
            is_transposed=True,
        )

        self.conv_block = UnetResBlock(
            spatial_dims,
            transformer_dim,
            out_channels,
            kernel_size=kernel_size,
            stride=1,
            norm_name=norm_name,
        )

        self.transformer_proj_layer = nn.ModuleList()
        self.unfold_proj_layer = nn.ModuleList()
        for i in range(2):
            layer = Convolution(
             2,
             np.prod(patch_size),
             transformer_dim,
             strides=1,
             kernel_size=1,
             adn_ordering="ADN",
             act=act_name,
             norm=norm_name,
             dropout=dropout_rate,
             dropout_dim=1,
             dilation=1,
             bias=True,
             conv_only=False,
             padding=0,
            )
            self.unfold_proj_layer.append(layer)

            layer = Convolution(
             3,
             in_channels//2,
             transformer_dim,
             strides=1,
             kernel_size=1,
             adn_ordering="ADN",
             act=act_name,
             norm=norm_name,
             dropout=dropout_rate,
             dropout_dim=1,
             dilation=1,
             bias=True,
             conv_only=False,
             padding=0,
            )
            self.transformer_proj_layer.append(layer)

        self.fold_proj_layer = Convolution(
             2,
             transformer_dim,
             np.prod(patch_size),
             strides=1,
             kernel_size=1,
             adn_ordering="ADN",
             act=act_name,
             norm=norm_name,
             dropout=dropout_rate,
             dropout_dim=1,
             dilation=1,
             bias=True,
             conv_only=False,
             padding=0,
            )

        self.transformer = CrossTransformerBlock(transformer_dim, hidden_dim, num_heads, dropout_rate)

    def forward(self, inp, skip):
        # number of channels for skip should equals to out_channels
        out = self.transp_conv(inp)

        out = self.transformer_proj_layer[0](out) #similar to local processing on 3dmobilevit, proj from cin to E
        skip = self.transformer_proj_layer[1](skip) #similar to local processing on 3dmobilevit, proj from cin to E

        out = unfold_proj(out, self.patch_size, self.unfold_proj_layer[0])
        skip = unfold_proj(skip, self.patch_size, self.unfold_proj_layer[1])

        out = self.transformer(skip, out)

        out = fold_proj(out, self.img_size, self.patch_size, self.fold_proj_layer, self.transformer_dim)

        out = self.conv_block(out)
        return out


import matplotlib
matplotlib.use('GTKAgg')
from matplotlib.backends.backend_gtk import DialogLineprops

import numpy as np
import matplotlib.pyplot as plt
import sys

def dataset(filename,x,y,splittype):
    fr = open(filename)
    xdata = []
    ydata = []
    for line in fr.readlines():
        data = line.strip().split(splittype)
        if x > 0:
            xdata.append(float(data[x-1]))
        ydata.append(float(data[y-1]))
    if x == 0:
        xdata = range(len(ydata))
    return xdata,ydata

def f(t):
    s1 = np.cos(2*np.pi*t)
    e1 = np.exp(-t)
    return np.multiply(s1, e1)

def drew(xdata,ydata):
    t1 = np.array(xdata)
    t1_1 = np.array(ydata)
    t2 = np.array(xdata)
    t2_2 = np.array(ydata)

    fig, ax = plt.subplots()
    l1, = ax.plot(t1, t1_1, 'bo', label='line 1')
    l2, = ax.plot(t2, t2_2, 'k--', label='line 2')

    dlg = DialogLineprops([l1,l2])
    dlg.show()
    plt.show()

if __name__ == "__main__":
    if sys.argv[4] == "\\t":
        xdata,ydata = dataset(sys.argv[1],int(sys.argv[2]),int(sys.argv[3]),"\t")
    else:
        xdata,ydata = dataset(sys.argv[1],int(sys.argv[2]),int(sys.argv[3]),sys.argv[4])
    drew(xdata,ydata)


import rpy2.robjects as robjects
from rpy2.robjects.packages import importr
from rpy2.robjects import pandas2ri

import config

pandas2ri.activate()

utils = importr('utils')
utils.chooseCRANmirror(ind=1)

try:
    forecast = importr('forecast')
    ts = robjects.r('ts')
    auto_order = robjects.r('arimaorder')
except Exception as e:
    utils.install_packages("forecast")

def auto_calibration_arima(df, col, freq=12):
    try:
        r_df = ts(df[col].values, frequency=freq)
        fit = forecast.auto_arima(r_df)
        return [x for x in auto_order(fit)]
    except Exception as e:
        print(e)

return None


import numpy as np
from PIL import Image

def apply_alpha(image, alpha):
    """Apply alpha to entire image, taking into account any existing alpha mask."""
    im = image.copy()
    if 'A' not in image.getbands():
        im.putalpha(int(round(alpha * 255)))
        return im
    elif image.mode == 'RGBA':
        a = bytes(bytearray((np.array(bytearray(image.split()[3].tobytes())) * alpha + 0.5).astype('int8')))
        im.putalpha(Image.frombytes('L', image.size, a))
        return im
    else:
        raise ValueError('incorrect image mode (was %s)' % image.mode)


import os
import time, pickle
import argparse
import numpy as np
import torch
import torch.nn as nn
from torch.cuda import amp
import torch.nn.functional as F
from torch.utils.data import DataLoader
from utils.dataloader import POI_Loader, MF_Loader, RNN_Dataset, STP_Dataloader
from utils.metric import *
from model.poiformer import TADGAT
from model.baseline import *
from model.hst_lstm import HSTLSTM_CLS
from model.stgcn import STGCN_CLS
from model.stp_udgat import STP_UDGAT

class FixNoiseStrategy():

    def __init__(self, hidden_size):
        self.hidden_size = hidden_size
        mu = 0
        sd = 1/self.hidden_size
        self.h0 = torch.randn(self.hidden_size, requires_grad=False) * sd + mu

    def on_init(self, user_len, device):
        hs = []
        for i in range(user_len):
            hs.append(self.h0)
        if not st_flag:
            return torch.stack(hs, dim=0).view(1, user_len, self.hidden_size).to(device)
        else:
            return torch.stack(hs, dim=0).to(device)

    def on_reset(self):
        return self.h0

class LstmStrategy():

    def __init__(self, hidden_size):
        self.hidden_size=hidden_size
        self.h_strategy = FixNoiseStrategy(hidden_size)
        self.c_strategy = FixNoiseStrategy(hidden_size)

    def on_init(self, user_len, device):
        h = self.h_strategy.on_init(user_len, device)
        c = self.c_strategy.on_init(user_len, device)
        return (h,c)

    def on_reset(self):
        h = self.h_strategy.on_reset()
        c = self.c_strategy.on_reset()
        return (h,c)

# class CELoss(nn.Module):
#     def __init__(self, label_smooth=None):
#         super().__init__()
#         self.label_smooth = label_smooth

#     def forward(self, pred, target):
#         eps = 1e-12
#         class_num = pred.size(-1)
#         if self.label_smooth is not None:
#             logprobs = F.log_softmax(pred, dim=1)
#             target = F.one_hot(target, pred.size(-1)).to(device)
#             target = torch.clamp(target.float(), min=self.label_smooth/(class_num-1), max=1.0-self.label_smooth)
#             loss = -1*torch.sum(target*logprobs, 1)

#         else:
#             loss = -1.*pred.gather(1, target.unsqueeze(-1)) + torch.log(torch.exp(pred+eps).sum(dim=1))
#         return loss.mean()

parser = argparse.ArgumentParser()
parser.add_argument('--mode', type=str, default='train', help='running mode, choose between [train, test]')
parser.add_argument('--dataset', type=str, help='dataset chosen for training')
parser.add_argument('--model', type=str, default='poiformer', help='model under verification in this run')
parser.add_argument('--checkpoint_dir', '-cd', type=str, default='checkpoint', help='path of checkpoint file')
parser.add_argument('--gpu_id', type=int, default=-1, help='gpu id chosen to run train/test')
parser.add_argument('--layer_num', type=int, default=1, help='number of stacked model layer')
parser.add_argument('--max_len', type=int, default=50, help='maximum length of item/user sequence')
parser.add_argument('--ubias_num', type=int, default=180, help='number of user ego-nets attention bias')
parser.add_argument('--ibias_num', type=int, default=360, help='number of item ego-nets attention bias')
parser.add_argument('--input_dim', type=int, default=32, help='dimension of input feature')
parser.add_argument('--hidden_channel', type=int, default=32, help='dimension of hidden layer')
parser.add_argument('--weight_decay', type=float, default=0, help='hyperparameter of weight decay')
parser.add_argument('--lr', type=float, default=0.01, help='learning rate')
parser.add_argument('--batch_size', type=int, default=256, help='batch size')
parser.add_argument('--test_bs', type=int, default=1024, help='batch size for test and validation')
parser.add_argument('--neg_sample_num', type=int, default=19, help='number of negative samples during training/validation')
parser.add_argument('--eval_per_n', type=int, default=1000, help='evaluate per n steps')
parser.add_argument('--epoch_num', type=int, default=30, help='training epoch number')
parser.add_argument('--ssl_temp', type=float, default=0.5, help='temperature hyperparameter for MTL')
parser.add_argument('--mtl_coef_1', type=float, default=0.5, help='coefficient of auxiliary task 1')
parser.add_argument('--mtl_coef_2', type=float, default=0.5, help='coefficient of auxiliary task 2')

args, _ = parser.parse_known_args()
checkpoint_path = '%s/%s.pth' % (args.checkpoint_dir, args.model)

# use_aux_init = args.add_aux_init == 1
# use_gt = args.add_gt == 1

# Print current experimental setting
print("Run %s" % time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time())))
print("Current process ID: %d" % os.getpid())
print("Current Model: %s" % args.model)
print("Current Dataset: %s" % args.dataset)
print("Running Mode: %s" % args.mode)
print("------HyperParameter Settings------")
print("Number of Stacking Layer: %d" % args.layer_num)
print("In Channel: %d" % args.input_dim)
print("Hidden Channel: %d" % args.hidden_channel)
# print("Use Graph Embedding for POI Init: %s" % str(use_aux_init))
# print('Use Graph Transformer Layer: %s' % str(use_gt))
if args.mode == 'train':
    print("Maximum Sequence Length: %d" % args.max_len)
    print("Learning Rate: %f" % args.lr)
    print("Weight Decay: %f" % args.weight_decay)
    print("Batch Size: %d" % args.batch_size)
    print("Epoch Number: %d" % args.epoch_num)
    print("Evaluate per %d step" % args.eval_per_n)
    print("Multi-task coefficient: %.2f, %.2f" % (args.mtl_coef_1, args.mtl_coef_2))
    print("\n\n")
else:
    print("Batch Size: %d" % args.batch_size)
    print("\n\n")

if args.gpu_id != -1:
    device = torch.device('cuda:%d' % args.gpu_id if torch.cuda.is_available() else 'cpu')
else:
    device = torch.device('cpu')
enable_amp = True if "cuda" in device.type else False
enable_amp = False

# TODO: ucenter need 80 day slot, icenter need 150 day slot for gowalla
if 'TADGAT' in args.model:
    if args.mode == 'train':
        train_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, dataset=args.dataset, mode='train', batch_size=args.batch_size, shuffle=True)    #, num_workers=1)
        print('train data loaded')
        val_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, dataset=args.dataset, mode='val', batch_size=args.batch_size, shuffle=False)   #, num_workers=1)
        print('valid data loaded')
        user_num, item_num = train_loader.user_num, train_loader.poi_num
        # test_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, mode='test', batch_size=args.batch_size, shuffle=False)  #, num_workers=1)
    else:
        test_loader = POI_Loader(args.max_len, args.layer_num, args.ubias_num, args.ibias_num, args.neg_sample_num, dataset=args.dataset, mode='test', batch_size=args.batch_size, shuffle=True)  #, num_workers=1)
        user_num, item_num = test_loader.user_num, test_loader.poi_num
        print('test data loaded')
elif args.model == 'MF' or args.model == 'LightGCN':
    if args.mode == 'train':
        train_loader = MF_Loader('train', batch_size=args.batch_size, dataset=args.dataset, shuffle=True)
        val_loader = MF_Loader('val', batch_size=args.batch_size, dataset=args.dataset, shuffle=False)
        user_num, item_num = train_loader.user_num, train_loader.poi_num
        edge_index = pickle.load(open('gowalla_processed/train_lgcn_eindex.pkl', 'rb'))
        edge_index = torch.tensor(edge_index, dtype=torch.int64).to(device)
    else:
        test_loader = MF_Loader('test', batch_size=args.batch_size, dataset=args.dataset, shuffle=False)
        user_num, item_num = test_loader.user_num, test_loader.poi_num
        edge_index = pickle.load(open('gowalla_processed/val_lgcn_eindex.pkl', 'rb'))
        edge_index = torch.tensor(edge_index, dtype=torch.int64).to(device)
    print('Data Loaded')
elif 'LSTM' in args.model or args.model == 'STGCN':
    st_flag = 'ST' in args.model and args.model != 'LSTM'
    if args.mode == 'train':
        train_dataset = RNN_Dataset(15, args.batch_size, st_flag, args.dataset, 'train')
        val_dataset = RNN_Dataset(15, args.batch_size, st_flag, args.dataset, 'val')
        train_loader, val_loader = DataLoader(train_dataset, batch_size=1), DataLoader(val_dataset, batch_size=1)
        user_num, item_num = train_dataset.user_num, train_dataset.poi_num
    else:
        test_dataset = RNN_Dataset(15, args.batch_size, st_flag, args.dataset, 'test')
        test_loader = DataLoader(test_dataset, batch_size=1)
        user_num, item_num = test_dataset.user_num, test_dataset.poi_num
elif args.model == 'STP_UDGAT':
    if args.mode == 'train':
        train_loader = STP_Dataloader('train', args.dataset, batch_size=args.batch_size, shuffle=True)
        val_loader = STP_Dataloader('val', args.dataset, batch_size=args.batch_size, shuffle=False)
        user_num, item_num = train_loader.user_num, train_loader.poi_num
    else:
        test_loader = STP_Dataloader('test', args.dataset, batch_size=args.batch_size, shuffle=False)
        user_num, item_num = test_loader.user_num, test_loader.poi_num

print(user_num, item_num)
neg_sample_num = args.neg_sample_num
if 'TADGAT' in args.model:
    model = TADGAT(args.input_dim, args.hidden_channel, args.layer_num, user_num, item_num, args.ubias_num, args.ibias_num, ssl_temp=args.ssl_temp, aux_feat=None)
elif args.model == 'MF':
    model = MF(args.input_dim, user_num, item_num)
elif 'LSTM' in args.model:
    if st_flag:
        model = HSTLSTM_CLS(args.input_dim, item_num, 168, 300)
    else:
        print('LSTM')
        model = LSTM_Basic(args.input_dim, item_num)
elif args.model == 'STGCN':
    model = STGCN_CLS(args.input_dim, item_num)
elif args.model == 'STP_UDGAT':
    model = STP_UDGAT(args.input_dim, args.hidden_channel, user_num, item_num)
elif args.model == 'LightGCN':
    model = LightGCN(args.input_dim, args.layer_num, user_num, item_num)
model = model.to(device)

if args.mode == 'train':
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.4)
    scaler = amp.GradScaler(enabled=enable_amp)

    # criterion = CELoss(label_smooth=0.05)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)

def train(best_ndcg, best_hr):
    model.train()
    step = 0
    train_loader.soc_edges = train_loader.soc_edges.to(device)
    for all_out in train_loader:
        optimizer.zero_grad()
        adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, sampled_edges = all_out
        # adjs, node_idx, soc_adjs, soc_node_idx, orig_seqs, time_diff, seq_lens, center_nid, soc_cnid, sampled_edges = all_out
        adjs = adjs.to(device)
        node_idx = node_idx.to(device)
        center_nid = center_nid.to(device)

        # soc_adjs = soc_adjs.to(device)
        # soc_node_idx = soc_node_idx.to(device)

        # orig_seqs = orig_seqs.to(device)
        # time_diff = time_diff.to(device)
        # seq_lens = seq_lens.to(device)

        # sampled_idx = sampled_edges[:, 1:]
        # sampled_idx = sampled_idx.to(device)
        # label = torch.zeros(sampled_edges.size(0), dtype=torch.int64)
        label = sampled_edges[:, 1]
        # sampled_idx = torch.arange(item_num) + user_num
        # sampled_idx = sampled_idx.repeat(label.size(0)).reshape(-1, item_num)
        # sampled_idx = sampled_idx.to(device)
        label = label.to(device)

        # lat_emb_idx, lon_emb_idx = lat_emb_idx.to(device), lon_emb_idx.to(device)

        with amp.autocast(enabled=enable_amp):
            logits = model(adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, device, train_loader.soc_edges)
            # logits = model(adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, device, soc_adjs, soc_node_idx, soc_cnid)

            # Code for ranking-based loss
            # loss = 0.
            # for k in range(logits.size(0)):
            #     pos_logit = logits[k, label[k]]
            #     neg_logit = torch.cat([logits[k, :label[k]], logits[k, (label[k]+1):]])
            #     # pos_loss = F.logsigmoid(pos_logit - pref_bound[k])
            #     # neg_loss = F.logsigmoid(pref_bound[k] - neg_logit).mean()
            #     pos_loss = F.logsigmoid(pos_logit)
            #     neg_loss = F.logsigmoid(-neg_logit).mean()
            #     loss += - pos_loss - neg_loss
            # loss /= logits.size(0)
            loss = criterion(logits, label) # + args.mtl_coef_1 * mtl_loss_1
            # new_end = time.time()
            # print(new_end - end)
        if np.isnan(loss.detach().cpu()):
            print("NaN detected!")
            exit(1)
        label = label.cpu()
        if step % 200 == 0:
            label = label.numpy()
            logits = logits.detach()
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            # hr_1 += hit_rate(sorted_idx, label, k=1)
            # hr_5 += hit_rate(sorted_idx, label, k=5)
            # hr_10 += hit_rate(sorted_idx, label, k=10)
            hr_1 = hit_rate(sorted_idx, label, k=1) / args.batch_size
            hr_5 = hit_rate(sorted_idx, label, k=5) / args.batch_size
            hr_10 = hit_rate(sorted_idx, label, k=10) / args.batch_size
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        # start = time.time()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        # end = time.time()
        # print(end - start)

        # if step % 20 == 0:
        #     hr_1 /= len(all_out) ; hr_5 /= len(all_out) ; hr_10 /= len(all_out)
        #     print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))

        # if step != 0 and step % args.eval_per_n == 0:
        #     start = time.time()
        #     hr_1, hr_5, hr_10, mrr, ndcg = test('val')
        #     end = time.time()
        #     print("time consumption: %.4f" % (end-start))
        #     print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG@10: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
        #     if ndcg > best_ndcg:
        #         best_ndcg = ndcg
        #         print("New best model saved")
        #         torch.save(model.state_dict(), checkpoint_path)
        step += 1
        # end = time.time()
        # print(end - all_start)
    scheduler.step()

    # if step % args.eval_per_n > 20:
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if mrr > best_ndcg: # or hr_10 > best_hr:
        # if ndcg > best_ndcg:
        #     best_ndcg = ndcg
        # else:
        #     best_hr = hr_10
        best_ndcg = mrr
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

def train_stp(best_ndcg, best_hr):
    model.train()
    step = 0
    train_loader.u_eindex = train_loader.u_eindex.to(device)
    for out in train_loader:
        optimizer.zero_grad()
        s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, anchors, label = out
        s_adj, t_adj, f_adj = s_adj.to(device), t_adj.to(device), f_adj.to(device)
        s_rw_adj, t_rw_adj, f_rw_adj = s_rw_adj.to(device), t_rw_adj.to(device), f_rw_adj.to(device)
        pp_adj = pp_adj.to(device)
        anchors = anchors.to(device)
        label = label.to(device)
        with amp.autocast(enabled=enable_amp):
            logits = model(s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, train_loader.u_eindex, anchors)
            loss = criterion(logits, label)
        label = label.cpu()
        if step % 200 == 0:
            label = label.numpy()
            logits = logits.detach()
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            hr_1 = hit_rate(sorted_idx, label, k=1) / args.batch_size
            hr_5 = hit_rate(sorted_idx, label, k=5) / args.batch_size
            hr_10 = hit_rate(sorted_idx, label, k=10) / args.batch_size
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        step += 1
    scheduler.step()
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test_stp('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if ndcg > best_ndcg or hr_10 > best_hr:
        if ndcg > best_ndcg:
            best_ndcg = ndcg
        else:
            best_hr = hr_10
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

def train_mf(best_ndcg, best_hr):
    model.train()
    step = 0
    for out in train_loader:
        optimizer.zero_grad()
        center_uid, seq = out
        with amp.autocast(enabled=enable_amp):
            if args.model == 'LightGCN':
                loss, logits = model(edge_index, center_uid, seq)
            else:
                loss, logits = model(center_uid, seq)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        if step % 5 == 0:
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            hr_1 = top_k_acc(sorted_idx, seq, k=1) / args.batch_size
            hr_5 = top_k_acc(sorted_idx, seq, k=5) / args.batch_size
            hr_10 = top_k_acc(sorted_idx, seq, k=10) / args.batch_size
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        step += 1
    scheduler.step()
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test_mf('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if hr_10 > best_hr:
        if ndcg > best_ndcg:
            best_ndcg = ndcg
        else:
            best_hr = hr_10
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

def train_rnn(best_ndcg, best_hr):
    model.train()
    step = 0
    h0_strategy = LstmStrategy(args.input_dim)
    h = h0_strategy.on_init(args.batch_size, device)
    for out in train_loader:
        optimizer.zero_grad()
        with torch.autograd.set_detect_anomaly(True):
            if not st_flag:
                seqs, lbls, reset_h = out
                seqs, lbls = seqs.to(device).squeeze(), lbls.to(device).squeeze()
            else:
                seqs, lbls, t_l, t_u, t_coef, s_l, s_u, s_coef, reset_h = out
                seqs, lbls = seqs.to(device).squeeze().T, lbls.to(device).squeeze().T
                # Note: the unsqueeze is for STGCN only
                t_l, t_u, t_coef = t_l.to(device).squeeze(), t_u.to(device).squeeze(), t_coef.to(device).squeeze()
                s_l, s_u, s_coef = s_l.to(device).squeeze(), s_u.to(device).squeeze(), s_coef.to(device).squeeze()
            lbls = lbls.reshape(-1)
            for j, reset in enumerate(reset_h):
                if reset:
                    hc = h0_strategy.on_reset()
                    if st_flag:
                        h[0][j] = hc[0]
                        h[1][j] = hc[1]
                    else:
                        h[0][0, j] = hc[0]
                        h[1][0, j] = hc[1]
            with amp.autocast(enabled=enable_amp):
                if not st_flag:
                    logits, h = model(seqs, h)
                else:
                    if args.model == 'ST_LSTM':
                        logits, h = model(seqs, t_l, t_u, s_l, s_u, t_coef, s_coef, h)
                    elif args.model == 'STGCN':
                        logits, h = model(seqs, t_l, s_l, h)
                loss = criterion(logits, lbls)
            scaler.scale(loss).backward(retain_graph=True)
            scaler.step(optimizer)
            scaler.update()
        h = (h[0].detach(), h[1].detach())
        if step % 100 == 0:
            sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
            hr_1 = hit_rate(sorted_idx, lbls, k=1) / logits.size(0)
            hr_5 = hit_rate(sorted_idx, lbls, k=5) / logits.size(0)
            hr_10 = hit_rate(sorted_idx, lbls, k=10) / logits.size(0)
            print("step %d : loss %.4f, HR@1 %.4f, HR@5 %.4f, HR@10 %.4f" % (step, float(loss), hr_1, hr_5, hr_10))
        step += 1
    scheduler.step()
    start = time.time()
    hr_1, hr_5, hr_10, mrr, ndcg = test_rnn('val')
    end = time.time()
    print("time consumption: %.4f" % (end - start))
    print("Test result on validation set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    if ndcg > best_ndcg or hr_10 > best_hr:
        if ndcg > best_ndcg:
            best_ndcg = ndcg
        else:
            best_hr = hr_10
        print("New best model saved")
        torch.save(model.state_dict(), checkpoint_path)
    return best_ndcg, best_hr

@torch.no_grad()
def test_mf(mode):
    print('Start Test...')
    model.eval()
    hr_1_tot, hr_5_tot, hr_10_tot = [], [], []
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    mrr = []
    for out in cur_loader:
        center_uid, seq = out
        if args.model == 'LightGCN':
            logits = model(edge_index, center_uid, seq)
        else:
            logits = model(center_uid, seq)
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        hr_1 = top_k_acc(sorted_idx, seq, k=1) / args.batch_size
        hr_5 = top_k_acc(sorted_idx, seq, k=5) / args.batch_size
        hr_10 = top_k_acc(sorted_idx, seq, k=10) / args.batch_size
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)

        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            for j in range(len(seq[i])):
                y_true[i, seq[i][j]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))
    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0

@torch.no_grad()
def test_rnn(mode):
    print('Start Test...')
    hr_1_tot, hr_5_tot, hr_10_tot = [], [], []
    mrr = []
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    h0_strategy = LstmStrategy(args.input_dim)
    h = h0_strategy.on_init(args.batch_size, device)
    for out in cur_loader:
        if not st_flag:
            seqs, lbls, reset_h = out
            seqs = seqs.to(device)
            seqs, lbls = seqs.squeeze(), lbls.squeeze()
        else:
            seqs, lbls, t_l, t_u, t_coef, s_l, s_u, s_coef, reset_h = out
            seqs, lbls = seqs.to(device).squeeze().T, lbls.squeeze().T
            t_l, t_u, t_coef = t_l.to(device).squeeze(), t_u.to(device).squeeze(), t_coef.to(device).squeeze()
            s_l, s_u, s_coef = s_l.to(device).squeeze(), s_u.to(device).squeeze(), s_coef.to(device).squeeze()
        lbls = lbls.reshape(-1)
        for j, reset in enumerate(reset_h):
            if reset:
                hc = h0_strategy.on_reset()
                if st_flag:
                    h[0][j] = hc[0]
                    h[1][j] = hc[1]
                else:
                    h[0][0, j] = hc[0]
                    h[1][0, j] = hc[1]
        if st_flag:
            if args.model == 'STGCN':
                logits, h = model(seqs, t_l, s_l, h)
            else:
                logits, h = model(seqs, t_l, t_u, s_l, s_u, t_coef, s_coef, h)
        else:
            logits, h = model(seqs, h)
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        label = lbls.numpy()
        hr_1 = hit_rate(sorted_idx, label, k=1) / logits.size(0)
        hr_5 = hit_rate(sorted_idx, label, k=5) / logits.size(0)
        hr_10 = hit_rate(sorted_idx, label, k=10) / logits.size(0)
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)

        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            y_true[i, label[i]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))
    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0

@torch.no_grad()
def test_stp(mode):
    print('Start Test...')
    model.eval()
    hr_1_tot, hr_5_tot, hr_10_tot = [], [], []
    mrr = []
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    cur_loader.u_eindex = cur_loader.u_eindex.to(device)
    for out in cur_loader:
        s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, anchors, label = out
        s_adj, t_adj, f_adj = s_adj.to(device), t_adj.to(device), f_adj.to(device)
        s_rw_adj, t_rw_adj, f_rw_adj = s_rw_adj.to(device), t_rw_adj.to(device), f_rw_adj.to(device)
        pp_adj = pp_adj.to(device)
        anchors = anchors.to(device)
        logits = model(s_adj, t_adj, f_adj, s_rw_adj, t_rw_adj, f_rw_adj, pp_adj, cur_loader.u_eindex, anchors)
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        label = label.numpy()
        hr_1 = hit_rate(sorted_idx, label, k=1) / logits.size(0)
        hr_5 = hit_rate(sorted_idx, label, k=5) / logits.size(0)
        hr_10 = hit_rate(sorted_idx, label, k=10) / logits.size(0)
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)
        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            y_true[i, label[i]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))

    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0

@torch.no_grad()
def test(mode):
    print('Start Test...')
    model.eval()
    hr_1_tot, hr_5_tot, hr_10_tot, mrr, ndcg_tot = [], [], [], [], []

    # TODO: uncomment the checkpoint loading operation during real test
    if mode == 'test':
        model.load_state_dict(torch.load(checkpoint_path))
    cur_loader = val_loader if mode == 'val' else test_loader
    cur_loader.soc_edges = cur_loader.soc_edges.to(device)
    # step = 0
    for n, all_out in enumerate(cur_loader):
        # adjs, node_idx, center_nid, sampled_edges = all_out
        # adjs, node_idx, soc_adjs, soc_node_idx, orig_seqs, time_diff, seq_lens, center_nid, soc_cnid, sampled_edges = all_out
        adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, sampled_edges = all_out

        # if mode == 'val':
        #     sampled_idx = sampled_edges[:, 1:]
        #     label = torch.zeros(sampled_edges.size(0), dtype=torch.int64)
        # else:
        label = sampled_edges[:, 1]
        # sampled_idx = torch.arange(item_num) + user_num
        # sampled_idx = sampled_idx.repeat(label.size(0)).reshape(-1, item_num)
        # adjs = [adj.to(device) for adj in adjs]
        adjs = adjs.to(device)
        node_idx = node_idx.to(device)
        center_nid = center_nid.to(device)

        # soc_adjs = soc_adjs.to(device)
        # soc_node_idx = soc_node_idx.to(device)

        # orig_seqs = orig_seqs.to(device)
        # time_diff = time_diff.to(device)
        # seq_lens = seq_lens.to(device)

        # logits = model(adjs, node_idx, center_nid, device)
        logits = model(adjs, node_idx, orig_seqs, time_diff, seq_lens, center_nid, device, cur_loader.soc_edges)# , soc_adjs, soc_node_idx, soc_cnid)

        # full_logits.append(logits)
        # full_labels.append(label)
        # if mode == 'test':
        #     full_logits.append(logits)

        # logits = torch.cat(full_logits, dim=0)
        # label = torch.cat(full_labels)
        # random_noise = torch.rand(logits.numpy().shape) * 1e-7
        # logits = logits + random_noise
        sorted_idx = torch.argsort(logits, dim=-1, descending=True).cpu().numpy()
        label = label.numpy()
        # if mode == 'test':
        hr_1 = hit_rate(sorted_idx, label, k=1) / logits.size(0)
        hr_5 = hit_rate(sorted_idx, label, k=5) / logits.size(0)
        hr_10 = hit_rate(sorted_idx, label, k=10) / logits.size(0)
        # else:
        #     hr_1 = hit_rate(sorted_idx, label, k=1)
        #     hr_5 = hit_rate(sorted_idx, label, k=5)
        #     hr_10 = hit_rate(sorted_idx, label, k=10)
        # hr_50 = hit_rate(sorted_idx, k=50)
        hr_1_tot.append(hr_1)
        hr_5_tot.append(hr_5)
        hr_10_tot.append(hr_10)
        y_true = np.zeros_like(sorted_idx)
        for i in range(sorted_idx.shape[0]):
            y_true[i, label[i]] = 1
        for i in range(sorted_idx.shape[0]):
            y_true[i] = y_true[i][sorted_idx[i]]

        rr_score = y_true / (np.arange(np.shape(y_true)[1]) + 1)
        mrr.append(np.sum(rr_score) / np.sum(y_true))

        # y_true = np.zeros_like(sorted_idx)
        # for i in range(len(label)):
        #     y_true[i, label[i]] = 1

        # mrr_true = np.take(y_true, sorted_idx)
        # rr_score = mrr_true / (np.arange(np.shape(mrr_true)[1]) + 1)
        # mrr += np.sum(np.sum(rr_score, axis=1) / np.sum(mrr_true, axis=1))

        # ndcg = ndcg_score(y_true, sorted_idx, k=10)
        # if np.nan not in ndcg:
        #     # ndcg_tot += np.mean(ndcg)
        #     ndcg_tot += np.sum(ndcg)
        # # del full_logits, full_labels
        # # full_logits, full_labels = [], []

        # if n % 100 == 0 and n != 0:
        #     print(n)
        # # step += 1

    return np.mean(hr_1_tot), np.mean(hr_5_tot), np.mean(hr_10_tot), np.mean(mrr), 0 # mrr / step, ndcg_tot / step

np.random.seed(1)
torch.manual_seed(1)
if args.mode == 'train':
    ndcg = 0.
    hr_10 = 0.
    # start = time.time()
    # hr_1, hr_5, hr_10, mrr, ndcg = test_stp('val')
    # end = time.time()
    # print("time consumption: %.4f" % (end-start))
    # print("Test result on val set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG@10: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))
    for epoch in range(1, args.epoch_num + 1):
        print("Currently epoch %d:" % epoch)
        print("Learning rate: %f" % scheduler.get_lr()[0])
        if 'TADGAT' in args.model:
            ndcg, hr_10 = train(ndcg, hr_10)
        elif args.model == 'MF' or args.model == 'LightGCN':
            ndcg, hr_10 = train_mf(ndcg, hr_10)
        elif 'LSTM' in args.model or args.model == 'STGCN':
            train_dataset.reset()
            ndcg, hr_10 = train_rnn(ndcg, hr_10)
        elif args.model == 'STP_UDGAT':
            ndcg, hr_10 = train_stp(ndcg, hr_10)
else:
    start = time.time()
    if 'TADGAT' in args.model:
        hr_1, hr_5, hr_10, mrr, ndcg = test('test')
    elif args.model == 'MF' or args.model == 'LightGCN':
        hr_1, hr_5, hr_10, mrr, ndcg = test_mf('test')
    elif 'LSTM' in args.model or args.model == 'STGCN':
        hr_1, hr_5, hr_10, mrr, ndcg = test_rnn('test')
    elif args.model == 'STP_UDGAT':
        hr_1, hr_5, hr_10, mrr, ndcg = test_stp('test')
    end = time.time()
    print("time consumption: %.4f" % (end-start))
    print("Test result on test set: HR@1 %.4f, HR@5 %.4f, HR@10 %.4f, MRR: %.4f, NDCG@10: %.4f" % (hr_1, hr_5, hr_10, mrr, ndcg))


import time
import math
from tqdm import tqdm  # tqdm_notebook as tqdm
import numpy as np
import os
import sys
import torch
from torch.utils import data
import utils
from torch.nn.utils import clip_grad_value_

# import torch.nn.functional as F
from transformer import Constants
from transformer.Generator import Generator

# from math_dataset import VOCAB_SZ, MAX_QUESTION_SZ, MAX_ANSWER_SZ, np_decode_string
from math_dataset import MAX_QUESTION_SZ, np_decode_string
from loss import compute_performance
from checkpoints import (
    rotating_save_checkpoint,
    build_checkpoint,
    save_checkpoint,
)
from math_dataset import np_encode_string, question_to_position_batch_collate_fn

def train(
    name,
    model,
    training_data,
    optimizer,
    device,
    epochs,
    validation_data=None,
    tb=None,
    log_interval=100,
    interpolate_interval=1,
    interpolate_data=None,
    start_epoch=0,
    start_batch=0,
    total_loss=0,
    n_char_total=0,
    n_char_correct=0,
    run_batches=0,
    best_valid_accu=0.0,
    best_valid_loss=float("Inf"),
    best_interpolate_accu=0.0,
    best_interpolate_loss=float("Inf"),
    run_max_batches=None,
    extrapolate_data=None,
    checkpoint=True,
    lr=None,
    warmup_lr=None,
    warmup_interval=None,
    smoothing=False,
):
    print("~~~ Beginning Training ~~~~")
    print(
        f"Start epoch: {start_epoch}, Start batch: {start_batch}, Max batch: {run_max_batches}"
    )

    for epoch_i in range(start_epoch, epochs):

        start = time.time()

        print(
            f"[ Epoch: {epoch_i} / {epochs}, Run Batch: {run_batches} / {run_max_batches}]"
        )

        train_loss, train_accu, new_batch_count, done = train_epoch(
            model=model,
            name=name,
            training_data=training_data,
            optimizer=optimizer,
            device=device,
            epoch=epoch_i,
            tb=tb,
            log_interval=log_interval,
            max_batches=run_max_batches,
            run_batch_count=run_batches,
            start_batch=start_batch,
            total_loss=total_loss,
            n_char_total=n_char_total,
            n_char_correct=n_char_correct,
            lr=lr,
            warmup_lr=warmup_lr,
            warmup_interval=warmup_interval,
            smoothing=smoothing,
        )

        run_batches = new_batch_count

        print(
            "[Training]  loss: {train_loss}, ppl: {ppl: 8.6f}, accuracy: {accu:3.3f} %, "
            "elapse: {elapse:3.3f}ms".format(
                train_loss=train_loss,
                ppl=math.exp(min(train_loss, 100)),
                accu=100 * train_accu,
                elapse=(time.time() - start) * 1000,
            )
        )

        if not utils.is_preempted():
            inference_datasets = {}
            if interpolate_data:
                inference_datasets["interpolate"] = interpolate_data
            if extrapolate_data:
                inference_datasets["extrapolate"] = extrapolate_data

            for group, dataset in inference_datasets.items():
                start = time.time()
                inference_loss, inference_acc = inference_epoch(
                    model, dataset, device, epoch_i, group, tb, log_interval,
                )
                print(
                    "[{group}]  loss: {inference_loss},  ppl: {ppl: 8.5f}, accuracy: {accu:3.3f} %, "
                    "elapse: {elapse:3.3f}ms".format(
                        group=group,
                        inference_loss=inference_loss,
                        ppl=math.exp(min(inference_loss, 100)),
                        accu=100 * inference_acc,
                        elapse=(time.time() - start) * 1000,
                    )
                )

        if done or checkpoint:
            print("Building checkpoint..")
            start = time.time()
            state = build_checkpoint(
                name=name,
                model=model,
                optimizer=optimizer,
                acc=train_accu,
                loss=train_loss,
                epoch=epoch_i if not done and checkpoint else epoch_i + 1,
                run_batches=run_batches,
                # is_preempted=utils.is_preempted(),
                start_batch=0,
                lr=lr,
            )

            if utils.is_cloud():
                print("Saving to google cloud..")
                checkpoint_name = "checkpoint"
                if done:
                    checkpoint_name = (
                        f"{checkpoint_name}_b{run_batches}_e{epoch_i}_complete"
                    )
                elif checkpoint:
                    checkpoint_name = f"{checkpoint_name}_b{run_batches}_e{epoch_i}"

                save_checkpoint(
                    state=state, name=checkpoint_name, path="./checkpoints",
                )
            else:
                rotating_save_checkpoint(
                    state,
                    prefix=f"{name}_{run_batches}_training",
                    path="./checkpoints",
                    nb=5,
                )
            print(f"Save checkpoint time: {(time.time() - start) * 1000}ms")
            # if utils.is_preempted():
            #     print("Completed preemption handling. Cleanly exiting.")
            #     sys.exit(0)

        if done:
            print(
                f"Reached max batch. Breaking out of training at the end of epoch {epoch_i}"
            )
            break

        start_batch = 0
        training_data.dataset.endEpoch()
        training_data.dataset.shuffleData()

    print("~~~~~~ Completed training ~~~~~~")

    if utils.is_cloud():
        print("Shutting down instance")
        os.system("sudo shutdown -h now")

def train_epoch(
    model,
    name,
    training_data,
    optimizer,
    device,
    epoch,
    tb=None,
    log_interval=100,
    max_batches=None,
    run_batch_count=0,
    start_batch=0,
    total_loss=0,
    n_char_total=0,
    n_char_correct=0,
    lr=None,
    warmup_lr=None,
    warmup_interval=None,
    smoothing=False,
):

    training_iter = iter(training_data)

    if start_batch > 0:
        last_question = np_encode_string(training_data.dataset.__getitem__(-1)["q"])
        print(f"Final question before checkpoint was {last_question}")

    model.train()
    # interrupted_batch = None
    done = False

    loss_per_char = 0
    accuracy = 0

    for batch_idx, batch in enumerate(training_iter, start=start_batch):
        if utils.is_preempted():
            print("Exiting...")
            sys.exit(0)

        if warmup_interval is not None and batch_idx == warmup_interval:
            print(f"End of warmup. Swapping learning rates from {warmup_lr} to {lr}")
            for param_group in optimizer.param_groups:
                warmup_lr = lr
                param_group["lr"] = lr

        batch_qs, batch_qs_pos, batch_as, batch_as_pos = map(
            lambda x: x.to(device), batch
        )

        gold_as = batch_as[:, 1:]

        optimizer.zero_grad()
        pred_as = model(batch_qs, batch_qs_pos, batch_as, batch_as_pos)

        loss, n_correct = compute_performance(pred_as, gold_as, smoothing=smoothing)

        loss.backward()

        # Clip gradients, paper uses 0.1
        clip_grad_value_(model.parameters(), 0.1)

        # update parameters
        optimizer.step()

        # note keeping
        total_loss += loss.item()

        non_pad_mask = gold_as.ne(Constants.PAD)
        n_char = non_pad_mask.sum().item()
        n_char_total += n_char
        n_char = n_char if n_char > 1 else 1

        batch_loss = loss / n_char
        loss_per_char = total_loss / n_char_total

        n_char_correct += n_correct
        batch_acc = n_correct / n_char
        accuracy = n_char_correct / n_char_total
        print(
            f"Batch: {batch_idx}. Acc: {accuracy:.6f}. Loss: {loss_per_char:.6f}. Batch_acc: {batch_acc:.6f}. Batch_loss: {batch_loss:.6f} "
        )

        # TODO: automatically trim the TB logs that go beyond the preempted checkpoint
        if tb is not None and batch_idx % log_interval == 0:
            tb.add_scalars(
                {
                    "loss_per_char": loss_per_char,
                    "accuracy": accuracy,
                    "batch_loss": batch_loss,
                    "batch_acc": batch_acc,
                },
                group="train",
                sub_group="batch",
                global_step=run_batch_count,
            )

        run_batch_count += 1

        if max_batches is not None and run_batch_count == max_batches:
            print(
                f"Reached {run_batch_count} batches on max_batches of {max_batches}. Breaking from epoch."
            )
            # interrupted_batch = batch_idx
            done = True
            break

        if batch_idx % 251 == 0 and batch_idx != 0:
            print(
                f"Checkpointing on batch: {batch_idx}. Accuracy: {accuracy}. Loss per char: {loss_per_char}. Time: {time.time()}"
            )
            print(f"Last question is {batch_qs[-1]}")

            state = build_checkpoint(
                name=name,
                model=model,
                optimizer=optimizer,
                acc=accuracy,
                loss=loss_per_char,
                epoch=epoch,
                run_batches=run_batch_count,
                start_batch=batch_idx + 1,
                total_loss=total_loss,
                n_char_total=n_char_total,
                n_char_correct=n_char_correct,
                lr=warmup_lr,
            )

            save_checkpoint(
                state=state, name=f"{name}_latest_checkpoint", path="./checkpoints"
            )

        # if utils.is_preempted():
        #     print(
        #         f"Preemption at end of Epoch batch: {batch_idx} and new Run batch: {run_batch_count}. Breaking from epoch."
        #     )
        #     interrupted_batch = batch_idx
        #     break

    if tb is not None and not utils.is_preempted():
        tb.add_scalars(
            {"loss_per_char": loss_per_char, "accuracy": accuracy},
            group="train",
            sub_group="epoch",
            global_step=epoch,
        )

    return loss_per_char, accuracy, run_batch_count, done

def inference_epoch(model, data, device, epoch, group, tb=None, log_interval=100):
    model.eval()

    total_loss = 0
    n_char_total = 0
    n_char_correct = 0

    with torch.no_grad():
        for batch_idx, batch in enumerate(tqdm(data, mininterval=2, leave=False)):
            # prepare data
            batch_qs, batch_qs_pos, batch_as, batch_as_pos = map(
                lambda x: x.to(device), batch
            )
            gold_as = batch_as[:, 1:]

            # forward
            pred_as = model(batch_qs, batch_qs_pos, batch_as, batch_as_pos)
            loss, n_correct = compute_performance(pred_as, gold_as, smoothing=False)

            # note keeping
            total_loss += loss.item()

            non_pad_mask = gold_as.ne(Constants.PAD)
            n_char = non_pad_mask.sum().item()
            n_char_total += n_char
            n_char_correct += n_correct

    loss_per_char = total_loss / n_char_total
    accuracy = n_char_correct / n_char_total

    if tb is not None:
        tb.add_scalars(
            {"loss_per_char": loss_per_char, "accuracy": accuracy},
            group=group,
            sub_group="epoch",
            global_step=epoch,
        )

    return loss_per_char, accuracy

def predict(generator, data, device, max_predictions=None):
    if max_predictions is not None:
        cur = max_predictions
    else:
        cur = len(data)

    resps = []
    for batch_idx, batch in enumerate(data):
        if cur == 0:
            break
        batch_qs, batch_qs_pos = map(lambda x: x.to(device), batch)
        all_hyp, all_scores = generator.generate_batch(batch_qs, batch_qs_pos)

        for i, idx_seqs in enumerate(all_hyp):
            for j, idx_seq in enumerate(idx_seqs):
                r = np_decode_string(np.array(idx_seq))
                s = all_scores[i][j].cpu().item()
                resps.append({"resp": r, "score": s})
        cur -= 1

    return resps

def predict_benchmark(generator, data, device, max_predictions=None):
    resps = []
    dataset_size = len(data)
    for batch_idx, (batch_qs, batch_qs_pos, batch_as) in enumerate(data):
        start = time.time()

        all_hyp, all_scores = generator.generate_batch(batch_qs, batch_qs_pos)

        for i, idx_seqs in enumerate(all_hyp):
            g = np_decode_string(np.array(idx_seqs[0]))
            s = all_scores[i][0].cpu().item()
            a = batch_as[i]
            c = g == a
            resp = {"correct": c, "guess": g, "answer": a, "score": s}
            print(resp)
            resps.append(resp)

        print(
            f"Batch {batch_idx + 1} of {math.ceil(dataset_size/len(idx_seqs))}, time {time.time() - start}s."
        )

    return resps

def predict_dataset(
    dataset,
    model,
    device,
    callback,
    max_batches=None,
    beam_size=5,
    max_token_seq_len=MAX_QUESTION_SZ,
    n_best=1,
    batch_size=1,
    num_workers=1,
):

    generator = Generator(
        model,
        device,
        beam_size=beam_size,
        max_token_seq_len=max_token_seq_len,
        n_best=n_best,
    )

    if max_batches is not None:
        cur = max_batches
    else:
        cur = len(dataset)

    resps = []
    for batch_idx, batch in enumerate(dataset):
        if cur == 0:
            break

        batch_qs, batch_qs_pos, _ = map(lambda x: x.to(device), batch)
        all_hyp, all_scores = generator.generate_batch(batch_qs, batch_qs_pos)

        callback(batch_idx, all_hyp, all_scores)

        cur -= 1
    return resps

def predict_multiple(
    questions,
    model,
    device,
    beam_size=5,
    max_token_seq_len=MAX_QUESTION_SZ,
    n_best=1,
    batch_size=1,
    num_workers=1,
):

    questions = list(map(lambda q: np_encode_string(q), questions))
    questions = data.DataLoader(
        questions,
        batch_size=1,
        shuffle=False,
        num_workers=1,
        collate_fn=question_to_position_batch_collate_fn,
    )

    generator = Generator(
        model,
        device,
        beam_size=beam_size,
        max_token_seq_len=max_token_seq_len,
        n_best=n_best,
    )

    return predict(generator, questions, device)

def predict_single(
    question, model, device, beam_size=5, max_token_seq_len=MAX_QUESTION_SZ, n_best=1
):

    generator = Generator(
        model,
        device,
        beam_size=beam_size,
        max_token_seq_len=max_token_seq_len,
        n_best=n_best,
    )

    qs = [np_encode_string(question)]
    qs, qs_pos = question_to_position_batch_collate_fn(qs)
    qs, qs_pos = qs.to(device), qs_pos.to(device)

    all_hyp, all_scores = generator.generate_batch(qs, qs_pos)
    # resp = np_decode_string(np.array(all_hyp[0][0]))

    resps = []
    for i, idx_seqs in enumerate(all_hyp):
        for j, idx_seq in enumerate(idx_seqs):
            r = np_decode_string(np.array(idx_seq))
            s = all_scores[i][j].cpu().item()
            resps.append({"resp": r, "score": s})

    return resps


from collections import deque

n = int(input())
quee = deque([i for i in range(1,n+1)])
while(len(quee)>1) :
    quee.popleft()
    num = quee.popleft()
    quee.append(num)

print(quee[0])


"""

    get_google_API_data.py

    Gets google API data, for a certain set of keywords, for a given date

    Takes the following args:

    DATE: date to scrape for (format: yyyy-mm-dd) (e.g., 2020-02-16)

"""
import os
import sys
import re
import boto3
import pytrends
from pytrends.request import TrendReq
import pandas as pd
import datetime

if __name__ == "__main__":

    # get consts, creds, args
    KEYWORDS = ["covid", "coronavirus", "lockdown", "quarantine", "vaccine"]

    AWS_BUCKET = os.environ["AWS_BUCKET"]
    AWS_ACCESS = os.environ["AWS_ACCESS"]
    AWS_SECRET = os.environ["AWS_SECRET"]

    FRESH_SCRAPES_FILENAME = ""
    INTEREST_OVER_TIME_FILENAME = "interest_over_time_"
    INTEREST_BY_REGION_FILENAME = "interest_by_region_"

    LOCAL_SCRAPES_DIR = "../../tweets/google_API_scrapes/"
    AWS_SCRAPES_DIR = "google_API_scrapes/"

    LOCAL_TIME_FILENAME = LOCAL_SCRAPES_DIR + INTEREST_OVER_TIME_FILENAME
    LOCAL_REGION_FILENAME = LOCAL_SCRAPES_DIR + INTEREST_BY_REGION_FILENAME

    AWS_TIME_FILENAME = AWS_SCRAPES_DIR + INTEREST_OVER_TIME_FILENAME
    AWS_REGION_FILENAME = AWS_SCRAPES_DIR + INTEREST_BY_REGION_FILENAME

    DATE = sys.argv[1]

    # check DATE format
    REGEX = re.compile("^[0-9]{4}-[0-9]{2}-[0-9]{2}$")

    if not REGEX.match(DATE):
        raise ValueError(
            "DATE has incorrect format (must be in yyyy-mm-dd format)")

    # log into S3
    s3 = boto3.client('s3',
                      aws_access_key_id=AWS_ACCESS,
                      aws_secret_access_key=AWS_SECRET)

    print(f"Getting Google API trends data for {DATE}")

    try:
        # start API call
        pytrends_API = TrendReq()

        pytrends_API.build_payload(
            kw_list=KEYWORDS, geo="US", timeframe=f"{DATE} {DATE}")

        # get trends over time and by region, and export locally:
        pytrends_API.interest_over_time().to_csv(LOCAL_TIME_FILENAME + DATE + ".csv")
        pytrends_API.interest_by_region().to_csv(LOCAL_REGION_FILENAME + DATE + ".csv")

        # upload to AWS
        s3.upload_file(LOCAL_TIME_FILENAME + DATE + ".csv",
                       AWS_BUCKET, AWS_TIME_FILENAME + DATE + ".csv")
        s3.upload_file(LOCAL_REGION_FILENAME + DATE + ".csv",
                       AWS_BUCKET, AWS_REGION_FILENAME + DATE + ".csv")

        # delete local versions of the file
        os.remove(LOCAL_TIME_FILENAME + DATE + ".csv")
        os.remove(LOCAL_REGION_FILENAME + DATE + ".csv")

        print(f"Finished getting Google API trends data for {DATE}")

    except Exception as e:
        print("Error in getting Google API trends data")
        print(e)
        print(f"Error occurred for the following date: {DATE}")

    finally:
        print(
            f"Finished running get_google_API_data.py on (in UTC time): {datetime.datetime.utcnow()}")


import numpy as np
import random
import torch
import os
import json
import argparse
import os.path

from utils.others import create_folder, dump_json, load_json, find_sub_string, find_all_metal_names
from utils.data_visualization import plot_latex
from utils.data_structure import DatasetTaskDesc, CompressMethod
from utils.data_preprocessing import load_transform_data, parse_series_time, load_metal_data, get_data
from experiments import algo_dict, gen_experiment, metal_desc
from statsmodels.tsa.arima.model import ARIMA

np.random.seed(48)
random.seed(48)
torch.manual_seed(48)
torch.random.manual_seed(48)

np.seterr(invalid='raise')

multi_task_algo = [
    "GPMultiTaskMultiOut", "GPMultiTaskIndex", "IndependentGP"
    "DeepGPMultiOut", "DSPPMultiOut", "SparseGPIndex",
    "SparseMaternGraphGP", "DeepGraphMultiOutputGP",
    "DeepGraphInfoMaxMultiOutputGP"
]
metric = ["MSE", "CRPS"]

default_len_dataset = 794
default_len_train_show = (274, 130)
optim_iter=5

def run_multi_task_gp(save_path, modifier, multi_task_desc, len_inp=10,
    len_dataset=default_len_dataset, len_train_show=default_len_train_show,
    kernel="Composite_1",
    graph_path="exp_result/graph_result/distance correlation_test_graph.npy",
    is_test=False, is_verbose=False,
    multi_task_algo=multi_task_algo, is_full_result=False
):
    base_multi_task_fast_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel=kernel,
        optim_iter=25,
        len_inp=len_inp,
        lr=0.1,
        graph_path=graph_path,
    )

    base_multi_task_slow_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel=kernel,
        optim_iter=optim_iter,
        len_inp=len_inp,
        lr=0.05,
        graph_path=graph_path
    )

    base_multi_task_slow_dspp = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel=kernel,
        optim_iter=3,
        len_inp=len_inp,
        lr=0.1,
        graph_path=graph_path
    )

    base_multi_task_slower_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel=kernel,
        optim_iter=500,
        # optim_iter=1,
        len_inp=len_inp,
        lr=0.005,
        graph_path=graph_path
    )

    default_config = {
        # "GPMultiTaskMultiOut": base_multi_task_fast_lr,
        # "GPMultiTaskIndex": base_multi_task_fast_lr,
        # "DeepGPMultiOut": base_multi_task_slow_lr,
        # Might be wrong.....
        # "DSPPMultiOut": base_multi_task_slow_dspp,
        "SparseGPIndex": base_multi_task_slow_lr,
        "SparseMaternGraphGP": base_multi_task_slow_lr,
        # "DeepGraphMultiOutputGP": base_multi_task_slower_lr,
        # "DeepGraphInfoMaxMultiOutputGP": base_multi_task_fast_lr,
        # Optim = 4 is enough....
        # "NonlinearMultiTaskGP": base_multi_task_slow_lr,
        # "NonlinearMultiTaskGSPP": base_multi_task_slow_dspp,
        "DeepGPGraphPropagate": base_multi_task_fast_lr,
        "DeepGPGraphInteract": base_multi_task_fast_lr,
        "DSPPGraphInteract": base_multi_task_fast_lr,
        "DSPPGraphPropagate": base_multi_task_fast_lr,
        "IndependentGP": algo_dict.encode_params(
            "gp", is_verbose=is_verbose,
            is_test=is_test,
            kernel="Matern",
            optim_iter=25,
            len_inp=10
        ),
        "IIDDataModel": algo_dict.encode_params(
            "iid", is_verbose=is_verbose,
            is_test=is_test, dist="gaussian"
        ),
        "ARIMAModel": algo_dict.encode_params(
            "arima", is_verbose=is_verbose,
            is_test=is_test, order=(11, 2, 11)
        ),
    }

    task = gen_experiment.gen_task_cluster(
        multi_task_algo, "metal", modifier,
        clus_metal_desc=multi_task_desc,
        clus_time_desc=None,
        algo_config=default_config,
        len_dataset=len_dataset,
        len_train_show=len_train_show
    )
    output = gen_experiment.run_experiments(task, save_path=save_path)

    if is_full_result:
        return output
    else:
        dump_json(f"{save_path}all_data.json", output)

        summary = {}

        for algo_name, result in output.items():
            data = {}
            for metric, r in result.items():
                mean_across_task = np.mean(np.array(r), axis=0)[0]
                data[metric] = mean_across_task
            summary[algo_name] = data

        return summary

def run_hyperparam_search(name, save_folder, multi_task_algo, kernel="Composite_1", is_test=False, is_verbose=False):
    num_inp = np.arange(2, 12, step=2)
    num_pca = np.arange(2, 7)

    all_results = {
        algo:{m: np.zeros((len(num_inp), len(num_pca))).tolist() for m in metric}
        for algo in multi_task_algo
    }
    all_metal_name = find_all_metal_names()

    for i, len_inp in enumerate(num_inp):
        for j, pca in enumerate(num_pca):
            pca_modifier = {
                metal: CompressMethod(int(pca), "pca", info={})
                for metal in all_metal_name
            }
            curr_result = run_multi_task_gp(
                f"{save_folder}/hyper_search_{name}/run_pca_{pca}_feat_{len_inp}/",
                pca_modifier,
                multi_task_desc=[all_metal_name],
                len_inp=int(len_inp),
                kernel=kernel,
                is_test=is_test,
                is_verbose=is_verbose,
                multi_task_algo=multi_task_algo
            )
            for algo in multi_task_algo:
                for met in metric:
                    all_results[algo][met][i][j] = curr_result[algo][met]

            dump_json(f"{save_folder}/hyper_search_{name}/final_result_{name}.json", all_results)

    return all_results

def general_testing_legacy(is_verbose, is_test):
    all_metal_name = find_all_metal_names()
    no_modifier = {
        metal: CompressMethod(0, "drop")
        for metal in all_metal_name
    }
    pca_modifier = {
        metal: CompressMethod(3, "pca")
        for metal in all_metal_name
    }

    all_algo = [
        "GPMultiTaskMultiOut", "IndependentGP", "GPMultiTaskIndex", "IIDDataModel", "ARIMAModel",
        "DeepGPMultiOut", "DSPPMultiOut", "SparseGPIndex", "SparseMaternGraphGP",
        "DeepGraphMultiOutputGP", "DeepGraphInfoMaxMultiOutputGP",
        "NonlinearMultiTaskGP", "NonlinearMultiTaskGSPP", "DeepGPGraphPropagate", "DeepGPGraphInteract", "DSPPGraphInteract", "DSPPGraphPropagate"
    ]
    display_name_to_algo = dict(zip(
        all_algo,[
            "Multi-Task Out", "Independent GP", "Multi-Task Index",
            "Mean", "ARIMA", "Deep GP", "DSPP",
            "Sparse Multi-Task Index", "Sparse Matern Graph GP",
            "Deep Graph Multi Output GP", "Deep Graph InfoMax",
            "Non-Linear Multi-Task GP", "Non-Linear Multi-Task DSPP",
            "Deep Gaussian Process Graph", "Deep Gaussian Process Interaction",
            "DSPP Interaction", "DSPP Graph Propagation"
        ],
    ))

    # all_algo = [
    #     "IndependentGP", "IIDDataModel",
    #     "DeepGPMultiOut", "DSPPMultiOut", "SparseGPIndex",
    #     "SparseMaternGraphGP", "NonlinearMultiTaskGSPP"
    # ]

    all_algo = [
        # "DeepGraphMultiOutputGP",
        # "DeepGraphInfoMaxMultiOutputGP",
        # "NonlinearMultiTaskGP",
        # "NonlinearMultiTaskGSPP",
        # "DeepGPGraphPropagate", "DeepGPGraphInteract",
        # "DSPPGraphInteract",
        "DSPPGraphPropagate"
    ]

    base_multi_task_fast_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel="Matern",
        optim_iter=optim_iter,
        len_inp=10,
        lr=0.1,
        graph_path="exp_result/graph_result/hsic_test_graph.npy"
    )

    base_multi_task_v_fast_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel="Matern",
        optim_iter=optim_iter,
        len_inp=10,
        lr=0.2,
        graph_path="exp_result/graph_result/hsic_test_graph.npy"
    )

    base_multi_task_slow_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel="Matern",
        optim_iter=optim_iter,
        len_inp=10,
        lr=0.05,
        graph_path="exp_result/graph_result/hsic_test_graph.npy"
    )

    base_multi_task_deep = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel="Matern",
        optim_iter=optim_iter,
        len_inp=10,
        lr=0.01,
        graph_path="exp_result/graph_result/hsic_test_graph.npy"
    )

    base_multi_task_slower_lr = algo_dict.encode_params(
        "gp_multi_task", is_verbose=is_verbose,
        is_test=is_test,
        kernel="Matern",
        optim_iter=500,
        len_inp=10,
        lr=0.005,
        graph_path="exp_result/graph_result/hsic_test_graph.npy"
    )

    default_config = {
        "GPMultiTaskMultiOut": base_multi_task_fast_lr,
        "GPMultiTaskIndex": base_multi_task_fast_lr,
        "DeepGPMultiOut": base_multi_task_slow_lr,
        "DSPPMultiOut": base_multi_task_slow_lr,
        "SparseGPIndex": base_multi_task_slow_lr,
        "SparseMaternGraphGP": base_multi_task_slow_lr,
        "DeepGraphMultiOutputGP": base_multi_task_slower_lr,
        "DeepGraphInfoMaxMultiOutputGP": base_multi_task_fast_lr,
        # Optim = 4 is enough....
        "NonlinearMultiTaskGP": base_multi_task_deep,
        "NonlinearMultiTaskGSPP": base_multi_task_deep,
        # "DeepGPGraphPropagate": base_multi_task_fast_lr,
        # "DeepGPGraphInteract": base_multi_task_fast_lr,
        # "DSPPGraphInteract": base_multi_task_v_fast_lr,
        # "DSPPGraphPropagate": base_multi_task_v_fast_lr,
        "IndependentGP": algo_dict.encode_params(
            "gp", is_verbose=is_verbose,
            is_test=is_test,
            kernel="Composite_1",
            optim_iter=optim_iter,
            len_inp=10
        ),
        "IIDDataModel": algo_dict.encode_params(
            "iid", is_verbose=is_verbose,
            is_test=is_test, dist="gaussian"
        ),
        "ARIMAModel": algo_dict.encode_params(
            "arima", is_verbose=is_verbose,
            is_test=is_test, order=(2, 0, 5)
        ),
    }

    def original_test():

        all_algo = [
            "GPMultiTaskMultiOut",
            "IndependentGP",
            "GPMultiTaskIndex",
            "IIDDataModel",
            "ARIMAModel",
            "DeepGPMultiOut",
            "DSPPMultiOut",
            "SparseGPIndex",
            # "NonlinearMultiTaskGP",
            # "NonlinearMultiTaskGSPP",
        ]

        # (time_al, time_cu, commodity, time_al_feat, time_cu_feat, commodity_feat) = [
        (time_al_feat, time_cu_feat, commodity_feat) = [
            gen_experiment.gen_task_cluster(
                all_algo, type_task, modifier,
                clus_metal_desc=metal_desc,
                clus_time_desc=time_desc,
                algo_config=default_config
            )
            for type_task, modifier, metal_desc, time_desc in [
                # ("time", no_modifier, "aluminium", [[22, 44, 66]]),
                # ("time", no_modifier, "copper", [[22, 44, 66]]),
                # ("metal", no_modifier, [["aluminium", "copper"]], None),
                ("time", pca_modifier, "aluminium", [[22, 44, 66]]),
                ("time", pca_modifier, "copper", [[22, 44, 66]]),
                ("metal", pca_modifier, [["aluminium", "copper"]], None),
            ]
        ]

        # task_train = [time_al, commodity, time_al_feat, commodity_feat]
        # task_names = ["Price", "Metal", "Price_Feat", "Metal_Feat"]

        task_train = [time_al_feat, commodity_feat]
        task_names = ["Price_Feat", "Metal_Feat"]

        def how_to_plot(super_task):
            # plot_latex(
            #     names=[all_algo, all_algo],
            #     results=[super_task["Price"], super_task["Metal"]],
            #     multi_task_name=[["Date (22)", "Date (44)", "Date (66)"], ["Aluminium", "Copper"]],
            #     display_name_to_algo=display_name_to_algo
            # )
            # print()
            # print()
            plot_latex(
                names=[all_algo, all_algo],
                results=[super_task["Price_Feat"], super_task["Metal_Feat"]],
                multi_task_name=[["Date (22)", "Date (44)", "Date (66)"], ["Aluminium", "Copper"]],
                display_name_to_algo=display_name_to_algo
            )

        return task_train, task_names, how_to_plot

    def test_all_metal():
        all_metal_name = find_all_metal_names()

        task = gen_experiment.gen_task_cluster(
            all_algo, "metal", pca_modifier,
            clus_metal_desc=[all_metal_name],
            clus_time_desc=None,
            algo_config=default_config
        )

        def how_to_plot(super_task):
            plot_latex(
                names=[all_algo],
                results=[super_task["All Metal"]],
                multi_task_name=[all_metal_name],
                display_name_to_algo=display_name_to_algo
            )

        return [task], ["All Metal"], how_to_plot

    task_train, task_names, how_to_plot = original_test()
    # task_train, task_names, how_to_plot = test_all_metal()

    super_task = {}
    for task, name in zip(task_train, task_names):
        all_out = gen_experiment.run_experiments(task)
        super_task.update({name: all_out})

    dump_json("save/all_data.json", super_task)
    # super_task = load_json("save/all_data.json")

    how_to_plot(super_task)

def general_testing(save_path, setting, is_verbose, is_test):
    print(save_path)
    all_metal_name = find_all_metal_names()

    for i, (curr_algo, len_inp, pca_dim, kernel, graph_path) in enumerate(setting):
        pca_modifier = {
            metal: CompressMethod(int(pca_dim), "pca", info={})
            for metal in all_metal_name
        }

        # try:
        error_results = run_multi_task_gp(
            save_path, pca_modifier,
            multi_task_desc=[all_metal_name],
            len_inp=len_inp,
            len_dataset=default_len_dataset,
            len_train_show=default_len_train_show,
            kernel=kernel,
            is_test=is_test, is_verbose=is_verbose,
            multi_task_algo=[curr_algo],
            is_full_result=True,
            graph_path=graph_path
        )
        # except RuntimeError:
        #     error_results = {curr_algo : {"CRPS": None}}

        if i > 0:
            result = load_json(f"{save_path}grid_result.json")
            result.update(error_results)
            dump_json(f"{save_path}grid_result.json", result)
        else:
            dump_json(f"{save_path}grid_result.json", error_results)

    result = load_json(f"{save_path}grid_result.json")
    plot_latex(
        names=[[data[0] for data in setting]],
        results=[result],
        multi_task_name=[all_metal_name],
        display_name_to_algo=algo_dict.class_name_to_display
    )

    pass

def grid_commodities_run(save_path, setting, is_test=False, is_verbose=False):
    all_metal_name = find_all_metal_names()

    num_metal = len(all_metal_name)
    # multi_task_algo = ["GPMultiTaskMultiOut", "GPMultiTaskIndex"]

    counter = 0

    # error_matrix = np.zeros((num_metal, num_metal))
    error_dict = {
        algo: np.zeros((num_metal, num_metal)).tolist()
        for algo, _, _, _ in setting
    }

    for i in range(num_metal):
        for j in range(i, num_metal):
            counter += 1
            if i != j:
                for curr_algo, len_inp, pca_dim, kernel in setting:
                    pca_modifier = {
                        metal: CompressMethod(int(pca_dim), "pca", info={})
                        for metal in all_metal_name
                    }

                    try:
                        error_results = run_multi_task_gp(
                            save_path, pca_modifier,
                            multi_task_desc=[[all_metal_name[j], all_metal_name[i]]],
                            len_inp=len_inp,
                            len_dataset=default_len_dataset,
                            len_train_show=default_len_train_show,
                            kernel=kernel,
                            is_test=is_test, is_verbose=is_verbose,
                            multi_task_algo=[curr_algo]
                        )
                    except RuntimeError:
                        error_results = {curr_algo : {"CRPS": None}}

                    error_dict[curr_algo][j][i] = 0 if i == j else error_results[curr_algo]["CRPS"]

                    dump_json(f"{save_path}/grid_result.json", error_dict)

def update_json(path, entry):
    if os.path.isfile(path):
        updated_version = {}
        current_json = load_json(path)
        for k in entry:
            if k in current_json:
                updated_version[k] = current_json[k].copy()
                updated_version[k].update(entry[k])
            else:
                updated_version[k] = entry[k]
    else:
        updated_version = entry

    dump_json(path, updated_version)

def grid_compare_clusters(setting, cluster_data_path, save_path, is_test=False, is_verbose=False, cluster_names=None):

    cluster_dict = load_json(cluster_data_path)
    all_metal_name = find_all_metal_names()

    all_results = {}

    compare_cluster_path = f"{save_path}/compare_cluster.json"

    for curr_algo, len_inp, pca_dim, kernel in setting:
        pca_modifier = {
            metal: CompressMethod(int(pca_dim), "pca", info={})
            for metal in all_metal_name
        }

        for test_name, cluster_index in cluster_dict.items():

            if not cluster_names is None:
                if not test_name in cluster_names:
                    continue

            curr_save_path = save_path + "/" + "-".join(test_name.split(" ")) + "/"
            try:
                error_results = run_multi_task_gp(
                    curr_save_path, pca_modifier,
                    multi_task_desc=gen_experiment.cluster_index_to_nested(cluster_index),
                    len_inp=len_inp,
                    len_dataset=default_len_dataset,
                    len_train_show=default_len_train_show,
                    kernel=kernel,
                    is_test=is_test, is_verbose=is_verbose,
                    multi_task_algo=[curr_algo]
                )
            except RuntimeError:
                error_results = {curr_algo : {"CRPS": None}}

            all_results.update({
                test_name: {
                    k: v["CRPS"]
                    for k, v in error_results.items()
                }
            })

            update_json(compare_cluster_path, all_results)

            # Running all
        curr_save_path = save_path + "/full_model/"

        try:
            error_results = run_multi_task_gp(
                curr_save_path, pca_modifier,
                multi_task_desc=[all_metal_name],
                len_inp=len_inp,
                len_dataset=default_len_dataset,
                len_train_show=default_len_train_show,
                kernel=kernel,
                is_test=is_test, is_verbose=is_verbose,
                multi_task_algo=[curr_algo]
            )
        except RuntimeError:
            error_results = {curr_algo : {"CRPS": None}}

        all_results.update({
            "full_model": {
                k: v["CRPS"]
                for k, v in error_results.items()
            }
        })

        update_json(compare_cluster_path, all_results)

def run_ARMA_param_search(save_folder="exp_result/hyper_param_arma_test"):
    metal_names = find_all_metal_names()
    all_output_data = [
        get_data(metal, is_price_only=False, is_feat=False)
        for metal in metal_names
    ]
    test = all_output_data[0]["Price"]

    result = {
        metal: np.zeros((10, 10, 10))
        for metal in metal_names
    }

    create_folder(save_folder)

    for metal in metal_names[:5]:
        for i, s1 in enumerate(np.arange(2, 12, step=1)):
            for j, s2 in enumerate(np.arange(2, 12, step=1)):
                for k, s3 in enumerate(np.arange(2, 12, step=1)):
                    print(f"At {i}, {j} and {k} with metal {metal}")
                    try:
                        model = ARIMA(test, order=(s1, s3, s2))
                        model_fit = model.fit(method="innovations_mle")
                        result[metal][i, j, k] = model_fit.aic
                    except ValueError:
                        print(f"Value Error At {s1}, {s2}")
                        result[metal][i, j, k] = 10

                    np.save(f"{save_folder}/{metal}.npy", result[metal])

def grid_compare_graph_run(save_path, setting, is_test, is_verbose):
    graph_path = [
        "exp_result/graph_result/distance correlation_test_graph.npy",
        "exp_result/graph_result/hsic_test_graph.npy",
        "exp_result/graph_result/kendell_test_graph.npy",
        "exp_result/graph_result/peason_test_graph.npy",
        "exp_result/graph_result/spearman_test_graph.npy"
    ]

    graph_path = graph_path[:2] if is_test else graph_path

    def get_graph_name(name):
        return name.split("/")[-1].split("_")[0]

    get_graph_name(graph_path[0])

    base_line_algo = ["SparseGPIndex", "DeepGPMultiOut"]
    all_metal_name = find_all_metal_names()

    all_results = {}
    save_graph_path = f"{save_path}/compare_graph.json"

    for curr_algo, len_inp, pca_dim, kernel in setting:
        pca_modifier = {
            metal: CompressMethod(int(pca_dim), "pca", info={})
            for metal in all_metal_name
        }

        for curr_graph_path in graph_path:
            test_name = get_graph_name(curr_graph_path)
            curr_save_path = save_path + test_name + "/"

            try:
                error_results = run_multi_task_gp(
                    curr_save_path, pca_modifier, multi_task_desc=[all_metal_name],
                    len_inp=len_inp,
                    len_dataset=default_len_dataset,
                    len_train_show=default_len_train_show,
                    kernel=kernel,
                    graph_path=curr_graph_path,
                    is_test=is_test, is_verbose=is_verbose,
                    multi_task_algo=[curr_algo]
                )
            except RuntimeError:
                error_results = {curr_algo : {"CRPS": None}}

            all_results.update({
                test_name : {
                    k: v["CRPS"]
                    for k, v in error_results.items()
                }
            })

            update_json(save_graph_path, all_results)

    curr_save_path = save_path + "full_model/"

    try:
        error_results = run_multi_task_gp(
            curr_save_path, pca_modifier, multi_task_desc=[all_metal_name],
            len_inp=len_inp,
            len_dataset=default_len_dataset,
            len_train_show=default_len_train_show,
            kernel="RBF",
            graph_path=None,
            is_test=is_test, is_verbose=is_verbose,
            multi_task_algo=base_line_algo
        )
    except RuntimeError:
        error_results = {curr_algo : {"CRPS": None}}

    all_results.update({
        "no_graph_model" : {
            k: v["CRPS"]
            for k, v in error_results.items()
        }
    })

    update_json(save_graph_path, all_results)

def argument_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", help="Are we testing", dest='is_test', action='store_true')
    parser.add_argument("--train", help="Are we training", dest='is_test', action='store_false')
    parser.set_defaults(is_test=True)

    parser.add_argument("--not-verbose", dest='is_verbose', action='store_false')
    parser.add_argument("--verbose", help="Are we training", dest='is_verbose', action='store_true')
    parser.set_defaults(is_verbose=False)

    args = parser.parse_args()
    return args

def hyperparameter_search():
    args = argument_parser()
    is_test=args.is_test
    is_verbose=args.is_verbose

    create_folder("save")

    # multi_task_algo=["GPMultiTaskIndex", "GPMultiTaskMultiOut"],

    # run_hyperparam_search(
    #     "matern", "exp_result/save_hyper",
    #     multi_task_algo=["SparseGPIndex", "GPMultiTaskMultiOut", "GPMultiTaskIndex"],
    #     kernel="Matern", is_test=is_test,
    #     is_verbose=is_verbose
    # )
    # run_hyperparam_search(
    #     "rbf", "exp_result/save_hyper",
    #     multi_task_algo=["SparseGPIndex", "GPMultiTaskMultiOut", "GPMultiTaskIndex"],
    #     kernel="RBF", is_test=is_test,
    #     is_verbose=is_verbose
    # )

    # run_hyperparam_search(
    #     "matern", "exp_result/save_hyper_deep_dspp_matern",
    #     multi_task_algo=["DSPPMultiOut"],
    #     # multi_task_algo=["DSPPMultiOut", "NonlinearMultiTaskGSPP"],
    #     kernel="Matern", is_test=is_test,
    #     is_verbose=is_verbose
    # )
    # run_hyperparam_search(
    #     "rbf", "exp_result/save_hyper_dspp_rbf",
    #     multi_task_algo=["DSPPMultiOut", "NonlinearMultiTaskGSPP"],
    #     kernel="RBF", is_test=is_test,
    #     is_verbose=is_verbose
    # )

    # run_hyperparam_search(
    #     "matern", "exp_result/save_hyper_deep_2",
    #     multi_task_algo=["DeepGPMultiOut", "NonlinearMultiTaskGP"],
    #     kernel="Matern", is_test=is_test,
    #     is_verbose=is_verbose
    # )
    # run_hyperparam_search(
    #     "rbf", "exp_result/save_hyper_deep_2",
    #     multi_task_algo=["DeepGPMultiOut", "NonlinearMultiTaskGP"],
    #     kernel="RBF", is_test=is_test,
    #     is_verbose=is_verbose
    # )

def general_test_run():
    args = argument_parser()
    is_test=args.is_test
    is_verbose=args.is_verbose

    best_setting = [
        # ("IIDDataModel",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("ARIMAModel",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("IndependentGP",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("SparseGPIndex",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("GPMultiTaskMultiOut",
        #     2, 2,
        #     "Matern", "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("GPMultiTaskIndex",
        #     8, 2, "RBF",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("DeepGPMultiOut",
        #     10, 6, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("DSPPMultiOut",
        #     10, 6, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("NonlinearMultiTaskGP",
        #     6, 3, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("NonlinearMultiTaskGSPP",
        #     6, 3, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        # ("DeepGraphMultiOutputGP",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/spearman_test_graph.npy"),
        # ("DeepGraphInfoMaxMultiOutputGP",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/hsic_test_graph.npy"),
        # ("SparseMaternGraphGP",
        #     2, 2, "Matern",
        #     "exp_result/graph_result/distance correlation_test_graph.npy"),
        ("DeepGPGraphPropagate",
            3, 10, "Matern",
            "exp_result/graph_result/hsic_test_graph.npy"),
        ("DeepGPGraphInteract",
            3, 10, "Matern",
            "exp_result/graph_result/hsic_test_graph.npy"),
        ("DSPPGraphInteract",
            3, 10, "Matern",
            "exp_result/graph_result/hsic_test_graph.npy"),
        ("DSPPGraphPropagate",
            3, 10, "Matern",
            "exp_result/graph_result/hsic_test_graph.npy"),
    ]

    save_path = "exp_result/general_running/"
    general_testing(save_path, best_setting, is_verbose, is_test)
    # general_testing_legacy(is_verbose, is_test)

def compare_cluster():
    args = argument_parser()
    is_test=args.is_test
    is_verbose=args.is_verbose

    best_setting = [
        # ("SparseGPIndex", 2, 2, "Matern"),
        # ("GPMultiTaskMultiOut", 2, 2, "Matern"),
        # ("GPMultiTaskIndex", 8, 2, "RBF"),
        ("DeepGPMultiOut", 10, 6, "Matern"),
        ("DSPPMultiOut", 10, 6, "Matern"),
        ("NonlinearMultiTaskGP", 6, 3, "RBF"),
        ("NonlinearMultiTaskGSPP", 6, 3, "RBF"),
    ]

    second_best_setting = [
        ("SparseGPIndex", 2, 3, "Matern"),
        ("GPMultiTaskMultiOut", 2, 4, "Matern"),
        ("GPMultiTaskIndex", 10, 2, "RBF"),
        ("DeepGPMultiOut", 10, 6, "RBF"),
        ("DSPPMultiOut", 10, 6, "RBF"),
        ("NonlinearMultiTaskGP", 6, 3, "Matern"),
        ("NonlinearMultiTaskGSPP", 6, 3, "Matern"),
    ]

    worst_setting = [
        ("SparseGPIndex", 4, 3, "RBF"),
        ("GPMultiTaskMultiOut", 10, 6, "RBF"),
        ("GPMultiTaskIndex", 2, 3, "RBF"),
        ("DeepGPMultiOut", 8, 2, "Matern"),
        ("DSPPMultiOut", 8, 2, "Matern"),
        ("NonlinearMultiTaskGP", 10, 4, "RBF"),
        ("NonlinearMultiTaskGSPP", 10, 4, "RBF"),
    ]

    grid_compare_clusters(
        best_setting,
        "exp_result/cluster_result/feat_data/cluster_4.json",
        "exp_result/cluster_compare",
        is_test=is_test, is_verbose=is_verbose
    )

def compare_many_clusters():
    args = argument_parser()
    is_test=args.is_test
    is_verbose=args.is_verbose

    best_setting = [
        # ("SparseGPIndex", 2, 2, "Matern"),
        # ("GPMultiTaskMultiOut", 2, 2, "Matern"),
        # ("GPMultiTaskIndex", 8, 2, "RBF"),
        # ("DeepGPMultiOut", 10, 6, "Matern"),
        # ("DSPPMultiOut", 10, 6, "Matern"),
        # ("NonlinearMultiTaskGP", 6, 3, "Matern"),
        ("NonlinearMultiTaskGSPP", 6, 3, "Matern"),
    ]

    worst_setting = [
        # ("SparseGPIndex", 4, 3, "RBF"),
        # ("GPMultiTaskMultiOut", 10, 6, "RBF"),
        # ("GPMultiTaskIndex", 2, 3, "RBF"),
        # ("DeepGPMultiOut", 8, 2, "Matern"),
        # ("DSPPMultiOut", 8, 2, "Matern"),
        ("NonlinearMultiTaskGP", 10, 4, "RBF"),
        # ("NonlinearMultiTaskGSPP", 10, 4, "Matern"),
    ]

    save_folder = "exp_result/range_cluster_test"

    num_cluster = [2, 3, 4, 5, 6, 7]
    create_folder(save_folder)

    missing_entries_best = {
        2: ["euclidean", "euclidean knn", "dtw knn"],
        # 3: ["soft-dtw divergence"],
        # 4: ["euclidean knn", "dtw knn", "softdtw knn", "expert"],
        # 5: ["peason", "softdtw knn"],
        # 6: ["euclidean"],
        # 7: ["dtw knn"]
    }

    missing_entries_worst = {
        2: ["dtw knn"],
        4: ["kendel"]
    }

    for n in num_cluster:
        if not n in missing_entries_best:
            continue

        grid_compare_clusters(
            best_setting,
            f"exp_result/cluster_result/feat_data/cluster_{n}.json",
            f"{save_folder}/cluster_compare_{n}",
            is_test=is_test, is_verbose=is_verbose,
            cluster_names=missing_entries_best[n]
        )

def grid_commodities():
    args = argument_parser()
    is_test=args.is_test
    is_verbose=args.is_verbose

    best_setting = [
        ("SparseGPIndex", 2, 2, "Matern"),
        ("GPMultiTaskMultiOut", 2, 2, "Matern"),
        ("GPMultiTaskIndex", 8, 2, "RBF"),
        # ("DeepGPMultiOut", 10, 6, "Matern"),
        # ("DSPPMultiOut", 10, 6, "Matern"),
        # ("NonlinearMultiTaskGP", 6, 3, "Matern"),
        # ("NonlinearMultiTaskGSPP", 6, 3, "Matern"),
    ]

    grid_commodities_run(
        "exp_result/grid_corr_plot/",
        best_setting,
        is_test=is_test, is_verbose=is_verbose
    )

def grid_compare_graph():
    args = argument_parser()
    is_test=args.is_test
    is_verbose=args.is_verbose

    best_graph_setting = [
        ("DeepGraphMultiOutputGP", 2, 2, "Matern"),
        ("DeepGraphInfoMaxMultiOutputGP", 2, 2, "Matern"),
        ("SparseMaternGraphGP", 2, 2, "Matern"),
    ]

    worst_graph_setting = [
        ("DeepGraphMultiOutputGP", 10, 6, "RBF"),
        ("DeepGraphInfoMaxMultiOutputGP", 10, 6, "RBF"),
        ("SparseMaternGraphGP", 4, 3, "RBF"),
    ]

    grid_compare_graph_run(
        "exp_result/graph_compare/",
        worst_graph_setting,
        is_test=is_test, is_verbose=is_verbose
    )

def main():
    # compare_cluster()
    # compare_many_clusters()
    # hyperparameter_search()
    run_ARMA_param_search()
    # general_test_run()
    # grid_commodities()
    # grid_compare_graph()

if __name__ == '__main__':
    main()


import pandas as pd, os
from bisect import bisect_left, bisect_right
from collections import defaultdict

class peak_table:
    def __init__(self, table_path):
        f = open(table_path, 'r')
        info = f.readlines()[1:]
        f.close()
        self.peaks = {}
        self.peaks_index = defaultdict(list)

        for line in info:
            cur_line = [x.strip() for x in line.split('\t')]
            chromosome, start, end = cur_line[0:3]
            width, total_signal, height = cur_line[4:7]
            cur_peak = peak(chromosome, int(start), int(end), int(width), float(total_signal), float(height))
            if chromosome in self.peaks:
                self.peaks[chromosome][start].add(cur_peak)
                self.peaks[chromosome][end].add(cur_peak)
            else:
                self.peaks[chromosome] = defaultdict(set)
                self.peaks[chromosome][start].add(cur_peak)
                self.peaks[chromosome][end].add(cur_peak)
        for chromosome in self.peaks.keys():
            self.peaks_index[chromosome] = sorted(self.peaks[chromosome].keys())

    def get_peaks(self, chromosome, start, end):
        indexes = self.peaks_index[chromosome]
        start_index = bisect_left(indexes, start)
        end_index = bisect_right(indexes, end)
        candidates = indexes[start_index:end_index]
        result_peaks = set()
        for c in candidates:
            for p in self.peaks[chromosome][c]:
                result_peaks.add(p)

        return result_peaks

class peak:
    def __init__(self, chromosome, start, end, width, total_signal, height):
        self.chromosome = chromosome
        self.start = start
        self.end = end
        self.width = width
        self.total_signal = total_signal
        self.height =height