import mne
import os
import glob
import pandas as pd
import gspread
from google.oauth2 import service_account
from datetime import datetime, timedelta
import shutil
import pytz
import numpy as np

# Load credentials and connect to Google Sheets API
scope = ['https://www.googleapis.com/auth/spreadsheets']
creds = service_account.Credentials.from_service_account_file(
    '/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/analysis/DeZambotti/GoogleSheet_api_key/api-access-381407-51bd2bb2d008.json', scopes=scope)
client = gspread.authorize(creds)

# Replace this with your Google Sheet key
with open('/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/analysis/DeZambotti/GoogleSheet_api_key/sheet_key.txt') as f:
    contents = f.readlines()
sheet_key = contents[0]
sheet_name = 'Devices/Session Tracking'
# Read the data from the Google Sheet
sheet = client.open_by_key(sheet_key).worksheet(sheet_name)

data = sheet.get_all_records()

# Convert the data to a pandas DataFrame

df1 = pd.DataFrame(data)

# add subject 3101 to the df structure
sheet_name2 = 'NUS Onsite Dry-Run Data'
# Read the data from the Google Sheet
sheet2 = client.open_by_key(sheet_key).worksheet(sheet_name2)
data2 = sheet2.get_all_records()

df2 = pd.DataFrame(data2)

df = pd.concat([df1, df2])

data_dir = '/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/data_raw/Zepplife/Zepp_Life_processed.csv'
HPB_BedwakeTime = pd.read_csv(data_dir)
subjects = (HPB_BedwakeTime['NUS_ID']).unique()

psg = '/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/data_raw/Sleep_Staging/Consensus_score/'

Saving_file_path = "/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/analysis/DeZambotti/Xiaomi"

# staging_file_path = "/Volumes/CSC5/SleepCognitionLab/Tera2b/Experiments/OuraValidation/Oura3/data_raw/Sleep_Staging"
# Now you can test your for loop using the `data` DataFrame

combined_data_AllDays = pd.DataFrame()
combined_data_AllDays_diff = pd.DataFrame()
i = -1

for subj in subjects:
    i = i+1

    match = df[(df['Night'] == int(subj[-1]))
               & (df['Participant ID'] == subj[:-3])]

    lightOff = match.iloc[0]['Session start local time']
    lightON = match.iloc[0]['Session end local time']

    lightOff = pd.Timestamp(lightOff)
    base_date = pd.Timestamp(match.iloc[0]['Session end date'])
    if lightOff.hour > 19:
        lightOff_day = base_date - pd.Timedelta(days=1)
        lightOff = lightOff.replace(year=lightOff_day.year,
                                    month=lightOff_day.month, day=lightOff_day.day).tz_localize(pytz.UTC)
    else:
        lightOff = lightOff.replace(year=base_date.year,
                                    month=base_date.month, day=base_date.day).tz_localize(pytz.UTC)
    lightON = pd.Timestamp(lightON).replace(
        year=base_date.year, month=base_date.month, day=base_date.day).tz_localize(pytz.UTC)

    HPB_BedwakeTime_tmp= HPB_BedwakeTime[HPB_BedwakeTime['NUS_ID']== subj]

    subject_psg=subj.replace('NUS3_','NUS3')
    psg_file = f'{psg}{subject_psg}*sleepstage_consensus.csv'
    PSG_data = glob.glob(psg_file)
    df_PSG = pd.read_csv(PSG_data[0])

    # first_sleep_ep = df_PSG.ne(0).idxmax().to_frame('pos')
    # last_sleep_ep = df_PSG.iloc[::-1, :].ne(0).idxmax().to_frame('pos')

    first_sleep_ep = df_PSG.ne(0).idxmax()
    last_sleep_ep = df_PSG.iloc[::-1, :].ne(0).idxmax()

    A = pd.to_datetime(
        HPB_BedwakeTime_tmp['start_stage'].iloc[0], utc=True)

    B = lightOff+pd.Timedelta(minutes= int(first_sleep_ep)/2)

    hour = int(A.strftime('%H'))
    minutes = int(A.strftime('%M'))
    seconds = int(A.strftime('%S'))
    hour_in_seconds_A1 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_A1 = hour_in_seconds_A1+hr

    hour = int(B.strftime('%H'))
    minutes = int(B.strftime('%M'))
    seconds = int(B.strftime('%S'))
    hour_in_seconds_B1 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_B1 = hour_in_seconds_B1+hr

    diff = hour_in_seconds_A1-hour_in_seconds_B1
    print(
        f"Bedtime start difference Based on Seconds  {subj}  NIGHT 1    {diff} minutes            OURA {A }  and  PSG {B}")

    # Calculate the time difference and include the sign
    time_diff = A - B
    diff_seconds1 = time_diff.total_seconds()

    # difference of Lights ON and Oura start of recording
    BedtimeStart_difference = diff_seconds1

    print(
        f"Bedtime start difference  {subj}  NIGHT 1   {diff_seconds1/60} minutes            OURA {A }  and  PSG {B}")

    # A2 = df_Oura['timestamp'][len(df_Oura)-1]

    A = pd.to_datetime(
        HPB_BedwakeTime_tmp['end_stage'].iloc[-1], utc=True)
    B = lightON-pd.Timedelta(minutes= (len(df_PSG)-int(last_sleep_ep))/2)

    hour = int(A.strftime('%H'))
    minutes = int(A.strftime('%M'))
    seconds = int(A.strftime('%S'))
    hour_in_seconds_A2 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_A2 = hour_in_seconds_A2+hr

    hour = int(B.strftime('%H'))
    minutes = int(B.strftime('%M'))
    seconds = int(B.strftime('%S'))
    hour_in_seconds_B2 = int(hour) * 60 + int(minutes) + int(seconds)/60

    if hour > 18:
        hr = -24*60

        hour_in_seconds_B2 = hour_in_seconds_B2+hr

    diff = hour_in_seconds_A2-hour_in_seconds_B2

    # Calculate the time difference and include the sign
    time_diff2 = A - B
    diff_seconds2 = time_diff2.total_seconds()

    BedtimeEnd_difference = diff_seconds2
    print(
        f"Waketime  difference    {subj}  NIGHT 1    {BedtimeEnd_difference} minutes            OURA {A }  and  PSG {B}")

    start_ep_diff = BedtimeStart_difference/30
    end_ep_diff = BedtimeEnd_difference/30

    combined_data_diff = pd.DataFrame({
        'subj': [subj],

        'StartTime_device': [hour_in_seconds_A1],
        'StartTime_ref': [hour_in_seconds_B1],
        'EndTime_device': [hour_in_seconds_A2],
        'EndTime_ref': [hour_in_seconds_B2]

    })
    if int(subj[-1])==1 :
        combined_data_AllDays_diff = combined_data_AllDays_diff.append(
            combined_data_diff, ignore_index=True
        )
        print(combined_data_AllDays_diff)

combined_data_AllDays_diff.to_csv(os.path.join(Saving_file_path,
                                               f"SOL_waketime_Diff_data_N1.csv"), index=False)

################################################################
################################################################
################################################################


from clutrr.actors.ancestry import Ancestry
from clutrr.relations.builder import RelationBuilder
from tqdm import tqdm
import random
import numpy as np
import json
import copy

from clutrr.args import get_args
from clutrr.store.store import Store
from clutrr.utils.utils import comb_indexes
import pandas as pd
from clutrr.relations.templator import *

#store = Store()

def generate_rows(args, store, task_name, split=0.8, prev_patterns=None):
    # pre-flight checks
    combination_length = min(args.combination_length, args.relation_length)
    if not args.use_mturk_template:
        if combination_length > 1:
            raise NotImplementedError("combination of two or more relations not implemented in Synthetic templating")
    else:
        if combination_length > 3:
            raise NotImplementedError("combinations of > 3 not implemented in AMT Templating")
    # generate
    print(args.relation_length)
    print("Loading templates...")
    all_puzzles = {}
    if args.template_split:
        train_templates = json.load(open(args.template_file + '.train.json'))
        test_templates = json.load(open(args.template_file + '.test.json'))
    else:
        train_templates = json.load(open(args.template_file + '.json'))
        test_templates = json.load(open(args.template_file + '.json'))
    if args.use_mturk_template:
        templatorClass = TemplatorAMT
    else:
        synthetic_templates_per_rel = {}
        for key, val in store.relations_store.items():
            for gender, gv in val.items():
                synthetic_templates_per_rel[gv['rel']] = gv['p']
        templatorClass = TemplatorSynthetic
        train_templates = synthetic_templates_per_rel
        test_templates = synthetic_templates_per_rel

    # Build a mapping from ANY relation to the SAME list of sentences for asking queries
    query_templates = {}
    for key, val in store.relations_store.items():
        for gender, gv in val.items():
            query_templates[gv['rel']] = store.question_store['relational']
    query_templator_class = TemplatorSynthetic

    pb = tqdm(total=args.num_rows)
    num_stories = args.num_rows
    stories_left = num_stories
    columns = ['id', 'story', 'query', 'text_query', 'target', 'text_target', 'clean_story', 'proof_state', 'f_comb',
               'task_name','story_edges','edge_types','query_edge','genders', 'syn_story', 'node_mapping', 'task_split']
    f_comb_count = {}
    rows = []
    anc_num = 0
    anc_num += 1
    anc = Ancestry(args, store)
    rb = RelationBuilder(args, store, anc)
    while stories_left > 0:
        status = rb.build()
        if not status:
            rb.reset_puzzle()
            rb.anc.next_flip()
            continue
        rb.add_facts()
        # keeping a count of generated patterns to make sure we have homogenous distribution
        if len(f_comb_count) > 0 and args.equal:
            min_c = min([v for k,v in f_comb_count.items()])
            weight = {k:(min_c/v) for k,v in f_comb_count.items()}
            rb.generate_puzzles(weight)
        else:
            rb.generate_puzzles()
        # if unique_test_pattern flag is set, and split is 0 (which indicates the task is test),
        # only take the same test patterns as before
        # also assert that the relation - test is present
        if args.unique_test_pattern and split == 0 and len(prev_patterns) > 0 and len(prev_patterns[args.relation_length]['test']) > 0:
            # if all these conditions met, prune the puzzles
            todel = []
            for pid,puzzle in rb.puzzles.items():
                if puzzle.relation_comb not in prev_patterns[args.relation_length]['test']:
                    todel.append(pid)
            for pid in todel:
                del rb.puzzles[pid]
        # now we have got the puzzles, assign the templators
        for pid, puzzle in rb.puzzles.items():
            if puzzle.relation_comb not in f_comb_count:
                f_comb_count[puzzle.relation_comb] = 0
            f_comb_count[puzzle.relation_comb] += 1
            pb.update(1)
            stories_left -= 1
        # store the puzzles
        all_puzzles.update(rb.puzzles)
        rb.reset_puzzle()
        rb.anc.next_flip()
    pb.close()
    print("Puzzles created. Now splitting train and test on pattern level")
    print("Number of unique puzzles : {}".format(len(all_puzzles)))
    pattern_puzzles = {}
    for pid, pz in all_puzzles.items():
        if pz.relation_comb not in pattern_puzzles:
            pattern_puzzles[pz.relation_comb] = []
        pattern_puzzles[pz.relation_comb].append(pid)
    print("Number of unique patterns : {}".format(len(pattern_puzzles)))
    train_puzzles = []
    test_puzzles = []
    sp = int(len(pattern_puzzles) * split)
    all_patterns = list(pattern_puzzles.keys())

    no_pattern_overlap = not args.holdout
    # if k=2, then set no_pattern_overlap=True
    if args.relation_length == 2:
        no_pattern_overlap = True

    if not no_pattern_overlap:
        # for case > 3, strict no pattern overlap
        train_patterns = all_patterns[:sp]
        pzs = [pattern_puzzles[p] for p in train_patterns]
        pzs = [s for p in pzs for s in p]
        train_puzzles.extend(pzs)
        test_patterns = all_patterns[sp:]
        pzs = [pattern_puzzles[p] for p in test_patterns]
        pzs = [s for p in pzs for s in p]
        test_puzzles.extend(pzs)
    else:
        # for case of 2, pattern overlap but templators are different
        # In this case, we have overlapping patterns, first choose the overlapping patterns
        # we directly split on puzzle level
        train_patterns = all_patterns
        test_patterns = all_patterns[sp:]
        pzs_train = []
        pzs_test = []
        for pattern in all_patterns:
            pz = pattern_puzzles[pattern]
            if pattern in test_patterns:
                # now split - hacky way
                sz = int(len(pz) * (split - 0.2))
                pzs_train.extend(pz[:sz])
                pzs_test.extend(pz[sz:])
            else:
                pzs_train.extend(pz)
        train_puzzles.extend(pzs_train)
        test_puzzles.extend(pzs_test)

    print("# Train puzzles : {}".format(len(train_puzzles)))
    print("# Test puzzles : {}".format(len(test_puzzles)))
    pb = tqdm(total=len(all_puzzles))
    # saving in csv
    for pid, puzzle in all_puzzles.items():
        task_split = ''
        if pid in train_puzzles:
            task_split = 'train'
            templator = templatorClass(templates=train_templates, family=puzzle.anc.family_data)
        elif pid in test_puzzles:
            task_split = 'test'
            templator = templatorClass(templates=test_templates, family=puzzle.anc.family_data)
        else:
            AssertionError("pid must be either in train or test")
        story_text = puzzle.generate_text(stype='story', combination_length=combination_length, templator=templator)
        fact_text = puzzle.generate_text(stype='fact', combination_length=combination_length, templator=templator)
        story = story_text + fact_text
        story = random.sample(story, len(story))
        story = ' '.join(story)
        clean_story = ' '.join(story_text)
        target_text = puzzle.generate_text(stype='target', combination_length=1, templator=templator)

        story_key_edges = puzzle.get_story_relations(stype='story') + puzzle.get_story_relations(stype='fact')
        # Build query text
        query_templator = query_templator_class(templates=query_templates, family=puzzle.anc.family_data)
        query_text = puzzle.generate_text(stype='query', combination_length=1, templator=query_templator)
        query_text = ' '.join(query_text)
        query_text = query_text.replace('?.', '?')  # remove trailing '.'
        puzzle.convert_node_ids(stype='story')
        puzzle.convert_node_ids(stype='fact')
        story_keys_changed_ids = puzzle.get_sorted_story_edges(stype='story') + puzzle.get_sorted_story_edges(stype='fact')
        query_edge = puzzle.get_sorted_query_edge()

        genders = puzzle.get_name_gender_string()

        rows.append([pid, story, puzzle.query_text, query_text, puzzle.target_edge_rel, target_text,
                     clean_story, puzzle.proof_trace, puzzle.relation_comb, task_name, story_keys_changed_ids,
                     story_key_edges, query_edge, genders, '', puzzle.story_sort_dict, task_split])
        pb.update(1)
    pb.close()

    print("{} ancestries created".format(anc_num))
    print("Number of unique patterns : {}".format(len(f_comb_count)))
    return columns, rows, all_puzzles, train_patterns, test_patterns

