import os
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import tqdm
import numpy as np
from AFSD.common.thumos_dataset import THUMOS_Dataset, get_video_info, \
    load_video_data, detection_collate, get_video_anno
from torch.utils.data import DataLoader
from AFSD.thumos14.BDNet import BDNet
from AFSD.thumos14.multisegment_loss import MultiSegmentLoss
from AFSD.common.config import config
import matplotlib.pyplot as plt

def get_basic_config(config):
    class cfg: pass
    cfg.batch_size = config['training']['batch_size']
    cfg.learning_rate = config['training']['learning_rate']
    cfg.weight_decay = config['training']['weight_decay']
    cfg.max_epoch = config['training']['max_epoch']
    cfg.checkpoint_path = config['training']['checkpoint_path']
    cfg.focal_loss = config['training']['focal_loss']
    cfg.edl_loss = config['training']['edl_loss'] if 'edl_loss' in config['training'] else False
    cfg.edl_config = config['training']['edl_config'] if 'edl_config' in config['training'] else None
    cfg.cls_loss_type = 'edl' if cfg.edl_loss else 'focal' # by default, we use focal loss
    cfg.os_head = config['model']['os_head'] if 'os_head' in config['model'] else False
    num_classes = config['dataset']['num_classes']
    cfg.num_classes = num_classes - 1 if cfg.os_head else num_classes
    cfg.random_seed = config['training']['random_seed']
    cfg.train_state_path = os.path.join(cfg.checkpoint_path, 'training')
    assert os.path.exists(cfg.train_state_path)
    cfg.resume = config['training']['resume']
    cfg.use_edl = config['model']['use_edl'] if 'use_edl' in config['model'] else False
    cfg.clip_length = config['dataset']['training']['clip_length']
    cfg.overlap_thresh = config['training']['piou']
    cfg.output_path = config['testing']['output_path']
    return cfg

def get_path(input_path):
    if os.path.lexists(input_path):
        fullpath = os.path.realpath(input_path) if os.path.islink(input_path) else input_path
        real_name = fullpath.split('/')[-1]
        real_full_path = os.path.join(os.path.dirname(input_path), real_name)
    else:
        raise FileNotFoundError
    return real_full_path

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

GLOBAL_SEED = 1

def worker_init_fn(worker_id):
    set_seed(GLOBAL_SEED + worker_id)

def set_rng_state(states):
    random.setstate(states[0])
    np.random.set_state(states[1])
    torch.set_rng_state(states[2])
    if torch.cuda.is_available():
        torch.cuda.set_rng_state(states[3])

def setup_model(cfg):
    net = BDNet(in_channels=config['model']['in_channels'],
                backbone_model=config['model']['backbone_model'], use_edl=cfg.use_edl)
    net = nn.DataParallel(net, device_ids=[0]).cuda()
    model_path = get_path(os.path.join(cfg.checkpoint_path, 'checkpoint-latest.ckpt'))
    net.module.load_state_dict(torch.load(model_path))
    train_path = get_path(os.path.join(cfg.train_state_path, 'checkpoint_latest.ckpt'))
    state_dict = torch.load(train_path)
    set_rng_state(state_dict['state'])
    net.eval()
    return net

def setup_dataloader(cfg):
    train_video_infos = get_video_info(config['dataset']['training']['video_info_path'])
    train_video_annos = get_video_anno(train_video_infos,
                                       config['dataset']['training']['video_anno_path'],
                                       config['dataset']['class_info_path'])
    train_data_dict = load_video_data(train_video_infos,
                                      config['dataset']['training']['video_data_path'])
    train_dataset = THUMOS_Dataset(train_data_dict,
                                   train_video_infos,
                                   train_video_annos)
    train_data_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True,
                                   num_workers=4, worker_init_fn=worker_init_fn,
                                   collate_fn=detection_collate, pin_memory=True, drop_last=True)
    epoch_step_num = len(train_dataset) // cfg.batch_size
    return train_data_loader, epoch_step_num

def compute_iou(pred, target):
    """
    jaccard: A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
    """
    pred_left = pred[:, 0]
    pred_right = pred[:, 1]
    target_left = target[:, 0]
    target_right = target[:, 1]

    pred_area = pred_left + pred_right
    target_area = target_left + target_right

    eps = torch.finfo(torch.float32).eps

    inter = torch.min(pred_left, target_left) + torch.min(pred_right, target_right)
    area_union = target_area + pred_area - inter
    ious = inter / area_union.clamp(min=eps)
    return ious

def get_matched_targets(targets, loc_data, priors, clip_length):
    num_batch = loc_data.size(0)
    num_priors = priors.size(0)
    # match priors and ground truth segments
    loc_t = torch.Tensor(num_batch, num_priors, 2).to(loc_data.device)
    conf_t = torch.LongTensor(num_batch, num_priors).to(loc_data.device)
    prop_loc_t = torch.Tensor(num_batch, num_priors, 2).to(loc_data.device)
    prop_conf_t = torch.LongTensor(num_batch, num_priors).to(loc_data.device)
    for idx in range(num_batch):
        truths = targets[idx][:, :-1]
        labels = targets[idx][:, -1]
        pre_loc = loc_data[idx]
        K = priors.size(0)
        N = truths.size(0)
        center = priors[:, 0].unsqueeze(1).expand(K, N)
        left = (center - truths[:, 0].unsqueeze(0).expand(K, N)) * clip_length
        right = (truths[:, 1].unsqueeze(0).expand(K, N) - center) * clip_length
        area = left + right
        maxn = clip_length * 2
        area[left < 0] = maxn
        area[right < 0] = maxn
        best_truth_area, best_truth_idx = area.min(1)

        loc_t[idx][:, 0] = (priors[:, 0] - truths[best_truth_idx, 0]) * clip_length
        loc_t[idx][:, 1] = (truths[best_truth_idx, 1] - priors[:, 0]) * clip_length
        conf = labels[best_truth_idx]
        conf[best_truth_area >= maxn] = 0
        conf_t[idx] = conf

        iou = compute_iou(pre_loc, loc_t[idx])  # [num_priors]
        prop_conf = conf.clone()
        prop_conf[iou < cfg.overlap_thresh] = 0
        prop_conf_t[idx] = prop_conf
        prop_w = pre_loc[:, 0] + pre_loc[:, 1]
        prop_loc_t[idx][:, 0] = (loc_t[idx][:, 0] - pre_loc[:, 0]) / (0.5 * prop_w)
        prop_loc_t[idx][:, 1] = (loc_t[idx][:, 1] - pre_loc[:, 1]) / (0.5 * prop_w)

    return loc_t, conf_t, prop_loc_t, prop_conf_t

def evidence_func(logit, evidence='exp'):
    if evidence == 'relu':
        return F.relu(logit)
    if evidence == 'exp':
        return torch.exp(torch.clamp(logit, -10, 10))
    if evidence == 'softplus':
        return F.softplus(logit)

def grad_edl(logit, target, num_cls=15, evidence='exp'):
    if logit.dim() > 2:
        # N,C,d1,d2 -> N,C,m (m=d1*d2*...)
        logit = logit.view(logit.size(0), logit.size(1), -1)
        logit = logit.transpose(1, 2).contiguous()  # [N,C,d1*d2..] -> [N,d1*d2..,C]
        logit = logit.view(-1, logit.size(-1))  # [N,d1*d2..,C]-> [N*d1*d2..,C]
    target = target.view(-1)  # [N,d1,d2,...]->[N*d1*d2*...,]
    # one-hot embedding for the target
    y = torch.eye(num_cls).to(logit.device, non_blocking=True)
    y = y[target]

    pred_alpha = evidence_func(logit, evidence=evidence) + 1  # (alpha = e + 1)
    uncertainty = num_cls / pred_alpha.sum(dim=-1, keepdim=True)  # (N, 1)
    # gradient length
    grad = (1 / pred_alpha - uncertainty) * y  # y_ij * (1/alpha_ij - u_i)
    grad_norm = torch.abs(grad)
    return grad, grad_norm

def get_grad_info(cfg):
    set_seed(cfg.random_seed)
    # Setup model
    net = setup_model(cfg)

    # Setup dataloader
    train_data_loader, epoch_step_num = setup_dataloader(cfg)

    # start loop
    all_grad_norms, all_grad_norms_prop = [], []
    all_grads, all_grads_prop = [], []
    with tqdm.tqdm(train_data_loader, total=epoch_step_num, ncols=0) as pbar:
        for n_iter, (clips, targets, scores, ssl_clips, ssl_targets, flags) in enumerate(pbar):
            clips = clips.cuda()
            targets = [t.cuda() for t in targets]
            with torch.no_grad():
                # run infernece
                output_dict = net(clips, ssl=False)
                # get the matched target
                loc_t, conf_t, prop_loc_t, prop_conf_t = get_matched_targets(targets, output_dict['loc'], output_dict['priors'], cfg.clip_length)

            # coarse stage
            conf_p = output_dict['conf'].view(-1, cfg.num_classes)
            targets_conf = conf_t.view(-1, 1)
            if cfg.cls_loss_type == 'focal':
                conf_p = F.softmax(conf_p, dim=1)
            if cfg.os_head:
                inds_keep = targets_conf > 0  # (N,1)
                targets_conf = targets_conf[inds_keep].unsqueeze(-1) - 1  # (M,1), starting from 0
                conf_p = conf_p[inds_keep.squeeze()]  # (M,15)
            if targets_conf.numel() > 0 and cfg.cls_loss_type == 'edl':
                # compute gradient norm (one-hot)
                grads, grad_norms = grad_edl(conf_p, targets_conf, num_cls=cfg.num_classes, evidence=cfg.edl_config['evidence'])

            # refined stage
            prop_conf_p = output_dict['prop_conf'].view(-1, cfg.num_classes)
            prop_conf_t = prop_conf_t.view(-1, 1)
            if cfg.cls_loss_type == 'focal':
                prop_conf_p = F.softmax(prop_conf_p, dim=1)
            if cfg.os_head:
                inds_keep = prop_conf_t > 0  # (N,1)
                prop_conf_t = prop_conf_t[inds_keep].unsqueeze(-1) - 1  # (M,1), starting from 0
                prop_conf_p = prop_conf_p[inds_keep.squeeze()]  # (M,15)
            if prop_conf_t.numel() > 0 and cfg.cls_loss_type == 'edl':
                # compute gradient norm
                grads_prop, grad_norms_prop = grad_edl(prop_conf_p, prop_conf_t, num_cls=cfg.num_classes, evidence=cfg.edl_config['evidence'])

            if grad_norms is not None:
                all_grad_norms.append(grad_norms.cpu().numpy())
                all_grads.append(grads.cpu().numpy())
            if grad_norms_prop is not None:
                all_grad_norms_prop.append(grad_norms_prop.cpu().numpy())
                all_grads_prop.append(grads_prop.cpu().numpy())
    return all_grad_norms, all_grad_norms_prop, all_grads, all_grads_prop

