import os
# import optuna
# import wandb
import mlflow
# from losses.label_smoothing import LabelSmoothingCrossEntropy
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from data_augmentations.mixup import mixup, cutmix
from optimizer.sam import SAM
from optimizer.adan import Adan
from optimizer.ranger21.ranger21 import Ranger21
from datetime import datetime
from data_augmentations.grid import GridMask
import json
from sklearn import metrics
# try:
#     from apex import amp
#     APEX_AVAILABLE = True
# except ModuleNotFoundError:
APEX_AVAILABLE = False
print(APEX_AVAILABLE)

# os.environ["CUDA_VISIBLE_DEVICES"]='3'

# importing the libraries
import pandas as pd
import numpy as np
import time
import torch
import random
import matplotlib.pyplot as plt
from PIL import Image
import torch.nn.functional as F

# for creating validation set
from torch.optim import lr_scheduler

# PyTorch libraries and modules
from datetime import datetime

import logging
import torchvision.transforms as transforms
from timm.utils import ModelEma
from model import Net, ArcFaceLossAdaptiveMargin, ArcNet

#import cv2

from torch.optim import Adam, SGD
from tqdm import tqdm

from utils import seed_torch, count_parameters
from dataset import ImageFolder
from torch.utils.data import DataLoader

import torch

from eval import eval
import argparse
import gc
import torch.nn.utils.prune as prune
from timm.utils import get_state_dict

weights_path = [None, None, None, None, None]

default_configs = {}

from torch.nn.modules.batchnorm import _BatchNorm

def disable_running_stats(model):
    def _disable(module):
        if isinstance(module, _BatchNorm):
            module.backup_momentum = module.momentum
            module.momentum = 0

    model.apply(_disable)

def enable_running_stats(model):
    def _enable(module):
        if isinstance(module, _BatchNorm) and hasattr(module, "backup_momentum"):
            module.momentum = module.backup_momentum

    model.apply(_enable)

def prune_model_global_unstructured(model, layer_type, proportion):
    module_tups = []
    for module in model.modules():
        if isinstance(module, layer_type):
            module_tups.append((module, 'weight'))

    prune.global_unstructured(
        parameters=module_tups, pruning_method=prune.L1Unstructured,
        amount=proportion
    )
    for module, _ in module_tups:
        prune.remove(module, 'weight')
    return model

def remove_folder(folder_path):
    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        try:
            if os.path.isfile(file_path) or os.path.islink(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as e:
            print('Failed to delete %s. Reason: %s' % (file_path, e))

def lr_func(epoch):
    LR_WARMUP_EPOCHS = 3
    LR_MAX = default_configs["lr"]
    LR_START = 1e-5
    # LR_SUSTAIN_EPOCHS = 65
    # LR_STEP_DECAY = 0.8
    if epoch < LR_WARMUP_EPOCHS:
        lr = (LR_MAX - LR_START) / LR_WARMUP_EPOCHS * epoch + LR_START
    else:
        lr = LR_MAX
    # else:
    #     lr = LR_MAX*pow(LR_STEP_DECAY, epoch - LR_WARMUP_EPOCHS - LR_SUSTAIN_EPOCHS)
    return lr

def train_one_fold(fold, train_loader, test_loader):
    print("FOLD: ", fold)
    grid = GridMask(96, 224, 360, 0.6, 1, 0.8)
    DATA_PATH = "train"
    now = datetime.now()
    weight_path = os.path.join("weights", default_configs["a_name"])
    os.makedirs(weight_path, exist_ok=True)
    weight_path = os.path.join(weight_path, str(fold))
    os.makedirs(weight_path, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    start_epoch = 0
    if default_configs["use_arcface"]:
        criterion_class = ArcFaceLossAdaptiveMargin(m=default_configs["arcface_m"], s=default_configs["arcface_s"])
    else:
        criterion_class = LabelSmoothingCrossEntropy(default_configs["smoothing_value"])
    # criterion_class = SoftTargetCrossEntropy()
    # criterion_class = torch.nn.CrossEntropyLoss()
    # early_stopping = EarlyStopping()
    # criterion_class = torch.nn.BCEWithLogitsLoss()
    # criterion_class = FocalCosineLoss()

    criterion_class.to(device)
    if default_configs["use_arcface"]:
        model = ArcNet(default_configs["backbone"], default_configs["n_frames"])
    else:
        model = Net(default_configs["backbone"], default_configs["n_frames"])

#     model = make_model(4)
    # base_optimizer = torch.optim.SGD  # define an optimizer for the "sharpness-aware" update
    # optimizer_model = torch.optim.SGD(model.parameters(), lr=default_configs["lr"], weight_decay=default_configs["weight_decay"])
    if default_configs["optimizer"] == "SAM":
        base_optimizer = torch.optim.SGD  # define an optimizer for the "sharpness-aware" update
        optimizer_model = SAM(model.parameters(), base_optimizer, lr=default_configs["lr"], momentum=0.9, weight_decay=default_configs["weight_decay"], adaptive=True)
    # optimizer_model = torch.optim.AdamW(model.parameters(), lr=default_configs["lr"], weight_decay=default_configs["weight_decay"])
    elif default_configs["optimizer"] == "Ranger21":
        optimizer_model = Ranger21(model.parameters(), lr=default_configs["lr"], weight_decay=default_configs["weight_decay"],
        num_epochs=default_configs["num_epoch"], num_batches_per_epoch=len(train_loader))
    elif default_configs["optimizer"] == "SGD":
        optimizer_model = torch.optim.SGD(model.parameters(), lr=default_configs["lr"], weight_decay=default_configs["weight_decay"], momentum=0.9)
    else:
        optimizer_model = Adan(model.parameters(), lr=default_configs["lr"], weight_decay=default_configs["weight_decay"])
    if default_configs["optimizer"] != "Ranger21":
        scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer_model, 5, T_mult=1, eta_min=1e-5, last_epoch=- 1, verbose=False)
    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer_model, T_max=args.num_epochs)
#     count_parameters(model.src_embed)
    model.to(device)

    scaler = torch.cuda.amp.GradScaler()

    if weights_path[fold] is not None:
        pretrained_dict = torch.load(weights_path[fold])
        model_dict = model.state_dict()
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.size() == model_dict[k].size()}
        model.load_state_dict(pretrained_dict, strict=False)
        #model.load_state_dict(pretrained_dict)
    # for param in model.parameters():
    #     param.requires_grad = False

    # model.model.last_linear.weight.requires_grad = True
    # model.model.last_linear.bias.requires_grad = True
    # model = torch.nn.DataParallel(model)
    if APEX_AVAILABLE:
        model, optimizer_model = amp.initialize(
            model, optimizer_model, opt_level="O2", loss_scale="dynamic"
        )

    best_metric = {"loss": 10000, "acc": 0, "roc_auc": 0, "tpr_fpr_5e-4": 0, "tpr_fpr_1e-4": 0, "eer": 10000}
    best_metric_ema = {"loss": {"score": 10000, "list": []}, "acc": {"score": 0, "list": []}, "roc_auc": {"score": 0, "list": []}, \
        "tpr_fpr_5e-4": {"score": 0, "list": []}, "tpr_fpr_1e-4": {"score": 0, "list": []}, "eer": {"score": 10000, "list": []}}
    best_model_path = ""

    input_list, output_list = [], []
    iter_size = 1

    for epoch in range(start_epoch, default_configs["num_epoch"]):
        print("\n-----------------Epoch: " + str(epoch) + " -----------------")
        grid.set_prob(epoch, default_configs["num_epoch"])
        for param_group in optimizer_model.param_groups:
            # logger.debug("LR = {}".format(param_group['lr']))
            # param_group['lr'] = lr_func(epoch)
            mlflow.log_metric("lr", param_group['lr'], step=epoch)
            print("LR: ", param_group['lr'])
        if epoch == default_configs['start_ema_epoch']:
            print("Start ema......................................................")
            model_ema = ModelEma(
                model,
                decay=default_configs["model_ema_decay"],
                device=device, resume='')
        running_loss = 0
        j = 0
        start = time.time()
        n_images = 0
        n_correct = 0
        optimizer_model.zero_grad()
        for batch_idx, (imgs, labels, img_paths) in enumerate(tqdm(train_loader)):
            model.train()
            imgs = imgs.to(device).float()
            labels = labels.to(device).long()
            # print(labels.shape)

            if torch.rand(1)[0] < 0.5 and (default_configs["use_mixup"] or default_configs["use_cutmix"]):
                rand_prob = torch.rand(1)[0]

                if default_configs["use_mixup"] == True and default_configs["use_cutmix"] == False:
                    mix_images, target_a, target_b, lam = mixup(imgs, labels, alpha=default_configs["mixup_alpha"])
                if default_configs["use_mixup"] == False and default_configs["use_cutmix"] == True:
                    mix_images, target_a, target_b, lam = cutmix(imgs, labels, alpha=default_configs["mixup_alpha"])
                if default_configs["use_mixup"] == True and default_configs["use_cutmix"] == True:
                    if rand_prob < 0.5:
                        mix_images, target_a, target_b, lam = mixup(imgs, labels, alpha=default_configs["mixup_alpha"])
                    else:
                        mix_images, target_a, target_b, lam = cutmix(imgs, labels, alpha=default_configs["mixup_alpha"])

                with torch.cuda.amp.autocast():
                    logits = model(mix_images)
                    loss = criterion_class(logits, target_a) * lam + \
                    (1 - lam) * criterion_class(logits, target_b)
                    loss /= default_configs["accumulation_steps"]

                scaler.scale(loss).backward()
                if ((batch_idx + 1) % default_configs["accumulation_steps"] == 0) or ((batch_idx + 1) == len(train_loader)):
                    # scaler.unscale_(optimizer_model)
                    # torch.nn.utils.clip_grad_norm_(model.parameters(), default_configs["max_norm"])
                    scaler.step(optimizer_model)
                    scaler.update()
                    if epoch >= default_configs['start_ema_epoch']:
                        model_ema.update(model)
                    optimizer_model.zero_grad()
            else:
                with torch.cuda.amp.autocast():
                    if default_configs["use_gridmask"] and torch.rand(1)[0] < 0.5:
                        imgs = grid(imgs)
                    logits = model(imgs)
                    loss = criterion_class(logits, labels)
                    loss /= default_configs["accumulation_steps"]
                scaler.scale(loss).backward()
                if ((batch_idx + 1) % default_configs["accumulation_steps"] == 0) or ((batch_idx + 1) == len(train_loader)):
                    # scaler.unscale_(optimizer_model)
                    # torch.nn.utils.clip_grad_norm_(model.parameters(), default_configs["max_norm"])
                    scaler.step(optimizer_model)
                    scaler.update()
                    if epoch >= default_configs['start_ema_epoch']:
                        model_ema.update(model)
                    optimizer_model.zero_grad()

            running_loss += loss.item() * imgs.size(0)
            j = j + imgs.size(0)
            n_images += len(labels)
            if default_configs["optimizer"] != "Ranger21":
                scheduler.step()

        # train_acc = n_correct/n_images
        end = time.time()
        mlflow.log_metric("train_loss_f{}".format(fold), running_loss/j, step=epoch)
        mlflow.log_metric("train_elapsed_time_f{}".format(fold), end - start, step=epoch)

        # start = time.time()
        # val_metric, ground_truths, scores = eval(test_loader, model, criterion_class, device, default_configs["n_eval"], fold, epoch, is_ema=False)
        # end = time.time()
        # mlflow.log_metric("val_elapsed_time_f{}".format(fold), end - start, step=epoch)

        val_metric_type_list = ["loss", "acc", "roc_auc", "tpr_fpr_5e-4", "tpr_fpr_1e-4", "eer"]
        # for val_metric_type in val_metric_type_list:
        #     print("Val {}: {}".format(val_metric_type, val_metric[val_metric_type]))
        #     mlflow.log_metric("val_{}".format(val_metric_type), val_metric[val_metric_type], step=epoch)
        #     flag = False
        #     if val_metric_type in ["loss", "eer"]:
        #         if(val_metric[val_metric_type] < best_metric[val_metric_type]):
        #             flag = True
        #     else:
        #         if(val_metric[val_metric_type] > best_metric[val_metric_type]):
        #             flag = True
        #     if flag == True:
        #         best_model_path = os.path.join(weight_path, 'checkpoint_{}_best.pt'.format(val_metric_type))
        #         torch.save(model.state_dict(), best_model_path)
        #         mlflow.log_artifact(best_model_path)
        #         best_metric[val_metric_type] = val_metric[val_metric_type]
        # print("\n=============================================================================")
        if epoch >= default_configs['start_ema_epoch']:
            val_metric, ground_truths, scores = eval(test_loader, model_ema.ema, criterion_class, device, default_configs["n_eval"], fold, epoch, True, default_configs)
            for val_metric_type in val_metric_type_list:
                print("Val ema {}: {}".format(val_metric_type, val_metric[val_metric_type]))
                mlflow.log_metric("val_{}_ema".format(val_metric_type), val_metric[val_metric_type], step=epoch)
                flag = False
                if val_metric_type in ["loss", "eer"]:
                    if(val_metric[val_metric_type] < best_metric_ema[val_metric_type]["score"]):
                        flag = True
                else:
                    if(val_metric[val_metric_type] > best_metric_ema[val_metric_type]["score"]):
                        flag = True
                if flag == True:
                    best_model_path = os.path.join(weight_path, 'checkpoint_{}_best_ema.pt'.format(val_metric_type))
                    try:
                        os.remove(best_model_path)
                    except Exception as e:
                        print(e)
                    torch.save(get_state_dict(model_ema), best_model_path)
                    mlflow.log_artifact(best_model_path)
                    best_metric_ema[val_metric_type] = {"score": val_metric[val_metric_type], "list": [ground_truths, scores]}
                    # print("Save best model ema: ", best_model_path, val_metric[val_metric_type])

    # model = prune_model_global_unstructured(model, nn.Conv2d, 0.1)
    # test_loss, test_score = eval(val_loader, model, criterion_class, device, False, fold, 15)
    # print("After pruning, test score: ", test_score)
    # torch.save(model.state_dict(), "model_prune_{}.pt".format(fold))
    # mlflow.log_params({"pruned_{}".format(fold): test_score})
    mlflow.log_artifact(best_model_path)
    del model
    torch.cuda.empty_cache()
    gc.collect()

    return best_metric_ema

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Train code")
    parser.add_argument("--exp", type=str, default="exp_1")
    args = parser.parse_args()

    f = open(os.path.join('code/configs', "{}.json".format(args.exp)))
    default_configs = json.load(f)
    f.close()
    print(default_configs)

    seed_torch()
    DATA_PATH = "train"
    avg_score = {"loss": 0, "acc": 0, "roc_auc": 0, "eer": 0}
    n_fold = 5

    train_loader_list = {}
    test_loader_list = {}
    for fold in range(n_fold):
        if default_configs["pseudo_label"]:
            train_df = pd.read_csv("code/data/train_fold{}.csv".format(fold))
            pseudo_df = pd.read_csv("code/data/pseudo_label.csv")
            train_df = pd.concat([train_df, pseudo_df])
        else:
            train_df = pd.read_csv("code/data/train_fold{}.csv".format(fold))

        val_df =  pd.read_csv("code/data/val_fold{}.csv".format(fold))
        train_data_1 = ImageFolder(train_df, DATA_PATH, default_configs, {default_configs["image_size"]: 9}, "train")
        test_data_1 = ImageFolder(val_df, DATA_PATH, default_configs, None, "submission")
        # test_data_2 = ImageFolder(val_df, DATA_PATH, default_configs["image_size"], 11, "test", None)
        if default_configs["batch_size"] <= 2:
            train_loader = DataLoader(train_data_1, batch_size=default_configs["batch_size"], shuffle=True, drop_last=True, pin_memory=True, num_workers=default_configs["num_workers"])
        else:
            train_loader = DataLoader(train_data_1, batch_size=default_configs["batch_size"], shuffle=True, drop_last=True, pin_memory=True, num_workers=default_configs["num_workers"])
        if default_configs["batch_size"] <= 2:
            test_loader = DataLoader(test_data_1, batch_size=8, shuffle=False, pin_memory=True, num_workers=default_configs["num_workers"])
        else:
            test_loader = DataLoader(test_data_1, batch_size=default_configs["batch_size"]*2, shuffle=False, pin_memory=True, num_workers=default_configs["num_workers"])
        train_loader_list[fold] = train_loader
        test_loader_list[fold] = test_loader

    # wandb.define_metric("fold")
    # wandb.define_metric("metric", step_metric="fold")
    avg_ground_truths = []
    avg_scores = []
    for fold in range(n_fold):
        with mlflow.start_run():
            mlflow.log_params({"Fold": fold})
            mlflow.log_params(default_configs)
            mlflow.log_artifacts("code")
            # mlflow.set_tags({"TRIAL": trial.number})
            score = train_one_fold(fold, train_loader_list[fold], test_loader_list[fold])
            for k, v in avg_score.items():
                avg_score[k] += score[k]["score"]
                mlflow.log_metric("metric_{}".format(k), score[k]["score"])
            ground_truths, scores = score["eer"]["list"]
            avg_ground_truths = avg_ground_truths + ground_truths.tolist()
            avg_scores = avg_scores + scores.tolist()

            # mlflow.log_params({"CV": avg_score/n_fold})
            mlflow.end_run()

    for k, v in avg_score.items():
        print("CV_{}: ".format(k), avg_score[k]/n_fold)

    roc_auc = metrics.roc_auc_score(avg_ground_truths, avg_scores)
    fpr, tpr, threshold = metrics.roc_curve(avg_ground_truths, avg_scores, pos_label=1)
    fnr = 1 - tpr
    eer_threshold = threshold[np.nanargmin(np.absolute((fnr - fpr)))]
    EER = fpr[np.nanargmin(np.absolute((fnr - fpr)))]
    print("Correction CV_EER: ", EER)
    print("Correction CV_ROC_AUC: ", roc_auc)