def test_run(args):
    store = Store(args)
    anc = Ancestry(args, store)
    rb = RelationBuilder(args, store, anc)
    rb.num_rel = 3
    all_patterns = set()
    while True:
        for j in range(len(anc.family_data.keys())):
            rb.build()
            up = rb.unique_patterns()
            all_patterns.update(up)
            print(len(all_patterns))
            rb.reset_puzzle()
        if not rb.anc.next_flip():
            break
    print("Number of unique puzzles : {}".format(len(all_patterns)))

    rb.add_facts()
    rb.generate_puzzles()
    print("Generated {} puzzles".format(len(rb.puzzles)))
    pid = random.choice(list(rb.puzzles.keys()))
    print(rb.puzzles[pid])

def main(args):
    store = Store(args)
    header, rows = generate_rows(args, store)
    df = pd.DataFrame(columns=header, data=rows)
    # split test train
    msk = np.random.rand(len(df)) > args.test
    train_df = df[msk]
    test_df = df[~msk]
    train_df.to_csv(args.output + '_train.csv')
    test_df.to_csv(args.output + '_test.csv')

if __name__ == '__main__':
    args = get_args()
    test_run(args)
    #main(args)


from PyQt5 import QtWidgets
import MainUI, RulesUI, StartUI, EditUI, GlobalManager

if __name__ == '__main__':
    app = QtWidgets.QApplication([])
    GlobalManager.app = app

    mainWindow = MainUI.MainPanel()
    GlobalManager.mainPanel = mainWindow
    rulesWindow = RulesUI.RulesPanel()
    GlobalManager.rulesPanel = rulesWindow
    editWindow = EditUI.EditPanel()
    GlobalManager.editPanel = editWindow
    startWindow = StartUI.StartPanel()
    startWindow.window.show()

    app.exec_()


import setuptools

with open("README.md", "r") as fh:
    long_description = fh.read()

exec(open('hydrotools/version.py').read())

setuptools.setup(
    name="hydrotools",
    version=__version__,
    author="Chas Egan",
    author_email="thecosmologist@gmail.com",
    description="A code torrent of washing sprawl which in watery worlds is able of all.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/chasegan/hydrotools",
    packages=setuptools.find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3",
        "Operating System :: OS Independent",
    ]
)


from torch.utils.data import DataLoader
from datasets import SEMDataset
from model_elmo import Model
import argparse
import torch
import torch.nn as nn
from tqdm import tqdm
import torch.utils.data as data
from evaluate import eval
import os
import warnings
warnings.filterwarnings("ignore")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def collate_fn(batch):
    # put question lengths in descending order so that we can use packed sequences later
    batch.sort(key=lambda x: x[-1], reverse=True)
    return data.dataloader.default_collate(batch)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=80)
    parser.add_argument('--num_hid', type=int, default=2048)#512,2048(elmo)
    parser.add_argument('--emb_dim', type=int, default=2048)#300,2048(elmo)
    parser.add_argument('--tag_emb_dim', type=int, default=100)
    parser.add_argument('--pos_emb_dim', type=int, default=512)#300,512(elmo)
    parser.add_argument('--wordG_out_dim', type=int, default=1024)  # 512,1024(transformer)
    parser.add_argument('--output', type=str, default='saved_models/')
    parser.add_argument('--idx2word', type=str, default='data/idx2word.pkl')
    parser.add_argument('--batch_size', type=int, default=16)
    parser.add_argument('--ntoken', type=int, default=5342)
    parser.add_argument('--ntag', type=int, default=46)
    parser.add_argument('--npos', type=int, default=74)
    parser.add_argument('--num_class', type=int, default=3)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--options_file', type=str, default="https://allennlp.s3.amazonaws.com/models/elmo/2x4096_512_2048cnn_2xhighway/elmo_2x4096_512_2048cnn_2xhighway_options.json")
    parser.add_argument('--weight_file', type=str, default="https://allennlp.s3.amazonaws.com/models/elmo/2x4096_512_2048cnn_2xhighway/elmo_2x4096_512_2048cnn_2xhighway_weights.hdf5")
    parser.add_argument('--alpha', type=float, default=0.2, help='Alpha for the leaky_relu.')
    parser.add_argument('--seed', type=int, default=1111, help='random seed')
    parser.add_argument('--model_name', type=str, default="model_elmo", help='saved models name')
    args = parser.parse_args()
    return args

def train(model,train_loader,dev_loader,args):
    optim = torch.optim.Adam(model.parameters(),lr=0.0001)
    log_softmax = nn.LogSoftmax(dim=2).to(device)
    div_loss = nn.KLDivLoss(reduction='none')
    softmax = nn.Softmax(dim=2).to(device)
    best_socre = 0
    best_epoch = -1
    for epoch in range(args.epochs):
        train_tq = tqdm(train_loader,desc='E{:03d}'.format(epoch),ncols=0)
        dev_tq = tqdm(dev_loader,desc='T{:03d}'.format(epoch),ncols=0)
        total_loss = 0
        model.eval()
        matchm,rankm ,score , acc, o_acc, b_acc,i_acc = evaluate(model, dev_tq,softmax,epoch)
        print(rankm,acc, o_acc, b_acc,i_acc)
        if score>best_socre:
            best_socre = score
            best_epoch = epoch
            best_match = matchm
            torch.save(model.state_dict(),args.model_name+'.pkl')

        model.train()
        for (s,a,p,a_s,s_len,mask,mask2,pos_len,graph,indicate,posid,distribution) in train_tq:
            s = s.to(device) #[B, 38]
            a = a.to(device) #[B, 38]
            p = p.to(device) #[B,38]
            a_s = a_s.to(device)#[B,9,38]
            s_len = s_len.to(device)#[B]
            mask = mask.to(device)
            mask2 = mask2.to(device)
            pos_len = pos_len.to(device)
            graph = graph.to(device)
            #indicate = indicate.to(device)
            posid = posid.to(device)
            distribution = distribution.to(device)

            prob = model(s,p,mask2,posid,graph,indicate) #[B,38,3]
            prob = log_softmax(prob)

            # div loss
            #distribution = distribution/(torch.sum(distribution,dim=2,keepdim=True)+0.00001)
            #div_l = div_loss(prob,distribution)
            _, p_max = prob.max(dim=2)

            #total = s_len.sum()
            #loss = div_l.sum()/total

            weight = torch.ones(prob.shape[0],38,3).to(device)
            weight[:,:,1:3]=5
            clamp_a = torch.clamp(a,0,3)
            #-----------------------
            #calculate loss
            new_a = torch.unsqueeze(clamp_a,-1)
            new_b = torch.zeros(prob.shape[0],38,args.num_class).to(device)
            new_b.scatter_(2, new_a, 1)
            total = s_len.sum()
            loss = (new_b*mask*prob).sum()
            loss = -loss/total
            #------------------------

            total_loss += loss.item()*total
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 2)
            optim.step()
            optim.zero_grad()
            fmt = '{:.4f}'.format
            train_tq.set_postfix(loss=fmt(loss.item()))
        print("total loss:",total_loss)

    print(best_socre)
    print(best_match)
    print(best_epoch)

def evaluate(model,dev_tq,softmax,epoch):

    total_o = 0
    total_b = 0
    total_i = 0
    o_acc = 0
    b_acc = 0
    i_acc = 0
    acc = 0
    total = 0
    with torch.no_grad():
        for (s,a,p,a_s,s_len,mask,mask2,pos_len,graph,indicate,posid,distribution) in dev_tq:

            s = s.to(device)  # [B, 38]
            p = p.to(device)
            #mask = mask.to(device)
            mask2 = mask2.to(device)

            pos_len = pos_len.to(device) #[64]
            graph = graph.to(device)    #[64,80,80]
            #indicate = indicate.to(device)#[64,80]
            posid = posid.to(device) #[64,80]

            prob = model(s,p,mask2,posid,graph,indicate) #[dev_set,38,3]
            prob = softmax(prob)
            prob = prob.cpu()
            ac, o, b, i, num_o, num_b, num_i, num = accuracy(prob,distribution,mask[:,:,0])
            acc+=ac
            o_acc+=o
            b_acc+=b
            i_acc+=i
            total_o+=num_o
            total_b+=num_b
            total_i+=num_i
            total += num
            prob = prob.numpy() #[B,38,3]

            #print(prob)
            with open("reslut"+str(epoch)+".txt",'a') as f:
                for i in range(prob.shape[0]):
                    for j in range(s_len[i]):
                        f.write('a \t b \t' +str(1-prob[i][j][0])+'\n')
                    f.write('\n')

    matchm,rankm = eval("result"+str(epoch)+".txt","train_dev_data/dev.txt")
    return matchm,rankm,sum(matchm.values()),acc/total, o_acc/total_o, b_acc/total_b, i_acc/total_i