def plot_grad_density(save_file, all_grad_norms, num_bins=30, momentum=0.75, fontsize=18):

    edges = [float(x) / num_bins for x in range(num_bins+1)]
    edges[-1] += 1e-6
    if momentum > 0:
        acc_sum = [0.0 for _ in range(num_bins)]
    grad_density = np.zeros((num_bins))
    weight_density = np.zeros((num_bins))
    grad_norm = np.concatenate(all_grad_norms, axis=0).sum(axis=-1)  # sum for one-hot grad_norms
    weights = np.zeros_like(grad_norm)
    for i in range(num_bins):
        # compute gradient density (number of gradients in each bin)
        inds = (grad_norm >= edges[i]) & (grad_norm < edges[i+1])
        num_in_bin = inds.sum()
        grad_density[i] = num_in_bin
        # compute the weights
        if num_in_bin > 0:
            if momentum > 0:
                acc_sum[i] = momentum * acc_sum[i] + (1 - momentum) * num_in_bin
                w = 1.0 / acc_sum[i]
            else:
                w = 1.0 / num_in_bin
            weights[inds] = w
            weight_density[i] = w

    # plot
    fig, ax1 = plt.subplots(1,1, figsize=(8,5))
    ax1.plot(edges[:-1], grad_density, 'r-', linewidth=2, label='Grad Density')
    # ax1.set_ylabel('fraction of samples', fontsize=fontsize)
    # ax1.set_yscale('log')
    ax1.legend(fontsize=fontsize, loc='upper left')

    ax2 = ax1.twinx()
    ax2.plot(edges[:-1], weight_density, 'b-', linewidth=2, label='weights')
    # ax2.set_ylabel('weights', fontsize=fontsize)
    # ax2.set_yscale('log')
    ax2.legend(fontsize=fontsize, loc='upper right')

    plt.xlabel('gradient norm', fontsize=fontsize)
    plt.xlim(0, grad_norm.max() - 1.0/num_bins)
    plt.tight_layout()
    plt.savefig(save_file)

def plot_grad_hist(save_file, all_grads, xlim=(-0.1, 0.1), ylim=(0, 100), fontsize=18):
    grads = np.concatenate(all_grads, axis=0).sum(axis=-1)
    fig, ax1 = plt.subplots(1,1, figsize=(8,5))
    plt.hist(grads, 200, density=True, alpha=0.8)
    plt.xlabel("gradient", fontsize=fontsize)
    plt.ylabel("probability density", fontsize=fontsize)
    plt.xlim(xlim) if xlim is not None else [min(grads), max(grads)]
    plt.ylim(ylim) if xlim is not None else None
    plt.tight_layout()
    plt.savefig(save_file)

def plot_gradnorm_weight(save_file, all_grad_norms, xlim=(0, 1), ylim=(0, 100), fontsize=18):
    grad_norm = np.concatenate(all_grad_norms, axis=0).sum(axis=-1)  # sum for one-hot grad_norms
    x_vals = np.linspace(grad_norm.min(), grad_norm.max(), num=200)
    weights = 1.0 / np.exp(10* x_vals)
    # plot figure
    fig, ax1 = plt.subplots(1,1, figsize=(8,5))
    ax1.hist(grad_norm, 200, density=True, alpha=0.8)
    ax1.set_yscale('log')
    ax1.set_ylabel("probability density", fontsize=fontsize)
    # ax1.set_ylim()
    ax1.legend(["frequency of gradient norm"], fontsize=fontsize, loc='upper left')

    ax2 = ax1.twinx()
    ax2.plot(x_vals, weights, 'r-', linewidth=2)
    ax2.set_yscale('log')
    ax2.set_ylabel('weights', fontsize=fontsize)
    ax2.legend(['weights'], fontsize=fontsize, loc='upper right')

    plt.xlim(xlim) if xlim is not None else [min(grad_norm), max(grad_norm)]
    plt.tight_layout()
    plt.savefig(save_file)

if __name__ == '__main__':

    cfg = get_basic_config(config)

    output_file = os.path.join(cfg.output_path, 'grad_norms.npz')
    if not os.path.exists(output_file):
        all_grad_norms, all_grad_norms_prop, all_grads, all_grads_prop = get_grad_info(cfg)
        np.savez(output_file[:-4], grad_norms=all_grad_norms,
                                   grad_norm_prop=all_grad_norms_prop,
                                   grads=all_grads,
                                   grad_prop=all_grads_prop)
    else:
        results = np.load(output_file, allow_pickle=True)
        all_grad_norms, all_grad_norms_prop = results['grad_norms'], results['grad_norm_prop']
        all_grads, all_grads_prop = results['grads'], results['grad_prop']
        print(f'Grad norm result file already exist at {output_file}!')

    plot_grad_density(os.path.join(cfg.output_path, 'grad_density.png'), all_grad_norms, num_bins=100, momentum=0.75)
    plot_grad_density(os.path.join(cfg.output_path, 'grad_prop_density.png'), all_grad_norms_prop, num_bins=100, momentum=0.75)

    plot_grad_hist(os.path.join(cfg.output_path, 'grad_hist.png'), all_grads, xlim=(-0.1,0.1), ylim=(0,100))
    plot_grad_hist(os.path.join(cfg.output_path, 'grad_prop_hist.png'), all_grads_prop, xlim=None, ylim=None)

    plot_gradnorm_weight(os.path.join(cfg.output_path, 'gradnorm_weight.png'), all_grad_norms, xlim=(0, 0.25))


import pytest