# Data model of the evolving crust (thermal, hydraulic and mechanical farfield)
# Physical units: kg, m, s, K

import numpy as np
import matplotlib.pyplot as plt

from glaciationBCs.constants_AREHS import gravity
from glaciationBCs.constants_AREHS import rho_wat
from glaciationBCs.constants_AREHS import c_p_wat
from glaciationBCs.constants_AREHS import beta_pw
from glaciationBCs.constants_AREHS import beta_Tw

class crust():
	# class variables:
	V_fluid_max = 1e-11	#m/s
	p_ref = 1e5 	#Pa TODO: should be equal to air.pressure
	T_ref = 298.15 	#K

	# constructor
	def __init__(self, q_geo, v_min, v_max, T_ini, T_bot):
		# instance variables: owned by instances of the class, can be different for each instance
		self.q_geo = q_geo
		self.v_min = v_min
		self.v_max = v_max
		self.T_bot = T_bot
		self.T_ini = T_ini

	def water_density(self, p ,T):
		return rho_wat * (1 + beta_Tw*(T-self.T_ref) + beta_pw*(p-self.p_ref))

	def geothermal_heatflux(self):
		return [0.0, self.q_geo, 0.0]

	def displacement_below(self):
		return [0.0, 0.0, 0.0]

	def displacement_aside(self):
		return [0.0, 0.0, 0.0]

	def geothermal_temperature(self, v, T_atm):
		# linear profile according to geothermal heatflux
		DT = self.T_bot - T_atm
		Dv = (self.v_min - self.v_max)
		return DT/Dv * (v - self.v_max) + T_atm

	def lateral_heatflux(self, v, T_atm):
		# only for salt models
		# if props.model_id == 3:
		# 	for i, lv in enumerate(props.south_layer_bounds[:-1]):
		# 		if lv >= v > props.south_layer_bounds[i+1]:
		# 			if props.H_deactivated[i]:
		# 				return 0.
		# 			break
		# linear profile according to decreasing fluid velocity
		Dv = (self.v_min - self.v_max)
		V_fluid = self.V_fluid_max * (v - self.v_max) / Dv
		q_heat = V_fluid * (self.T_ini - T_atm) * rho_wat * c_p_wat
		return q_heat

	def hydrostatic_pressure(self, v):
		#return self.hydrostatic_pressure_lin(v)
		return self.hydrostatic_pressure_exp(v)

	def hydrostatic_pressure_lin(self, v):
		# linear profile according to gravity
		p_pore = rho_wat * gravity * (self.v_max - v)
		return p_pore

	def hydrostatic_pressure_exp(self, v):
		# exponential profile according to pressure dependence
		arg = rho_wat * gravity * (self.v_max - v) * beta_pw
		p_pore = (np.exp(arg) - 1) / beta_pw
		return p_pore

	def hydrostatic_pressure_qua(self, v, T_atm):
		# prescribe estimated constant temperature gradient
		gradT_geo = 0.02 # K/m, ~3./100
		# quadratic profile according to density dependending on temperature
		linear = rho_wat * gravity * (self.v_max - v)
		factor = 1 + beta_Tw * ((T_atm-self.T_ref) - 0.5 * gradT_geo * (self.v_max - v))
		p_pore = linear * factor
		return p_pore

	# def lithostatic_stresses(self, v, props):
	# 	heights = np.abs(np.diff(props.south_layer_bounds))
	# 	rho_eff = (1. - props.poro_array) * props.rho_array + \
	# 				(props.poro_array - props.biot_array) * 1000.
	# 	layer_stress = rho_eff * gravity * heights
	# 	total_stress = np.append(0, np.add.accumulate(layer_stress[:-1]))
	# 	stress = [0., 0., 0.]
	# 	for i, (ls, lh, lv, ts, nu, ab) in enumerate(zip(layer_stress, heights, props.south_layer_bounds[:-1], total_stress, props.nu_array, props.biot_array)):
	# 		if lv >= v >= props.south_layer_bounds[i+1]:
	# 			stress = np.array([nu / (1. - nu), 1., nu / (1. - nu)]) * ls/lh * (lv - v) + ts
	# 			break
	# 	return -stress

	def plot_profile(self, T_atm):
		vRange = np.linspace(self.v_min,self.v_max,20)
		#fRange = self.hydrostatic_pressure(vRange)
		#fRange = self.lateral_heatflux(vRange, T_atm)
		fRange1 = self.hydrostatic_pressure_lin(vRange)
		fRange2 = self.hydrostatic_pressure_exp(vRange)
		fRange3 = self.hydrostatic_pressure_qua(vRange, T_atm)
		fRange4 = fRange2 - fRange1
		fRange5 = fRange3 - fRange1
		fig,ax = plt.subplots()
		ax.set_title('Vertical profile')
		ax.plot(fRange1, vRange, label='linear')
		ax.plot(fRange2, vRange, label='exponential')
		ax.plot(fRange3, vRange, label='quadratic')
		#ax.plot(fRange4, vRange, label='exponential - linear')
		#ax.plot(fRange5, vRange, label='quadratic - linear')
		ax.set_xlabel('$p$ / Pa')
		ax.set_ylabel('$y$ / m')
		ax.grid()
		fig.legend()
		plt.show()

	# def plot_lithostatic_stress(self, props):
	# 	vRange = np.linspace(0,-1000, 100)
	# 	fRange = [1e-6*self.lithostatic_stresses(v, props)[0] for v in vRange]
	# 	fig,ax = plt.subplots()
	# 	ax.set_title('Vertical profile')
	# 	ax.plot(vRange, fRange)
	# 	ax.set_ylabel('$sigma$ / MPa')
	# 	ax.set_xlabel('$y$ / m')
	# 	ax.grid()
	# 	plt.show()

	def plot_profile_evolution(self):
		vRange = np.linspace(self.v_min,self.v_max,20)
		TRange = np.linspace(self.T_ini,self.T_ini-10,10)
		fig,ax = plt.subplots()
		for T_atm in TRange:
			fRange = self.lateral_heatflux(vRange, T_atm)
			ax.plot(fRange, vRange, label='T_atm=$%.2f $ ' %(T_atm))
		ax.set_title('Vertical profile')
		ax.set_xlabel('$q_x$ / W/m²')
		ax.set_ylabel('$v$ / m')
		ax.grid()
		#fig.legend()
		plt.legend(loc='upper left', bbox_to_anchor = (1.05, 1.0))
		plt.show()


"""
@Project: Calibrate Cameras
@File: validation.py
@Function: validate the results.
@Date: 23/06/17
"""
import numpy as np
import cv2
from calibrate_initial import *
from match_points import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def get_pts(file):
    with open(file, 'r') as file:
        lines = file.readlines()
    # 提取匹配点坐标
    pts = []

    for line in lines:
        # 解析每一行的坐标
        x1, y1, x2, y2, x3, y3, x4, y4 = map(float, line.strip().split(','))
        # 存储到pts1和pts2数组中
        pts.append([x1, y1])
        pts.append([x2, y2])
        pts.append([x3, y3])
        pts.append([x4, y4])
    # 转换为NumPy数组
    pts = np.array(pts)

    return pts

def get_2dto3dpoints(K1,K2,R,T,pts1,pts2):

    #获得相机1的投影矩阵
    P1 = K1@np.array([[1, 0, 0, 0],
                   [0, 1, 0, 0],
                   [0, 0, 1, 0]])
    # print(K1.shape)
    #获得相机2的投影矩阵
    # print(K2.shape)
    # print(R.shape)
    # print(T.shape)
    # print(P1)
    P2 = K2@np.hstack((R, T))
    # print(P2)
    #利用匹配点和投影矩阵，采用三角测量算法可以获得三维坐标
    points_4d_homogeneous = cv2.triangulatePoints(P1, P2, pts1.T,pts2.T)
    # print(points_4d_homogeneous)
    #转为齐次坐标
    points_3d_homogeneous = points_4d_homogeneous / points_4d_homogeneous[3]
    points_3d = points_3d_homogeneous[:3].T

    return points_3d , points_3d_homogeneous