def accuracy(prob,label,mask):
    #prob [B,38,3] label [B,38,3],mask:[B,38]

    prob_max = torch.argmax(prob,dim=2)
    label_max = torch.argmax(label,dim=2)
    acc = torch.sum((prob_max==label_max).float()*mask)
    o_prob = prob_max == 0
    o_label = label_max == 0
    b_prob = prob_max == 1
    b_label = label_max == 1
    i_prob = prob_max == 2
    i_label = label_max == 2
    o_acc = torch.sum((o_prob==o_label).float()*mask*o_label.float())
    b_acc = torch.sum((b_prob==b_label).float()*mask*b_label.float())
    i_acc = torch.sum((i_prob==i_label).float()*mask*i_label.float())
    total_o = torch.sum(o_label.float()*mask)
    total_b = torch.sum(b_label.float()*mask)
    total_i = torch.sum(i_label.float()*mask)
    return acc, o_acc, b_acc, i_acc, total_o,total_b,total_i,torch.sum(mask)

if __name__ == '__main__':
    args = parse_args()
    torch.manual_seed(args.seed)

    train_set = SEMDataset("train",'data/word2idx.pkl','data/pos2idx.pkl')
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)

    dev_set = SEMDataset('dev','data/word2idx.pkl','data/pos2idx.pkl')
    dev_loader = DataLoader(dev_set, batch_size=args.batch_size, shuffle=False)

    m = Model(args).to(device)

    #m = nn.DataParallel(m).to(device)
    train(m,train_loader,dev_loader,args)


import numpy as np
from sklearn import svm
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier, GradientBoostingClassifier
from sklearn.tree import DecisionTreeClassifier

def get_labeled_windowed_data(observations, window_size=7):
    """
    Split up the observations into windowed chunks. Each windowed chunk of
    observations is associated with a label vector of what the price change is
    per market *immediately after* the windowed chunk (+1 for price goes up,
    0 for no change, and -1 for price goes down). Thus, a classifier's task
    for the data is given a windowed chunk, to predict what its label is
    (i.e., given recent percent changes in all the markets, predict the
    directions of the next price changes per market).

    Inputs
    ------
    - observations: 2D array; each column is a percent-change time series data
        for a specific market
    - window_size: how large the window is (in number of time points)

    Outputs
    -------
    - windows: 3D array; each element of the outermost array is a 2D array
        of the same format as `observations` except where the number of time
        points is exactly `window_size`
    - window_labels: 2D array; `window_labels[i]` is a 1D vector of labels
        corresponding to the time point *after* the window specified by
        `windows[i]`; `window_labels[i]` says what the price change is for
        each market (+1 for going up, 0 for staying the same, and -1 for going
        down)

    *WARNING*: Note that the training data produced here is inherently not
    i.i.d. in that `windows[0]` and `windows[1]`, for instance, will largely
    overlap!
    """
    num_time_points, num_markets = observations.shape
    windows = []
    window_labels = []
    for start_idx in range(num_time_points-window_size):
        windows.append(observations[start_idx:start_idx+window_size])
        window_labels.append(1*(observations[start_idx+window_size] > 0)
                             -1*(observations[start_idx+window_size] < 0))
    windows = np.array(windows)
    window_labels = np.array(window_labels)
    return windows, window_labels

# global variables to be saved for the trained classifier
clfs = None
clfs1 = None