def test_quick_pipeline_with_fragmented_compute():

	"""
	How can I used tcrdist3 to test for TCRs that may HLA restricted.

	"""

	import os
	import pandas as pd
	import numpy as np
	from scipy import sparse
	from tcrdist.repertoire import TCRrep
	from tcrdist.rep_funcs import  compute_pw_sparse_out_of_memory

	f = 'mira_epitope_67_382_APHGVVFL_APHGVVFLHV_GVVFLHVTY_VVFLHVTYV.tcrdist3.csv'
	f = os.path.join('tcrdist','data','covid19',f)
	assert os.path.isfile(f)

	df = pd.read_csv(f)
	df = df[['subject', 'cell_type', 'v_b_gene', 'j_b_gene', 'cdr3_b_aa', 'cdr3_b_nucseq',  'cohort', 'hla-a', 'hla-a_1','hla-b', 'hla-b_1']]
	tr = TCRrep(cell_df = df,
				organism = 'human',
				chains = ['beta'],
				db_file = 'alphabeta_gammadelta_db.tsv',
				compute_distances = False,
				store_all_cdr = False)

	from tcrdist.rep_funcs import  compute_pw_sparse_out_of_memory

	S, fragments = compute_pw_sparse_out_of_memory(	tr = tr,
													row_size      = 100,
													pm_processes  = 2,
													pm_pbar       = True,
													max_distance  = 1000,
													matrix_name   = 'rw_beta',
													reassemble    = True,
													cleanup       = False)

	tr.clone_df['B07'] = (tr.clone_df['hla-b'].str.startswith("B*07") | tr.clone_df['hla-b_1'].str.startswith("B*07"))
	tr.clone_df['B07'] = ["B*07" if (x) else "NOTB*07 " for x in tr.clone_df['B07']]

	#sparse.save_npz("S.npz", S)
	from tcrdist.rep_funcs import  compute_n_tally_out_of_memory
	nn_tally_df_cohort = compute_n_tally_out_of_memory(fragments,
												matrix_name = "rw_beta",
												pm_processes  = 6,
												to_file = False,
												to_memory = True,
												knn_radius = 25,
												x_cols = ['B07'])

	from hierdiff.association_testing import cluster_association_test
	nn_associations = cluster_association_test(res = nn_tally_df_cohort, y_col='cmember', method='fishers')
	nn_associations = nn_associations.sort_values('pvalue', ascending = True)
	import ast
	nn_associations['neighbors_i'] = nn_associations.neighbors.apply(lambda x: ast.literal_eval(x))

	from tcrdist.summarize import test_for_almost_subsets, filter_is, filter_gt
	nn_associations['mostly_unique'] = test_for_almost_subsets(nn_associations['neighbors_i'], thr = 5)
	nr_nn_associations = filter_is(nn_associations, 'mostly_unique', 1).copy()

	#nr_nn_associations = filter_gt(nr_nn_associations, 'K_neighbors', 25).copy()
	nr_nn_associations

	# MOTIF GENERATION
	from tcrsampler.sampler import TCRsampler
	t = TCRsampler()
	if  'olga_human_beta_t.sampler.tsv' not in t.currently_available_backgrounds():
		t.download_background_file('olga_sampler.zip')
	#t.download_background_file('olga_sampler.zip') # ONLY IF NOT ALREADY DONE
	tcrsampler_beta = TCRsampler(default_background = 'olga_human_beta_t.sampler.tsv')
	tcrsampler_beta.build_background(max_rows = 1000)

	"""SEE PALMOTIF DOCS (https://github.com/agartland/palmotif)"""
	from palmotif import compute_pal_motif, svg_logo
	from tcrdist.summarize import _select

	"""GENERATE SVG GRAPHIC FOR EACH NODE OF THE TREE"""
	#pwmat_str = 'pw_beta'
	cdr3_name = 'cdr3_b_aa'
	gene_names = ['v_b_gene','j_b_gene']
	svgs_beta = list()
	svgs_beta_raw = list()
	info_list = list()

	from tcrdist.rep_diff import member_summ
	summary = member_summ(  res_df = nr_nn_associations,
							clone_df = tr.clone_df,
							addl_cols=['cohort','hla-a', 'hla-a_1', 'hla-b', 'hla-b_1', 'subject'])

	nr_nn_associations = pd.concat([nr_nn_associations, summary], axis = 1).reset_index()

	for i,r in nr_nn_associations.head(25).iterrows():
		dfnode  = tr.clone_df.iloc[r['neighbors_i'],:].copy()
		# <pwnode> Pairwise Matrix for node sequences
		pwnode = S[r['neighbors_i'],:] [:,r['neighbors_i']].todense()
		if dfnode.shape[0] > 2:
			iloc_idx = pwnode.sum(axis = 0).argmin()
			centroid = dfnode[cdr3_name].to_list()[iloc_idx]
		else:
			centroid = dfnode[cdr3_name].to_list()[0]

		print(f"CENTROID: {centroid}")

		gene_usage_beta = dfnode.groupby(gene_names).size()
		sampled_rep = tcrsampler_beta.sample( gene_usage_beta.reset_index().to_dict('split')['data'],
			flatten = True, depth = max(100, 1000 // dfnode.shape[0]))

		sampled_rep  = [x for x in sampled_rep if x is not None]

		motif, stat = compute_pal_motif(
						seqs = _select(df = tr.clone_df,
									   iloc_rows = r['neighbors_i'],
									   col = cdr3_name),
						refs = sampled_rep,
						centroid = centroid)

		svgs_beta.append(svg_logo(motif, return_str= True))

		sampled_rep = sampled_rep.append(centroid)
		motif_raw, _ = compute_pal_motif(
					 seqs =_select(df = tr.clone_df,
									iloc_rows = r['neighbors_i'],
									col = cdr3_name),
					 centroid = centroid)
		svgs_beta_raw.append(svg_logo(motif_raw, return_str= True))
		info_list.append(r)

	def row_to_string(r, vals = ['ct_columns', 'val_0', 'ct_0', 'val_1', 'ct_1', 'val_2', 'ct_2','val_3', 'ct_3', 'levels', 'K_neighbors', 'R_radius', 'RR', 'OR', 'pvalue', 'FWERp','FDRq']):
		#d = {v:r[v] for v in vals}
		return "

".join([f"\t{v} : {r[v]}" for v in vals])

	def to_html_table(r, vals = ['ct_columns', 'hla-a', 'hla-a_1', 'hla-b', 'hla-b_1', 'val_0', 'ct_0', 'val_2', 'ct_2', 'K_neighbors', 'R_radius', 'pvalue', 'FDRq','cdr3_b_aa','v_b_gene', 'j_b_gene', 'cohort','subject']):
		return pd.DataFrame(r[vals]).transpose().to_html()

	def shrink(html_str):
		return html_str.replace('height="100%"',  'height="10%"').\
			replace('width="100%"', 'width="10%"')

	with open('svgs_in_line.html', 'w') as fh:
		fh.write(f"\n")

		for svg, svg_raw, details in zip(svgs_beta, svgs_beta_raw, info_list):
			fh.write(f"{shrink(svg_raw)}{shrink(svg)}")
			try:
				fh.write(to_html_table(details))
			except:
				print("F")
			fh.write("

")
		fh.write(f"\n")


from rdkit import Chem
from rdkit.Chem import AllChem
import re
import os
from subprocess import Popen
import matplotlib.pyplot as plt
from scipy.stats import linregress
import numpy as np
import argparse
from multiprocessing import Pool

# path_to_orca = "/usr/local/orca_4_1_1_linux_x86-64/orca"
# path_to_orca = "/usr/local/orca_4_1_1/orca"
path_to_orca = "/usr/local/orca_4_1_2_linux_x86-64/orca"

conversion_factor = 627.5  # (kcal/mol)/Ha^-1

method_dict = {"Default": ("! PBE0 def2-SVP RIJCOSX def2/J TIGHTSCF TightOpt Freq D3BJ",
                           "! wB97X-D3 def2-TZVPP RIJCOSX def2/J TIGHTSCF Grid6 GridX6"),
               "High-level": ("! RI-MP2 def2-TZVP RIJCOSX def2/J def2-TZVP/C TIGHTSCF TightOpt NumFreq",
                              "! DLPNO-CCSD(T) def2-QZVPP RIJCOSX AutoAux TIGHTSCF Grid6 GridX6"),
               "Cheap": ("! PBE def2-SVP RIJCOSX def2/J Opt Freq D3BJ", None),
               "Proposed": ("! wB97X-D3 def2-TZVP RIJCOSX def2/J TightOpt TIGHTSCF Freq Grid6 GridX6 CPCM",
                            "! RI-B2GP-PLYP def2-TZVPP def2-TZVPP/C D3BJ RIJCOSX AutoAux TIGHTSCF Grid6 GridX6 SP CPCM"),
               "Reference": ("! RI-SCS-MP2 def2-TZVP RIJCOSX def2/J def2-TZVP/C TIGHTSCF TightOpt NumFreq CPCM",
                             "!DLPNO-CCSD(T) def2-TZVPP RIJCOSX AutoAux TIGHTSCF Grid6 GridX6 SP CPCM"),
               "maProposed": ("! wB97X-D3 ma-def2-TZVP RIJCOSX AutoAux TightOpt TIGHTSCF Freq Grid6 GridX6 CPCM",
                              "! RI-B2GP-PLYP ma-def2-TZVPP D3BJ RIJCOSX AutoAux TIGHTSCF Grid6 GridX6 SP CPCM"),
               "maReference": ("! RI-SCS-MP2 ma-def2-TZVP RIJCOSX AutoAux TIGHTSCF TightOpt NumFreq CPCM",
                               "! DLPNO-CCSD(T) ma-def2-TZVPP RIJCOSX AutoAux TIGHTSCF Grid6 GridX6 SP CPCM"),
               #   "PBE_def2-SVP": ("! PBE def2-SVP RI def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE_def2-TZVP": ("! PBE def2-TZVP RI def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE_def2-TZVPP": ("! PBE def2-TZVPP RI def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE_def2-QZVP": ("! PBE def2-QZVP RI def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE0_def2-SVP": ("! PBE0 def2-SVP RIJCOSX def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE0_def2-TZVP": ("! PBE0 def2-TZVP RIJCOSX def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE0_def2-TZVPP": ("! PBE0 def2-TZVPP RIJCOSX def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "PBE0_def2-QZVP": ("! PBE0 def2-QZVP RIJCOSX def2/J TightOpt TIGHTSCF Freq D3BJ Grid6 GridX6",
               #             None),
               #   "M062X_def2-SVP": ("! M062X def2-SVP RIJCOSX def2/J TightOpt TIGHTSCF NumFreq Grid6 GridX6",
               #             None),
               #   "M062X_def2-TZVP": ("! M062X def2-TZVP RIJCOSX def2/J TightOpt TIGHTSCF NumFreq Grid6 GridX6",
               #             None),
               #   "M062X_def2-TZVPP": ("! M062X def2-TZVPP RIJCOSX def2/J TightOpt TIGHTSCF NumFreq Grid6 GridX6",
               #             None),
               #   "M062X_def2-QZVP": ("! M062X def2-QZVP RIJCOSX def2/J TightOpt TIGHTSCF NumFreq Grid6 GridX6",
               #             None),
               #   "wB97X-D3_def2-SVP": ("! wB97X-D3 def2-SVP RIJCOSX def2/J TightOpt TIGHTSCF Freq Grid6 GridX6",
               #             None),
               #   "wB97X-D3_def2-TZVP": ("! wB97X-D3 def2-TZVP RIJCOSX def2/J TightOpt TIGHTSCF Freq Grid6 GridX6",
               #             None),
               #   "wB97X-D3_def2-TZVPP": ("! wB97X-D3 def2-TZVPP RIJCOSX def2/J TightOpt TIGHTSCF Freq Grid6 GridX6",
               #             None),
               #   "wB97X-D3_def2-QZVP": ("! wB97X-D3 def2-QZVP RIJCOSX def2/J TightOpt TIGHTSCF Freq Grid6 GridX6",
               #             None),
               #   "RI-MP2_def2-SVP": ("! RI-MP2 def2-SVP RIJCOSX def2/J def2-SVP/C TIGHTSCF TightOpt NumFreq",
               #             None),
               #   "RI-MP2_def2-TZVP": ("! RI-MP2 def2-TZVP RIJCOSX def2/J def2-TZVP/C TIGHTSCF TightOpt NumFreq",
               #             None),
               #   "RI-MP2_def2-TZVPP": ("! RI-MP2 def2-TZVPP RIJCOSX def2/J def2-TZVPP/C TIGHTSCF TightOpt NumFreq",
               #             None),
               #   "RI-MP2_def2-QZVP": ("! RI-MP2 def2-QZVP RIJCOSX def2/J def2-QZVP/C TIGHTSCF TightOpt NumFreq",
               #             None),
               #   "SCS-MP2_def2-SVP": ("! RI-SCS-MP2 def2-SVP RIJCOSX def2/J def2-SVP/C TIGHTSCF TightOpt NumFreq",
               #             None),
               #   "SCS-MP2_def2-TZVP": ("! RI-SCS-MP2 def2-TZVP RIJCOSX def2/J def2-TZVP/C TIGHTSCF TightOpt NumFreq",
               #             None),
               #   "SCS-MP2_def2-TZVPP": ("! RI-SCS-MP2 def2-TZVPP RIJCOSX def2/J def2-TZVPP/C TIGHTSCF TightOpt NumFreq",
               #             None)
               }

probe_dict = {"SiH3": ("[H][Si]([H])[H]", "[H][Si]([H])([H])[H]", ".[Si]%99([H])([H])[H]"),
              "SeH": ("[Se][H]", "[H][Se][H]", ".[Se]%99[H]"),
              "Br": ("[Br]", "Br", ".[Br]%99"),
              "Cl": ("[Cl]", "Cl", ".[Cl]%99"),
              "F": ("[F]", "F", ".[F]%99"),
              "I": ("[I]", "I", ".[I]%99"),
              "OH": ("[O][H]", "[H]O[H]", ".[O]%99[H]"),
              "SH": ("[S][H]", "[H]S[H]", ".[S]%99[H]"),
              "TeH": ("[Te][H]", "[H][Te][H]", ".[Te]%99[H]"),
              "NH2": ("[N]([H])[H]", "[H]N([H])[H]", ".[N]%99([H])[H]"),
              "PH2": ("[P]([H])[H]", "[H]P([H])[H]", ".[P]%99([H])[H]"),
              "AsH2": ("[As]([H])[H]", "[H][As]([H])[H]", ".[As]%99([H])[H]"),
              "CH3": ("[H][C]([H])[H]", "C", ".C%99"),
              "GeH3": ("[H][Ge]([H])[H]", "[H][Ge]([H])([H])[H]", ".[Ge]%99([H])([H])[H]"),
              "SnH3": ("[H][Sn]([H])[H]", "[H][Sn]([H])([H])[H]", ".[Sn]%99([H])([H])[H]")
              }

def get_args():

    parser = argparse.ArgumentParser()
    parser.add_argument("strained", action='store', type=str, help='SMILES string of strained molecule')
    parser.add_argument("adduct", action='store', type=str, help='SMILES string of adduct')
    parser.add_argument("-np", '--number_processors', action='store', type=int, default=1, help='Number of processors')
    parser.add_argument("-l", '--level', action='store', type=str, default="Default",
                        help='Level of theory for calculations', choices=['Default', 'High-level', 'Cheap', 'Proposed',
                                                                          'Reference', 'maProposed', 'maReference'])
    parser.add_argument("-chg", '--charge_on_probe', action='store', type=int, default=0,
                        help='Charge on probe for addition to strained molecule')
    parser.add_argument("-m", '--max_core', action='store', type=float, default=4000, help='Maximum memory per core')

    return parser.parse_args()

def make_mol_obj(smiles_string):
    """
    Make an RDKit molecule object from a SMILES string (e.g. generated from ChemDraw)
    :param smiles_string: (str) SMILES of a molecule
    :return: (object) RDKit Mol object
    """
    obj = Chem.MolFromSmiles(smiles_string)
    obj = Chem.AddHs(obj)
    AllChem.EmbedMultipleConfs(obj, numConfs=1, params=AllChem.ETKDG())
    return obj

def gen_conformer_xyzs(mol_obj, conf_ids):
    """
    Generate xyz lists for all the conformers in mol.conf_ids
    :param mol_obj: rdkit object
    :param conf_ids: (list) list of conformer ids to convert to xyz
    :return: (list) of xyz lists
    """
    xyzs = []

    for i in range(len(conf_ids)):
        mol_block_lines = Chem.MolToMolBlock(mol_obj, confId=conf_ids[i]).split('\n')
        mol_file_xyzs = []

        for line in mol_block_lines:
            split_line = line.split()
            if len(split_line) == 16:
                atom_label, x, y, z = split_line[3], split_line[0], split_line[1], split_line[2]
                mol_file_xyzs.append([atom_label, float(x), float(y), float(z)])

        xyzs.append(mol_file_xyzs)

    if len(xyzs) == 0:
        exit('Could not generate xyzs from RDKit object')

    return xyzs

def modify_adduct_smiles(smiles_string):

    # e.g, [*]C([H])([H])[C]([H])[H]
    if smiles_string.startswith('[*]'):

        # If the 3rd and 4th characters are letters..
        if smiles_string[3].isalpha() and smiles_string[4].isalpha():
            smiles_string = smiles_string[3:5] + '[*]' + smiles_string[5:]

        # If the 3rd character is a letter..
        elif smiles_string[3].isalpha():
            smiles_string = smiles_string[3:4] + '[*]' + smiles_string[4:]

        else:
            exit('Failed to modify the adduct SMILES string')

    return smiles_string.replace("[*]", "%99")

def xyzs2xyzfile(xyzs, filename=None, basename=None, title_line=''):
    """
    For a list of xyzs in the form e.g [[C, 0.0, 0.0, 0.0], ...] convert create a standard .xyz file

    :param xyzs: List of xyzs
    :param filename: Name of the generated xyz file
    :param basename: Name of the generated xyz file without the file extension
    :param title_line: String to print on the title line of an xyz file
    :return: The filename
    """

    if basename:
        filename = basename + '.xyz'

    if filename is None:
        return 1

    if filename.endswith('.xyz'):
        with open(filename, 'w') as xyz_file:
            if xyzs:
                print(len(xyzs), '\n', title_line, sep='', file=xyz_file)
            else:
                return 1
            [print('{:<3}{:^10.5f}{:^10.5f}{:^10.5f}'.format(*line), file=xyz_file) for line in xyzs]

    return filename

def add_h_to_adduct(adduct_smiles):
    """
    For a SMILES string of and adduct

    i.e.

    C

    :param adduct_smiles:
    :return:
    """

    pattern1 = "\[.\]"
    pattern2 = "\[..\]"
    atom_labels_sq_brackets = re.findall(pattern1, adduct_smiles)
    atom_labels_sq_brackets += re.findall(pattern2, adduct_smiles)

    len_strained_smiles = len(adduct_smiles.split('.', 1)[0])

    for atom_label_sq_brackets in atom_labels_sq_brackets:
        if atom_label_sq_brackets != "[H]":
            if atom_label_sq_brackets in adduct_smiles[:len_strained_smiles]:
                adduct_smiles = adduct_smiles.replace(atom_label_sq_brackets, atom_label_sq_brackets[1:-1])

    return adduct_smiles

def did_orca_calculation_terminate_normally(out_filename):

    out_lines = [line for line in open(out_filename, 'r', encoding="utf-8")]
    for n_line, line in enumerate(reversed(out_lines)):
        if 'ORCA TERMINATED NORMALLY' or 'The optimization did not converge' in line:
            return True
        if n_line > 50:
            # The above lines are pretty close to the end of the file – there's no point parsing it all
            break

    return False

def gen_orca_inp(mol, name, opt=False, sp=False, pal=1):
    inp_filename = name + ".inp"
    with open(inp_filename, "w") as inp_file:
        if opt:
            keyword_line = method_dict[level][0]
            if len(mol.xyzs) == 1:
                if "TightOpt" in keyword_line:
                    keyword_line = keyword_line.replace("TightOpt", "")
                if "Opt" in keyword_line:
                    keyword_line = keyword_line.replace("Opt", "")
            print(keyword_line, file=inp_file)
        if sp:
            print(method_dict[level][1], file=inp_file)
        print("%maxcore", maxcore, file=inp_file)
        print("%pal nprocs", pal, "end", file=inp_file)
        print("*xyz", mol.charge, mol.mult, file=inp_file)
        [print('{:<3}{:^12.8f}{:^12.8f}{:^12.8f}'.format(*line), file=inp_file) for line in mol.xyzs]
        print('*', file=inp_file)
    return inp_filename

def run_orca(inp_filename, out_filename):
    """
    Run the ORCA calculation given the .inp file as a subprocess
    :param inp_filename:
    :param out_filename:
    :return:
    """

    if os.path.exists(os.path.join("Library", level)):
        if os.path.exists(os.path.join("Library", level, out_filename)):
            return [line for line in open(os.path.join("Library", level, out_filename), 'r', encoding="utf-8")]

    orca_terminated_normally = False

    if os.path.exists(out_filename):
        orca_terminated_normally = did_orca_calculation_terminate_normally(out_filename)

    if not orca_terminated_normally:
        with open(out_filename, 'w') as orca_out:
            orca_run = Popen([path_to_orca, inp_filename], stdout=orca_out)
        orca_run.wait()

    return [line for line in open(out_filename, 'r', encoding="utf-8")]

def get_orca_opt_xyzs_energy(out_lines):
    """
    For a lost of ORCA output file lines find the optimised xyzs and energy
    :param out_lines:
    :return:
    """

    opt_converged, geom_section = False, False
    opt_xyzs, energy, gibbs_corr = [], None, None

    for line in out_lines:

        if 'THE OPTIMIZATION HAS CONVERGED' in line:
            opt_converged = True
        if 'CARTESIAN COORDINATES' in line and opt_converged:
            geom_section = True

        if geom_section and len(line.split()) == 0:
            geom_section = False

        if geom_section and len(line.split()) == 4:
            atom_label, x, y, z = line.split()
            opt_xyzs.append([atom_label, float(x), float(y), float(z)])

        if 'FINAL SINGLE POINT ENERGY' in line:
            energy = float(line.split()[4])             # e.g. line = 'FINAL SINGLE POINT ENERGY     -4143.815610365798'

        if 'G-E(el)' in line:
            gibbs_corr = float(line.split()[2])

    return opt_xyzs, energy, gibbs_corr

def get_orca_gibbs_corr_energy_single_atom(out_lines):

    s_trans, h_total, energy = None, None, None

    for line in out_lines:

        if 'Translational entropy' in line:
            s_trans = float(line.split()[3])

        if 'Total enthalpy' in line:
            h_total = float(line.split()[3])

        if 'FINAL SINGLE POINT ENERGY' in line:
            energy = float(line.split()[4])

    # If any of the energies are not found return None
    if any([e is None for e in [s_trans, h_total, energy]]):
        return None, None

    gibbs_corr = (h_total - s_trans) - energy

    return energy, gibbs_corr

def get_orca_sp_energy(out_lines):

    for line in out_lines[::-1]:
        if 'FINAL SINGLE POINT ENERGY' in line:
            return float(line.split()[4])             # e.g. line = 'FINAL SINGLE POINT ENERGY     -4143.815610365798'

def get_atoms_in_smiles_string(smiles):

    atoms = []

    smiles_str_list = list(smiles)
    for i, char in enumerate(smiles_str_list):
        if i < len(smiles_str_list) - 1:
            if char.isupper() and smiles_str_list[i+1].islower():
                atoms.append(''.join(smiles_str_list[i:i+2]))
            if char.isupper() and not smiles_str_list[i+1].islower():
                atoms.append(char)
        else:
            if char.isupper():
                atoms.append(char)

    return atoms

def print_output(process, name, state):
    return print("{:<30s}{:<50s}{:>10s}".format(process, name, state))

class Molecule(object):

    def calc_gibbs(self):
        self.optimise()
        self.single_point()
        self.set_gibbs()

    def optimise(self):
        print_output('Optimisation of', self.name, 'Running')
        inp_filename = gen_orca_inp(mol=self, name=self.name + "_opt", opt=True, pal=self.pal)
        orca_output_lines = run_orca(inp_filename, out_filename=inp_filename.replace(".inp", ".out"))
        if len(self.xyzs) == 1:
            self.energy, self.gibbs_corr = get_orca_gibbs_corr_energy_single_atom(out_lines=orca_output_lines)
        else:
            self.xyzs, self.energy, self.gibbs_corr = get_orca_opt_xyzs_energy(out_lines=orca_output_lines)
        print_output('', '', 'Done')

    def single_point(self):
        print_output('Single point of of', self.name, 'Running')
        if method_dict[level][1] is not None:
            inp_filename = gen_orca_inp(mol=self, name=self.name + "_sp", sp=True, pal=self.pal)
            orca_output_lines = run_orca(inp_filename, out_filename=inp_filename.replace(".inp", ".out"))
            self.energy = get_orca_sp_energy(out_lines=orca_output_lines)
        print_output('', '', 'Done')

    def set_gibbs(self):
        if self.energy is None or self.gibbs_corr is None:
            self.gibbs = None
        else:
            self.gibbs = self.energy + self.gibbs_corr

    def __init__(self, smiles, charge=0, mult=1, name="strained", pal=1):
        print_output('Molecule object for', name, 'Generating')

        self.smiles = smiles
        self.charge = charge
        self.mult = mult
        self.name = name
        self.energy = None
        self.gibbs_corr = None
        self.gibbs = None
        self.obj = make_mol_obj(self.smiles)
        self.xyzs = gen_conformer_xyzs(mol_obj=self.obj, conf_ids=[0])[0]
        self.pal = pal

        self.calc_gibbs()
        print_output('', '', '')

def calc_dG_addition(strained, probe, adduct):
    if any(gibbs is None for gibbs in [adduct.gibbs, strained.gibbs, probe.gibbs]):
        return None
    return (adduct.gibbs - (strained.gibbs + probe.gibbs))*conversion_factor

def calc_dG_stabilisation(probeH, adduct, probe, adductH):
    if any(gibbs is None for gibbs in [probeH.gibbs, adduct.gibbs, probe.gibbs, adductH.gibbs]):
        return None
    return ((probe.gibbs + adductH.gibbs) - (probeH.gibbs + adduct.gibbs))*conversion_factor

def get_xs_ys_not_none(xs, ys):
    xs_not_none, ys_not_none = [], []
    for i in range(len(xs)):
        if xs[i] is not None and ys[i] is not None:
            xs_not_none.append(xs[i])
            ys_not_none.append(ys[i])
    return xs_not_none, ys_not_none

def get_xs_to_zero(xs):

    if all([x < 0 for x in xs]):
        return list(sorted(xs)) + [0]
    else:
        return list(sorted(xs))

def calc_dGs(general_adduct_smiles, charge_on_probe, probe_name, mult, strained, pal):
    probe = Molecule(smiles=probe_dict[probe_name][0], name=probe_name + str(charge_on_probe),
                     charge=charge_on_probe, mult=mult, pal=pal)
    probeH = Molecule(smiles=probe_dict[probe_name][1], name=probe_name + "H", pal=pal)
    adduct = Molecule(smiles=general_adduct_smiles + probe_dict[probe_name][2],
                      name=strained.name + "_" + probe.name, charge=charge_on_probe, mult=mult, pal=pal)
    adductH = Molecule(smiles=add_h_to_adduct(adduct_smiles=adduct.smiles),
                       name=strained.name + "_" + probe.name + "H", pal=pal)
    dG_addition = calc_dG_addition(strained, probe, adduct)
    dG_stabilisation = calc_dG_stabilisation(probeH, adduct, probe, adductH)

    return dG_addition, dG_stabilisation

def calc_strain_graph(strained_smiles, general_adduct_smiles, charge_on_probe):
    mult = 1
    if charge_on_probe == 0:
        mult = 2
    strained = Molecule(smiles=strained_smiles, pal=n_procs)
    general_adduct_smiles = modify_adduct_smiles(smiles_string=general_adduct_smiles)

    n_processes = len(probe_dict)
    n_core_per_process = int(n_procs/n_processes) if int(n_procs/n_processes) > 0 else 1
    processes = int(n_processes/n_core_per_process)
    with Pool(processes=processes) as pool:
        results = [pool.apply_async(calc_dGs, (general_adduct_smiles, charge_on_probe, probe_name, mult, strained,
                                               n_core_per_process))
                   for probe_name in probe_dict.keys()]

        dGs = [res.get(timeout=None) for res in results]

    xs = [val[1] for val in dGs]
    ys = [val[0] for val in dGs]

    return xs, ys

def plot_strain_graph(strained_smiles, general_adduct_smiles, charge_on_probe):

    xs, ys = calc_strain_graph(strained_smiles, general_adduct_smiles, charge_on_probe)
    plt.scatter(xs, ys)

    xs_not_None, ys_not_None = get_xs_ys_not_none(xs, ys)
    m, c, r, p, err = linregress(xs_not_None, ys_not_None)
    xs_sq = []
    for i in xs_not_None:
        xs_sq.append(i**2)
    intercept_err = err*np.sqrt(sum(xs_sq)/len(xs_sq))
    plt.annotate("gradient = " + str(np.round(m,2)) + "$\pm$" + str(np.round(err, 2)) +
                 "\nstrain relief = " + str(np.round(c,1)) + "$\pm$" + str(np.round(intercept_err, 2)) +
                 "\n$r^2$ = " + str(np.round(np.square(r),3)), (0.6*min(xs_not_None), 0.8*min(ys_not_None)),
                 ha='center', va='center')

    xs_to_zero = get_xs_to_zero(xs=xs_not_None)
    plt.plot(xs_to_zero, np.array(xs_to_zero)*m + c, color = 'black', linestyle = 'dashed')
    plt.xlabel("$\Delta G_{stab}$ / kcal mol$^{-1}$")
    plt.ylabel("$\Delta G_{add}$ / kcal mol$^{-1}$")
    plt.axhline(y=0, color='k', linewidth = '0.5')
    plt.axvline(x=0, color='k', linewidth = '0.5')

    return plt.savefig("_strain_graph.png", dpi=1200)

if __name__ == "__main__":
    args = get_args()
    level = args.level
    maxcore = args.max_core
    n_procs = args.number_processors

    plot_strain_graph(strained_smiles=args.strained, general_adduct_smiles=args.adduct,
                      charge_on_probe=args.charge_on_probe)


import os
import urllib
import imageio
import cv2
import numpy as np
import tarfile
import shutil
import onnx
import onnxruntime

from google_drive_downloader import GoogleDriveDownloader as gdd
from imread_from_url import imread_from_url

model_url = "https://github.com/Kazuhito00/Skin-Clothes-Hair-Segmentation-using-SMP/blob/main/02.model/DeepLabV3Plus(timm-mobilenetv3_small_100)_452_2.16M_0.8385/best_model_simplifier.onnx?raw=true"
model_path = "models/hair_segmentation.onnx"

class HairSegmentation():

    def __init__(self, webcam_width, webcam_height):

        # Initialize model
        self.model = self.initialize_model()

        # Read fire gif image
        self.num_fire_imgs, self.fire_imgs = get_fire_gif(webcam_width, webcam_height)
        self.fire_img_num = 0

    def __call__(self, image):

        return self.segment_hair(image)

    def initialize_model(self):

        # Donwload model if not available
        download_github_model(model_url, model_path)

        # Create interpreter for the model
        self.session = onnxruntime.InferenceSession(model_path)

        # Get model info
        self.getModel_input_details()
        self.getModel_output_details()

    def segment_hair(self, image):

        input_tensor = self.prepare_input(image)

        # Perform inference on the image
        outputs = self.inference(input_tensor)

        # Process output data
        hair_mask = self.process_output(outputs)

        return hair_mask

    def prepare_input(self, image):

        self.img_height, self.img_width, self.img_channels = image.shape

        input_image  = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        input_image  = cv2.resize(input_image , (self.input_width,self.input_height))

        mean = [0.485, 0.456, 0.406]
        std = [0.229, 0.224, 0.225]
        input_image  = (input_image  / 255 - mean) / std

        input_image = input_image.transpose(2, 0, 1)
        input_tensor = input_image[np.newaxis,:,:,:]

        return input_tensor.astype(np.float32)

    def inference(self, input_tensor):
        return self.session.run(self.output_names, {self.input_name: input_tensor})

    def process_output(self, outputs):

        hair_mask = np.squeeze(outputs[0])
        hair_mask = hair_mask.transpose(1, 2, 0)
        hair_mask = hair_mask[:,:,2]
        hair_mask = cv2.resize(hair_mask, (self.img_width,self.img_height))

        return np.round(hair_mask).astype(np.uint8)

    def get_output_tensor(self, index):

        tensor = np.squeeze(self.interpreter.get_tensor(self.output_details[index]['index']))
        return tensor

    def getModel_input_details(self):

        self.input_name = self.session.get_inputs()[0].name

        self.input_shape = self.session.get_inputs()[0].shape
        self.input_height = self.input_shape[2]
        self.input_width = self.input_shape[3]

    def getModel_output_details(self):

        model_outputs = self.session.get_outputs()
        self.output_names = [model_outputs[0].name]

        self.output_shape = model_outputs[0].shape
        self.output_height = self.output_shape[2]
        self.output_width = self.output_shape[3]

    def draw_fire_hair(self, img, hair_mask):

        fire_img = np.zeros(img.shape,dtype=np.uint8)

        # Find the rectangle that surrounds all the contours
        left, top, right, bottom = find_contours_rectangle(hair_mask)

        fire_img[top:bottom, left:right] = cv2.resize(self.fire_imgs[self.fire_img_num][self.input_height//4:,:,:],(right-left, bottom-top))
        img[hair_mask>0] = fire_img[hair_mask>0]

        self.fire_img_num += 1
        if self.fire_img_num >= self.num_fire_imgs:
            self.fire_img_num = 0

        return img

def get_fire_gif(img_width, img_height):
    # Read fire image
    fire_path = "images/fire.gif"
    fire_image_url = "https://thumbs.gfycat.com/ShrillCooperativeBobwhite-max-1mb.gif"
    imdata = urllib.request.urlopen(fire_image_url).read()
    imbytes = bytearray(imdata)
    open(fire_path,"wb+").write(imdata)

    ## Read the gif from disk to `RGB`s using `imageio.miread`
    gif = imageio.mimread(fire_path)

    # convert form RGB to BGR and resize
    fire_imgs = [cv2.resize(cv2.cvtColor(img, cv2.COLOR_RGB2BGR),(img_width, img_height), interpolation = cv2.INTER_AREA) for img in gif]
    num_fire_imgs = len(fire_imgs)

    return num_fire_imgs, fire_imgs

def download_gdrive_model(gdrive_id, model_path):

    if not os.path.exists(model_path):
        gdd.download_file_from_google_drive(file_id=gdrive_id,
                                    dest_path='./tmp/tmp.tar.gz')
        tar = tarfile.open("tmp/tmp.tar.gz", "r:gz")
        tar.extractall(path="tmp/")
        tar.close()

        shutil.move("tmp/saved_model_512x512/model_float32_opt.onnx", model_path)
        shutil.rmtree("tmp/")

def download_github_model(model_url, model_path):

    if not os.path.exists(model_path):
        model_data = urllib.request.urlopen(model_url).read()
        model_bytes = bytearray(model_data)
        open(model_path,"wb+").write(model_bytes)

def find_contours_rectangle(mask):

    contours,hierarchy = cv2.findContours(mask*255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    if len(contours) != 0:
        min_left = mask.shape[1]
        min_top  = mask.shape[0]
        max_right = 0
        max_bottom = 0
        for contour in contours:
            left, top, rect_width, rect_height = cv2.boundingRect(contour)
            bottom = top + rect_height
            right = left + rect_width

            min_left = min([min_left,left])
            min_top = min([min_top,top])
            max_right = max([max_right,right])
            max_bottom = max([max_bottom,bottom])

        contour_rectangle = [min_left, min_top, max_right, max_bottom]
    else:
        contour_rectangle = [0, 0, mask.shape[1], mask.shape[0]]

    return contour_rectangle

if __name__ == '__main__':

    image = imread_from_url("https://thispersondoesnotexist.com/image")

    hair_segmentation = HairSegmentation(image.shape[1], image.shape[0])

    hair_mask = hair_segmentation(image)
    fire_hair_image = hair_segmentation.draw_fire_hair(image, hair_mask)

    cv2.namedWindow("Hair Segmentation", cv2.WINDOW_NORMAL)
    cv2.imshow("Hair Segmentation", fire_hair_image)
    cv2.waitKey(0)


from concurrent.futures import ProcessPoolExecutor
from functools import partial
import numpy as np
import os
import audio
from nnmnkwii.datasets import vctk
from nnmnkwii.io import hts
from hparams import hparams
from os.path import exists
import librosa
import re
import glob

def build_from_path(in_dir, out_dir, num_workers=1, tqdm=lambda x: x):
    executor = ProcessPoolExecutor(max_workers=num_workers)
    futures = []

    # speakers = vctk.available_speakers
    #
    # td = vctk.TranscriptionDataSource(in_dir, speakers=speakers)
    # transcriptions = td.collect_files()
    # speaker_ids = td.labels
    # speaker_ids_unique = np.unique(speaker_ids)
    # speaker_to_speaker_id = {}
    # for i, j in zip(speakers, speaker_ids_unique):
    #     speaker_to_speaker_id[i] = j
    # wav_paths = vctk.WavFileDataSource(
    #     in_dir, speakers=speakers).collect_files()
    #
    # _ignore_speaker = hparams.not_for_train_speaker.split(", ")
    # ignore_speaker = [speaker_to_speaker_id[i] for i in _ignore_speaker]
    wav_paths = glob.glob(os.path.join(in_dir, "*", "*", "*.wav"))

    # Create dictionary of sequence speaker id and original speaker id
    speakers = os.listdir(in_dir)
    speakers.sort()
    speaker_id_to_speaker = {}
    speaker_to_speaker_id = {}
    for i in range(len(speakers)):
        speaker_id_to_speaker[904 + i] = speakers[i]
        speaker_to_speaker_id[speakers[i]] = 904 + i
    meta_path1 = os.path.join(out_dir, "speaker_id_to_speaker.txt")
    with open(meta_path1, 'w', encoding='utf-8') as m:
        for i in range(len(speaker_id_to_speaker)):
            key = i + 904
            m.write('{}|{}\n'.format(str(key), speaker_id_to_speaker[key]))

    for index, wav_path in enumerate(wav_paths):
        speaker = os.path.dirname(wav_path)
        speaker = os.path.dirname(speaker)
        speaker = os.path.basename(speaker)

        text_path = wav_path.replace(".wav", ".normalized.txt")
        with open(text_path, 'r') as f:
            text = f.readline().rstrip('\n')
        futures.append(executor.submit(
            partial(_process_utterance, out_dir, index + 1, speaker_to_speaker_id[speaker], wav_path, text)))
    return [future.result() for future in tqdm(futures)]

def start_at(labels):
    has_silence = labels[0][-1] == "pau"
    if not has_silence:
        return labels[0][0]
    for i in range(1, len(labels)):
        if labels[i][-1] != "pau":
            return labels[i][0]
    assert False

def end_at(labels):
    has_silence = labels[-1][-1] == "pau"
    if not has_silence:
        return labels[-1][1]
    for i in range(len(labels) - 2, 0, -1):
        if labels[i][-1] != "pau":
            return labels[i][1]
    assert False

def _process_utterance(out_dir, index, speaker_id, wav_path, text):
    sr = hparams.sample_rate
    filename = os.path.basename(wav_path).replace('.wav', '')

    # Load the audio to a numpy array:
    wav = audio.load_wav(wav_path)

    # Librosa trim seems to cut off the ending part of speech
    wav, _ = librosa.effects.trim(wav, top_db=60, frame_length=2048, hop_length=512)

    if hparams.rescaling:
        wav = wav / np.abs(wav).max() * hparams.rescaling_max

    # Save trimmed wav
    save_wav_path = hparams.save_preprocessed_wav
    env_dir = os.path.dirname(wav_path)
    env = os.path.basename(env_dir)

    speaker = os.path.dirname(env_dir)
    speaker = os.path.basename(speaker)
    save_wav_path = os.path.join(save_wav_path, speaker, env, os.path.basename(wav_path))
    dir = os.path.dirname(save_wav_path)
    if not os.path.exists(dir):
        os.system('mkdir {} -p'.format(dir))
    audio.save_wav(wav, save_wav_path)

    # Compute the linear-scale spectrogram from the wav:
    spectrogram = audio.spectrogram(wav).astype(np.float32)
    n_frames = spectrogram.shape[1]

    # Compute a mel-scale spectrogram from the wav:
    mel_spectrogram = audio.melspectrogram(wav).astype(np.float32)

    # Write the spectrograms to disk:
    spectrogram_filename = '{}{}spec.npy'.format(filename, hparams.spec_cmp_separator)
    mel_filename = '{}{}mel.npy'.format(filename, hparams.spec_cmp_separator)
    np.save(os.path.join(out_dir, spectrogram_filename), spectrogram.T, allow_pickle=False)
    np.save(os.path.join(out_dir, mel_filename), mel_spectrogram.T, allow_pickle=False)

    # Return a tuple describing this training example:
    return (spectrogram_filename, mel_filename, n_frames, text, speaker_id)


import os
from dataclasses import dataclass
from functools import cached_property, lru_cache
from Bio.SeqRecord import SeqRecord
from Bio.SeqUtils.ProtParam import ProteinAnalysis
from abnumber.chain import Chain, Position
import pandas as pd
from abnumber.germlines import get_imgt_v_chains, get_imgt_j_chains
from sqlalchemy import create_engine
from typing import Dict, Union, List, Optional, Tuple
import numpy as np
from sqlalchemy.engine import Engine
from biophi.humanization.methods.stats import get_oasis_percentile, \
    get_germline_family_residue_frequency, get_chain_type_residue_frequency

OASIS_MIN_SUBJECTS_THRESHOLDS = {
    'loose': 0.01,
    'relaxed': 0.1,
    'medium': 0.5,
    'strict': 0.9
}

@dataclass
class OASisParams:
    oasis_db_path: str
    min_fraction_subjects: float

    @property
    def min_fraction_subjects_label(self):
        label = f'{self.min_fraction_subjects:.0%}'
        for label, threshold in OASIS_MIN_SUBJECTS_THRESHOLDS.items():
            if threshold == self.min_fraction_subjects:
                break
        return label

@dataclass
class PeptideHumanness:
    seq: str
    num_oas_subjects: int
    fraction_oas_subjects: float
    num_oas_occurrences: int

    def is_human(self, min_fraction_subjects):
        assert 0 < min_fraction_subjects < 1, 'min_fraction_subjects should be between 0 and 1'
        if self.fraction_oas_subjects is None:
            return np.nan
        return self.fraction_oas_subjects >= min_fraction_subjects

    def to_series(self):
        return pd.Series({
            'Seq': self.seq,
            'Num OAS Subjects': self.num_oas_subjects,
            'Fraction OAS Subjects': self.fraction_oas_subjects,
            'Num OAS Hits': self.num_oas_occurrences
        })

@dataclass
class ChainHumanness:
    chain: Chain
    imgt_chain: Chain
    peptides: Dict[Position, PeptideHumanness]
    num_germline_residues: int
    v_germline_names: List[str]
    j_germline_names: List[str]
    v_germline_family: str
    v_germline_suffix: str
    germline_family_residue_frequency: Dict[Position, Dict[str, float]]
    chain_type_residue_frequency: Dict[Position, Dict[str, float]]

    @cached_property
    def pos_imgt_mapping(self) -> Dict[Position, Position]:
        return {pos: imgt_pos for pos, imgt_pos in zip(self.chain.positions, self.imgt_chain.positions)}

    def get_oasis_curve(self, frequency=True, cumulative=True) -> pd.Series:
        bins = [get_fraction_subjects_bin(peptide.fraction_oas_subjects) for peptide in self.peptides.values()]
        curve = pd.Series(bins, name=self.chain.name).value_counts().reindex(FRACTION_SUBJECTS_BINS).fillna(0)

        if frequency:
            curve = curve / curve.sum()
        else:
            curve = curve.astype(np.int)

        if cumulative:
            curve = curve[::-1].cumsum()[::-1]
        return curve[FRACTION_SUBJECTS_BINS]

    def get_oasis_identity(self, min_fraction_subjects) -> float:
        return self.get_num_human_peptides(min_fraction_subjects) / self.get_num_peptides()

    def get_oasis_percentile(self, min_fraction_subjects) -> float:
        return get_oasis_percentile(
            chain_type=self.chain.chain_type,
            oasis_identity=self.get_oasis_identity(min_fraction_subjects),
            min_fraction_subjects=min_fraction_subjects
        )

    def get_v_germline_chains(self, limit=None) -> List[Chain]:
        germlines = get_imgt_v_chains(self.chain.chain_type)
        return [germlines[name] for name in self.v_germline_names[:limit]]

    def get_j_germline_chains(self, limit=None) -> List[Chain]:
        germlines = get_imgt_j_chains(self.chain.chain_type)
        return [germlines[name] for name in self.j_germline_names[:limit]]

    def get_num_human_peptides(self, min_fraction_subjects) -> int:
        return sum(p.is_human(min_fraction_subjects) for p in self.peptides.values())

    def get_num_nonhuman_peptides(self, max_fraction_subjects) -> int:
        return sum(not p.is_human(max_fraction_subjects) for p in self.peptides.values())

    def get_num_peptides(self) -> int:
        return len(self.peptides)

    def has_position(self, position) -> bool:
        return position in self.peptides

    def get_peptide(self, position, edges=True) -> PeptideHumanness:
        if self.has_position(position):
            return self.peptides[position]
        if not edges:
            raise KeyError(str(position))
        positions = list(self.peptides.keys())
        return self.peptides[positions[0] if position < positions[0] else positions[-1]]

    def get_peptide_length(self) -> int:
        lengths = [len(peptide.seq) for peptide in self.peptides.values()]
        assert len(set(lengths)) == 1, f'Peptides should have same lengths, got: {set(lengths)}'
        return lengths[0]

    def get_positional_humanness(self, min_fraction_subjects) -> List[Tuple[Position, str, List[PeptideHumanness]]]:
        chain_positions = list(self.chain.positions)
        peptide_len = self.get_peptide_length()
        assert peptide_len % 2 == 1, 'Peptide needs to have odd length'
        annots = []
        for raw_pos, (pos, aa) in enumerate(self.chain):
            window = [chain_positions[i] for i in range(max(0, raw_pos-peptide_len+1), raw_pos+1)]
            peptides = [self.peptides[peptide_pos]
                         for peptide_pos in window if peptide_pos in self.peptides]
            non_human_peptides = [p for p in peptides if not p.is_human(min_fraction_subjects)]
            annots.append((pos, aa, non_human_peptides))
        return annots

    def to_peptide_dataframe(self) -> pd.DataFrame:
        scheme = list(self.peptides)[0].scheme
        imgt_positions = self.imgt_chain.positions.keys()
        germline = {
            **self.get_j_germline_chains(1)[0].positions,
            **self.get_v_germline_chains(1)[0].positions
        }
        scheme_col = f'{scheme.title()}'
        df = pd.DataFrame([{
            'Antibody': self.chain.name,
            'Chain type': self.chain.chain_type,
            'Region': pos.get_region(),
            scheme_col: pos,
            'AA': aa,
            'Germline': germline.get(imgt_pos, '-'),
            'Chain Freq': self.chain_type_residue_frequency[pos].get(aa, 0)
                                if self.chain_type_residue_frequency.get(pos) else None,
            'Family Freq': self.germline_family_residue_frequency[pos].get(aa, 0)
                                if self.germline_family_residue_frequency.get(pos) else None
        } for (pos, aa), imgt_pos in zip(self.chain, imgt_positions)])
        df = df.merge(pd.DataFrame(
            [pep.to_series() for pep in self.peptides.values()],
            index=list(self.peptides.keys())
        ).add_prefix('Peptide '), left_on=scheme_col, right_index=True, how='left')
        return df

    @classmethod
    def to_sequence_dataframe(cls, humanness_list: List['ChainHumanness'], species=True) -> pd.DataFrame:
        chains = [humanness.chain for humanness in humanness_list]
        df = Chain.to_dataframe(chains)
        df.insert(0, 'scheme', [chain.scheme for chain in chains])
        if not species and 'species' in df.columns:
            df.drop(columns=['species'], inplace=True)
        df.index.name = 'Antibody'
        return df

    def get_germline_content(self):
        return self.num_germline_residues / len(self.chain)

    def get_top_freqs(self, n):
        top_freqs = []
        for i in range(n):
            top_n_freqs = []
            for pos, aa_freqs in self.germline_family_residue_frequency.items():
                if aa_freqs is not None and len(aa_freqs) > i:
                    aa, freq = sorted(aa_freqs.items(), key=lambda d: -d[1])[i]
                else:
                    aa, freq = 'X', 0
                top_n_freqs.append((pos, aa, freq))
            top_freqs.append(top_n_freqs)
        return top_freqs

@dataclass
class AntibodyHumanness:
    vh: ChainHumanness
    vl: ChainHumanness

    def get_oasis_identity(self, min_fraction_subjects) -> float:
        return self.get_num_human_peptides(min_fraction_subjects) / self.get_num_peptides()

    def get_oasis_percentile(self, min_fraction_subjects) -> float:
        return get_oasis_percentile(
            chain_type='mean',
            oasis_identity=self.get_oasis_identity(min_fraction_subjects),
            min_fraction_subjects=min_fraction_subjects
        )

    def get_oasis_curve(self, frequency=True) -> pd.Series:
        if self.vh and self.vl:
            counts_curve = self.vh.get_oasis_curve(frequency=False, cumulative=False) \
                           + self.vl.get_oasis_curve(frequency=False, cumulative=False)
            if frequency:
                counts_curve = counts_curve / counts_curve.sum()
            return counts_curve[::-1].cumsum()[::-1]
        if self.vh:
            return self.vh.get_oasis_curve(frequency=frequency)
        if self.vl:
            return self.vl.get_oasis_curve(frequency=frequency)

    def get_num_peptides(self) -> int:
        return (self.vh.get_num_peptides() if self.vh else 0) + (self.vl.get_num_peptides() if self.vl else 0)

    def get_num_human_peptides(self, min_fraction_subjects) -> int:
        return (self.vh.get_num_human_peptides(min_fraction_subjects) if self.vh else 0) \
               + (self.vl.get_num_human_peptides(min_fraction_subjects) if self.vl else 0)

    def get_num_nonhuman_peptides(self, min_fraction_subjects) -> int:
        return (self.vh.get_num_nonhuman_peptides(min_fraction_subjects) if self.vh else 0) \
               + (self.vl.get_num_nonhuman_peptides(min_fraction_subjects) if self.vl else 0)

    def to_peptide_dataframe(self) -> pd.DataFrame:
        dfs = []
        if self.vh:
            dfs.append(self.vh.to_peptide_dataframe())
        if self.vl:
            dfs.append(self.vl.to_peptide_dataframe())
        return pd.concat(dfs)

    def get_germline_content(self):
        num_germline_residues = 0
        num_total_residues = 0
        if self.vh:
            num_germline_residues += self.vh.num_germline_residues
            num_total_residues += len(self.vh.chain)
        if self.vl:
            num_germline_residues += self.vl.num_germline_residues
            num_total_residues += len(self.vl.chain)
        if num_total_residues == 0:
            return None
        return num_germline_residues / num_total_residues

def chop_seq_peptides(seq: Union[SeqRecord, Chain], peptide_length):
    if isinstance(seq, SeqRecord):
        seq = str(seq.seq)
        # numbered list from 1 to length of sequence
        positions = list(range(1, len(seq) + 1))
    elif isinstance(seq, Chain):
        # numbers are coming from Chain Position objects
        positions = list(seq.positions.keys())
        seq = ''.join(seq.positions.values())
    else:
        raise ValueError(f'Unsupported sequence type: {type(seq)}')

    return [(pos, seq[i:i + peptide_length]) for i, pos in enumerate(positions[:-peptide_length+1])]

def get_antibody_humanness(vh: Optional[Chain], vl: Optional[Chain], params: OASisParams) -> AntibodyHumanness:
    return AntibodyHumanness(
        vh=get_chain_humanness(vh, params=params) if vh else None,
        vl=get_chain_humanness(vl, params=params) if vl else None
    )

def get_chain_oasis_peptides(chain, params: OASisParams):
    pos_peptides = chop_seq_peptides(chain, peptide_length=9)

    if params.oasis_db_path is None:
        return {pos: PeptideHumanness(
                    seq=peptide,
                    num_oas_subjects=None,
                    fraction_oas_subjects=None,
                    num_oas_occurrences=None
                ) for pos, peptide in pos_peptides}

    if params.oasis_db_path.endswith('.gz'):
        raise ValueError('The OASis DB file needs to be unzipped (use "gunzip DB_PATH.db.gz")')

    if not os.path.exists(params.oasis_db_path):
        raise FileNotFoundError(f'The OASis DB path does not exist: {params.oasis_db_path}')

    oas_engine = create_engine('sqlite:///' + os.path.abspath(params.oasis_db_path), echo=False)

    oas_filter_chain = "Heavy" if chain.is_heavy_chain() else "Light"
    result = oas_engine.execute(f'SELECT COUNT(*) FROM subjects WHERE Complete{oas_filter_chain}Seqs >= 10000')
    num_total_oas_subjects = int(result.fetchall()[0][0])

    oas_hits = get_oas_hits(
        [peptide for pos, peptide in pos_peptides],
        engine=oas_engine,
        filter_chain=oas_filter_chain
    )
    oas_grouped = oas_hits.groupby('peptide')

    return {
        pos: parse_peptide_humanness(
            pep,
            oas_hits=oas_grouped.get_group(pep) if pep in oas_grouped.groups else oas_hits[:0],
            num_total_oas_subjects=num_total_oas_subjects
        ) for pos, pep in pos_peptides
    }

def get_chain_humanness(chain: Chain, params: OASisParams) -> ChainHumanness:

    peptides = get_chain_oasis_peptides(chain, params=params)

    imgt_chain = chain if chain.scheme == 'imgt' and chain.cdr_definition == 'imgt' else chain.renumber('imgt')

    v_germline_chains, j_germline_chains = chain.find_human_germlines(limit=10)
    top_v, top_j = v_germline_chains[0], j_germline_chains[0]
    num_germline_residues = sum(top_v.positions.get(pos) == aa or top_j.positions.get(pos) == aa
                                    for pos, aa in imgt_chain)

    v_germline_family = top_v.name.split('-')[0].split('/')[0]
    v_germline_suffix = top_v.name.replace(v_germline_family, '')
    return ChainHumanness(
        chain=chain,
        imgt_chain=imgt_chain,
        num_germline_residues=num_germline_residues,
        v_germline_names=[chain.name for chain in v_germline_chains],
        j_germline_names=[chain.name for chain in j_germline_chains],
        peptides=peptides,
        v_germline_family=v_germline_family,
        v_germline_suffix=v_germline_suffix,
        germline_family_residue_frequency=get_germline_family_residue_frequency(chain, imgt_chain, v_germline_family),
        chain_type_residue_frequency=get_chain_type_residue_frequency(chain, imgt_chain)
    )

def parse_peptide_humanness(peptide: str, oas_hits: pd.DataFrame, num_total_oas_subjects: int = None) -> PeptideHumanness:
    return PeptideHumanness(
        seq=peptide,
        num_oas_subjects=len(oas_hits),
        fraction_oas_subjects=len(oas_hits) / num_total_oas_subjects if num_total_oas_subjects else None,
        num_oas_occurrences=oas_hits['count'].sum()
    )

def get_oas_hits(peptides: Union[str, List[str]], engine: Engine, filter_chain=None):
    if not isinstance(peptides, list):
        peptides = list(peptides)
    filter_chain_statement = ""
    if filter_chain:
        assert filter_chain in ['Heavy', 'Light']
        filter_chain_statement = f"AND Complete{filter_chain}Seqs >= 10000"

    statement = "SELECT peptides.* FROM peptides " \
                "LEFT JOIN subjects ON peptides.subject=subjects.id " \
                "WHERE peptide IN (" + ",".join("?" * len(peptides)) + ") AND subjects.StudyPath <> 'Corcoran_2016' " \
                + filter_chain_statement
    return pd.read_sql(statement, params=peptides, con=engine)

FRACTION_SUBJECTS_THRESHOLDS = np.arange(0, 0.91, 0.01)
FRACTION_SUBJECTS_BINS = [f'{v:.0%}' for v in FRACTION_SUBJECTS_THRESHOLDS]

def get_fraction_subjects_bin(fraction):
    for bin_threshold, bin_label in zip(FRACTION_SUBJECTS_THRESHOLDS[::-1], FRACTION_SUBJECTS_BINS[::-1]):
        if bin_threshold is None or fraction >= bin_threshold:
            return bin_label


#! /usr/bin/env python
# coding=utf-8
#================================================================
#   Copyright (C) 2019 * Ltd. All rights reserved.
#
#   Editor      : VIM
#   File name   : freeze_graph.py
#   Author      : YunYang1994
#   Created date: 2019-03-20 15:57:33
#   Description :
#
#================================================================

import tensorflow as tf
import numpy as np
import base64
from core.yolov3 import YOLOV3
import core.utils as utils
from PIL import Image

input_size      = 416

def preprocess_image(image_buffer):
  """Preprocess JPEG encoded bytes to 3D float Tensor."""

  # Decode the string as an RGB JPEG.
  # Note that the resulting image contains an unknown height and width
  # that is set dynamically by decode_jpeg. In other words, the height
  # and width of image is unknown at compile-time.
  image = tf.image.decode_jpeg(image_buffer, channels=3)
  image.set_shape((None, None, 3))
  # After this point, all image pixels reside in [0,1)
  # until the very end, when they're rescaled to (-1, 1).  The various
  # adjust_* ops all require this range for dtype float.
  # image = tf.image.convert_image_dtype(image, dtype=tf.float32)
  # Crop the central region of the image with an area containing 87.5% of
  # the original image.
  # image = tf.image.central_crop(image, central_fraction=0.875)
  # # Resize the image to the original height and width.
  image = tf.expand_dims(image, 0)
  # image = tf.image.resize_bilinear(
  #     image, [FLAGS.image_size, FLAGS.image_size], align_corners=False)
  # image = tf.squeeze(image, [0])
  # # Finally, rescale to [-1,1] instead of [0, 1)
  # image = tf.subtract(image, 0.5)
  # image = tf.multiply(image, 2.0)
  image_paded = tf.image.resize_image_with_pad(image, input_size, input_size)
  return image_paded/255

def _encoded_image_string_tensor_input_placeholder():
  """Returns input that accepts a batch of PNG or JPEG strings.

  Returns:
    a tuple of input placeholder and the output decoded images.
  """
  batch_image_str_placeholder = tf.placeholder(
      dtype=tf.string,
      shape=[None],
      name='input_data')

  def decode(encoded_image_string_tensor):
    image_tensor = tf.image.decode_image(encoded_image_string_tensor,
                                         channels=3)
    image_tensor.set_shape((None, None, 3))
    return image_tensor

  return (batch_image_str_placeholder,
          tf.map_fn(
              decode,
              elems=batch_image_str_placeholder,
              dtype=tf.uint8,
              parallel_iterations=32,
              back_prop=False))

pb_file = "./yolov3_coco.pb"
ckpt_file = "./checkpoint/yolov3_test_loss=18.1357.ckpt-30"
output_node_names = ["input/input_data", "pred_sbbox/concat_2", "pred_mbbox/concat_2", "pred_lbbox/concat_2"]

with tf.name_scope('input'):

    input_data = tf.placeholder(dtype=tf.string, name='input_data')
    # image_tensor = tf.image.decode_image(input_data,
    #                                      channels=3)
    # image_tensor.set_shape((None, None, 3))
    #
    # frame_resize = tf.image.resize_images(image_tensor, [416, 416], method=0)
    #
    # image_input = frame_resize[np.newaxis, ...] / 255
    images = preprocess_image(input_data)
    # images = preprocess_image(serialized_tf_example)

    # images = tf.cast(images, tf.float32)

model = YOLOV3(images, trainable=False)
print(model.conv_sbbox, model.conv_mbbox, model.conv_lbbox)

sess  = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
saver = tf.train.Saver()
saver.restore(sess, ckpt_file)

converted_graph_def = tf.graph_util.convert_variables_to_constants(sess,
                            input_graph_def  = sess.graph.as_graph_def(),
                            output_node_names = output_node_names)

with tf.gfile.GFile(pb_file, "wb") as f:
    f.write(converted_graph_def.SerializeToString())


import utils
import torch
import torch.nn as nn
from tqdm import tqdm
import numpy as np
import string

def loss_fn(o1,o2,t1,t2):
    l1 = nn.BCEWithLogitsLoss()(o1,t1)
    l2 = nn.BCEWithLogitsLoss()(o2,t2)
    return l1 + l2

def train_fn(data_loader,model,optimizer,device,scheduler):
    model.train()
    losses = utils.AverageMeter()
    tk0 = tqdm(data_loader,total=len(data_loader))
    for bi,d, in enumerate(tk0):
        ids = d["ids"]
        token_type_ids = d["token_type_ids"]
        mask = d["mask"]
        targets_start = d["targets_start"]
        targets_end = d["targets_end"]

        ids = ids.to(device, dtype=torch.long)
        token_type_ids = token_type_ids.to(device, dtype=torch.long)
        mask = mask.to(device, dtype=torch.long)
        targets_start = targets_start.to(device, dtype=torch.float)
        targets_end = targets_end.to(device, dtype=torch.float)

        optimizer.zero_grad()
        o1,o2 = model(
            ids=ids,
            mask=mask,
            token_type_ids=token_type_ids
        )

        loss = loss_fn(o1,o2, targets_start,targets_end)
        loss.backward()
        optimizer.step()
        scheduler.step()
        losses.update(loss.item(),ids.size(0))
        tk0.set_postfix(loss=losses.avg)

def eval_fn(data_loader,model,device):
    model.eval()
    fin_output_start = []
    fin_output_end = []
    fin_padding_lens = []
    fin_tweet_tokens = []
    fin_orig_selected = []
    fin_orig_sentiment = []
    fin_orig_tweet = []

    for bi,d, in enumerate(data_loader):
        ids = d["ids"]
        token_type_ids = d["token_type_ids"]
        mask = d["mask"]
        tweet_tokens = d["tweet_tokens"]
        padding_len = d["padding_len"]
        orig_sentiment = d["orig_sentiment"]
        orig_selected = d["orig_selected"]
        orig_tweet = d["orig_tweet"]

        ids = ids.to(device, dtype=torch.long)
        token_type_ids = token_type_ids.to(device, dtype=torch.long)
        mask = mask.to(device, dtype=torch.long)

        o1,o2 = model(
            ids=ids,
            mask=mask,
            token_type_ids=token_type_ids
        )

        fin_output_start.append(torch.sigmoid(o1).cpu().detach().numpy())
        fin_output_end.append(torch.sigmoid(o2).cpu().detach().numpy())
        fin_padding_lens.extend(padding_len.cpu().detach().numpy().tolist())

        fin_tweet_tokens.extend(tweet_tokens)
        fin_orig_sentiment.extend(orig_sentiment)
        fin_orig_selected.extend(orig_selected)
        fin_orig_tweet.extend(orig_tweet)

    fin_output_start = np.vstack(fin_output_start)
    fin_output_end = np.vstack(fin_output_end)

    threshold = 0.2
    jaccards = []
    for j in range(len(fin_tweet_tokens)):
        target_string = fin_orig_selected[j]
        tweet_tokens = fin_tweet_tokens[j]
        padding_len = fin_padding_lens[j]
        original_tweet = fin_orig_tweet[j]
        sentiment = fin_orig_sentiment[j]

        if padding_len >0:
            mask_start = fin_output_start[j,:][:-padding_len] >= threshold
            mask_end = fin_output_end[j,:][:-padding_len] >= threshold
        else:
            mask_start = fin_output_start[j,:] >= threshold
            mask_end = fin_output_end[j,:] >= threshold

        mask = [0] * len(mask_start)
        idx_start = np.nonzero(mask_start)[0]
        idx_end = np.nonzero(mask_end)[0]

        if len(idx_start) > 0:
            idx_start = idx_start[0]
            if len(idx_end) >0:
                idx_end = idx_end[0]
            else:
                idx_end = idx_start

        else:
            idx_start=0
            idx_end=0

        for mj in range(idx_start,idx_end+1):
            mask[mj] =1

        output_tokens = [x for p,x in enumerate(tweet_tokens,split()) if mask[p]==1]
        output_tokens = [x for x in output_tokens if x not in ("[CLS]","[SEP]")]

        final_output = ' '
        for ot in output_tokens:
            if ot.startswith("##"):
                final_output = final_output + ot[2:]
            elif len(ot) == 1 and ot in string.punctuation:
                final_output = final_output + ot
            else:
                final_output = final_output+" " + ot
        final_output = final_output.strip()
        if sentiment == 'neutal' or len(original_tweet.split()) < 4:
            final_output = original_tweet
        jac = utils.jaccard(target_string.strip(),final_output.strip())
        jaccards.append(jac)
    mean_jac = np.mean(jaccards)
    return mean_jac


import json
import os

import dotenv
import joblib
import numpy as np
import torch
from azureml.core.model import Model
from torch_geometric.data import Batch, Data

from src.models.model import GCN

# ws = Workspace.from_config("src/cloud/config.json")
# print('Ready to use Azure ML {} to work with {}'.format(azureml.core.VERSION, ws.name))
# model = ws.models['diabetes_model']

# Called when the service is loaded
def init():

    project_dir = os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)
    dotenv_path = os.path.join(project_dir, ".env")
    dotenv.load_dotenv(dotenv_path)
    model_dir = os.path.join(os.getenv('AZUREML_MODEL_DIR'))
    global model

    for f in files:
        if "ckpt" in f:
            model_path = os.path.join(model_dir, f)
            break
    else:
        model_path=os.path.join(model_dir, '22-10-12_190dze4r.ckpt')

    model = GCN.load_from_checkpoint(model_path)

    model.eval()
    # torch.no_grad()
    # global model
    # model = 2
    # model_path = os.path.join(os.getenv('AZUREML_MODEL_DIR'), 'model.pkl')
    # model = joblib.load(model_path)

# Called when a request is received
def run(raw_data):
    # return json.dumps([model])
    d = json.loads(raw_data)
    x = torch.from_numpy(np.array(d["x"])).float()
    edge_index = torch.from_numpy(np.array(d["edge_index"]))
    edge_attr = torch.from_numpy(np.array(d["edge_attr"])).float()

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
    batch = Batch.from_data_list([data])
    predictions = model.predict_step(batch)

    pred_list = predictions.data.numpy().tolist()
    return json.dumps(pred_list)

if __name__ == "__main__":
    init()
    x = [
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
    ]

    edge_index = [
        [
            0,
            0,
            1,
            1,
            2,
            2,
            2,
            3,
            3,
            3,
            4,
            4,
            5,
            5,
            6,
            6,
            7,
            7,
            8,
            8,
            8,
            9,
            9,
            10,
            10,
            10,
            11,
            12,
        ],
        [
            1,
            9,
            0,
            2,
            1,
            3,
            7,
            2,
            4,
            8,
            3,
            5,
            4,
            6,
            5,
            7,
            2,
            6,
            3,
            9,
            10,
            0,
            8,
            8,
            11,
            12,
            10,
            10,
        ],
    ]

    edge_attr = [
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
    ]
    data = json.dumps({"x": x, "edge_index": edge_index, "edge_attr": edge_attr})

    a = run(data)
    print(a)


import numpy as np
import h5py
import pandas as pd
import glob
import ipdb
import cPickle as pkl
import hickle as hkl

DATA_PATH = '/data1/yj/tgif/vid_c3d_concat/'
data_list = glob.glob(DATA_PATH + '*/*.c3d')
dataframe = pd.DataFrame().from_csv('/data1/yj/tgif/DataFrame/tgif-new-group.tsv', sep='\t')
#ipdb.set_trace()
'''
with h5py.File('/data1/yj/tgif/features/POOL5.h5','w') as hf:
    for items in dataframe.iterrows():
        try:
            dat =np.squeeze(np.load(open('/data1/yj/tgif/vid_resnet2/'+ items[1]['web_url'].split('/')[-1].split('.')[0]+'/pool5.npy')))
        except:
            print items[1]['web_url'].split('/')[-1]
            continue
        if len(dat) > 1000:
            #ipdb.set_trace()
            dat = dat[:10]
            print dat.shape
            print items[1]['web_url'].split('/')[-1]
        hf.create_dataset(str(items[0]), data=dat)
'''
'''
with h5py.File('/data1/yj/tgif/features/TGIF_C3D.h5','w') as hf:
    for items in dataframe.iterrows():
        try:
            dat =np.squeeze(pkl.load(open('/data1/yj/tgif/vid_c3d_concat/'+ items[1]['web_url'].split('/')[-1].split('.')[0]+'/fc.c3d')))
        except:
            print items[1]['web_url'].split('/')[-1]
            continue
        if len(list(dat.shape)) > 1 and len(dat) > 1000:
            print dat.shape
            dat = dat[:10]
            #print dat.shape
            print items[1]['web_url'].split('/')[-1]
        hf.create_dataset(str(items[0]), data=dat)
'''
with h5py.File('/data1/yj/tgif/features/TGIF_VGG_fc7.h5','w') as hf:
    for items in dataframe.iterrows():
        try:
            dat =np.squeeze(np.load('/data1/yj/tgif/vgg_feat/'+ items[1]['web_url'].split('/')[-1].split('.')[0]+'/vgg_fc7.npy'))
        except:
            print items[1]['web_url'].split('/')[-1]
            continue
        if len(list(dat.shape)) > 1 and len(dat) > 3000:
            print dat.shape
            dat = dat[:10]
            #print dat.shape
            print items[1]['web_url'].split('/')[-1]
        hf.create_dataset(str(items[0]), data=dat)