if __name__ == '__main__':
    #待验证的相机数量
    nums_cameras=3

    ptss=[]
    #获取各相机的L型匹配点
    for i in range(nums_cameras):
        pt_cam="validation/"+str(i+1)+".txt"
        pts=get_pts(pt_cam)
        # print(pts)
        ptss.append(pts)

    #读取获得的内外参数
    Ks=np.loadtxt("Result\IntrinsicMatrix.txt",delimiter=',')
    Rs=np.loadtxt("Result\RotationMatrix.txt",delimiter=',')
    Ts=np.loadtxt("Result\Translation.txt",delimiter=',')

    ACss=[]
    ABss=[]
    DCss=[]

    points_3ds=[]

    for i in range(nums_cameras-1):
        dot_ADC=[]
        dot_BDC=[]

        points_3d,points_3d_homogeneous=get_2dto3dpoints(Ks[0].reshape(3,3),Ks[i+1].reshape(3,3),Rs[i+1].reshape(3,3),Ts[i+1].reshape(3,1),ptss[0],ptss[i+1])
        ACs=[]
        ABs=[]
        DCs=[]
        points_3ds.append(points_3d)
        # np.savetxt("pts"+str(i)+".txt",points_3d)
        ZCs=[]
        AB_weighted_averages=[]
        AC_weighted_averages=[]
        DC_weighted_averages=[]
        for j in range(0,points_3d.shape[0],4):

            #读取A,B,C,D点
            A=points_3d[j]
            B=points_3d[j+1]
            C=points_3d[j+2]
            D=points_3d[j+3]

            # 按照点C距离原点的距离进行加权平均
            ZCs.append(C[2])
        #     #对三维的A,B,C,D进行可视化
        #     x=points_3d[j:j+4,0]
        #     y=points_3d[j:j+4,1]
        #     z=points_3d[j:j+4,2]
        #     # 创建三维坐标系
        #     fig = plt.figure()
        #     ax = fig.add_subplot(111, projection='3d')

        #     # 绘制四个点
        #     ax.scatter(x,y,z, c=['red', 'green', 'blue', 'purple'])

        #     # 设置坐标轴标签
        #     ax.set_xlabel('X')
        #     ax.set_ylabel('Y')
        #     ax.set_zlabel('Z')

        #     # 显示图形
        #     plt.show()

            vector_DC = D - C
            vector_BC = B - C
            vector_AC = A - C

            ACs.append(np.linalg.norm(A-C))
            ABs.append(np.linalg.norm(A-B))
            DCs.append(np.linalg.norm(D-C))

            ACss.append(np.linalg.norm(A-C))
            ABss.append(np.linalg.norm(A-B))
            DCss.append(np.linalg.norm(D-C))

            # print(vector_AC,vector_AD)
            #保存内积，理论上两个向量垂直，即内积为0，归一化后即为cos角度
            dot_ADC.append(np.dot(vector_AC/np.linalg.norm(A-C), vector_DC/np.linalg.norm(D-C)))
            dot_BDC.append(np.dot(vector_BC/np.linalg.norm(B-C), vector_DC/np.linalg.norm(D-C)))
        #输出内积均值

        #C点小的值权重大
        weights = 1 / np.array(ZCs)
        AC_weighted_average = np.average(np.array(ACs), weights=weights)
        AB_weighted_average = np.average(np.array(ABs), weights=weights)
        DC_weighted_average = np.average(np.array(DCs), weights=weights)

        AC_weighted_averages.append(AC_weighted_average)
        AB_weighted_averages.append(AB_weighted_average)
        DC_weighted_averages.append(DC_weighted_average)

        # print("AC:",AC_weighted_average)
        # print("AB:",AB_weighted_average)
        # print("DC:",DC_weighted_average)
        print("第",i+1,"对相机的余弦输出及其角度：")
        print("cosADC:",np.mean(dot_ADC),"∠ADC=",np.arccos(np.mean(dot_ADC))*180/np.pi)
        print("cosBDC:",np.mean(dot_BDC),"∠BDC=",np.arccos(np.mean(dot_BDC))*180/np.pi)

        # print(np.mean(ACs))
        # print(np.mean(ABs))
        # print(np.mean(DCs))
    ACss=np.array(ACss).reshape(-1,1)
    ABss=np.array(ABss).reshape(-1,1)
    DCss=np.array(DCss).reshape(-1,1)
    np.savetxt("distance.csv",np.concatenate([ACss,ABss,DCss],axis=1),header="AC,AB,DC")
    # np.savetxt("distanceAB.csv",ABss)
    # np.savetxt("distanceDC.csv",DCss)

    print("AC_weights:",np.mean(AC_weighted_average))
    print("AB_weights:",np.mean(AB_weighted_average))
    print("DC_weights:",np.mean(DC_weighted_average))

    print("AC:",np.mean(ACss))
    print("AB:",np.mean(ABss))
    print("DC:",np.mean(DCss))


import snscrape.modules.twitter as sntwitter
import pandas as pd
import csv
from fetch_cur_tweets_dict import get_cur_usernames_dict

def scrape_single_page(username):
    username = "from:" + username
    scraper = sntwitter.TwitterSearchScraper(username)
    return scraper.get_items()
    # for tweet in scraper.get_items():
    #     print(tweet.content)
    #     count += 1
    #     if count > 100:
    #         break

def main():
    with open('ScrapedCurrentTweets.csv', 'w', newline='\n') as file:
        writer = csv.writer(file)
        username_dict = get_cur_usernames_dict()
        #Republicans
        count = 0
        for username in username_dict["Republican"]:
            tweets = scrape_single_page(username)
            tweet_count = 0
            for tweet in tweets:
                writer.writerow(["Republican", tweet.content])
                tweet_count += 1
                if tweet_count == 200:
                    break
            count += 1
            print(count)
        #Democrats
        count = 0
        for username in username_dict["Democrat"]:
            tweets = scrape_single_page(username)
            tweet_count = 0
            for tweet in tweets:
                writer.writerow(["Democrat", tweet.content])
                tweet_count += 1
                if tweet_count == 200:
                    break
            count += 1
            print(count)

main()


'''
2021-06-10 17:02:38
Write a function to find the longest common prefix string amongst an array of strings.

If there is no common prefix, return an empty string "".

Example 1:

Input: strs = ["flower","flow","flight"]
Output: "fl"
Example 2:

Input: strs = ["dog","racecar","car"]
Output: ""
Explanation: There is no common prefix among the input strings.

Constraints:

1 <= strs.length <= 200
0 <= strs[i].length <= 200
strs[i] consists of only lower-case English letters.

'''

strs = ["flower","flow","flight", "fles", "flor"]
strs = ["dog","racecar","car","star"]
strs = ["a"]
strs=["ab", "a"]
strs = ["dog","racecar","car","star"]
strs = ["flower","flow","flight", "fles", "flor"]

def longestCommonPrefix(strs):

    common=""

    # check if there is a null
    if "" in strs:
        common=""

    # check if there is only one element in the array
    elif len(strs)==1:
        common=strs[0]

    else:
        # get the length of the shortest word(s)
        shortest_length=len(min((word for word in strs if word), key=len))
        print('shortest length word is of size:', shortest_length)

        # get all words of a certain length
        shortest_length_strs=[word for word in strs if len(word) == shortest_length]

        print('printing all words with shortest len', shortest_length_strs)

        # chosing a random shortest word
        shortest_length_word=shortest_length_strs[0]
        print('chosing a random shortest word to play around with:', shortest_length_word)

        # loop over the shortest word
        while shortest_length>0:

            word_to_test=shortest_length_word[:shortest_length]
            print('word is now:', word_to_test)

            # keep total matches found
            total_matches_found=0

            # looping over all the words
            for i, n in enumerate(strs):

                print('checking word: ', n)
                if word_to_test in n:
                    print('full match found, going to the next word')
                    common=word_to_test
                    if total_matches_found<len(strs):
                        total_matches_found=total_matches_found+1
                        print('total matches found:', total_matches_found)
                    else:
                        print('total matches found:', total_matches_found)
                        #return common
                else:
                    print('didnt find matching, breaking here: ', n)
                    common=""
                    total_matches_found=0
                    break
            shortest_length=shortest_length-1
    print('in the end:', common)

longestCommonPrefix(strs)


import helpers.helper as helper
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import scipy
import helpers.pre_processing as pre
import math
from scipy.signal import hilbert

#plots all regions for mouse group in 3 by 3 plot
def plot_regions(data, mouse, region1, region2, shift=0, xlab="Region 1", ylab="Region 2", title="Region 1 vs Region 2"):
    x_axis = np.arange(data.shape[-1]) + shift
    plt.plot(x_axis, data[mouse, region1, :], label=xlab, color='blue')
    plt.plot(x_axis, data[mouse, region2, :], label=ylab, color='red')

    plt.xlabel('Time(sec) post ketamine')
    plt.ylabel('Signal')
    plt.title(title)
    plt.legend()
    plt.show()

#plot corr, can reorder if necessary (if not using Tommasso's data)
def plot_corr(data, title="Correlations", reorder=False):
    corr = helper.get_corr_matrix(data)[0]
    if reorder:
        corr = reorder_corr(corr)
    plt.figure()
    sns.heatmap(corr, cmap='plasma', vmin=0.3, vmax=0.8)
    plt.xlabel("Region 1")
    plt.ylabel("Region 2")
    locs = [i + 0.5 for i in range(len(corr[0]))]
    labels = ['AI L','GIDI L','S1 L','M1 L','M2 L','Cg1 L','PrL L','IL L','CPu L','NAcC L','NAcSh L',
    'NAcSh R','NAcC R','CPu R','IL R','PrL R','Cg1 R','M2 R','M1 R','S1 R','GIDI R','AI R']
    plt.xticks(locs, labels, rotation=90)
    plt.yticks(locs, labels)
    plt.title(title)
    plt.show()

#gets indices for splicing, takes in center/length in minutes
#return left, right, shift
def get_splicing_info(center, length):
    right = center * 60 + length * 30 + 1200
    left = center * 60 - length * 30 + 1200
    shift = left - 1200
    return left, right, shift

#maps label to index
#allows for errors with capitalization and spacing
def label_to_index(label):
    label_dict = {
    'ail': 0,
    'gidil': 1,
    's1l': 2,
    'm1l': 3,
    'm2l': 4,
    'cg1l': 5,
    'prll': 6,
    'ill': 7,
    'cpul': 8,
    'naccl': 9,
    'nacshl': 10,
    'nacshr': 11,
    'naccr': 12,
    'cpur': 13,
    'ilr': 14,
    'prlr': 15,
    'cg1r': 16,
    'm2r': 17,
    'm1r': 18,
    's1r': 19,
    'gidir': 20,
    'air': 21
    }
    key = label.replace(" ", "").lower()
    return label_dict[key]

#plotting subplot, helper function
def plot_subplot(ax, data, mouse, region1, region2, shift=100, title="Region 1 vs Region 2", ylabel="Signal"):
    x_axis = (np.arange(data.shape[-1]) + shift) / 60

    labels = ['AI L','GIDI L','S1 L','M1 L','M2 L','Cg1 L','PrL L','IL L','CPu L','NAcC L','NAcSh L',
    'NAcSh R','NAcC R','CPu R','IL R','PrL R','Cg1 R','M2 R','M1 R','S1 R','GIDI R','AI R']

    ax.plot(x_axis, data[mouse, region1, :], label=labels[region1], color='blue')
    ax.plot(x_axis, data[mouse, region2, :], label=labels[region2], color='red')

    ax.set_xlabel('Time(min) post ketamine')
    ax.set_ylabel(ylabel)
    ax.set_title(title, fontsize=8)
    ax.legend(fontsize='small')
    # ax.legend(loc='upper left', bbox_to_anchor=(1, 1))
    # ax.legend()

    return ax

def plot_subplot_phase(ax, data, mouse, region1, region2, shift=100):
    x_axis = (np.arange(data.shape[-1]) + shift) / 60
    phase_difference = np.arccos(np.cos(data[mouse, region2, :] - data[mouse, region1, :])) * 180/math.pi
    labels = ['AI L','GIDI L','S1 L','M1 L','M2 L','Cg1 L','PrL L','IL L','CPu L','NAcC L','NAcSh L',
    'NAcSh R','NAcC R','CPu R','IL R','PrL R','Cg1 R','M2 R','M1 R','S1 R','GIDI R','AI R']

    ax.plot(x_axis, phase_difference, color='blue')
    r1 = labels[region1]
    r2 = labels[region2]
    title = "Phase Difference: " + r2 + " - " + r1
    ax.set_xlabel('Time(min) post ketamine')
    ax.set_ylabel("Phase Difference")
    ax.set_title(title, fontsize=8)
    # ax.legend(loc='upper left', bbox_to_anchor=(1, 1))
    # ax.legend()

    return ax

def plot_subplots_phase(data, region1, region2, center, length, labels=False):
    if labels:
        region1 = label_to_index(region1)
        region2 = label_to_index(region2)
    left, right, shift = get_splicing_info(center, length)
    data = data[:, :, left:right]
    data = pre.hamming(data)
    hilbert_transform = hilbert(data, axis=-1)
    data = np.angle(hilbert_transform)
    num_mice = data.shape[0]
    fig, axes = plt.subplots(3, 3)
    for i in range(3):
        for j in range(3):
            if 3*i + j > num_mice - 1:
                continue
            mouse_num = 3*i + j
            plot_subplot_phase(axes[i, j], data, mouse_num, region1,
                         region2, shift)
    title_big="Phase Differences"
    fig.suptitle(title_big)
    plt.tight_layout()
    plt.show()