def train(windows, window_labels):
    """
    Your training procedure goes here! It should train a classifier where you
    store whatever you want to store for the trained classifier as *global*
    variables. `train` will get called exactly once on the exact same training
    data you have access to. However, you will not get access to the mystery
    test data.

    Inputs
    ------
    - windows, window_labels: see the documentation for the output of
        `get_labeled_windowed_data`
    """

    # -------------------------------------------------------------------------
    # YOUR CODE HERE
    #

    # The autograder wants you to explicitly state which variables are global
    # and are supposed to thus be saved after training for use with prediction.
    global clfs, clfs1, clfs2

    #guess = 0

    #from sklearn.metrics import log_loss

    n, m = 7, 5
    X = np.concatenate((windows[:,-n:,0], windows[:,-n:,1], windows[:,-n:,2], windows[:,-n:,3]), axis=1)
    for i in range(4):
        for k in range(2,8):
            X = np.concatenate((X, ((windows[:,-k:,i] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            X = np.concatenate((X, ((windows[:,-k:,i] > 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            X = np.concatenate((X, ((windows[:,-k:,i] < 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            #X = np.concatenate((X, ((windows[:,-k:,1] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            #X = np.concatenate((X, ((windows[:,-k:,2] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
            #X = np.concatenate((X, ((windows[:,-k:,3] == 0).sum(axis=1)==k).reshape(X.shape[0],1)), axis=1)
        X = np.concatenate((X, ((windows[:,-m:,i] == 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)
        X = np.concatenate((X, ((windows[:,-m:,i] > 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)
        #X = np.concatenate((X, ((windows[:,-m:,0] >= 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)
        X = np.concatenate((X, ((windows[:,-m:,i] < 0).sum(axis=1)).reshape(X.shape[0],1)), axis=1)

    clfs = [RandomForestClassifier(n_estimators=10) for i in range(4)]
    #clfs = [ AdaBoostClassifier(DecisionTreeClassifier(max_depth=1),
    #                     algorithm="SAMME",
    #                    n_estimators=20) for i in range(4)]
    clfs1 = [GradientBoostingClassifier(n_estimators=100, learning_rate=1.0, max_depth=1, random_state=0) for i in range(4)]
    clfs2 = [svm.SVC(C=10**6, kernel='rbf') for i in range(4)]

    for i in range(4):
       clfs[i] = clfs[i].fit(X, window_labels[:,i])
       clfs1[i] = clfs1[i].fit(X, window_labels[:,i])
       clfs2[i] = clfs2[i].fit(X, window_labels[:,i])
       #clf_probs = clf.predict_proba(X)
       #score = log_loss(y, clf_probs)
    #
    # END OF YOUR CODE
    # -------------------------------------------------------------------------

def forecast(window):
    """
    Your forecasting method goes here! You may assume that `train` has already
    been called on training data and so any global variables you stored as a
    result of running `train` are available to you here for prediction
    purposes.

    Input
    -----
    - window: 2D array; each column is 7 days worth of percent changes in
        price for a specific market

    Output
    ------
    1D array; the i-th entry is a prediction for whether the percentage
    return will go up (+1), stay the same (0), or go down (-1) for the i-th
    market
    """

    # -------------------------------------------------------------------------
    # YOUR CODE HERE
    #
    #from collections import Counter

    global clfs, clfs1, clfs2
    from scipy.stats import mode

    n, m = 7, 5
    X = window[-n:,:]
    X = X.flatten('F')
    for i in range(4):
        for k in range(2,8):
            X = np.append(X, (window[-k:,i] == 0).sum()==k)
            X = np.append(X, (window[-k:,i] > 0).sum()==k)
            X = np.append(X, (window[-k:,i] < 0).sum()==k)
            #X = np.append(X, (window[-k:,1] == 0).sum()==k)
            #X = np.append(X, (window[-k:,2] == 0).sum()==k)
            #X = np.append(X, (window[-k:,3] == 0).sum()==k)
        X = np.append(X, (window[-m:,i] == 0).sum())
        X = np.append(X, (window[-m:,i] > 0).sum())
        #X = np.append(X, (window[-m:,0] >= 0).sum())
        X = np.append(X, (window[-m:,i] < 0).sum())

    predicted_labels = np.array([mode([clfs[idx].predict(X), clfs1[idx].predict(X), clfs2[idx].predict(X)], axis=0).mode.flatten() for idx in range(window.shape[1])])

    #
    # END OF YOUR CODE
    # -------------------------------------------------------------------------

    return predicted_labels

def main():
    # get coconut oil challenge training data
    observations = []
    with open('coconut_challenge.csv', 'r') as f:
        for line in f.readlines():
            pieces = line.split(',')
            if len(pieces) == 5:
                observations.append([float(pieces[1]),
                                     float(pieces[2]),
                                     float(pieces[3]),
                                     float(pieces[4])])
    observations = np.array(observations)
    train_windows, train_window_labels = \
        get_labeled_windowed_data(observations, window_size=7)

    train(train_windows, train_window_labels)

    # figure out accuracy of the trained classifier on predicting labels for
    # the training data
    train_predictions = []
    for window, window_label in zip(train_windows, train_window_labels):
        train_predictions.append(forecast(window))
    train_predictions = np.array(train_predictions)

    train_prediction_accuracy_plus1 = \
        np.mean(train_predictions[train_window_labels == 1]
                == train_window_labels[train_window_labels == 1])
    train_prediction_accuracy_minus1 = \
        np.mean(train_predictions[train_window_labels == -1]
                == train_window_labels[train_window_labels == -1])
    train_prediction_accuracy_0 = \
        np.mean(train_predictions[train_window_labels == 0]
                == train_window_labels[train_window_labels == 0])
    print('Training accuracy for prediction +1:',
          train_prediction_accuracy_plus1)
    print('Training accuracy for prediction -1:',
          train_prediction_accuracy_minus1)
    print('Training accuracy for prediction 0:',
          train_prediction_accuracy_0)

if __name__ == '__main__':
    main()

courses/MITx/6.008.1x - Computational Probability & Inference (2017)/Week11/final_proj_challenge.py


# Scrape list of YA authors from wikipedia and save to CSV file
# Author list is broken into sections alphabetically
#
# Key HTML Structure:
#

#

#       : List of books...
#

#

#
# Note: The "section" with a class="mf-section-26" has list items we don't want
#       to capture.

import requests
from bs4 import BeautifulSoup
import pandas as pd

url = 'https://en.m.wikipedia.org/wiki/List_of_young_adult_fiction_writers'
output_file = "data/raw/authors.csv"

data = requests.get(url)
soup = BeautifulSoup(data.content, "html.parser")
sections = soup.find_all("section", {"class":"collapsible-block"})

authors = []
author_index = []
index_val = 1

for section in sections:
    if "mf-section-26" not in section["class"]: # don't capture the last block
        author_list = section.find_all("li")
        for author in author_list:
            authors.append(author.find_next("a").get("title"))
            author_index.append(index_val)
            index_val += 1

df = pd.DataFrame(authors, index= author_index, columns=['Name'])

remove_list = [
    '(author)',
    '(writer)',
    '(novelist)',
    '(author, born 1954)',
    '(TV producer)',
    '(Irish author)',
    '(young adult author)',
    '(historian)']

for item in remove_list:
    df['Name'] = df['Name'].str.replace(item, '', regex= False)

df.to_csv(output_file)


import torch
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
import argparse
import torchvision.models as models
from classification_dataset import DVTDataset
from torch import nn
from sklearn.metrics import roc_auc_score, precision_score, recall_score
import numpy as np
import re
from torch.utils.tensorboard import SummaryWriter
import json

class Identity(nn.Module):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x

def evaluate(model, dataloader, criterion, writer=None, epoch=0, val=False):
    auc = 0
    prec = 0
    recall = 0
    model = model.eval()
    total_preds = np.array([])
    total_labels = np.array([])
    total_hard_preds = np.array([])

    for i, data in enumerate(dataloader):
        inputs, labels, access_number = data['image'], data['label'], data['access_number']
        inputs = inputs.repeat(1, 3, 1, 1).float()
        inputs, labels = inputs.cuda(), labels.cuda()
        outputs = model(inputs)
        if val:
            # print(torch.argmax(outputs, 1))
            loss = criterion(outputs, labels)
            preds = torch.argmax(outputs, 1)
            # print(preds)
            # auc = auc + roc_auc_score(labels.cpu().numpy(), preds.cpu().numpy())
            # prec = prec + precision_score(labels.cpu().numpy(), preds.cpu().numpy())
            # recall = recall + recall_score(labels.cpu().numpy(), preds.cpu().numpy())
            # print("AUC: ", roc_auc_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("prec: ", precision_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("recall: ", recall_score(labels.cpu().numpy(), preds.cpu().numpy()))
            if i % args.log_interval == 0:
                batch_number = i + epoch * len(dataloader)
                print("Val Batch Number: {} | Val Loss: {} ".format(batch_number, loss))
                writer.add_scalar("Validation Loss", loss, batch_number)

        else:
            # Compute AUC/ Precision
            preds = torch.argmax(outputs, 1)
            # soft = nn.Softmax()
            # out = soft(outputs)
            # print(preds)
            try:
                # print("Labels", labels)
                # print("Preds", preds)
                out_preds = torch.max(outputs, 1)[0]
                out_hard_preds = torch.argmax(outputs, 1)
                total_hard_preds = np.concatenate((total_hard_preds, out_hard_preds.cpu().detach().numpy()))
                total_preds = np.concatenate((total_preds, out_preds.cpu().detach().numpy()))
                total_labels = np.concatenate((total_labels, labels.cpu().numpy()))
                current_prec = precision_score(labels.cpu().numpy(), preds.cpu().numpy())
                current_recall = recall_score(labels.cpu().numpy(), preds.cpu().numpy())
                # print("Current Prec", current_prec)
                # print("Current Recall", current_recall)
                # print("Equal", current_prec == 1.0)
                auc = auc + roc_auc_score(labels.cpu().numpy(), preds.cpu().detach().numpy())
                prec = prec + precision_score(labels.cpu().numpy(), preds.cpu().numpy())
                recall = recall + recall_score(labels.cpu().numpy(), preds.cpu().numpy())

            except ValueError:
                pass
            # print("AUC: ", roc_auc_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("prec: ", precision_score(labels.cpu().numpy(), preds.cpu().numpy()))
            # print("recall: ", recall_score(labels.cpu().numpy(), preds.cpu().numpy()))

    auc = float(auc) / i
    prec = float(prec) / i
    recall = float(recall) / i
    print("Average AUC: ", auc)
    print("Average Prec: ", prec)
    print("Average Recall: ", recall)
    d = {'preds': list(total_preds), 'labels': list(total_labels), 'hard_preds': list(total_hard_preds)}
    jn = json.dumps(d)
    f = open("dict_2.json", "w")
    f.write(jn)
    f.close()
    # writer.add_scalar("Validation Loss", loss, i)
    # writer.add_scalar("Validation AUC", auc, i)
    # writer.add_scalar("Validation Precision", prec, i)
    # writer.add_scalar("Validation Recall", recall, i)
    return auc

def main(args):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    train_dataset = DVTDataset(csv_file=args.train, root_dir=args.root_dir, device=device)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)

    model = models.densenet161()
    print("Loading model onto CUDA")
    state_dict = torch.load("densenet161-8d451a50.pth")
    pattern = re.compile(
        r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')

    for key in list(state_dict.keys()):
        res = pattern.match(key)
        if res:
            new_key = res.group(1) + res.group(2)
            state_dict[new_key] = state_dict[key]
            del state_dict[key]
    model.load_state_dict(state_dict)
    model.classifier = nn.Linear(2208, 2)

    model = model.cuda()
    model.train()

    print("Finished loading model onto CUDA")

    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    # optimizer = optim.Adam(model.parameters(), lr = 0.0001)
    # optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9)
    # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
    writer = SummaryWriter("classification_lowlr_runs")
    # weights = np.zeros(1000)
    # weights[0] = 1.0
    # weights[1] = 1.0
    criterion = nn.CrossEntropyLoss()
    if not args.evaluate:
        for epoch in range(args.epochs):
            for i, data in enumerate(train_loader):
                # get the inputs; data is a list of [inputs, labels]
                inputs, labels = data['image'], data['label']
                # print(inputs.shape)
                inputs = inputs.repeat(1, 3, 1, 1).float()
                inputs, labels = inputs.cuda(), labels.cuda()
                # print(inputs.shape)
                # zero the parameter gradients
                optimizer.zero_grad()
                # print("Forward pass")
                # forward + backward + optimize

                outputs = model(inputs)
                # print(outputs)
                # print(torch.argmax(outputs, 1))
                # print("Forward pass complete.")
                loss = criterion(outputs, labels)
                loss.backward()
                # print("Backward pass.")
                optimizer.step()
                if i % args.log_interval == 0:
                    batch_number = i + epoch * len(train_loader)
                    # print(torch.argmax(outputs, 1))
                    writer.add_scalar("Training Loss", loss, batch_number)
                    print("Batch Number: {} | Loss: {}".format(i, loss))
            torch.save(model.state_dict(), "classification_lowlr/model_{}.pt".format(epoch))
            val_dataset = DVTDataset(csv_file=args.val, root_dir=args.root_dir, device=device)
            val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True,
                                                     num_workers=0)
            evaluate(model, val_loader, criterion, writer, epoch, val=True)
            # scheduler.step()
            # test_dataset = DVTDataset(csv_file = args.test, root_dir = args.root_dir, device = device)
            # test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = args.batch_size, shuffle = True, num_workers = 0)
            # evaluate(model, test_loader, criterion)

    model.load_state_dict(torch.load(args.checkpoint_name + "classification_dropout/model_9.pt"))
    model.eval()
    test_dataset = DVTDataset(csv_file=args.test, root_dir=args.root_dir, device=device)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0)
    evaluate(model.eval(), test_loader, criterion)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--train', default='train_femoral.csv')
    parser.add_argument('--val', default='validation_femoral.csv')
    parser.add_argument('--test', default='train_femoral.csv')
    parser.add_argument('--root_dir', default='')
    parser.add_argument('--batch_size', default=10)
    parser.add_argument('--epochs', default=10)
    parser.add_argument('--lr', default=0.001)
    parser.add_argument('--log_interval', default=10)
    parser.add_argument('--val_interval', default=100)
    parser.add_argument('--evaluate', default=False)
    parser.add_argument('--checkpoint_name', default="")

    args = parser.parse_args()
    main(args)


'''Class which contains canvas (self.scene) for inserting nodes and creating pipeline.
Functions->
addObjectModel: assigns model to self.shape
addObject: assigns non-model objects to self.shape
mouseDoubleClickEvent: inserts object in self.shape to diagram at position of double-click
addArrow: inserts arrow from first object selected to second object
getTotalCost: returns total cost of models in diagram
keyPressEvent: deletes selected object from diagram
calcCloudCost: returns total cost of deploying models in diagram to cloud
updateConfigList: updates list of configurations in each model
updateTfrList: updates list of transformers in each non-model object
removeTfrFromList: removes transformer from list
getNetLatency: returns latency of pipeline
'''

from PyQt5 import QtCore
from UniqueModel import UniqueModel
from CloudCost import *
from Arrow import Arrow
from Shape import Shape
from ModelShape import ModelShape
from PyQt5.QtWidgets import QGraphicsEllipseItem, QGraphicsLineItem, QGraphicsPixmapItem, QGraphicsRectItem, QGraphicsTextItem, QGraphicsView, QGraphicsScene
from PyQt5.QtCore import QPointF, Qt
CustomObjectRole = QtCore.Qt.UserRole + 1

class GraphicView(QGraphicsView):  #Canvas for inserting nodes and creating pipeline
    def __init__(self):
        super().__init__()

        self.scene = QGraphicsScene()
        self.setScene(self.scene)
        self.setSceneRect(0, 0, 1200, 1000)
        self.shape = None
        self.prevShape = None
        self.connectorActive = False
        self.setViewportUpdateMode(QGraphicsView.FullViewportUpdate)
        self.configs = []
        self.tfrs = []

    def addObjectModel(self,model):
        self.shape = ModelShape(model)

    def addObject(self,text:str,a,b,cost,latency,type1):
        self.shape = Shape(text,a,b,cost,latency,type1)

    def mouseDoubleClickEvent(self, event): #add object in self.shape to location clicked, if location is empty
        position = self.mapToScene(event.pos())
        if(self.shape is not None):
            if(self.itemAt(position.x(),position.y()) is None):
                if(self.prevShape is None or self.prevShape != self.shape):
                    self.scene.addItem(self.shape)
                    self.shape.setPos(QPointF(position.x(), position.y()))
                    if(self.shape.type() == QGraphicsRectItem().type()):
                        self.shape.configs = self.configs
                    elif(self.shape.type() == QGraphicsEllipseItem().type()):
                        self.shape.tfrs = self.tfrs
                    self.prevShape = self.shape

    def addArrow(self): #add arrow from first object clicked to second object clicked, connected to 'Connector' button
        items = self.scene.selectedItems()
        if(len(items) == 2):
            if(items[0].selectedTime < items[1].selectedTime):
                start_item = items[0]
                end_item = items[1]
            else:
                start_item = items[1]
                end_item = items[0]
            if((items[0].type() == QGraphicsRectItem().type()) & (items[1].type() == QGraphicsRectItem().type())):
                selector = Shape("",20,20,0,0,"")
                self.scene.addItem(selector)
                selector.tfrs = self.tfrs
                arrow1 = Arrow(start_item,selector)
                arrow2 = Arrow(selector,end_item)
                selector.setPos((start_item.scenePos()+end_item.scenePos())/2)
                self.scene.addItem(arrow1)
                self.scene.addItem(arrow2)
            else:
                arrow = Arrow(start_item,end_item)
                self.scene.addItem(arrow)

            items[0].setSelected(False)
            items[1].setSelected(False)

    def getTotalCost(self):
        items = self.scene.items()
        totalcost = 0
        for item in items:
            if(item.type() == QGraphicsRectItem().type() or item.type() == QGraphicsEllipseItem().type()):
                totalcost += item.Cost
        return totalcost

    def keyPressEvent(self, event): #for deleting object from scene. Error: object gets removed but not deallocated
        if event.key() == Qt.Key_Backspace:
            items = self.scene.selectedItems()
            for item in items:
                self.scene.removeItem(item)
                if(item.type() != QGraphicsLineItem().type()):
                    for arrow in item.in_arrows:
                        self.scene.removeItem(arrow)
                    for arrow in item.out_arrows:
                        self.scene.removeItem(arrow)

    def calcCloudCost(self): #cost for cloud for all item in scene
        totalsum = 0
        items = self.scene.items()
        for item in items:
            if(item.type() == QGraphicsRectItem().type()):
                    totalsum += item.model.CloudCost
        return totalsum

    def updateConfigList(self,config):
        self.configs.append(config)
        for item in self.scene.items():
            if(item.type() == QGraphicsRectItem().type()):
                item.configs = self.configs

    def updateTfrList(self,tfr):
        self.tfrs.append(tfr)
        for item in self.scene.items():
            if(item.type() == QGraphicsEllipseItem().type()):
                item.tfrs = self.tfrs

    def removeTfrFromList(self,tfr):
        self.tfrs.remove(tfr)
        for item in self.scene.items():
            if(item.type() == QGraphicsEllipseItem().type()):
                item.tfrs = self.tfrs

    def getNetLatency(self):
        def longestPath(graph):
            ans = 0
            n = len(graph)
            table = [-1] * n

            def dfs(u):
                if table[u] != -1:
                    return table[u]
                p_len = 0
                for v in graph[u]:
                    p_len = max(p_len, v[1] + dfs(v[0]))
                table[u] = p_len
                return p_len

            for i in range(n):
                ans = max(ans, dfs(i))    #graph[u] contains arrays of outgoing nodes; graph is an adjacency list
            return ans

        source = Shape("",10,10,0,0,"") #virtual node connected to all sources
        nodelist = [source]
        items = self.scene.items()
        for item in items:
            if((item.type() == QGraphicsRectItem().type()) | (item.type() == QGraphicsEllipseItem().type())):
                nodelist.append(item)
                if((item.in_arrows == []) and (item is not source)):
                    Arrow(source,item)

        graph = [[(nodelist.index(arrow.end_item),arrow.end_item.Latency) if arrow.end_item.type() == QGraphicsRectItem().type() or arrow.end_item.type() == QGraphicsEllipseItem().type() else (nodelist.index(arrow.end_item),0) for arrow in node.out_arrows] for node in nodelist]
        for arrow in source.out_arrows:
            arrow.end_item.in_arrows = []
        source.out_arrows = []
        return longestPath(graph)


"""empty message

Revision ID: f205d13cfa6d
Revises: 11551aecc7cb
Create Date: 2018-03-19 14:24:20.894749

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = 'f205d13cfa6d'
down_revision = '11551aecc7cb'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index('ix_customers_hash', table_name='customers')
    op.create_index(op.f('ix_customers_hash'), 'customers', ['hash'], unique=True)
    op.create_unique_constraint(None, 'customers', ['id'])
    op.drop_index('customer_hash', table_name='openned')
    op.drop_index('customer_id', table_name='openned')
    op.create_unique_constraint(None, 'openned', ['id'])
    op.drop_constraint('openned_ibfk_2', 'openned', type_='foreignkey')
    op.drop_constraint('openned_ibfk_1', 'openned', type_='foreignkey')
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_foreign_key('openned_ibfk_1', 'openned', 'customers', ['customer_hash'], ['hash'])
    op.create_foreign_key('openned_ibfk_2', 'openned', 'customers', ['customer_id'], ['id'])
    op.drop_constraint(None, 'openned', type_='unique')
    op.create_index('customer_id', 'openned', ['customer_id'], unique=True)
    op.create_index('customer_hash', 'openned', ['customer_hash'], unique=True)
    op.drop_constraint(None, 'customers', type_='unique')
    op.drop_index(op.f('ix_customers_hash'), table_name='customers')
    op.create_index('ix_customers_hash', 'customers', ['hash'], unique=False)
    # ### end Alembic commands ###


import numpy as np

from brl_gym.wrapper_envs import BayesEnv
from brl_gym.wrapper_envs.env_sampler import DiscreteParamEnvSampler, UniformSampler
from brl_gym.estimators.param_env_discrete_estimator import ParamEnvDiscreteEstimator
from brl_gym.envs.classic_control.cartpole import CartPoleEnv
from brl_gym.wrapper_envs.explicit_bayes_env import ExplicitBayesEnv
from gym.spaces import Box, Dict

class BayesCartPoleEnv(BayesEnv):
    # Wrapper envs for mujoco envs
    def __init__(self):
        env = CartPoleEnv()
        self.estimator = ParamEnvDiscreteEstimator(env, discretization=2)

        self.env_sampler = DiscreteParamEnvSampler(env, 2)
        super(BayesCartPoleEnv, self).__init__(env, self.estimator)
        self.nominal_env = env

    def reset(self):
        self.env = self.env_sampler.sample()
        return super().reset()

    def step(self, action):
        prev_belief = self.estimator.get_belief()
        prev_state = self.env.get_state()
        obs, reward, done, info = self.env.step(action)
        info['prev_state'] = prev_state
        info['curr_state'] = self.env.get_state()

        # Estimate
        self.estimator.estimate(action, obs, **info)
        belief = self.estimator.get_belief()
        info['belief'] = belief

        obs = np.concatenate([obs, belief], axis=0)
        return obs, reward, done, info

class ExplicitBayesCartPoleEnv(ExplicitBayesEnv):
    def __init__(self):
        env = CartPoleEnv()
        self.estimator = ParamEnvDiscreteEstimator(env, discretization=2)

        self.env_sampler = DiscreteParamEnvSampler(env, 2)
        super(ExplicitBayesCartPoleEnv, self).__init__(env, self.estimator)
        self.nominal_env = env

        self.observation_space = Dict(
            {"obs": env.observation_space, "zbel": self.estimator.belief_space})
        self.internal_observation_space = env.observation_space
        self.env = env

    def _update_belief(self,
                             action,
                             obs,
                             **kwargs):
        # Estimate
        self.estimator.estimate(
                action, obs, **kwargs)
        belief = self.estimator.get_belief()
        return belief, kwargs

    def step(self, action):
        prev_state = self.env.get_state()
        obs, reward, done, info = self.env.step(action)
        info['prev_state'] = prev_state
        info['curr_state'] = self.env.get_state()

        bel, info = self._update_belief(
                                        action,
                                        obs,
                                        **info)
        true_param = self.env.get_params()
        # mass = true_param['masspole']
        length = true_param['length']

        exp1 = np.argwhere(self.env_sampler.param_sampler_space['length'] == length)[0,0]
        # exp2 = np.argwhere(self.env_sampler.param_sampler_space['masspole'] == mass)[0,0]
        exp_id = exp1
        info['expert'] = exp_id

        return {'obs':obs, 'zbel':bel}, reward, done, info

    def reset(self):
        self.env = self.env_sampler.sample()
        obs = self.env.reset()
        self.estimator.reset()
        bel, _ = self._update_belief(action=None, obs=obs)
        self.last_obs = (obs, bel)
        return {'obs':obs, 'zbel':bel}

if __name__ == "__main__":
    env = ExplicitBayesCartPoleEnv()
    print(env.reset())
    for _ in range(200):
        print(env.step(env.action_space.sample())[0])

    import IPython; IPython.embed()


# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

from django.db import migrations, models
import django.utils.timezone

class Migration(migrations.Migration):
    dependencies = []

    operations = [
        migrations.CreateModel(
            name="StatusMessage",
            fields=[
                (
                    "id",
                    models.AutoField(
                        verbose_name="ID",
                        serialize=False,
                        auto_created=True,
                        primary_key=True,
                    ),
                ),
                ("message", models.TextField()),
                ("severity", models.CharField(max_length=20)),
                ("enabled", models.BooleanField(default=True)),
                ("created_at", models.DateTimeField(default=django.utils.timezone.now)),
            ],
        )
    ]


'''
Author: lmio 2091319361@qq.com
Date: 2023-08-12 20:53:28
LastEditors: lmio 2091319361@qq.com
Description: 23. 合并 K 个升序链表
'''

import heapq
from typing import List, Optional

from utils.node import ListNode

ListNode.__lt__ = lambda a, b: a.val < b.val

class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        cur = ListNode()
        dummy = cur
        h = [node for node in lists if node]
        heapq.heapify(h)
        while h:
            node = heapq.heappop(h)
            cur.next = node
            if node.next:
                heapq.heappush(h, node.next)
            cur = cur.next
        return dummy.next


__author__ = "Andrew DiFabbio"
__email__ = "avd38@cornell.edu"

import argparse
import numpy as np
import pandas as pd
import os
from matplotlib import pyplot as plt
import exp_params as exp
from scipy.stats import chi2

np.seterr(divide='ignore', invalid='ignore')

# COMBINING REDUCED DATA----------------------------------------------------------------------------------------
def sum_over_results(path,num_ch,num_bins,verbose=False):
    """Sum the coincidences, accidentals, and fluorescence counts for a thickness within a
    folder, returning the result in the format <fluorescence counts>,<coincidences>,
    <accidentals>.

    If verbose, these results are returned with respect to every permuation of fluorescence
    and scattered detector. Otherwise they are simply returned according to fluorescence
    detector."""

    fl_ct=np.zeros((num_ch,1))
    coin=np.zeros((num_ch*num_ch,num_bins))
    acoin=np.zeros((num_ch*num_ch,num_bins))

    for l in os.listdir(path):
        if '.txt' in l: continue

        if 'fluor_count' in l:
            fl_ct+=np.genfromtxt(os.path.join(path,l),delimiter=',').reshape((4,1))
        elif 'coincidence' in l:
            coin+=np.genfromtxt(os.path.join(path,l),delimiter=',')
        elif 'accidentals' in l:
            acoin+=np.genfromtxt(os.path.join(path,l),delimiter=',')

    if not verbose:
        coin=np.sum(coin.reshape((num_ch,num_ch,num_bins)),axis=1)
        acoin=np.sum(acoin.reshape((num_ch,num_ch,num_bins)),axis=1)

    return fl_ct,coin,acoin

def get_ang_corr(num_ch):
    """Calculate the factor for angular dependence in IAB as it relates to the detectors. This
    is essentially an integral over the solid angle of the angular coefficient."""

    ang_corr=np.zeros((num_ch,1))
    for r in range(num_ch):
        for ch in range(num_ch):
            if r==ch: continue
            ang_corr[r]+=exp.CHANNEL_ANGULAR_COEFF[ch]*exp.CHANNEL_SOLID_ANGLE[ch]

    return np.power(ang_corr,-1)

def combine_summed_results(data,num_ch,num_orb,e_bin,n_bins,atten=True,thick=None):
    """Combine the already summed counts by applying angular and attenuation corrections and
    handling error propagation as necessary. Returns <normalized excess>,<normalized error>.

    If atten, attenuation correcion is applied, and a thickness must be given in the form
    'data/40nm/', for example."""

    if atten and thick == None: raise Exception('combine_summed_results: attenuation correction needs thickness.')
    ang_corr=get_ang_corr(num_ch)
    fl_ct,coin,acoin=data

    excesses=ang_corr*(coin-acoin/num_orb)/fl_ct
    variances=ang_corr**2*((coin+acoin/num_orb**2)/fl_ct**2 + (coin-acoin/num_orb)**2/fl_ct**3)
    if atten: excesses,variances=correct_atten(excesses,variances,e_bin,n_bins,thick)

    norm_vars=np.power(np.sum(np.power(variances,-1),axis=0),-1)

    return (norm_vars*np.sum(excesses/variances,axis=0))/(e_bin*exp.XMAP_ENERGY_UNIT), \
        np.sqrt(norm_vars)/(e_bin*exp.XMAP_ENERGY_UNIT)

def combine_folder(path,num_ch,num_bins,num_orb,e_bin,atten=True,thick=None):
    """Read and combine the data for a given thickness."""

    if atten and thick == None: raise Exception('combine_folder: attenuation correction needs thickness.')
    return combine_summed_results(sum_over_results(path,num_ch,num_bins),num_ch,num_orb,e_bin,num_bins,atten,thick)

def combine_data(wdir,num_ch,num_bins,num_orb,e_bin,atten=True,write=True):
    """Fully combine the reduced data. If atten, attenuation correction is applied. If write, data
    will be written to CSV files afterward, otherwise the result is returned. The results consist
    of normalized excess and corresponding errors with rows representing thickness in the order 40nm,
    80nm, 160nm, 320nm, empty, and columns representing energy bin."""

    result_excesses=np.zeros((5,num_bins))
    result_errors=np.zeros((5,num_bins))

    for r,t in enumerate(['data/40nm/','data/80nm/','data/160nm/','data/320nm/','data/empty/']):
        exc,err=combine_folder(os.path.join(wdir,t),num_ch,num_bins,num_orb,e_bin,atten,t)
        result_excesses[r,:]+=exc
        result_errors[r:]+=err

    if write:
        np.savetxt(os.path.join(wdir,'combined_result_excesses.csv'),result_excesses,delimiter=',')
        np.savetxt(os.path.join(wdir,'combined_result_errors.csv'),result_errors,delimiter=',')
    else:
        return result_excesses,result_errors

# PLOTTING COMBINED RESULTS AND PREDICTIONS---------------------------------------------------------------------
def k(t):
    return exp.K_COEFF*(t/exp.D)**0.75

def W_secondaries(t,Es):
    return 2*k(t)*exp.Z*(exp.E0-exp.EB-Es)/(4*np.pi*Es)

def W_LET(Es):
    return 2*exp.ALPHA*((exp.E0-exp.EB)/exp.MC2)/(4*np.pi**2*Es)

def plot_thickness(ax,exc,err,Emin,Emax,e_bin,t):
    """Plot the combined results for a given thickness and energy range along with the predicted rates
    for that thickness. Blue for secondaries, red for LET, and green for secondaries+LET."""

    start,stop=int(Emin/(e_bin*exp.XMAP_ENERGY_UNIT)),int(Emax/(e_bin*exp.XMAP_ENERGY_UNIT)+0.5)
    Erange=np.arange((start+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT), \
                     (stop+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT),e_bin*exp.XMAP_ENERGY_UNIT)
    Erange_=np.linspace((start+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT), \
                      (stop+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT),1000)
    # plot predictions
    ax.plot(Erange_,0*Erange_,color='k')
    ax.plot(Erange_,W_secondaries(t,Erange_),color='b')
    ax.plot(Erange_,W_LET(Erange_),color='r')
    ax.plot(Erange_,W_LET(Erange_)+W_secondaries(t,Erange_),color='g')

    # plot excesses
    ax.errorbar(Erange,exc[start:stop],yerr=err[start:stop],ls='none')
    ax.scatter(Erange,exc[start:stop])
    ax.set_title(f'{t}nm')
    ax.set_xlabel('Energy (keV)')
    ax.set_ylabel('Normalized Excess Rate')

def plot(wdir,Emin,Emax,e_bin,save=False):
    """Plot the combined results and predictions for 40nm,80nm,160nm,320nm thicknesses in a given energy
    range."""

    excess=np.genfromtxt(os.path.join(wdir,'combined_result_excesses.csv'),delimiter=',')
    errors=np.genfromtxt(os.path.join(wdir,'combined_result_errors.csv'),delimiter=',')

    fig,((ax0,ax1),(ax2,ax3))=plt.subplots(2,2)
    plot_thickness(ax0,excess[0,:],errors[0,:],Emin,Emax,e_bin,40)
    plot_thickness(ax1,excess[1,:],errors[1,:],Emin,Emax,e_bin,80)
    plot_thickness(ax2,excess[2,:],errors[2,:],Emin,Emax,e_bin,160)
    plot_thickness(ax3,excess[3,:],errors[3,:],Emin,Emax,e_bin,320)
    plt.tight_layout()

    if save:
        if not os.path.isdir(os.path.join(wdir,'figures/')): os.mkdir(os.path.join(wdir,'figures/'))
        plt.savefig(os.path.join(os.path.join(wdir,'figures/'), \
                                 f'normalized_excess_plots_(Emin={Emin},Emax={Emax}).png'))
    else:
        plt.show()

# CALCULATE AND REPORT CHI SQUARE VALUES------------------------------------------------------------------------
def calc_chisq(obs,err,pred,rnd=True):
    """Calculate the chi square and probability given observed results, errors, and predicted results.
    If rnd, round the returned values accordingly."""

    chi_sq=np.sum((obs-pred)**2/err**2)
    nu=obs.shape
    p=(1-chi2.cdf(chi_sq,nu))[0]
    if rnd: return round(chi_sq,3),round(p,5)
    return chi_sq,p

def excess_chisq(wdir,Emin,Emax,e_bin,save=False):
    """Compute the chi square and probability between observed and predicted results of the combined data
    given an energy range. Print the results as a table to the command line and write them to a TXT"""

    excess=np.genfromtxt(os.path.join(wdir,'combined_result_excesses.csv'),delimiter=',')
    errors=np.genfromtxt(os.path.join(wdir,'combined_result_errors.csv'),delimiter=',')
    start,stop=int(Emin/(e_bin*exp.XMAP_ENERGY_UNIT)),int(Emax/(e_bin*exp.XMAP_ENERGY_UNIT)+0.5)
    Erange=np.arange((start+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT), \
                     (stop+0.5)*(e_bin*exp.XMAP_ENERGY_UNIT),e_bin*exp.XMAP_ENERGY_UNIT)

    table=[('','df','SEC chi2','SEC p','LET chi2','LET p','TOT chi2','TOT p')]
    df=(Erange.shape)[0]
    for r,t in enumerate([40,80,160,320]):
        SECchi2,SECp=calc_chisq(excess[r,start:stop],errors[r,start:stop],W_secondaries(t,Erange))
        LETchi2,LETp=calc_chisq(excess[r,start:stop],errors[r,start:stop],W_LET(Erange))
        TOTchi2,TOTp=calc_chisq(excess[r,start:stop],errors[r,start:stop],W_LET(Erange)+W_secondaries(t,Erange))
        table.append((str(s) for s in (f'{t}nm',df,SECchi2,SECp,LETchi2,LETp,TOTchi2,TOTp)))

    if save:
        if not os.path.isdir(os.path.join(wdir,'chisq/')): os.mkdir(os.path.join(wdir,'chisq/'))
        with open(os.path.join(os.path.join(wdir,'chisq/'),f'chisq_table_(Emin={Emin},Emax={Emax}).txt'),'w+') as f:
            for row in table: f.write('| {:>6} | {:>2} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} |\n'.format(*row))
    for row in table: print('| {:>6} | {:>2} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10} |'.format(*row))

# CALCULATE AND PRESENT CORRECTIONS FOR ATTENUATION-------------------------------------------------------------
def get_atten(Erange,n_ch=4,verbose=False):
    """Compute the attentuation from all sources for a given energy range. If verbose, the result includes the
    raw transmittance for every source of attenuation in the format <He>,<kapton>,<1-elem detector QE>,<4-elem
    detector QE>,<320nm>,<160nm>,<80nm>,<40nm>.
    Otherwise, it returns transmittance results for each mode (based on the detector channel that saw the
    fluorescence event) and for each target thickness relative to the transmittance at fluorescence energy.
    These are returned in the format <average mode transmittance>,(<320nm>,<160nm>,<80nm>,<40nm>).

    This function, along with other functions related to attenuation correction, will likely need to be modified
    if any of the experimental conditions change. This is far less generalizable than everything else."""

    atten_raw=pd.read_csv(exp.ATTENUATION_DOC_PATH)
    He_T=np.interp(Erange,atten_raw['He_E'].values,atten_raw['He_T'].values)
    Kap_T=np.interp(Erange,atten_raw['Kap_E'].values,atten_raw['Kap_T'].values)
    D1_T=np.interp(Erange,atten_raw['D_E'].values,atten_raw['D1_T'].values)
    D4_T=np.interp(Erange,atten_raw['D_E'].values,atten_raw['D4_T'].values)
    M320_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M320_T'].values)
    M160_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M160_T'].values)
    M80_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M80_T'].values)
    M40_T=np.interp(Erange,atten_raw['M_E'].values,atten_raw['M40_T'].values)

    if verbose: return He_T,Kap_T,D1_T,D4_T,M320_T,M160_T,M80_T,M40_T

    D1_Tot=He_T*Kap_T*D1_T
    D4_Tot=He_T*Kap_T*D4_T
    Mode_T=np.zeros((n_ch,len(Erange)))

    for r in range(n_ch):
        S=0
        for ch in range(n_ch):
            if r==ch: continue
            S+=exp.CHANNEL_SOLID_ANGLE[ch]*exp.CHANNEL_ANGULAR_COEFF[ch]
            if ch == 0:
                Mode_T[r,:]+=exp.CHANNEL_SOLID_ANGLE[ch]*exp.CHANNEL_ANGULAR_COEFF[ch]*D1_Tot
            else:
                Mode_T[r,:]+=exp.CHANNEL_SOLID_ANGLE[ch]*exp.CHANNEL_ANGULAR_COEFF[ch]*D4_Tot

        Mode_T[r,:]/=S
        if r == 0:
            Mode_T[r,:]/=np.interp(exp.E_FLUOR,Erange,D1_Tot)
        else:
            Mode_T[r,:]/=np.interp(exp.E_FLUOR,Erange,D4_Tot)

    M320_T/=np.interp(exp.E_FLUOR,Erange,M320_T)
    M160_T/=np.interp(exp.E_FLUOR,Erange,M160_T)
    M80_T/=np.interp(exp.E_FLUOR,Erange,M80_T)
    M40_T/=np.interp(exp.E_FLUOR,Erange,M40_T)

    return Mode_T,(M320_T,M160_T,M80_T,M40_T)

def correct_atten(exc,err,e_bin,n_bins,t):
    """Similarly to how corrections are applied for the angular component, this function corrects the excesses
    and energies for attenuation."""

    Mode_T,M_T=get_atten(exp.XMAP_ENERGY_UNIT*e_bin*np.arange(n_bins),4)
    th={'data/40nm/':3,'data/80nm/':2,'data/160nm/':1,'data/320nm/':0}

    if t == 'data/empty/': return exc/Mode_T,err/Mode_T
    return exc/(Mode_T*M_T[th[t]]),err/(Mode_T*M_T[th[t]])

def plot_atten(Emin,Emax,n_ch=4,save=False,wdir=None):
    """Plot the transmittance for all sources of attenuation over a given energy range. If save, save the result
    as a figure."""

    if save and wdir == None: raise Exception('Instructed to save with no wdir provided.')

    Erange=np.linspace(Emin,Emax,1000)
    He_T,Kap_T,D1_T,D4_T,M320_T,M160_T,M80_T,M40_T=get_atten(Erange,n_ch,True)

    plt.plot(Erange,M40_T,label='Material Transmittance 40nm')
    plt.plot(Erange,M80_T,label='Material Transmittance 80nm')
    plt.plot(Erange,M160_T,label='Material Transmittance 160nm')
    plt.plot(Erange,M320_T,label='Material Transmittance 320nm')
    plt.plot(Erange,He_T,label='Helium Transmittance')
    plt.plot(Erange,Kap_T,label='Kapton Transmittance')
    plt.plot(Erange,D1_T,label='1-Elem Detector QE')
    plt.plot(Erange,D4_T,label='4-Eleme Detector QE')

    plt.legend(loc='lower right')
    plt.title('Sources of Attenuation')
    plt.xlabel('Energy (keV)')
    plt.ylabel('Transmittance')
    plt.tight_layout()

    if save:
        if not os.path.isdir(os.path.join(wdir,'figures/')): os.mkdir(os.path.join(wdir,'figures/'))
        plt.savefig(os.path.join(os.path.join(wdir,'figures/'), \
                                 f'atten_sources_(Emin={Emin},Emax={Emax}).png'))
    else:
        plt.show()

# HANDLING COMMAND LINE ARGUMENTS-------------------------------------------------------------------------------
if __name__=='__main__':
    # Parsing command line arguments
    p = argparse.ArgumentParser()
    sp=p.add_subparsers(dest='command')

    # combine command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py combine <PATH TO WORKING DIRECTORY> <ENERGY BIN> \\
    <NUMBER OF BINS> <NUMBER OF CHANNELS> <NUMBER OF ORBITS> <--skipwrite> <--noatten>"""

    combine_p=sp.add_parser('combine',help='combine reduced data')
    combine_p.add_argument(help='working directory',dest='wdir')
    combine_p.add_argument(help='length of energy bin (xMAP units)',type=int,dest='e_bin')
    combine_p.add_argument(help='number of bins in final histogram',type=int,dest='n_bins')
    combine_p.add_argument(help='number of detector channels',type=int,dest='n_ch')
    combine_p.add_argument(help='number of orbits over which to average accidentals',type=int,dest='n_orb')
    combine_p.add_argument('--skipwrite',help='do not write the result to a CSV',action='store_false')
    combine_p.add_argument('--noatten',help='do not correct for attenuation',action='store_false')

    # plot command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py plot <PATH TO WORKING DIRECTORY> <ENERGY BIN> \\
    <LOWER ENERGY BOUND> <UPPER ENERGY BOUND> <--save>"""

    plot_p=sp.add_parser('plot',help='plot excesses command')
    plot_p.add_argument(help='working directory',dest='wdir')
    plot_p.add_argument(help='length of energy bin (xMAP units)',type=int,dest='e_bin')
    plot_p.add_argument(help='lower bound of scattered energies to plot (keV)',type=float,dest='Emin')
    plot_p.add_argument(help='upper bound of scattered energies to plot (keV)',type=float,dest='Emax')
    plot_p.add_argument('--save',help='save the plot generated',action='store_true')

    # chisq command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py chisq <PATH TO WORKING DIRECTORY> <ENERGY BIN> \\
    <LOWER ENERGY BOUND> <UPPER ENERGY BOUND> <--save>"""

    chisq_p=sp.add_parser('chisq',help='calculate chi square results')
    chisq_p.add_argument(help='working directory',dest='wdir')
    chisq_p.add_argument(help='length of energy bin (xMAP units)',type=int,dest='e_bin')
    chisq_p.add_argument(help='lower bound of scattered energies to plot (keV)',type=float,dest='Emin')
    chisq_p.add_argument(help='upper bound of scattered energies to plot (keV)',type=float,dest='Emax')
    chisq_p.add_argument('--save',help='save the plot generated',action='store_true')

    # pltatten command
    """USAGE: python <PATH TO fl_coin>/fl_coin/calc_exc.py pltatten <LOWER ENERGY BOUND> <UPPER ENERGY BOUND> \\
    <--wdir PATH TO WORKING DIRECTORY> <--save>"""

    pltatten_p=sp.add_parser('pltatten',help='plot sources of attenuation')
    pltatten_p.add_argument(help='lower bound of scattered energies to plot (keV)',type=float,dest='Emin')
    pltatten_p.add_argument(help='upper bound of scattered energies to plot (keV)',type=float,dest='Emax')
    pltatten_p.add_argument('--wdir',help='working directory for saving the plot',default=None)
    pltatten_p.add_argument('--save',help='save the plot generated',action='store_true')

    a=p.parse_args()
    if a.command=='plot':
        plot(a.wdir,a.Emin,a.Emax,a.e_bin,a.save)
    elif a.command=='combine':
        combine_data(a.wdir,a.n_ch,a.n_bins,a.n_orb,a.e_bin,a.noatten,a.skipwrite)
    elif a.command=='chisq':
        excess_chisq(a.wdir,a.Emin,a.Emax,a.e_bin,a.save)
    elif a.command=='pltatten':
        plot_atten(a.Emin,a.Emax,4,a.save,a.wdir)


import os
import pickle as plk
import numpy as np

from tqdm import tqdm

# load tập dữ liệu đã chọn chứa các audio và video, đồng thời tính toán độ dài của từng mẫu theo số frames (đối với video)
# hoặc số lượng audio. Sau đó, độ dài được thêm vào tập dữ liệu và tập dữ liệu cập nhật được chọn lại.

def get_lengths(mode, feature_name):
    fd = data[mode][feature_name]
    max_len = fd.shape[1]

    c_sum = np.sum(fd, axis=-1)
    lengths = []
    for i in tqdm(range(fd.shape[0])):
        null = True
        zeros = np.zeros([fd.shape[1], fd.shape[2]])
        cur_length = max_len
        for j in range(max_len):
            if c_sum[i][j] == 0:
                cur_length = j
                null = False
                break
        if cur_length == 0:
            cur_length = 1
        lengths.append(cur_length)
    return lengths

with open('/home/sharing/disk3/dataset/multimodal-sentiment-dataset/ALL/mosei/unaligned_50.pkl', 'rb') as lf:
    data = plk.load(lf)

def handleData(mode):
    # data[mode]['audio_lengths'], _ = get_lengths(mode, 'feature_A')
    # data[mode]['vision_lengths'], _ = get_lengths(mode, 'feature_V')
    data[mode]['audio_lengths'] = get_lengths(mode, 'audio')
    data[mode]['vision_lengths'] = get_lengths(mode, 'vision')

handleData('train')
handleData('valid')
handleData('test')

with open('/home/sharing/disk3/dataset/multimodal-sentiment-dataset/ALL/mosei/unaligned_50.pkl', 'wb') as df:
    plk.dump(data, df, protocol = 4)


from collections import OrderedDict, defaultdict
import torch.nn.functional as F
import torch
import settings
from typing import Iterator, List, Dict
from allennlp.data.vocabulary import Vocabulary
from allennlp.models import Model
from penalty import BasePenalty, SpanImplicationMax, TerminalTag, Transition
from allennlp.modules.conditional_random_field import allowed_transitions
from IPython.core.debugger import Pdb

@Model.register("srl_constraints_cagg_choice")
class SRLConstraintsCustomAggChoice(Model):
    def __init__(self,
                vocab: Vocabulary,
                config: Dict):
        super().__init__(vocab)
        self.constraint_dict = OrderedDict()
        #Pdb().set_trace()
        #implication constraint
        bl_config = config.get('span_implication_bl',None)
        lb_config = config.get('span_implication_lb',None)
        b_config = config.get('terminal_tag_b',None)
        l_config = config.get('terminal_tag_l',None)

        if bl_config is not None:
            self.constraint_dict['BL'] = SpanImplicationMax(vocab, bl_config, "BL")

        if lb_config is not None:
            self.constraint_dict['LB'] = SpanImplicationMax(vocab, lb_config, "LB")

        if b_config is not None:
            self.constraint_dict['B'] = TerminalTag(vocab, b_config, "B")

        if l_config is not None:
            self.constraint_dict['L'] = TerminalTag(vocab, l_config, "L")

        config_transition_constraints = config.get('transition', None)
        if config_transition_constraints is not None:
            self.constraint_dict['transition'] = Transition(vocab, config_transition_constraints)

        self.id2key = list(self.constraint_dict.keys())
        self.key2id = dict((k,i) for i,k in enumerate(self.id2key))
        self.dd_constant_lambda = config.get('dd_constant_lambda',False)
        if self.dd_constant_lambda:
            #self.lambda_list = torch.nn.ParameterList([torch.nn.Parameter(torch.ones(self.constraint_dict[k].num_constraints())) for k in self.id2key],requires_grad = False)
            self.lambda_list = torch.nn.ParameterList([torch.nn.Parameter(torch.ones(self.constraint_dict[k].num_constraints()),requires_grad=False) for k in self.id2key])
        else:
            self.lambda_list = torch.nn.ParameterList([torch.nn.Parameter(torch.zeros(self.constraint_dict[k].num_constraints())) for k in self.id2key])
        #self.lambda_dict = torch.nn.ParameterDict([(k,torch.nn.Parameter(torch.zeros(v.num_constraints()))) for k,v in self.constraint_dict.items()])
        if settings.cuda:
            self.lambda_list = self.lambda_list.cuda()
        #for k, v in self.constraint_dict.items():
        #    self.lambda_dict[k] = torch.nn.Parameter(
        #        torch.zeros(v.num_constraints()))

    def forward(self, scores, mask, span_matrix,debug=False):
        loss = 0.0
        penalties = {}
        for k in self.constraint_dict:
            #Pdb().set_trace()
            violations,debug_out = self.constraint_dict[k].get_penalty(scores,mask.float(),span_matrix.float(),debug = debug)
            penalty = (self.lambda_list[self.key2id[k]]*violations.mean(dim=0)).sum()
            if debug:
                penalties[k+'_violations'] = debug_out.sum(dim=-1).detach()
            #
            loss += self.constraint_dict[k].weight*penalty
            penalties[k+'_pen'] = penalty.item()

        penalties['loss'] = loss
        return penalties

    def violation_stats(self,scores,indicator=0):
        rs = OrderedDict()
        for k in self.constraint_dict:
            penalty = self.constraint_dict[k].get_penalty(scores,indicator)
            if indicator == 0:
                k1 = k+'_v'
                k2 = k+'_tc'
                k3 = k+'_p'
            else:
                k1 = k + '_pv'
                k2 = k + '_ptc'
                k3 = k + '_pp'

            rs[k1] = (penalty > 0).sum().item()
            rs[k2] = penalty.numel()
            rs[k3] = penalty.sum().item()
        return rs

    def lambda_stats(self):
        ls = OrderedDict()
        for k in self.constraint_dict:
            ls[k+'_nzl'] = (self.lambda_list[self.key2id[k]] > 0).float().sum().item()
            ls[k+'_tl'] = self.lambda_list[self.key2id[k]].numel()
            ls[k+'_apl'] = 0 if (ls[k+'_nzl'] == 0) else (self.lambda_list[self.key2id[k]].sum().item()/ls[k+'_nzl'])
            assert ((self.lambda_list[self.key2id[k]] < 0).float().sum().item() == 0)
        return ls

# The mean - NOT per word BUT per sentence
"""
# Span constraints
class SpanImplicationMax(BasePenalty):
    def __init__(self, vocab, config, spanType: str):
        super(SpanImplicationMax,self).__init__()
        # TODO - is the weight supposed to be default value 1?
        self.weight = 1
        self.spanType = spanType
        try:
            self.weight = config['weight']
        except:
            pass
        if (spanType not in ["BL", "LB"]):
            raise Exception("SpanImplication constraint: spanType should be BL or LB")
        self.populate_dd_vars(vocab, self.spanType, config)

        if settings.cuda:
            self.Btags_tensor = self.Btags_tensor.cuda()
            self.Ltags_tensor = self.Ltags_tensor.cuda()

        # Also initialize tensor of B and L tags

    # In config, add an option - "aggregate_type" - can be max, sum etc
    def populate_dd_vars(self, vocab, spanType, config): #NOTE

        Btag = "B"
        Ltag = "L"
        if(spanType == "LB"):
            Btag = "L"
            Ltag = "B"
        self.Btags_tensor = None
        self.Ltags_tensor = None
        Btags_dict = {}
        Ltags_dict = {}

        index2tagtype = [] # List the tag types

        index2labels = vocab.get_index_to_token_vocabulary(namespace="labels")

        self.nconstraints = 0

        for idx in index2labels:
            lbl = index2labels[idx]
            if(lbl != 'O' and lbl[2:] != 'V'):
                l = lbl[2:]
                posn = lbl[0]
                if (l not in index2tagtype):
                    index2tagtype.append(l)

                if posn == Btag:
                    Btags_dict[l] = idx
                elif posn == Ltag:
                    Ltags_dict[l] = idx

        # Now form Btags_tensor and Ltags_tensor
        self.Btags_tensor, self.Ltags_tensor = [], []
        for idx, tagtype in enumerate(index2tagtype):
            self.Btags_tensor.append(Btags_dict[tagtype])
            self.Ltags_tensor.append(Ltags_dict[tagtype])

        self.Btags_tensor = torch.LongTensor(self.Btags_tensor)
        self.Ltags_tensor = torch.LongTensor(self.Ltags_tensor)

        self.nconstraints = len(index2tagtype)

        # aggregate_type
        aggregate_type = config.get("aggregate_type", "max")
        # max penalty over all expressions in  the logic formula (Yatin's formula)
        self.use_maxp = config.get("use_maxp", False)

        if self.use_maxp == True:
	        self.max_penalty = config.get("max_penalty", "log")
	        if self.max_penalty not in ["log", "negate"]:
	        	raise Exception("Penalty type for max should be log or negate")

        if aggregate_type.lower() not in ["sum", "max"]:
        	raise Exception("SpanImplicationMax aggregation type invalid")

        if aggregate_type.lower() == "sum":
        	self.aggregate_mat = lambda x: x.sum(dim=2)
        elif aggregate_type.lower() == "max":
        	self.aggregate_mat = lambda x: x.max(dim=2)[0]

    def num_constraints(self):
        return self.nconstraints

    # Pass in class PROBABILITIES, NOT logits!
    def get_penalty(self, prob, mask, span_matrix, indicator=0):
        # prob is class PROBABILITIES, NOT logits(keep the option?)

        # prob - bsize x senlen x ntags
        # span_matrix - bsize x senlen x senlen

        # NOTE - the summation can be replaced by max later if we want
        # Here we take NEGATIVE of expression so we can do hinge on it
        # constr_mat - bsize x senlen x |Btags_tensor|
        #mask out the span matrix

        if settings.cuda:
            span_matrix = span_matrix.cuda()
        #
        span_matrix = span_matrix*mask.unsqueeze(-1).expand_as(span_matrix)
        span_matrix = span_matrix*mask.unsqueeze(1).expand_as(span_matrix)

        if(self.spanType == "LB"):
            local_span_matrix = span_matrix.transpose(1,2)
        else:
            local_span_matrix = span_matrix

        if settings.cuda:
            local_span_matrix = local_span_matrix.cuda()

        # local_span_matrix - bsize x senlen x senlen
        # prob[:, :, self.Ltags_tensor] - bsize x senlen x ntags
        bsize,senlen,_ = prob.size()
        ntags = self.nconstraints
        #Pdb().set_trace()
        expanded_span_mat = local_span_matrix.unsqueeze(-1).expand(bsize, senlen, senlen, ntags)
        expanded_ltags_mat = prob[:, :, self.Ltags_tensor].unsqueeze(1).expand(bsize, senlen, senlen, ntags)

        #constr_mat.shape  = B x S x S x nconstraints
        constr_mat = expanded_span_mat * expanded_ltags_mat

        active_span_mat = local_span_matrix.sum(dim=2)
        active_span_mat[active_span_mat > 0] = 1

        if self.use_maxp == False:
            constr_mat = prob[:, :, self.Btags_tensor] - self.aggregate_mat(constr_mat)
        else:
            constr_mat = torch.max(1-prob[:, :, self.Btags_tensor], (constr_mat.max(dim=2)[0]))
            constr_mat[active_span_mat.unsqueeze(-1).expand_as(constr_mat) == 0] = 1
            if self.max_penalty == "log":
            	constr_mat = -1.0 * torch.log(constr_mat)
            elif self.max_penalty == "negate":
            	constr_mat = 1.0 - constr_mat

        # constr_mat = prob[:, :, self.Btags_tensor] - torch.bmm(local_span_matrix, prob[:, :, self.Ltags_tensor])
        # active_span_mat - bsize x senlen - Check which spans are inactive
        constr_mat = constr_mat * active_span_mat.unsqueeze(-1).expand_as(constr_mat) # Shapes match?
        #mask it again
        constr_mat =  constr_mat*mask.unsqueeze(-1).expand_as(constr_mat)
        F.relu(constr_mat, inplace=True)
        constr_mat = constr_mat.sum(dim=1)
        return constr_mat

# B/L constraint - can't be one if no span is present
class TerminalTag(BasePenalty):
    def __init__(self, vocab, config, terminalType: str):
        super(TerminalTag,self).__init__()
        # TODO - is the weight supposed to be default value 1?
        self.weight = 1
        self.terminalType = terminalType
        try:
            self.weight = config['weight']
        except:
            pass

        if terminalType not in ["B", "L"]:
            raise Exception("TerminalTag constraint: Tag must be B- or L- type.")

        self.populate_dd_vars(vocab, terminalType)
        # Also initialize tensor of B and L tags
        if settings.cuda:
            self.Ttags_tensor = self.Ttags_tensor.cuda()

    def populate_dd_vars(self, vocab, terminalType): #NOTE

        self.Ttags_tensor = None

        Ttags_dict = {}
        index2tagtype = [] # List the tag types

        index2labels = vocab.get_index_to_token_vocabulary(namespace="labels")

        self.nconstraints = 0

        for idx in index2labels:
            lbl = index2labels[idx]
            if(lbl != 'O' and lbl[2:] != 'V'):
                l = lbl[2:]
                posn = lbl[0]
                if (l not in index2tagtype):
                    index2tagtype.append(l)

                if posn == terminalType:
                    Ttags_dict[l] = idx

        # Now form Ttags_tensor and Ltags_tensor
        self.Ttags_tensor = []
        for idx, tagtype in enumerate(index2tagtype):
            self.Ttags_tensor.append(Ttags_dict[tagtype])

        self.Ttags_tensor = torch.LongTensor(self.Ttags_tensor)

        self.nconstraints = len(index2tagtype)

    def num_constraints(self):
        return self.nconstraints

    # Pass in class PROBABILITIES, NOT logits!
    def get_penalty(self, prob, mask, span_matrix, indicator=0):
        # prob is class PROBABILITIES, NOT logits(keep the option?)

        # prob - bsize x senlen x ntags
        # span_matrix - bsize x senlen x senlen
        # constr_mat - bsize x sellen

        # active_span_mat - bsize x senlen - Check which spans are inactive
        if (self.terminalType == "B"):
            local_span_matrix = span_matrix
        else:
            local_span_matrix = span_matrix.transpose(1,2)

        if settings.cuda:
            local_span_matrix = local_span_matrix.cuda()

        active_span_mat = local_span_matrix.sum(dim=2)
        active_span_mat[active_span_mat > 0] = 1
        constr_mat = prob[:, :, self.Ttags_tensor]
        # IMPORTANT - inactive spans are active for terminal constraint
        constr_mat = constr_mat * (1 - active_span_mat.unsqueeze(-1))
        constr_mat = constr_mat*mask.unsqueeze(-1).expand_as(constr_mat)
        constr_mat = constr_mat.sum(dim=1)

        return constr_mat

# Transition constraints
class Transition(BasePenalty):
    def __init__(self, vocab, config):
        super(Transition,self).__init__()
        # TODO - is the weight supposed to be default value 1?
        self.weight = config.get("weight", 1)

        self.config = config

        self.populate_dd_vars(vocab)

        # Put tensors on CUDA
        if settings.cuda:
            self.transition_matrix = self.transition_matrix.cuda()

    def populate_dd_vars(self, vocab): #NOTE

        # Set up allowed transitions
        all_labels = vocab.get_index_to_token_vocabulary(namespace="labels")
        num_labels = len(all_labels)
        constraints = allowed_transitions("BIOUL", all_labels)
        self.transition_matrix = torch.zeros([num_labels+2, num_labels+2]).fill_(0.0)

        for c in constraints:
            # if (c[0] < num_labels) and (c[1] < num_labels):
            self.transition_matrix[c[0],c[1]] = 1.0

        self.nconstraints = num_labels + 2

        aggregate_type = self.config.get("aggregate_type", "max")
        if aggregate_type.lower() not in ["sum", "max"]:
            raise Exception("Transition aggregation type invalid")
        if aggregate_type.lower() == "sum":
            self.aggregate_mat = lambda x: x.sum(dim=3)
        elif aggregate_type.lower() == "max":
            self.aggregate_mat = lambda x: x.max(dim=3)[0]

    def num_constraints(self):
        return self.nconstraints

    # Pass in class PROBABILITIES, NOT logits!
    def get_penalty(self, prob0, mask, span_matrix, indicator=0):
        # prob is class PROBABILITIES, NOT logits(keep the option?)

        # prob - bsize x senlen x ntags

        prob = F.pad(prob0, (0,2,1,1,0,0))
        bsize, senlen, ntags = prob.shape

        prob[:,0,ntags-2] = 1
        prob[:,-1,ntags-1] = 1

        m1 = prob[:, :-1, :]
        m2 = prob[:, 1:, :]

        m2 = m2.unsqueeze(-1).expand(bsize, senlen-1, ntags, ntags).transpose(2,3)
        t_mat = self.transition_matrix.unsqueeze(0).unsqueeze(0).expand(bsize, senlen-1, ntags, ntags)

        constr_mat = m2 * self.transition_matrix
        constr_mat = m1 - self.aggregate_mat(constr_mat)

        # Apply mask
        # Pdb().set_trace()
        constr_mat = constr_mat * F.pad(mask, (1,0,0,0), value=1.0).unsqueeze(-1).expand_as(constr_mat)

        F.relu(constr_mat, inplace=True)
        constr_mat = constr_mat.sum(dim=1)

        return constr_mat

"""


'''
Given a sorted list of integers of length N, determine if an element x is in the list without performing any multiplication, division, or bit-shift operations.

Do this in O(log N) time.
'''

from random import randint
from typing import List

# In place of a standard binary search which would involve using division, select an index in the range at random as the pivot on each iteration.
# On average, this will work in O(log N) time, the same as a binary search, just with higher volatility.
# int_list is assumed to be pre-sorted.
def hasElement(x: int, int_list: List[int]) -> bool:
    # Initialize
    lo, hi = 0, len(int_list)-1

    while lo <= hi:
        selection = randint(lo, hi)
        # Move a boundary based on the selected index's value if it is not x. Return True if the selected index's value is x.
        if int_list[selection] < x:
            lo = selection + 1
        elif int_list[selection] > x:
            hi = selection - 1
        else: # int_list[selection] == x
            return True
    # x not found.
    return False

int_list = list(range(2, 22, 2))

print(hasElement(10, int_list))
print(hasElement(11, int_list))
print(hasElement(2, int_list))
print(hasElement(20, int_list))


'''
1125. Smallest Sufficient Team
https://leetcode.com/problems/smallest-sufficient-team/
'''
from collections import defaultdict
from typing import List

from test_tool import assert_value

class Solution:
    def smallestSufficientTeam(self, req_skills: List[str], people: List[List[str]]) -> List[int]:
        group = set()
        roster = defaultdict(list)
        for idx, skills in enumerate(people):
            for skill in skills:
                roster[skill].append(idx)

        return list(self.helper(set(req_skills), people, roster, group))

    def helper(self, req_skills, people, roster, group):
        if not req_skills:
            return set()

        skill = list(req_skills)[0]
        min_combo = set()
        for person in roster[skill]:
            if person in group:
                continue

            combo = self.helper(req_skills - set(people[person]), people, roster, group | {person}) | {person}
            if not min_combo or len(min_combo) > len(combo):
                min_combo = combo

        return min_combo

assert_value([0, 2], Solution().smallestSufficientTeam,
             req_skills=["java", "nodejs", "reactjs"],
             people=[["java"], ["nodejs"], ["nodejs", "reactjs"]])
assert_value([1, 2], Solution().smallestSufficientTeam,
             req_skills=["algorithms", "math", "java", "reactjs", "csharp", "aws"],
             people=[["algorithms", "math", "java"],
                     ["algorithms", "math", "reactjs"],
                     ["java", "csharp", "aws"],
                     ["reactjs", "csharp"],
                     ["csharp", "math"],
                     ["aws", "java"]])