def phase_diff(data, mouse, region1, region2, center, length, labels=False):

    if labels:
        region1 = label_to_index(region1)
        region2 = label_to_index(region2)
    left, right, shift = get_splicing_info(center, length)
    data = data[:, :, left:right]
    data = pre.hamming(data)
    hilbert_transform = hilbert(data, axis=-1)
    data = np.angle(hilbert_transform)
    x_axis = (np.arange(data.shape[-1]) + shift) / 60
    phase_difference = np.arccos(np.cos(data[mouse, region2, :] - data[mouse, region1, :])) * 180/math.pi
    labels = ['AI L','GIDI L','S1 L','M1 L','M2 L','Cg1 L','PrL L','IL L','CPu L','NAcC L','NAcSh L',
    'NAcSh R','NAcC R','CPu R','IL R','PrL R','Cg1 R','M2 R','M1 R','S1 R','GIDI R','AI R']

    plt.plot(x_axis, phase_difference, color='blue')
    r1 = labels[region1]
    r2 = labels[region2]
    title = "Phase Difference: " + r2 + " - " + r1
    plt.xlabel('Time(min) post ketamine')
    plt.ylabel("Phase Difference")
    plt.title(title, fontsize=8)
    # ax.legend(loc='upper left', bbox_to_anchor=(1, 1))
    # ax.legend()
    plt.show()

#plot several subplots
def plot_subplots(data, region1, region2, center, length, labels=False, title_big = "", ylabel="Signal", hilb=False):
    if labels:
        region1 = label_to_index(region1)
        region2 = label_to_index(region2)
    left, right, shift = get_splicing_info(center, length)
    data = data[:, :, left:right]
    data = pre.hamming(data)
    if hilb:
        hilbert_transform = hilbert(data, axis=-1)
        data = (np.angle(hilbert_transform))*180/math.pi
    num_mice = data.shape[0]
    fig, axes = plt.subplots(3, 3)
    for i in range(3):
        for j in range(3):
            if 3*i + j > num_mice - 1:
                continue
            mouse_num = 3*i + j
            title = "Mouse" + str(mouse_num)
            plot_subplot(axes[i, j], data, mouse_num, region1,
                         region2, shift, title=title, ylabel=ylabel)
    fig.suptitle(title_big)
    plt.tight_layout()
    plt.show()

#plot several subplots
def plot_phases_both(data, mouse_num, region1, region2, center, length, labels=False, unwrap = False, title_big = ""):
    if labels:
        region1 = label_to_index(region1)
        region2 = label_to_index(region2)
    left, right, shift = get_splicing_info(center, length)
    data = data[:, :, left:right]
    data = pre.hamming(data)
    hilbert_transform = hilbert(data, axis=-1)
    if unwrap:
        data = np.unwrap(np.angle(hilbert_transform)) * 180/math.pi
    else:
         data = np.angle(hilbert_transform) * 180/math.pi
    fig, axes = plt.subplots(figsize=(12, 8))
    plot_subplot(axes, data, mouse_num, region1,
                    region2, shift, title="Phase Plots", ylabel="Phase Angle")
    fig.suptitle(title_big)
    plt.tight_layout()
    plt.show()

#reorder correlation matrix to be more brain oriented, only when going from my data to Tommasso's
def reorder_corr(corr):
    brain_ordering = ['AI L','GIDI L','S1 L','M1 L','M2 L','Cg1 L','PrL L','IL L','CPu L','NAcC L','NAcSh L',
    'NAcSh R','NAcC R','CPu R','IL R','PrL R','Cg1 R','M2 R','M1 R','S1 R','GIDI R','AI R']
    old_ordering = [
    'DP L', 'DP R', 'IL L', 'IL R', 'PrL L', 'PrL R', 'Cg1 L', 'Cg1 R',
    'M2 L', 'M2 R', 'M1 L', 'M1 R', 'S1 L', 'S1 R', 'GIDI L', 'GIDI R',
    'Cl L', 'Cl R', 'CPu L', 'CPu R', 'NAcC L', 'NAcC R', 'NAcSh L', 'NAcSh R',
    'S1DZ L', 'S1DZ R', 'S1J L', 'S1J R', 'AI L', 'AI R'
    ]
    old_dict = {old_ordering[i]:i for i in range(len(old_ordering))}

    num_reg = len(brain_ordering)
    result = np.zeros((num_reg, num_reg))
    for i in range(num_reg):
        for j in range(num_reg):
            first_region = brain_ordering[i]
            sec_region = brain_ordering[j]
            first_index = old_dict[first_region]
            sec_index = old_dict[sec_region]
            result[i,j] = corr[first_index, sec_index]
    return result

#normalize based off baseline data using standard deviations
def normalize_to_baseline(baseline, sample):
    baseline_means = np.mean(baseline, axis=2)
    baseline_std = np.std(baseline, axis=2)
    baseline_means = baseline_means[:, :, np.newaxis]
    baseline_std = baseline_std[:, :, np.newaxis]
    normalized_data = (sample - baseline_means) / baseline_std
    return normalized_data

#normalize sample to itself
def normalize_to_sample(sample):
    sample_mean = np.mean(sample, axis=2)
    sample_std = np.std(sample, axis=2)
    sample_mean = sample_mean[:, :, np.newaxis]
    sample_std = sample_std[:, :, np.newaxis]
    return (sample-sample_mean) / sample_std

#plotting subplots of Dianni data

def load_dianni():
    data = helper.load_data_np('matlab_files/dianni_data/nalket_dianni.mat') #di ianni data
    hamming = helper.load_data_np('matlab_files/dianni_data/hamm_win.mat')
    windowed_data = hamming.T * data
    plot_corr(windowed_data)

    # data_baseline = helper.load_data_np('matlab_files/dianni_data/nalket_dianni_baseline.mat')
    # hamming_baseline = scipy.signal.windows.hamming(data_baseline.shape[-1])

    # plot_subplots(windowed_data, 2, 10, 150, title_big="Low Correlation: 2.5-17.5min")
# load_dianni()

#my data corr and plots
def plot_mine():
    data = helper.load_data_np('matlab_files/Time_Series_Data/ROI_CBV_Data/salket_m_roi.mat')
    data = pre.bandpass_filter(data, 0.01, 0.1, 5)[:, :, 2500:3500]
    hamming = scipy.signal.windows.hamming(data.shape[-1])
    data = hamming.T * data
    plot_corr(data, reorder=True)
    # plot_corr(data, reorder=True)
    # corr = helper.get_corr_matrix(data)[0]
    # corr = reorder_corr(corr)
    # helper.plot_correlation_ROI(corr)
# plot_regions(windowed_data, 7, 2, 14)

# plot_mine()
# def random_corr_matrix():
#     corr = helper.get_corr_matrix(windowed_data)[0]
#     plot_corr(windowed_data)
#     high 8/12, low 2/14

def normalize_test():
    data_baseline = helper.load_data_np('matlab_files/dianni_data/nalket_dianni_baseline.mat')
    data = helper.load_data_np('matlab_files/dianni_data/nalket_dianni.mat')
    hamming_baseline = scipy.signal.windows.hamming(data_baseline.shape[-1])
    hamming_data = scipy.signal.windows.hamming(data.shape[-1])
    data_baseline = hamming_baseline.T * data_baseline
    data = hamming_data.T * data
    # data = normalize_to_baseline(data_baseline, data)
    data = normalize_to_sample(data)
    plot_subplots(data, 8, 13, 150)
    # plot_corr(data, title="Correlations 2.5-17.5min")

# normalize_test()


def get_input() -> list:
    with open(r'input.txt', 'r') as file:
        file_list = file.read().splitlines()
    return file_list

def part_one(bag_list:list, looking_for:str) -> list:
    found_list = []
    for bag in bag_list:
        sep = bag.find(' bags contain')
        if looking_for in bag[sep:]:
            found_list.append(bag[:sep])
            found_list.extend(part_one(bag_list, bag[:sep]))
    return found_list
print(f'\nPart 1:  The shiny gold bag can go in {len(set(part_one(get_input(), "shiny gold")))} bags.')


import pandas as pd
import requests

# Function to find the category for an article based on its content
def get_category(article, category_ids):
    for category_id, category_description in category_ids.items():
        if category_description.lower() in article.lower():
            return category_description
    return None

# URLs for the dataset and categories files
dataset_url = 'https://drive.google.com/uc?id=1mAkm62rINY_l5ZOhtp8fJGPp6nRm1jWt'
categories_url = 'https://drive.google.com/uc?id=1m3zfhMlJmQFwwza7r6Dga04wpKPljhG9'

# Download dataset and categories files using pandas
dataset = pd.read_csv(dataset_url)
categories = pd.read_csv(categories_url)

# Create a dictionary that maps category IDs to their descriptions
category_ids = {}
for index, row in categories.iterrows():
    category_id = row['id']
    category_name = row['name']
    category_description = row['description']
    category_ids[category_id] = category_description

# Initialize a list to store the output data
output_data = []

# Classify the articles and populate the output list with categorized data
for index, row in dataset.iterrows():
    article_parts = [str(row['news']), str(row['title']), str(row['snippet']), str(row['content'])]
    article = ' '.join(article_parts)
    category = get_category(article, category_ids)
    output_data.append({'Article': article, 'Category': category})  # Store category in the 'Category' column

# Create a DataFrame from the output list
output_df = pd.DataFrame(output_data)

# Save the classified dataset to a CSV file
classified_dataset_file = 'classified_dataset.csv'
output_df.to_csv(classified_dataset_file, index=False)

# Print completion message
print("Classification completed. Classified dataset saved as", classified_dataset_file)


from os import chdir as os_chdir
from scipy import constants
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches  # for custom legends

os_chdir("./organized_code_modelling")

if_vienna_avgtypes = True
if_pharm_avgtypes = False

##### celltype parameters #####
if if_vienna_avgtypes:   # cell types in the already-published Vienna study differ from our own generated dataset
    df_cell_info = pd.read_csv("cell_freq_table_vienna.csv", index_col = 'Celltype')
elif if_pharm_avgtypes:
    df_cell_info = pd.read_csv("cell_freq_table_zurich.csv", index_col = 'Celltype')
else:
    print("Warning: have not picked cell type set")
#end

v_cell_names = df_cell_info.index.to_series()
v_cell_pcts = df_cell_info['Percent_PBMC']
v_cell_rad = df_cell_info['Diameter_um'] / 2

def calculate_SA(cell_radius):
    ''' input cell radius and outputs surface area in the same units squared'''
    return 4 * constants.pi * (cell_radius)**2
#end

cellTotal_conc = 2.5E6                         # cells / mL total cell concentration for whole PBMCs  (ends up not mattering in this model formulation)
v_cell_concs = cellTotal_conc * v_cell_pcts    # cells / mL (not used elsewhere in this version of the model)
v_cell_SA = calculate_SA(v_cell_rad)           # um^2

m_proportions = v_cell_pcts.to_numpy()[:, None] * v_cell_pcts.to_numpy()
m_proportions = pd.DataFrame(m_proportions, columns = v_cell_names.rename(None), index = v_cell_names.rename(None))

##### molecule copy numbers #####
prot_condition = "steady.state"
proteomics_csv = pd.read_csv("Rieckmann_proteomics_summary_mean_leuk.csv")
proteomics_csv.set_index('Uniprot', inplace = True)

def molecule_matrices(proteomics_csv,
                      condition = "steady.state",     # "steady.state"   or    "activated"
                      pick_mostabundant = False,      # match proteomics data cell type labales to microscopy labels by picking most abundant single subtype within the match category
                      pick_vienna_avgtypes = True,    # match proteomics data cell type labales to microscopy labels by averaging the subtypes
                      pick_pharm_avgtypes = False,    # match based on averaging for when using our pharmacoscopy cell types
                      pick_null_celltypes = False,        # null model where randomize cell type labels on expression matrices, but keep everything else constant
                      inp_null_celltypes_vec = np.array([]),      # if pick_null_celltypes = True, optionally can manually specify the randomized vector here, which is useful given the small number of possible permutations
                      cd = {'pT4':0.70, 'pBn':0.66, 'pBm':0.34, 'pMc':0.84, 'pMn':0.09, 'pMi':0.07,
                            'pDm':0.65, 'pDp':0.35, 'pT4n':0.41, 'pT4c':0.39, 'pT4e':0.19, 'pT4r':0.01,
                            'pT8n':0.39, 'pT8c':0.07, 'pT8e':0.22, 'pT8r':0.32, 'pNd':0.90, 'pNb':0.10,
                            'pTh1':0.17, 'pTh2':0.01, 'pTh17':0.01, 'pTreg':0.04,
                            'fTh1em':0.45 , 'fTh1cm':0.55, 'fTh2em':0.45, 'fTh2cm':0.55, 'fTregem':0.37, 'fTregcm':0.63, 'fTregmem': 0.80}):
    '''
    Converts matrix of celltype expression in absolute protein counts (such as Rieckmann data)
    to 2D concentrations formatted for vienna celltypes. \n
    Involves a translation between proteomics data's labels and vienna labels, for which options given.  \n
    Returns matrix of molecule 2D concentrations and vector of molecule rownames
    '''

    if pick_mostabundant and pick_vienna_avgtypes:
        m_molecules = proteomics_csv.filter(regex = (("_"+condition+"|").join(df_cell_info.Rieckmann_mostabundant))+"_"+condition).copy()
        renamingcells = dict(zip(df_cell_info.Rieckmann_mostabundant + "_" + condition, df_cell_info.index))
        m_molecules.rename(columns = renamingcells, inplace = True)
        m_molecules = m_molecules.reindex(columns = renamingcells.values(), copy = True)
    elif pick_vienna_avgtypes:
        m_molecules = pd.DataFrame()
        m_molecules['Bcell'] = cd['pBn']*proteomics_csv['B.naive'+"_"+condition] + cd['pBm']*proteomics_csv['B.memory'+"_"+condition]
        m_molecules['Mono'] = cd['pMc']*proteomics_csv['MO.classical'+"_"+condition] + cd['pMn']*proteomics_csv['MO.nonclassical_steady.state'] + cd['pMi']*proteomics_csv['MO.intermediate_steady.state']
        m_molecules['DC'] = cd['pDm']*proteomics_csv['mDC'+"_"+condition] + cd['pDp']*proteomics_csv['pDC'+"_"+condition]
        m_molecules['Tcell'] = (cd['pT4']*(cd['pT4n']*proteomics_csv['T4.naive'+"_"+condition] + cd['pTh1']*proteomics_csv['Th1_steady.state'] + cd['pTh2']*proteomics_csv['Th2_steady.state'] + cd['pTreg']*cd['fTregmem']*proteomics_csv["mTregs"+"_"+condition] + cd['pTreg']*(1-cd['fTregmem'])*proteomics_csv["nTregs"+"_"+condition] +
                                 (cd['pT4c'] - (cd['pTh1']*cd['fTh1cm']) - (cd['pTh2']*cd['fTh2cm']) - (cd['pTreg']*cd['fTregcm']))*proteomics_csv['T4.CM'+"_"+condition] + (cd['pT4e'] - (cd['pTh1']*cd['fTh1em']) - (cd['pTh2']*cd['fTh2em']) - (cd['pTreg']*cd['fTregem']))*proteomics_csv['T4.EM'+"_"+condition] + cd['pT4r']*proteomics_csv['T4.EMRA'+"_"+condition]) +
                                (1-cd['pT4'])*(cd['pT8n']*proteomics_csv['T8.naive'+"_"+condition] + cd['pT8c']*proteomics_csv['T8.CM'+"_"+condition] + cd['pT8e']*proteomics_csv['T8.EM'+"_"+condition] + cd['pT8r']*proteomics_csv['T8.EMRA'+"_"+condition]))
        m_molecules['NKcell'] = cd['pNd']*proteomics_csv['NK.dim'+"_"+condition] + cd['pNb']*proteomics_csv['NK.bright'+"_"+condition]
    elif pick_pharm_avgtypes:
        m_molecules = pd.DataFrame()
        m_molecules['B'] = cd['pBn']*proteomics_csv['B.naive'+"_"+condition] + cd['pBm']*proteomics_csv['B.memory'+"_"+condition]
        m_molecules['M14'] = cd['pMc']/(1-cd['pMn'])*proteomics_csv['MO.classical'+"_"+condition] + cd['pMi']/(1-cd['pMn'])*proteomics_csv['MO.intermediate_steady.state']
        m_molecules['M16'] = proteomics_csv['MO.nonclassical_steady.state']
        m_molecules['DC'] = cd['pDm']*proteomics_csv['mDC'+"_"+condition] + cd['pDp']*proteomics_csv['pDC'+"_"+condition]
        m_molecules['T4'] = (cd['pT4n']*proteomics_csv['T4.naive'+"_"+condition] + cd['pTh1']*proteomics_csv['Th1_steady.state'] + cd['pTh2']*proteomics_csv['Th2_steady.state'] + cd['pTreg']*cd['fTregmem']*proteomics_csv["mTregs"+"_"+condition] + cd['pTreg']*(1-cd['fTregmem'])*proteomics_csv["nTregs"+"_"+condition] +
                                 (cd['pT4c'] - (cd['pTh1']*cd['fTh1cm']) - (cd['pTh2']*cd['fTh2cm']) - (cd['pTreg']*cd['fTregcm']))*proteomics_csv['T4.CM'+"_"+condition] + (cd['pT4e'] - (cd['pTh1']*cd['fTh1em']) - (cd['pTh2']*cd['fTh2em']) - (cd['pTreg']*cd['fTregem']))*proteomics_csv['T4.EM'+"_"+condition] + cd['pT4r']*proteomics_csv['T4.EMRA'+"_"+condition])
        m_molecules['T8'] = cd['pT8n']*proteomics_csv['T8.naive'+"_"+condition] + cd['pT8c']*proteomics_csv['T8.CM'+"_"+condition] + cd['pT8e']*proteomics_csv['T8.EM'+"_"+condition] + cd['pT8r']*proteomics_csv['T8.EMRA'+"_"+condition]
        m_molecules['NK'] = cd['pNd']*proteomics_csv['NK.dim'+"_"+condition] + cd['pNb']*proteomics_csv['NK.bright'+"_"+condition]
    #end
    v_molec_names = m_molecules.index.to_series()

    m_molec_conc = m_molecules.divide(v_cell_SA, axis = 1)

    if pick_null_celltypes:

        if inp_null_celltypes_vec.any():
            shuffle_vec_indxs = inp_null_celltypes_vec

        else:
            shuffle_vec_indxs = np.random.permutation(np.arange(m_molec_conc.shape[1]))
        #end
        m_molec_conc_original = np.array(m_molec_conc)
        m_molec_conc = pd.DataFrame(m_molec_conc_original[:, shuffle_vec_indxs],
                                    index = m_molec_conc.index, columns = m_molec_conc.columns)
    #end

    return m_molec_conc, v_molec_names
#end

cellavg_proportions = {'pT4':0.70, 'pBn':0.66, 'pBm':0.34, 'pMc':0.84, 'pMn':0.09, 'pMi':0.07,
                       'pDm':0.65, 'pDp':0.35, 'pT4n':0.41, 'pT4c':0.39, 'pT4e':0.19, 'pT4r':0.01,
                       'pT8n':0.39, 'pT8c':0.07, 'pT8e':0.22, 'pT8r':0.32, 'pNd':0.90, 'pNb':0.10,
                       'pTh1':0.17, 'pTh2':0.01, 'pTh17':0.01, 'pTreg':0.04,
                       'fTh1em':0.45 , 'fTh1cm':0.55, 'fTh2em':0.45, 'fTh2cm':0.55, 'fTregem':0.37, 'fTregcm':0.63, 'fTregmem': 0.80}

m_molec_conc, v_molec_names = molecule_matrices(proteomics_csv,
                                                condition = prot_condition,
                                                pick_mostabundant = False,
                                                pick_vienna_avgtypes = if_vienna_avgtypes, pick_pharm_avgtypes = if_pharm_avgtypes,
                                                cd = cellavg_proportions)

if prot_condition == "steady.state":  # non-activated cells do not put HLA-F on their surface (https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3867582/)
    m_molec_conc.loc[m_molec_conc.index == "P30511", :] = np.zeros(m_molec_conc.iloc[0, :].shape)
#end

##### interaction matrix ####
interactions_csv = pd.read_csv("Leuk_interaction_affinities.csv",
                               usecols =[i for i in range(8)])

interactions_csv.Gene_L.replace(r'-[^-]+$', '', regex = True, inplace = True)
interactions_csv.Gene_R.replace(r'-[^-]+$', '', regex = True, inplace = True)

def pair_labeler(str1, str2, sep = " + "):
    ''' Combines two strings into one in alphabetical order (so unique). Returns str'''
    return np.where(str1 < str2, (str1 + sep + str2), (str2 + sep + str1))
#end
interactions_csv['gene_pair'] = interactions_csv.apply(lambda row: pair_labeler(row.Gene_L, row.Gene_R), axis = 1).astype(str)
interactions_csv['pair_label'] = interactions_csv.apply(lambda row: pair_labeler(row.Uniprot_L, row.Uniprot_R), axis = 1).astype(str)

def interaction_matrices(interactions_csv,
                         v_molec_names,
                         pick_excludeintxns = [],         # any interactions to exclude (can pass [] as blank)
                         pick_nullintxn = False,          # randomly permute interaction network to generate a null distribution
                         rand_seed = None):               # if randomly permuting, can optionally provide seed for RNG, otherwise defaults to python's standard rand seed
    '''
    Reformats interaction table in list format from csv into matrices of all protein-protein pairs and their affinities  \n
    Includes options to create "null models" where interactions or affinities are scrambled, to see how much the specific interaction network
    contributes to the results (and how much may be explainable by general factors like a celltype's average expression magnitude). \n
    Returns an updated version of the input csv table and a binary intxn matrix and a Kd matrix
    '''
    if len(pick_excludeintxns) > 0:
        interactions_csv = interactions_csv[~interactions_csv['pair_label'].isin(pick_excludeintxns)]
    #end

    m_interaction_names = pd.DataFrame((v_molec_names.values[:,None] + " + " + v_molec_names.values),
                                  columns = v_molec_names, index = v_molec_names)

    if pick_nullintxn:
        uniprots_with_intxn = [*interactions_csv['Uniprot_L'].to_list(), *interactions_csv['Uniprot_R'].to_list()]
        uniprots_with_intxn = pd.Series(list(set(uniprots_with_intxn)))
        uniprots_intxn_names_pairs = pd.DataFrame((uniprots_with_intxn.values[:,None] + " + " + uniprots_with_intxn.values),
                              columns = uniprots_with_intxn, index = uniprots_with_intxn)
        pairlabels_rand = uniprots_intxn_names_pairs.melt().sample(n = interactions_csv.shape[0], random_state = rand_seed)
        interactions_csv.loc[:, 'pair_label'] = pairlabels_rand['value'].values
    #end

    remappingkd = dict(zip(interactions_csv[interactions_csv.Kd_nM.notnull()].pair_label, interactions_csv[interactions_csv.Kd_nM.notnull()].Kd_nM))
    remappingkd_recip = dict(zip([" + ".join(intxn_name_i.split(" + ")[::-1]) for intxn_name_i in list(remappingkd.keys())],  # reverses ordering of protein names
                                 remappingkd.values()))
    remappingkd.update(remappingkd_recip)
    m_interactions_Kd3D = m_interaction_names.applymap(lambda x: remappingkd.get(x, 0))

    m_interactions = m_interactions_Kd3D > 0

    return interactions_csv, m_interactions, m_interactions_Kd3D
#end

(interactions_csv,
 m_interactions, m_interactions_Kd3D) = interaction_matrices(interactions_csv,
                                                            v_molec_names,
                                                            pick_excludeintxns = ['C10orf54 + HLA-E', 'C10orf54 + HLA-F', 'TNF + TNR1A'],  # non-surface
                                                            pick_nullintxn = False)

def calculate_3Dto2D(Kd_3D, height = 0.010):
    ''' Converts 3D Kd in Molar units to 2D Kd in units moelcules/um2 (height is in um units)  \n
        uses constant height factor to get rough estimates since very few are empirically measured'''
    return constants.Avogadro * 1E-15 * height * Kd_3D
#end

m_interactions_Kd2D = calculate_3Dto2D(m_interactions_Kd3D * 1E-9,
                                       height = 0.010)

m_interactions_Kd2D.loc["P06729", "P19256"] = m_interactions_Kd2D.loc["P19256", "P06729"] = 1.1 # rare case of litterature-published 2D Kd can use instead of inferred value (doi:10.1007/s10439-005-2504-5)

def bound_concentration(conc_A_total, conc_B_total, Kd_2D):
    '''
    Calculates amount of bound A-B given the total amount of A and B and their 2D affinity \n
     dervied from the quadratic solution the system of equations \n
     Kd = Rec * Lig / Bound \n
     Rec = Rec_tot - Bound  \n
     Lig = Lig_tot - Bound  \n
    Where all concentrations and Kd are in the same units (e.g. molecules/um2)
    Beware this equation does not handle case if 0 is input as a concentration or Kd (because smaller numbers = stronger binding so zero is perfect conversion to bound state)
    '''
    conc_AB = ((Kd_2D + conc_A_total + conc_B_total) -
               ((Kd_2D + conc_A_total + conc_B_total)**2 -
               (4 * conc_A_total * conc_B_total))**(1/2)) / 2
    return conc_AB
#end

# extract from binary intxn matrix pairs of molecule indexes
v_interaction_tuples = list(zip(*np.asarray(m_interactions_Kd2D).nonzero()))

def bound_matrices(v_interaction_tuples, v_molec_names,
                   m_molec_conc, m_interactions_Kd2D, v_cell_names):
    '''
    Iterates through all the binary interactions with Kd values, calculates bound concentrations for each.
     Sums up total bound concentration for each cell-cell pair across all the protein interactions. \n
    Returns matrix of cell pairs with sum of bound concentrations, and list of bound concentrations per protein interaction that fed into the sum
    '''
    t_bound = []
    sum_bound = pd.DataFrame(np.zeros((len(v_cell_names),len(v_cell_names))), columns = v_cell_names, index = v_cell_names)
    bound_rownames = pd.Series(["null"] * len(v_interaction_tuples), dtype = "str")

    for n, intxn_n in enumerate(v_interaction_tuples):
        intxn_name = "%s + %s"%(v_molec_names[intxn_n[0]], v_molec_names[intxn_n[1]])
        bound_rownames[n] = intxn_name

        conc_A_n = m_molec_conc.iloc[intxn_n[0], :]
        conc_B_n = m_molec_conc.iloc[intxn_n[1], :]

        conc_A = conc_A_n.to_frame()
        conc_B = conc_B_n.to_frame().transpose()

        m_bound = bound_concentration(conc_A.to_numpy()[:, None],
                                       conc_B.to_numpy(),
                                       Kd_2D = m_interactions_Kd2D.iloc[intxn_n])

        m_bound = np.minimum(m_bound, np.minimum(conc_A.to_numpy()[:, None], conc_B.to_numpy()))
        m_bound = m_bound.squeeze(axis = 1)

        sum_bound += pd.DataFrame(m_bound, columns = v_cell_names, index = v_cell_names)
        t_bound.append(m_bound)
    #end
    return sum_bound, t_bound, bound_rownames
#end

(sum_bound, t_bound, bound_rownames) = bound_matrices(v_interaction_tuples, v_molec_names,
                                                      m_molec_conc, m_interactions_Kd2D, v_cell_names)

##### fitting kinetic model's matrix of bound sums to vienna data #####

from sklearn import linear_model, metrics
from scipy import stats as stats
import statsmodels.api as sm

m_contact_obs =  pd.DataFrame({'DC':    pd.Series([2.4,     1.4,     np.nan,    0.2]),  # data : https://pubmed.ncbi.nlm.nih.gov/28437395/
                               'Mono':  pd.Series([1.3,     1.5,     0.4,       np.nan]),
                               'Bcell': pd.Series([np.nan,  0.3,     0.2,       0.0]),
                               'Tcell': pd.Series([0.0,     np.nan,  0.0,       0.0])})
m_contact_obs.set_index(m_contact_obs.columns, inplace = True)
m_contact_obs = m_contact_obs[v_cell_names[0:4]]  # re-order columns to match other matrices
m_contact_obs = m_contact_obs.reindex(v_cell_names[0:4], axis = 'index')

x_obs = sum_bound.drop('NKcell', axis=1).drop('NKcell', axis=0).to_numpy().flatten()
y_obs = m_contact_obs.to_numpy().flatten()

linreg = linear_model.LinearRegression()  # import model (ordinary least squares)
mask = np.isfinite(x_obs) & np.isfinite(y_obs)
linreg.fit(X = x_obs[mask, None],
           y = y_obs[mask, None])

beta = linreg.coef_
intercept = linreg.intercept_
fit_r2 = metrics.r2_score(y_pred = linreg.predict(x_obs[mask, None]),
                          y_true = y_obs[mask, None])  # calculate a few different ways to check correct
fit_r2 = pd.DataFrame({'sumbound':x_obs[mask], 'obscontact':y_obs[mask]}).corr(method = "pearson") ** 2
fit_r2 = linreg.score(X = x_obs[mask, None],
                      y = y_obs[mask, None])

fit_stats = sm.OLS(y_obs[mask], sm.add_constant(x_obs[mask], prepend = False))
fit_stats = fit_stats.fit()
fit_pval = fit_stats.pvalues[0]  # fit_stats.summary()

colors_dict = {'Bcell':'#de412c', 'Mono':'#ba91bd', 'DC':'#618f59', 'Tcell':'#619fcf'}
color_m_contact_obs = m_contact_obs.rename(columns = colors_dict).rename(index = colors_dict)
m_labels_flat = pd.Series((color_m_contact_obs.columns.values[:, None] + "," + color_m_contact_obs.index.values).flatten())
m_labels_flat.replace(colors_dict, inplace = True)
m_labels = pd.DataFrame({'Cell1': m_labels_flat.replace(r',.*$', '', regex = True),
                         'Cell2': m_labels_flat.replace(r'^.*,', '', regex = True)})

xval = x_obs[mask]   # manually define 95% confidence interval
yval = y_obs[mask]
yfit = x_obs[mask]*fit_stats.params[0] + fit_stats.params[1]
dof = xval .size - 2
t = stats.t.ppf(0.975, df = dof)
resid = yval - yfit
chi2 = np.sum((resid / yfit)**2)
chi2_red = chi2 / dof
s_err = np.sqrt(np.sum(resid**2) / dof)
x2 = np.linspace(xval.min(),xval.max(),500)
y2 = np.linspace(xval.min(),xval.max(),500)*fit_stats.params[0] + fit_stats.params[1]
ci = t * s_err * np.sqrt(1/xval.size + (x2 - np.mean(xval))**2 / np.sum((xval - np.mean(xval))**2))

fig, ax = plt.subplots(dpi = 300)  # dpi = 150 for console printing  ; else 300 to save
ax.fill_between(x2, y2 + ci, y2 - ci, facecolor ="#918e86", edgecolor = "none", alpha = 0.12)
plt.plot(np.linspace(x_obs[mask].min(),x_obs[mask].max(),500), np.linspace(x_obs[mask].min(),x_obs[mask].max(),500)*fit_stats.params[0] + fit_stats.params[1],
         linestyle = ":", color = "grey")
ax.scatter(x = x_obs[mask],
           y = y_obs[mask],
           s = 50,           # size of points
           c = m_labels.Cell1[mask])
ax.scatter(x = x_obs[mask],
           y = y_obs[mask],
           s = 10,           # size of points
           c = m_labels.Cell2[mask])
fig.text(0.16, 0.56, r'$r^2 = %.3f$' '\n' r'$p = %.3f$'%(fit_r2, fit_pval), fontsize=10)
ax.set_xlabel('Kinetic model interaction score')
ax.set_ylabel('Vienna microscopy contact score')
ax.legend(handles = [mpatches.Patch(color=colors_dict[key], label=key) for key in colors_dict],
          loc = 'upper left')
#fig.savefig('kineticmodel_v8_viennafig_cellavg_linreg.png', bbox_inches='tight')

#### export values to csv ####

names_cellpair = (sum_bound.columns.values[:, None] + " + " + sum_bound.index.values).flatten()  # make square matrix of cell-cell pairs in score table, then flatten to vector as above for 'y_obs' etc
names_intxnpair = [v_molec_names[tuple_i[0]] + " + " + v_molec_names[tuple_i[1]] if v_molec_names[tuple_i[0]] < v_molec_names[tuple_i[1]] else v_molec_names[tuple_i[1]] + " + " + v_molec_names[tuple_i[0]] for tuple_i in v_interaction_tuples]

table_bound_raw = pd.DataFrame([t_sub.flatten() for t_sub in t_bound])  # convert each adhesion matrix into vector
table_bound_raw.columns = names_cellpair   # insert cell pair (col) and protein intxn (row) names
table_bound_raw.index = names_intxnpair
table_bound_raw.index.name = "Interaction"

table_bound_raw['Pair_label'] = table_bound_raw.index.map(dict(zip(interactions_csv.pair_label, interactions_csv.gene_pair)))
#table_bound_raw.to_csv("kineticmodel_v8_zurich_resting_leukmanual_tablebound.csv")


import enum
import zlib
from llpdf.types.PDFName import PDFName

class Filter(enum.IntEnum):
	Uncompressed = 0
	FlateDecode = 1
	DCTDecode = 2
	RunLengthDecode = 3
	CCITTFaxDecode = 4
	ASCIIHexDecode = 5
	ASCII85Decode = 6

class PNGPredictor(enum.IntEnum):
	No = 0
	Sub = 1
	Up = 2
	Average = 3
	Paeth = 4

class Predictor(enum.IntEnum):
	NoPredictor = 1
	TIFFPredictor2 = 2
	PNGPredictionNone = 10
	PNGPredictionSub = 11
	PNGPredictionUp = 12
	PNGPredictionAverage = 13
	PNGPredictionPaeth = 14
	PNGPredictionOptimum = 15

class EncodedObject(object):
	_FILTER_MAP = {
		Filter.FlateDecode:		PDFName("/FlateDecode"),
		Filter.RunLengthDecode:	PDFName("/RunLengthDecode"),
		Filter.DCTDecode:		PDFName("/DCTDecode"),
		Filter.CCITTFaxDecode:	PDFName("/CCITTFaxDecode"),
		Filter.ASCIIHexDecode:	PDFName("/ASCIIHexDecode"),
		Filter.ASCII85Decode:	PDFName("/ASCII85Decode"),
	}
	_REV_FILTER_MAP = { value: key for (key, value) in _FILTER_MAP.items() }

	def __init__(self, encoded_data, filtering, columns = 1, predictor = Predictor.NoPredictor):
		self._encoded_data = encoded_data
		self._filtering = filtering
		self._columns = columns
		self._predictor = predictor

	@property
	def decompressible(self):
		return self._filtering != Filter.DCTDecode

	@property
	def compressed(self):
		return self._filtering != Filter.Uncompressed

	@property
	def encoded_data(self):
		return self._encoded_data

	@property
	def filtering(self):
		return self._filtering

	@property
	def columns(self):
		return self._columns

	@property
	def predictor(self):
		return self._predictor

	@property
	def lossless(self):
		return self._filtering != Filter.DCTDecode

	@property
	def meta_dict(self):
		meta = {
			PDFName("/Length"):		len(self),
		}
		if self._filtering != Filter.Uncompressed:
			meta[PDFName("/Filter")] = self._FILTER_MAP[self._filtering]
		if self._predictor != Predictor.NoPredictor:
			meta[PDFName("/DecodeParms")] = {
				PDFName("/Columns"): self.columns,
				PDFName("/Predictor"): int(self.predictor),
			}
		return meta

	def update_meta_dict(self, content_object):
		if PDFName("/Filter") in content_object:
			del content_object[PDFName("/Filter")]
		if PDFName("/DecodeParms") in content_object:
			del content_object[PDFName("/DecodeParms")]
		content_object.update(self.meta_dict)

	@staticmethod
	def _rle_decode(rle_data):
		result = bytearray()
		index = 0
		while index < len(rle_data):
			length = rle_data[index]
			index += 1
			if 0 <= length <= 127:
				bytecnt = 1 + length
				result += rle_data[index : index + bytecnt]
				index += bytecnt
			elif length == 128:
				# EOD
				break
			else:
				bytecnt = 257 - length
				value = rle_data[index]
				result += bytecnt * bytes([ value ])
				index += 1
		return result

	def _decompress(self):
		"""Only decompress filter, but do not de-predict."""
		if self._filtering == Filter.Uncompressed:
			return self._encoded_data
		elif self._filtering == Filter.FlateDecode:
			return zlib.decompress(self._encoded_data)
		elif self._filtering == Filter.RunLengthDecode:
			return self._rle_decode(self._encoded_data)
		else:
			raise Exception(NotImplemented, self._filtering)

	@staticmethod
	def _paeth_predictor(a, b, c):
		# a = left, b = above, c = upper left
		p = a + b - c
		pa = abs(p - a)
		pb = abs(p - b)
		pc = abs(p - c)
		if (pa <= pb) and (pa <= pc):
			return a
		elif pb <= pc:
			return b
		else:
			return c

	def _depredict(self, deencoded_data):
		if self.predictor == Predictor.NoPredictor:
			return deencoded_data
		else:
			result = bytearray()
			previous_scanline = bytes(self.columns)
			for i in range(0, len(deencoded_data), self.columns + 1):
				png_filter = PNGPredictor(deencoded_data[i])
				scanline = deencoded_data[i + 1 : i + self.columns + 1]
				if png_filter == PNGPredictor.No:
					decompressed_scanline = scanline
				elif png_filter == PNGPredictor.Sub:
					previous_value = 0
					decompressed_scanline = bytearray()
					for (index, value) in enumerate(scanline):
						new_value = (value + previous_value) & 0xff
						decompressed_scanline.append(new_value)
						previous_value = new_value
				elif png_filter == PNGPredictor.Up:
					decompressed_scanline = bytes((prev + cur) & 0xff for (prev, cur) in zip(previous_scanline, scanline))
				elif png_filter == PNGPredictor.Average:
					previous_value = 0
					decompressed_scanline = bytearray()
					for (index, value) in enumerate(scanline):
						new_value = (value + ((previous_value + previous_scanline[index]) // 2)) & 0xff
						decompressed_scanline.append(new_value)
						previous_value = new_value
				elif png_filter == PNGPredictor.Paeth:
					previous_value = 0
					decompressed_scanline = bytearray()
					for (index, value) in enumerate(scanline):
						new_value = (value + self._paeth_predictor(previous_value, previous_scanline[index], previous_scanline[index - 1] if (index > 0) else 0)) & 0xff
						decompressed_scanline.append(new_value)
						previous_value = new_value
				else:
					raise Exception(NotImplemented, png_filter)
				result += decompressed_scanline
#				print("%2d %2d: %s" % (int(png_filter), i // (self.columns + 1), scanline.hex()))
#				print("   %2d: %s" % (i // (self.columns + 1), decompressed_scanline.hex()))
				previous_scanline = decompressed_scanline
			return result

	def decode(self):
		return self._depredict(self._decompress())

	@classmethod
	def _predict(cls, unpredicted_data, columns):
		assert(len(unpredicted_data) % columns == 0)
		rows = len(unpredicted_data) // columns

		predicted_data = bytearray()
		if rows == 1:
			predictor = Predictor.PNGPredictionSub
			predicted_data.append(PNGPredictor.Sub)
			predicted_data += bytearray((x - y) & 0xff for (x, y) in zip(unpredicted_data[1:], unpredicted_data))
		else:
			predictor = Predictor.PNGPredictionUp
			previous_scanline = bytes(columns)
			for row in range(rows):
				scanline = unpredicted_data[row * columns : (row + 1) * columns]
				predicted_data.append(PNGPredictor.Up)
				predicted_data += bytearray((x - y) & 0xff for (x, y) in zip(previous_scanline, scanline))

		return (predicted_data, predictor, columns)

	@classmethod
	def create(cls, unencoded_data, compress = True, predict = False, columns = None):
		if predict:
			if columns is None:
				columns = len(unencoded_data)
			(predicted_data, used_predictor, predictor_columns) = cls._predict(unencoded_data, columns)
		else:
			(predicted_data, used_predictor, predictor_columns) = (unencoded_data, Predictor.NoPredictor, 1)

		if compress:
			encoded_data = zlib.compress(predicted_data)
			filtering = Filter.FlateDecode
		else:
			encoded_data = predicted_data
			filtering = Filter.Uncompressed
		encoded_object = cls(encoded_data = encoded_data, filtering = filtering, predictor = used_predictor, columns = predictor_columns)
		return encoded_object

	@classmethod
	def from_object(cls, obj):
		if PDFName("/Filter") in obj.content:
			pdf_filter = obj.content[PDFName("/Filter")]
			if isinstance(pdf_filter, list):
				if len(pdf_filter) == 1:
					pdf_filter = pdf_filter[0]
				else:
					raise Exception("Cannot create EncodedObject from object that has multiple filters applied: %s" % (str(pdf_filter)))
			filtering = cls._REV_FILTER_MAP[pdf_filter]
		else:
			filtering = Filter.Uncompressed

		if (PDFName("/DecodeParms") in obj.content) and (PDFName("/Predictor") in obj.content[PDFName("/DecodeParms")]):
			predictor = Predictor(obj.content[PDFName("/DecodeParms")][PDFName("/Predictor")])
			columns = obj.content[PDFName("/DecodeParms")][PDFName("/Columns")]
		else:
			predictor = Predictor.NoPredictor
			columns = 1

		return cls(encoded_data = obj.raw_stream, filtering = filtering, predictor = predictor, columns = columns)

	def __len__(self):
		return len(self._encoded_data)

	def __str__(self):
		details = [ ]
		details.append("%d bytes" % (len(self)))
		details.append(self._filtering.name)
		if self.predictor != Predictor.NoPredictor:
			details.append(self.predictor.name)
			details.append("%d columns" % (self.columns))
		return "EncodedObject<%s>" % (", ".join(details))

if __name__ == "__main__":
	with open("xref.bin", "rb") as f:
		data = f.read()
	enc = EncodedObject(data, filtering = Filter.Uncompressed, columns = 5, predictor = Predictor.PNGPredictionUp)
	print(enc)
	print(enc.decode().hex())

	data = b"Foobar Bar Foobar " * 25
	print(EncodedObject.create(data, predict = True, columns = len(data) // 25))
	print(EncodedObject.create(data, predict = False))
	print(data)


from google.colab import files

uploaded_file = files.upload()
uploaded_file_name = next(iter(uploaded_file))

import cv2
import matplotlib.pyplot as plt

img = cv2.imread(uploaded_file_name)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

height = img.shape[0]
width = img.shape[1]

fix = 0.1
img = cv2.resize(img,None,fx = fix , fy = fix)
img_next = cv2.resize(img , (width,height),interpolation = cv2.INTER_NEAREST)

plt.imshow(img_next)


import pandas as pd

def main():
        df = pd.read_csv("./finalProj/avoPrices.csv")
        df2015 = df[df['year'] == 2015]
        df2016 = df[df['year'] == 2016]
        df2017 = df[df['year'] == 2017]
        df2018 = df[df['year'] == 2017]

        out2015 = open("./finalProj/avo2015.csv", "w+")
        out2016 = open("./finalProj/avo2016.csv", "w+")
        out2017 = open("./finalProj/avo2017.csv", "w+")
        out2018 = open("./finalProj/avo2018.csv", "w+")

        df2015.to_csv(out2015, index=False)
        df2016.to_csv(out2016, index=False)
        df2017.to_csv(out2017, index=False)
        df2018.to_csv(out2018, index=False)

if __name__ == '__main__':
    main()


import argparse
import logging
from typing import List

from ..ui.utils import alphanumeric
from ...foundation.cli.command.Command import Command
from ...manager.Kathara import Kathara
from ...model.Lab import Lab
from ...strings import strings, wiki_description

class VconfigCommand(Command):
    def __init__(self) -> None:
        Command.__init__(self)

        self.parser: argparse.ArgumentParser = argparse.ArgumentParser(
            prog='kathara vconfig',
            description=strings['vconfig'],
            epilog=wiki_description,
            add_help=False
        )

        self.parser.add_argument(
            '-h', '--help',
            action='help',
            default=argparse.SUPPRESS,
            help='Show a help message and exit.'
        )
        self.parser.add_argument(
            '-n', '--name',
            metavar='DEVICE_NAME',
            required=True,
            help='Name of the device to be connected on desired collision domains.'
        )

        group = self.parser.add_mutually_exclusive_group(required=True)

        group.add_argument(
            '--add',
            type=alphanumeric,
            dest='to_add',
            metavar='CD',
            nargs='+',
            help='Specify the collision domain to add.'
        )
        group.add_argument(
            '--rm',
            type=alphanumeric,
            dest='to_remove',
            metavar='CD',
            nargs='+',
            help='Specify the collision domain to remove.'
        )

    def run(self, current_path: str, argv: List[str]) -> None:
        self.parse_args(argv)
        args = self.get_args()

        lab = Lab("kathara_vlab")

        machine_name = args['name']
        device = lab.get_or_new_machine(machine_name)
        device.api_object = Kathara.get_instance().get_machine_api_object(machine_name, lab_name=lab.name)

        if args['to_add']:
            for cd in args['to_add']:
                logging.info(
                    "Adding interface to device `%s` on collision domain `%s`..." % (machine_name, cd)
                )
                link = lab.get_or_new_link(cd)
                Kathara.get_instance().connect_machine_to_link(device, link)

        if args['to_remove']:
            for cd in args['to_remove']:
                logging.info(
                    "Removing interface on collision domain `%s` from device `%s`..." % (cd, machine_name)
                )
                (_, link) = lab.connect_machine_to_link(machine_name, cd)
                link.api_object = Kathara.get_instance().get_link_api_object(cd, lab_name=lab.name)

                Kathara.get_instance().disconnect_machine_from_link(device, link)


# -*- coding: utf-8 -*-
"""
Created on Sat Nov 25 19:51:39 2017

@author: Madan
"""
from __future__ import division
def calculate_emd(y_pred, y_test):
    y_pred = list(y_pred)
    y_test = list(y_test)
    d_pred_prob = {}
    d_test_prob = {}
    total_instances = len(y_test)
    #print total_instances
    #all_classes = set(y_test)
    all_classes = [2, 1, 0, -1, -2]
    C = len(all_classes)

    for i in all_classes:
        count1 = y_pred.count(i)
        count2 = y_test.count(i)
        d_pred_prob[i] = count1/total_instances
        d_test_prob[i] = count2/total_instances
    #print d_pred_prob
    #print d_test_prob
    score = 0;
    for i in range(0,C):
        sum1 = 0;
        sum2 = 0;
        for j in range(0,i+1):
            sum1 += d_pred_prob[all_classes[j]]
            sum2 += d_test_prob[all_classes[j]]
        score += abs(sum1-sum2)
    return score


# -*- coding: utf-8 -*-
"""
Created on Thu Nov  8 12:28:08 2018

@author: Tomas
"""
import matplotlib.pyplot as plt
import numpy as np

# Import this using below:
#import sys
#sys.path.append(r'C:\Users\Tomas\Imperial College London\Battle, Ronan - Shared MSci project\Scripts') # "/Users/Tomas/Imperial\ College London/Battle,\ Ronan\ -\ Shared\ MSci\ project/Scripts")
#sys.path.append(r'C:\Users\rbatt\OneDrive - Imperial College London\Shared MSci project\Scripts')
#from Shot_Class import Shot

# import signals dictionaries
from signal_dict_10_NOV_11 import signals
# import pickle loading function
from data_access_funcs import load_signal_data

class Shot():
    def __init__(self, ShotNumber, transitions = None):
#       must define signals dictionary here based on shot number/session series
        self.data = {} # initialise data dictionary
        self.ShotNumber = ShotNumber
        self._transitions = transitions

        self.unloaded = []
        for sig in signals:
            try:
                fn = str(ShotNumber) + '_' + sig + '.p'
                self.data[sig] = load_signal_data(fn)
            except FileNotFoundError:
                self.unloaded.append(sig)

    def signals_present(self):
        keys = self.data.keys()
        return len(keys), keys

    def signal_has_errors(self, SignalName):
        if self.data[SignalName]['errors'] is None:
            return False
        else:
            return True

    def plot_signal(self, SignalName):
        # this should go into another file anyway with plotting stuff
        x,y = self.data[SignalName]['time'], self.data[SignalName]['data']
        units = self.data[SignalName]['units']
        try:
            plt.plot(x,y)

            plt.ylabel('%s %s' % (SignalName, units))
            plt.annotate(self.ShotNumber, (0.9,0.9), xycoords = 'axes fraction')
            plt.xlim(-0.1)
        except:
            print('Can\'t deal with this signal yet')

    def plot_compare(self, signalslist):
        n = len(signalslist)
        fig, ax = plt.subplots(n, sharex=True, figsize=(11,7))
        ax[-1].set_xlim([0,1])
        fig.suptitle('Shot %s' % (self. ShotNumber))

        for i, signal in enumerate(signalslist):
            self._plot_ax_sig(ax, signal, i, signame= signal)

    def plot_JP(self, tlim = (0,1), ip = 'IP', wmhd = 'WMHD', coreTe = 'AYC_TE0',
                ne = 'ANE_DENSITY', Dalpha = 'AIM_DA_TO', Bt = 'BT',
                Ploss = 'Ploss', PINJ = 'PINJ', POHM = 'POHM'
                ):
        """
        Plot some signals together on single figure
        """
        n_signals = 7 # number of signals to be plotted

        # make figure and adjust boundaries
        fig, ax = plt.subplots(n_signals, sharex=True, figsize=(11,7))
        fig.subplots_adjust(top=0.935,bottom=0.09,left=0.08,right=0.975,hspace=0.0,wspace=0.2)

        # add figure title, labels, set time range
        fig.suptitle('Shot %s' % (self. ShotNumber))
        ax[-1].set_xlabel('$time \ [s]$')
        ax[-1].set_xlim(tlim)

        # mark LH, HL transition points
        if self._transitions:
            pass
        else:
            print('No transitions known')

        # plot plasma current, ax = 0
        self._plot_ax_sig(ax, ip, 0, signame = 'I_{p}')
        # plot n_e
        self._plot_ax_sig(ax, ne, 1, signame = 'n_{e}', plot_errors='fill') # , units='m^{-2}')
        # plot WMHD
        self._plot_ax_sig(ax, wmhd,2 , signame = 'W_{MHD}')
        # plot core T_e
        self._plot_ax_sig(ax, coreTe, 3 , signame = 'T_{e}^{core}',plot_errors='fill')
        # plot Dalpha of some kind
        self._plot_ax_sig(ax, Dalpha, 4, signame = 'D_{\\alpha}')
        # plot B_T
        self._plot_ax_sig(ax, Bt, 5, signame = 'B_{T}')
        # plot powers
        #self._plot_ax_sig(ax, POHM, panel = 6, signame = 'Pohm', label='POHM', errors=True) #Ploss = 'Ploss', PINJ = 'PINJ', POHM = 'POHM'
        self._plot_ax_sig(ax, Ploss, panel = 6, signame = 'Ploss', label='Ploss', plot_errors='fill')
        self._plot_ax_sig(ax, PINJ, panel = 6, signame = 'PINJ', label='PINJ', plot_errors='fill')

    def _plot_ax_sig(self, ax, sig, panel, signame='', plot_errors=False, units=None, label=None):
        """ for use in plot_JP
        plots signal (sig) on subplot axes (ax[panel]) and labels with signame (and sig units)
        need to add error handeling
        Note: use units with care, as not longer automatic, maybe just use for presentation plots
        """
        # type(ax[panel]) might want to make sure its correct
        shape = self.data[sig]['data'].shape
        if len(shape) == 2 and shape[0] != 1:
            print('Data format of {} has another dimension -> taking mean'.format(sig))
            time = self.data[sig]['time']
            data = np.nanmean(self.data[sig]['data'],axis=1)
            try:
                errors = np.nanmean(self.data[sig]['errors'],axis=1)
            except:
                print('No errors found in {}'.format(sig))
                errors = self.data[sig]['errors']
        #elif len(shape) == 2 and shape[0] == 1:
        #    time = self.data[sig]['time']
         #   data = self.data[sig]['data'][0]
         #   errors = self.data[sig]['errors'][0]

        else:
            time = self.data[sig]['time']
            data = self.data[sig]['data']
            errors = self.data[sig]['errors'] # may be None, deal with later

        if not units:
            units = self.data[sig]['units']
        else:
            pass
        if units == 'MW':
            data=data*10e6/5 #what?

        if signame!='':
            label=signame

        # Decide how to plot depending on plot_error call
        if plot_errors == True:
            #linestyle = '.'
            ax[panel].errorbar(time,data, yerr=errors,ecolor="red",label=label)
        elif plot_errors == 'fill':
            #print(sig)
            #print(data)                                        # why not just
            # fix if errors is nothing then make it zero        # if errors != None: ax[panel].fill_between(...)
            try:                                                # else: pass      #just goes on to plot instead
                if errors ==None: errors = np.zeros(len(data))
            except:
                pass
            #print(errors)

            ax[panel].plot(time,data,label=label)
            ax[panel].fill_between(time, data-np.nan_to_num(errors), data + np.nan_to_num(errors), alpha=0.3)

        else:
            ax[panel].plot(time,data,label=label)

        ax[panel].annotate(r'$%s \ [%s]$' %(signame, units), xy=(0.01,0.7), xycoords='axes fraction', fontsize=11)
        ax[panel].legend()


from scripts.helpful_scripts import get_account
from scripts.deploy_and_pull import deploy, deployMock, fundWithLink, approveToken, pull
from brownie import network, config, exceptions

from web3 import Web3
import pytest

# For global marker
pytestmark = [pytest.mark.require_network("mainnet-fork")]

# deploy erc-20 and mint to account
def token(ERC20Token):
    account = get_account(index=-2)
    token = account.deploy(ERC20Token, Web3.toWei(1000, "ether"))
    return token

# module wide deploy and mock.
@pytest.fixture(scope="module", autouse=True)
def default_deploy_and_mock(VRFCoordinatorMock, RussianRoulette, interface):
    account = get_account(index=-2)
    mockVRF = deployMock(account, VRFCoordinatorMock)

    RR = deploy(
        account,
        RussianRoulette,
        mockVRF.address,
        config["networks"][network.show_active()]["link"],
        config["networks"][network.show_active()]["keyhash"],
        Web3.toWei(config["networks"][network.show_active()]["fee"], "ether"),
    )

    txFund = fundWithLink(
        account,
        RR,
        Web3.toWei(config["networks"][network.show_active()]["fee"], "ether"),
    )

    DAI = config["networks"][network.show_active()]["dai"]
    amount = Web3.toWei(config["networks"][network.show_active()]["amount_in"], "ether")
    token_contract = interface.IERC20(DAI)

    txApprove = approveToken(account, DAI, RR.address, amount)

    starting_token_balance = Web3.fromWei(
        token_contract.balanceOf(account.address), "ether"
    )

    return account, mockVRF, RR, DAI, amount, token_contract, starting_token_balance

@pytest.fixture(autouse=True)
def isolation(fn_isolation):
    pass

def test_pull(default_deploy_and_mock):
    # Arrange

    (
        account,
        mockVRF,
        RR,
        DAI,
        amount,
        token_contract,
        starting_token_balance,
    ) = default_deploy_and_mock

    # Act
    txPull = pull(account, RR, amount, DAI)

    assert "requestId" in list(txPull.events["RequestRandomness"].keys())

    # request_id = txPull.events["RequestRandomness"]["requestId"]

    # Assert contract balance == amount in. Assert user balance is amount in less.
    assert (
        Web3.fromWei(token_contract.balanceOf(RR.address), "ether")
        == config["networks"][network.show_active()]["amount_in"]
    )
    assert (
        starting_token_balance
        - Web3.fromWei(token_contract.balanceOf(account.address), "ether")
    ) == config["networks"][network.show_active()]["amount_in"]

def test_return_token(default_deploy_and_mock, chain):
    # Arrange

    (
        account,
        mockVRF,
        RR,
        DAI,
        amount,
        token_contract,
        starting_token_balance,
    ) = default_deploy_and_mock

    # Act
    txPull = pull(account, RR, amount, DAI)

    request_id = txPull.events["RequestRandomness"]["requestId"]

    # has not been 50 blocks yet
    with pytest.raises(exceptions.VirtualMachineError):
        txReturn = RR.returnToken(request_id, {"from": account})

    # > 50 blocks waited
    chain.mine(51)
    txReturn = RR.returnToken(request_id, {"from": account})

    assert Web3.fromWei(token_contract.balanceOf(RR.address), "ether") == 0
    assert (
        Web3.fromWei(token_contract.balanceOf(account.address), "ether")
        == starting_token_balance
    )

def test_return_token_2(default_deploy_and_mock, chain):
    # Arrange

    (
        account,
        mockVRF,
        RR,
        DAI,
        amount,
        token_contract,
        starting_token_balance,
    ) = default_deploy_and_mock

    # Act
    txPull = pull(account, RR, amount, DAI)

    request_id = txPull.events["RequestRandomness"]["requestId"]

    RNG = 601
    txRNG = mockVRF.callBackWithRandomness(
        request_id, RNG, RR.address, {"from": account}
    )

    chain.mine(51)

    # RNG was successful, so requestIdToRandomnessZero is False
    with pytest.raises(exceptions.VirtualMachineError):
        txReturn = RR.returnToken(request_id, {"from": account})

# test a random number that is divisible by 6. User should lose amount in, contract should burn amount in.
# mainnet-fork with index=-2 unlocked
# tested token is DAI
# @pytest.mark.require_network("mainnet-fork")
@pytest.mark.parametrize("RNG", [600, 601, 0])
def test_full(default_deploy_and_mock, RNG, chain):
    # Arrange

    (
        account,
        mockVRF,
        RR,
        DAI,
        amount,
        token_contract,
        starting_token_balance,
    ) = default_deploy_and_mock

    # Assert default setting is false
    # assert RR.requestIdToRandomnessZero(request_id) == False

    # Act
    txPull = pull(account, RR, amount, DAI)
    # STATIC_RNG = 601
    request_id = txPull.events["RequestRandomness"]["requestId"]

    # Assert account is mapped
    assert RR.requestIdToSender(request_id) == account.address
    assert RR.requestIdToToken(request_id) == token_contract.address
    assert RR.requestIdToAmount(request_id) == amount
    assert RR.requestIdToBlockNumber(request_id) == txPull.block_number

    # Assert default setting is false
    assert RR.requestIdToRandomnessZero(request_id) == True

    # Assert contract balance == amount in. Assert user balance is amount in less.
    assert (
        Web3.fromWei(token_contract.balanceOf(RR.address), "ether")
        == config["networks"][network.show_active()]["amount_in"]
    )
    assert (
        starting_token_balance
        - Web3.fromWei(token_contract.balanceOf(account.address), "ether")
    ) == config["networks"][network.show_active()]["amount_in"]

    # can mock since vrfCoordinator is the only address that can call rawFullfillRandomness
    # mockVRF is deployed before RR, so mockVRF address can be input for RR constructor
    txRNG = mockVRF.callBackWithRandomness(
        request_id, RNG, RR.address, {"from": account}
    )

    txRNG.wait(1)

    # Assert contract balance is 0, token returned. Assert user balance == starting balance, tokens returned. Due to randomness not found
    if RNG == 0:
        chain.mine(51)
        txReturn = RR.returnToken(request_id, {"from": account})
        assert Web3.fromWei(token_contract.balanceOf(RR.address), "ether") == 0
        assert (
            Web3.fromWei(token_contract.balanceOf(account.address), "ether")
            == starting_token_balance
        )
        pass

    # Assert contract balance is 0, token returned. Assert user balance == starting balance, tokens returned.
    elif RNG % 6 != 0:
        assert Web3.fromWei(token_contract.balanceOf(RR.address), "ether") == 0
        assert (
            Web3.fromWei(token_contract.balanceOf(account.address), "ether")
            == starting_token_balance
        )

    # Assert contract balance is 0, token burned. Assert user balance == starting balance less amount in.
    else:
        assert Web3.fromWei(token_contract.balanceOf(RR.address), "ether") == 0
        assert (
            starting_token_balance
            - Web3.fromWei(token_contract.balanceOf(account.address), "ether")
        ) == config["networks"][network.show_active()]["amount_in"]

    zeroAddress = "0x0000000000000000000000000000000000000000"
    assert RR.requestIdToSender(request_id) == zeroAddress
    assert RR.requestIdToToken(request_id) == zeroAddress
    assert RR.requestIdToAmount(request_id) == zeroAddress
    assert RR.requestIdToBlockNumber(request_id) == zeroAddress

# test the fn_isolation works
def test_transfer(accounts):
    starting_balance = accounts[1].balance()
    accounts[0].transfer(accounts[1], "10 ether")
    assert (accounts[1].balance() - starting_balance) == Web3.toWei(10, "ether")

def test_chain_reverted(accounts):
    assert accounts[0].balance() == accounts[1].balance()

# def test_mapping_initialization(default_deploy_and_mock):

#     (
#         account,
#         mockVRF,
#         RR,
#         DAI,
#         amount,
#         token_contract,
#         starting_token_balance,
#     ) = default_deploy_and_mock

#     # RR.


import string

data = []

with open("C:\\Dev\\projects\\advent-of-code\\python\\day04\\input.txt") as f:
    record = []
    for line in f.read().split("\n"):
        if line != '':
            for element in line.split(" "):
                record.append(element)
        else:
            data.append(record)
            record = []

valid = 0
fields = ["byr", "iyr", "eyr", "hgt", "hcl", "ecl", "pid", "cid"]

for record in data:
    found = [0, 0, 0, 0, 0, 0, 0, 1]

    for element in record:
        found[fields.index(element.split(":")[0])] += 1
    if 0 not in found:
        # use for part 1        valid += 1
        good = True
        for element in record:
            left = element.split(":")[0]
            right = element.split(":")[1]
            if left == "byr":
                if((int(right) < 1920) or (int(right) > 2002)):
                    good = False
            elif left == "iyr":
                if ((int(right) < 2010) or (int(right) > 2020)):
                    good = False
            elif left == "eyr":
                if ((int(right) < 2020) or (int(right) > 2030)):
                    good = False
            elif left == "hgt":
                measure = right[-2::]
                if measure == "cm":
                    value = int(right.split("cm")[0])
                    if value < 150 or value > 193:
                        good = False
                elif measure == "in":
                    value = int(right.split("in")[0])
                    if value < 59 or value > 76:
                        good = False
                else:
                    good = False
            elif left == "hcl":
                if len(right) != 7 or right[0] != "#":
                    good = False
                for x in right[1::]:
                    if x not in string.hexdigits:
                        good = False
            elif left == "ecl":
                if right not in ["amb", "blu", "brn", "gry", "grn", "hzl", "oth"]:
                    good = False
            elif left == "pid":
                if len(right) != 9 or right.isnumeric() == False:
                    good = False

        if good == True:
            valid += 1

print(str(valid))


#
# Advent of Code 2017
# Bryan Clair
#
# Day 24
#

# Picking up again some years later, switching to Python3
import sys
sys.path.append('../../Advent 2022')
from aocutils import *

args = parse_args()

inputlines = [x.strip() for x in open(args.file).readlines()]

loops = set()
used = set()
edges = {}

for i,line in enumerate(inputlines):
    v1,v2 = [int(x) for x in line.split('/')]
    if v1 < v2:
        v1,v2 = v2,v1
    e = (v1,v2)

    if v1 == v2:
        loops.add(e)

    if v1 not in edges:
        edges[v1] = []
    if v2 not in edges:
        edges[v2] = []
    edges[v1].append(e)
    if v1 != v2:
        edges[v2].append(e)

def weight(edge):
    v1,v2 = edge
    return v1+v2

maxdepth = 0
best_at_maxdepth = 0

def best_bridge(start,depth=0,sofar=0):
    """Return the best bridge you can build starting at vertext start.
    depth tracks how many edges used.
    sofar tracks how much strenght we've built so far."""

    # print(' '*depth+'visiting',start,end=':')
    # find unused edges
    myedges = []
    for e in edges[start]:
        if e not in used:
            myedges.append(e)

    #print(' '*depth,myedges)

    if not myedges:
        global maxdepth
        global best_at_maxdepth
        if depth > maxdepth:
            maxdepth = depth
            best_at_maxdepth = sofar
        elif depth == maxdepth:
            if sofar > best_at_maxdepth:
                best_at_maxdepth = sofar

    # use any loops - not really worth the effort, runs a smidge faster
    for e in myedges:
        if e in loops:
            used.add(e)
            best = weight(e) + best_bridge(start, depth+1, sofar+weight(e))
            used.remove(e)
            return best

    # try each edge
    best = 0
    for e in myedges:
        v1,v2 = e
        next = v2 if (v1 == start) else v1
        used.add(e)
        s = weight(e) + best_bridge(next, depth+1, sofar+weight(e))
        used.remove(e)
        if s > best:
            best = s

    return best

print('part 1:',best_bridge(0))
print('part 2:',best_at_maxdepth)