import argparse
import pickle
import json

from copy import deepcopy

from src.config import conf
# diayn with model-free RL
from src.diayn import DIAYN
# diayn with evolution stratigies
from src.diayn_es import DIAYN_ES
# diayn with evolution model-based RL
# from src.diayn_mb import DIAYN_MB
from src.utils import seed_everything, evaluate_pretrained_policy_ext, evaluate_adapted_policy, evaluate_state_coverage, evaluate_pretrained_policy_intr, evaluate, save_final_results, plot_learning_curves

import torch
import torch.multiprocessing as mp
torch.set_num_threads(1)
torch.multiprocessing.set_sharing_strategy('file_system')

import numpy as np
import pandas as pd
import gym

import time
import os
import random
from collections import defaultdict

import matplotlib.pyplot as plt
import seaborn as sns
sns.set()
sns.set_style('whitegrid')
sns.set_context("paper", font_scale = conf.font_scale)

# Asymptotic performance of the expert
asymp_perofrmance = {
    'HalfCheetah-v2': 8070.205213363435,
    'Walker2d-v2': 5966.481684037183,
    'Ant-v2': 6898.620003217143,
    "MountainCarContinuous-v0": 95.0

    }

# shared parameters
params = dict( n_skills = 30,
           pretrain_steps = int(20e3),
           finetune_steps = int(1e5),
           buffer_size = int(1e6),
           min_train_size = int(1e4),
             )

# list for multip-processing
envs_mp = [
    # 1
#     { # 1 dof
#      'ppo':{
#         'MountainCarContinuous-v0': dict(
#            pretrain_steps = int(100e6),
#             n_skills = 10
#              ),
#         },
#     },

#     {
#         'sac':{
#         'MountainCarContinuous-v0': dict(
#            pretrain_steps = int (7e5),
#             n_skills = 10
#              ),
#         },
#     },

#     {
#         'es':{
#         'MountainCarContinuous-v0': dict(
#            pretrain_iterations = 20000,
#             n_skills = 10
#              ),
#         }
#     },

#     {
#         'pets':{
#         'MountainCarContinuous-v0': dict(
#            pretrain_steps = int (3e4),
#             n_skills = 10
#              ),
#         },
#     },
    # 2
#     { # 6 dof
#      'ppo':{
#         'HalfCheetah-v2': dict(
#            pretrain_steps = int(100e6),
#             n_skills = 30
#              ),
#         },
#     },

#     {
#       'sac':{
#         'HalfCheetah-v2': dict(
#            pretrain_steps = int(2e6),
#             n_skills = 30
#              ),
#         },
#     },

#     {
#       'es':{
#         'HalfCheetah-v2': dict(
#            pretrain_iterations = 25000,
#             n_skills = 30
#              ),
#         },
#     },

#     {
#        'pets':{
#         'HalfCheetah-v2': dict(
#            pretrain_steps = int (3.5e4),
#             n_skills = 30
#              ),
#         },
#     },
    # 3
#     { # 6 dof
#      'ppo':{
#         'Walker2d-v2': dict(
#            pretrain_steps = int(150e6),
#             n_skills = 30
#              ),
#         },

#     },

    {
      'sac':{
        'Walker2d-v2': dict(
           pretrain_steps = int(7e5),
            n_skills = 30
             ),
        },
    },

#     {
#       'es':{
#         'Walker2d-v2': dict(
#            pretrain_iterations = 25000,
#             n_skills = 30
#              ),
#         },
#     },

#     {
#      'pets':{
#         'Walker2d-v2': dict(
#            pretrain_steps = int (4e4),
#             n_skills = 30
#              ),
#         },
#     },

    # 4
#     { # 8 dof
#      'ppo':{
#         'Ant-v2': dict(
#            pretrain_steps = int(150e6),
#             n_skills = 30
#              ),
#         },

#     },

#     {
#     'sac':{
#         'Ant-v2': dict(
#            pretrain_steps = int(2e6),
#             n_skills = 30
#              ),
#         },

#     },

#     {
#      'es':{
#         'Ant-v2': dict(
#            pretrain_iterations = 30000,
#             n_skills = 30
#              ),
#         },
#     },

#     {
#      'pets':{
#         'Ant-v2': dict(
#            pretrain_steps = int (5e4),
#             n_skills = 30
#              ),
#         },
#     },

]

# ppo hyperparams
ppo_hyperparams = dict(
    n_steps = 2048,
    learning_rate = 1e-4,
    n_epochs = 10,
    batch_size = 4096,
    gamma = 0.99,
    gae_lambda = 0.95,
    clip_range = 0.1,
    ent_coef=0.1,
    n_actors = 32,
    algorithm = "ppo",

)

# SAC Hyperparameters
sac_hyperparams = dict(
    learning_rate = 3e-4,
    gamma = 0.99,
    buffer_size = int(1e7),
    batch_size = 128,
    tau = 0.001,
    gradient_steps = 1,
    ent_coef=0.1,
    learning_starts = 10000,
    algorithm = "sac"
)

# PETS Hyperparameters
pets_hyperparams = dict(
    learning_rate = 1e-3,
    batch_size = 256,
    ensemble_size = 15,
    trial_length = 200,
    population_size = 150,
    planning_horizon = 12,
    elite_ratio = 0.05,
    num_particles = 15,
    weight_decay = 5e-5,
    algorithm = "pets"
)

# Evolution Stratigies Hyperparameters after hyperparameters search
es_hyperparams = dict(
    lr = 1e-2, # learning rate
    iterations = 100, # iterations
    pop_size = 48, # population size testing
    algorithm = "es"
)

hyperparams = {
    'ppo': ppo_hyperparams,
    'sac': sac_hyperparams,
    'pets': pets_hyperparams,
    'es' : es_hyperparams
}

# Discriminator Hyperparameters
discriminator_hyperparams = dict(
    learning_rate = 3e-4,
    batch_size = 128,
    n_epochs = 1,
    weight_decay = 0,
    dropout = None, # The dropout probability
    label_smoothing = 0.0, # usually 0.2
    gp = None, # the weight of the gradient penalty term
    mixup = False,
    gradient_clip = None,
    temperature = 1,
    parametrization = "MLP", # TODO: add this as a CMD argument MLP, Linear, Separable, Concat
    lower_bound = "ba", # ba, tuba, nwj, nce, interpolate
    log_baseline = None,
    alpha_logit = -5., # small value of alpha => reduce the variance of nwj by introduce some nce bias
)

# Extract arguments from terminal
def cmd_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', type=str, default="MountainCarContinuous-v0")
    parser.add_argument("--alg", type=str, default="ppo")
    parser.add_argument("--skills", type=int, default=6)
    parser.add_argument("--presteps", type=int, default=int(1e6))
    parser.add_argument("--pm", type=str, default="MLP")
    parser.add_argument("--lb", type=str, default="ba")
    parser.add_argument("--run_all", type=bool, default=False)
    args = parser.parse_args()
    return args

# save the hyperparameters
def save_params(alg, directory, params_copy, algo_params):
    # save the hyperparams in a csv files
    alg_hyperparams = algo_params
    alg_hyperparams_df = pd.DataFrame(alg_hyperparams, index=[0])
    alg_hyperparams_df.to_csv(f"{directory}/{alg}_hyperparams.csv")

    discriminator_hyperparams_df = pd.DataFrame(discriminator_hyperparams, index=[0])
    discriminator_hyperparams_df.to_csv(f"{directory}/discriminator_hyperparams.csv")

    # convert the config namespace to a dictionary
    exp_params_df = pd.DataFrame(params_copy, index=[0])
    exp_params_df.to_csv(f"{directory}/params.csv")

    # general configrations
    config_d = vars(conf)
    dictionary_data = config_d
    a_file = open(f"{directory}/configs.json", "w")
    json.dump(config_d, a_file)
    a_file.close()
    print(f"{alg} hyperparame`ters: {alg_hyperparams}\n" )
    print(f"discriminator hyperparameters: {discriminator_hyperparams}\n" )
    print(f"shared experiment parameters: {params}\n" )
    print(f"configurations: {config_d }\n" )

def train_all(env_params, results_df_list, plots_d_list, seed , time, params_copy=deepcopy(params)):
    plots_dict = {}
    columns = ['Algorithm', "Reward After Adaptation", "State Entropy", "Intrinsic Reward" , "Reward Before Adaptation", "Environment"]
    results_df = pd.DataFrame(columns=columns)
    print(f"{env_params}")
    #####################
    for alg in env_params:
        for env in env_params[alg]:
            if env not in plots_dict:
                plots_dict[env] = {}
            # save a timestamp
            timestamp = time
            # Environment directory
            env_dir = main_exper_dir + f"env: {env}, alg:{alg}, stamp:{timestamp}/"
            os.makedirs(env_dir, exist_ok=True)
            alg_params = deepcopy(hyperparams[alg])
            if alg in ['ppo', 'sac', 'pets']:
                params_copy['n_skills'] = env_params[alg][env]['n_skills']
                params_copy['pretrain_steps'] = env_params[alg][env]['pretrain_steps']
                print(f"stamp: {timestamp}, alg: {alg}, env: {env}, n_skills: {params_copy['n_skills']}, pretrain_steps: {params_copy['pretrain_steps']}")
                if alg == 'pets':
                    params_copy['min_train_size'] = int(1e4)
            elif alg == "es":
                params_copy['n_skills'] = env_params[alg][env]['n_skills']
                alg_params['iterations'] = env_params[alg][env]['pretrain_iterations']
                params_copy['min_train_size'] = int(1e4)
                print(f"stamp: {timestamp}, alg: {alg}, env: {env}, n_skills: {params_copy['n_skills']}, pretrain_iterations: {alg_params['iterations']}")
            save_params(alg, env_dir, params_copy, alg_params)
            if alg == "ppo":
                params_copy['min_train_size'] = int(1e4)
            if (env_params[alg][env]).get('clip_range'):
                # clip_range
                alg_params['clip_range'] = (env_params[alg][env]).get('clip_range')
            seed_results = []
            seeds = [seed]
            for i in range(len(seeds)):
                seed_everything(seeds[i])
                seed_dir = env_dir + f"seed:{seeds[i]}"
                os.makedirs(seed_dir, exist_ok=True)
                if alg in ("sac", "ppo"):
                    diayn = DIAYN(params_copy, alg_params, discriminator_hyperparams, env, alg, seed_dir, seed=seeds[i], conf=conf, timestamp=timestamp, adapt_params=sac_hyperparams)
                    # pretraining step
                    pretrained_policy, discriminator = diayn.pretrain()
                    # fine-tuning step
                    # adapted_policy, best_skill = diayn.finetune()
                    # Evaluate the policy
                    intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = evaluate(env, params_copy['n_skills'], pretrained_policy, adapted_policy, discriminator,
                                                                                discriminator_hyperparams['parametrization'], best_skill, alg)
                    # for testing
                    # intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = (1 + np.random.randn()*0.25,1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25)

                # elif alg == "pets":
                #     diayn = DIAYN_MB(params_copy, alg_params, discriminator_hyperparams, env, alg, seed_dir, seed=seeds[i], conf=conf, timestamp=timestamp, adapt_params=sac_hyperparams)
                #     # pretraining step
                #     pretrained_policy, discriminator = diayn.pretrain()
                #     # fine-tune step
                #     # adapted_policy, best_skill = diayn.finetune()
                #     # Evaluate the policy
                #     intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = evaluate(env, params_copy['n_skills'], pretrained_policy, adapted_policy, discriminator,
                #                                                                                discriminator_hyperparams['parametrization'], best_skill, alg)
                #     # for testing
                #     # intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = (1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25)

                elif alg == "es":
                    diayn = DIAYN_ES(params_copy, alg_params, discriminator_hyperparams, env, "es", seed_dir, seed=seeds[i], conf=conf, timestamp=timestamp, adapt_params=sac_hyperparams)
                    # pretraining step
                    pretrained_policy, discriminator = diayn.pretrain()
                    # adaptation step
                    # adapted_policy, best_skill = diayn.finetune()
                    # Evaluate the policy
                    intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = evaluate(env, params_copy['n_skills'], pretrained_policy, adapted_policy, discriminator,
                                                                                discriminator_hyperparams['parametrization'], best_skill, alg)
                    # for testing
                    # intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = (1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25, 1 + np.random.randn()*0.25)

                asym = asymp_perofrmance[env]
                # asym = 1
                # save results
                seed_results.append([intrinsic_reward_mean, reward_beforeFinetune_mean/asym*100, reward_mean/asym*100, entropy_mean])
                d = {'Algorithm': alg.upper(),
                     "Reward After Adaptation": reward_mean/asym*100,
                     "Intrinsic Reward": intrinsic_reward_mean,
                     "State Entropy": entropy_mean,
                     "Reward Before Adaptation": reward_beforeFinetune_mean/asym*100,
                     "Environment": env[: env.index('-')] if env != "MountainCarContinuous-v0" else "MountainCar"
                    }
                results_df = results_df.append(d, ignore_index=True)
            # add to the plot dict
            if "algs" not in plots_dict[env]:
                plots_dict[env]["algs"] = []
            if "stamps" not in plots_dict[env]:
                plots_dict[env]["stamps"] = []
            if "skills" not in plots_dict[env]:
                plots_dict[env]["skills"] = []
            if "pms" not in plots_dict[env]:
                plots_dict[env]["pms"] = []
            if "lbs" not in plots_dict[env]:
                plots_dict[env]["lbs"] = []

            plots_dict[env]["algs"].append(alg)
            plots_dict[env]["stamps"].append(timestamp)
            plots_dict[env]["skills"].append(params_copy['n_skills'])
            plots_dict[env]["pms"].append(discriminator_hyperparams['parametrization'])
            plots_dict[env]["lbs"].append(discriminator_hyperparams['lower_bound'])

            save_final_results(seed_results, env_dir)
            results_df_list.append(results_df)
            plots_d_list.append(plots_dict)

def df_from_list(df_list, columns):
    results_df = pd.concat(df_list,ignore_index=True)
    return results_df

if __name__ == "__main__":
    args = cmd_args()
    run_all = args.run_all
    discriminator_hyperparams['lower_bound'] = args.lb
    discriminator_hyperparams['parametrization'] = args.pm
    main_exper_dir = None
    main_exper_dir = conf.log_dir_finetune + f"cls:{discriminator_hyperparams['parametrization']}, lb:{discriminator_hyperparams['lower_bound']}/"

    columns = ['Algorithm', "Environment", "Reward After Adaptation", "State Entropy", "Intrinsic Reward" , "Reward Before Adaptation"]

    if run_all:
        # for mp
        manager = mp.Manager()
        results_df_list = manager.list()
        plots_d_list = manager.list()
        n_processes = len(envs_mp)
        # print(n_processes)
        # print(envs_mp[:1])
        # input()
        processes_list = []
        # [[1645398097.9415703, 1645398097.927781], [1645375267.8525789, 1645312624.386389], [1645375267.87822, 1645398097.9635289], [1645375267.904899, 1645312624.4639032]]
        # stamps_list = [1645398097.927781, 1645398097.9415703, 1645312624.386389, 1645375267.8525789, 1645398097.9635289, 1645375267.87822, 1645312624.4639032, 1645375267.904899 ]
        seeds = conf.seeds
        for i in range(n_processes):
            timestamp = time.time()
            for seed in seeds:
                p = mp.Process(target=train_all, args=(envs_mp[i], results_df_list, plots_d_list, seed, timestamp))
                p.start()
                processes_list.append(p)
        for p in processes_list:
            p.join()
        print(f"results_df_list: {results_df_list}")
        print(f"plots d: {plots_d_list}")
        input()
        for plot_dict in plots_d_list:
            env_name = list(plot_dict.keys())[0]
            algs = plot_dict[env_name]['algs']
            stamps = plot_dict[env_name]['stamps']
            skills_list = plot_dict[env_name]['skills']
            pms = plot_dict[env_name]['pms']
            lbs = plot_dict[env_name]['lbs']
            asym = asymp_perofrmance[env_name]
            plot_learning_curves(env_name, algs, stamps, skills_list, pms, lbs, asym)
        results_df = df_from_list(results_df_list, columns)
        results_df.to_csv(f"{main_exper_dir}/results_df.csv")
        print("results df")
        print(f"{results_df}")
        for col in columns[2:]:
            plt.figure(figsize=conf.barplot_figsize)
            y_axis = col
            # ax = sns.barplot(x="Algorithm", y=y_axis, data=results_df, hue="Environment")
            if col == "Reward After Adaptation":
                ax = sns.barplot(x="Environment", y=y_axis, data=results_df, hue="Algorithm", hue_order=['ppo'.upper(), 'sac'.upper(), 'es'.upper(), 'pets'.upper()])
                sns.move_legend(ax, "lower left", title='Algorithm')
            else:
                ax = sns.barplot(x="Environment", y=y_axis, data=results_df, hue="Algorithm", hue_order=['ppo'.upper(), 'sac'.upper(), 'es'.upper(), 'pets'.upper()])
                ax.legend_.remove()
            filename = f"adaptation_experiment_final_{y_axis}.png"
            files_dir = f"Vis/adaptation_experiment_cls:{discriminator_hyperparams['parametrization']}, lb:{discriminator_hyperparams['lower_bound']}/"
            os.makedirs(files_dir, exist_ok=True)
            plt.savefig(files_dir + filename)
        # save the results
        with open(f"{main_exper_dir}/results.txt", "w") as o:
            for i in plots_d_list:
                o.write(f"{i}\n")
    else:
        # save a timestamp
        timestamp = time.time()
        print(f"Experiment timestamp: {timestamp}")
        params['pretrain_steps'] = args.presteps
        params['n_skills'] = args.skills
        alg_params = hyperparams[args.alg]
        if args.alg in ['ppo', 'sac', 'pets']:
                params['n_skills'] = args.skills
                params['pretrain_steps'] = args.presteps
                print(f"stamp: {timestamp}, alg: {args.alg}, env: {args.env}, n_skills: {params['n_skills']}, pretrain_steps: {params['pretrain_steps']}")
        elif args.alg == "es":
            params['n_skills'] = args.skills
            alg_params['iterations'] = args.presteps
            # alg_params['iterations_finetune'] = env_params[alg][env]['adaptation_iterations']
            print(f"stamp: {timestamp}, alg: {args.alg}, env: {args.env}, n_skills: {params['n_skills']}, pretrain_iterations: {alg_params['iterations']}")
        env_dir = main_exper_dir + f"env: {args.env}, alg:{args.alg}, stamp:{timestamp}/"
        os.makedirs(env_dir, exist_ok=True)
        save_params(args.alg, env_dir, params, alg_params)
        seed_results = []
        for i in range(len(conf.seeds)):
            seed_everything(conf.seeds[i])
            seed_dir = env_dir + f"seed:{conf.seeds[i]}"
            os.makedirs(seed_dir, exist_ok=True)
            if args.alg in ("sac", "ppo"):
                # TODO add ES, when it is ready
                diayn = DIAYN(params, alg_params, discriminator_hyperparams, args.env, args.alg, seed_dir, seed=conf.seeds[i], conf=conf, timestamp=timestamp, adapt_params=sac_hyperparams)
                # pretraining step
                pretrained_policy, discriminator = diayn.pretrain()
                # fine-tuning step
                adapted_policy, best_skill = diayn.finetune()
                # Evaluate the policy
                intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = evaluate(args.env, params['n_skills'], pretrained_policy, adapted_policy, discriminator,
                                                                                discriminator_hyperparams['parametrization'], best_skill, args.alg)
            # elif args.alg == "pets":
            #     diayn = DIAYN_MB(params, alg_params, discriminator_hyperparams, args.env, args.alg, seed_dir, seed=conf.seeds[i], conf=conf, timestamp=timestamp, adapt_params=sac_hyperparams)
            #     # pretraining step
            #     pretrained_policy, discriminator = diayn.pretrain()
            #     print("Finished")
            #     # fine-tune step
            #     adapted_policy, best_skill = diayn.finetune()
            #     # Evaluate the policy
            #     # Evaluate the policy
            #     intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = evaluate(args.env, params['n_skills'], pretrained_policy, adapted_policy, discriminator,
                                                                                # discriminator_hyperparams['parametrization'], best_skill, args.alg)
            elif args.alg == "es":
                diayn = DIAYN_ES(params, alg_params, discriminator_hyperparams, args.env, "es", seed_dir, seed=conf.seeds[i], conf=conf, timestamp=timestamp, args=args, adapt_params=sac_hyperparams)
                # pretraining step
                pretrained_policy, discriminator = diayn.pretrain()
                # adaptation step
                adapted_policy, best_skill = diayn.finetune()
                # Evaluate the policy
                intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean = evaluate(args.env, params['n_skills'], pretrained_policy, adapted_policy, discriminator,
                                                                                discriminator_hyperparams['parametrization'], best_skill, args.alg)
            seed_results.append([intrinsic_reward_mean, reward_beforeFinetune_mean, reward_mean, entropy_mean])
        save_final_results(seed_results, env_dir)


#modcom.co.ke/datascience
#modcom.co.ke/flask
#modcom.co.ke/flask/datascience/banks
#modcom.co.ke/bank
import pandas as pd

df = pd.read_csv("school.csv") #df is the dataframe
# print(df)
# print(df["StudyTime"])
print(df[["StudyTime", "Gender"]])
#narrow down
subset = df[["StudyTime", "Gender"]] #example of extracting different columns in a dataframe , maybe for plotting.
#below code shows rows and columns
print(df.shape)
#below checks and maps empty collumns
# print(df.isnull().sum())
# df.dropna(inplace = True ) #removing empties
# print(df.isnull().sum())
# print(df.shape)

#imputation, filling the empties,
df["Gender"].fillna(2, inplace=True)
df["Athlete"].fillna(2, inplace=True)
df["Smoking"].fillna(3, inplace=True)
df["Rank"].fillna(5, inplace=True)

#lets do for math
medianmath = df["Math"].median()
df["Math"].fillna(medianmath, inplace=True)

mediansprint = df["Sprint"].median()
df["Sprint"].fillna(mediansprint, inplace=True)

medianreading = df["Reading"].median()
df["Reading"].fillna(medianreading, inplace=True)

mediansleeptime = df["SleepTime"].median()
df["SleepTime"].fillna(mediansleeptime, inplace=True)

medianstudytime = df["StudyTime"].median()
df["StudyTime"].fillna(medianstudytime, inplace=True)

medianwriting = df["Writing"].median()
df["Writing"].fillna(medianwriting, inplace=True)

print(df.isnull().sum())

#===========PLOTTING===========

import matplotlib.pyplot as plt
fig, ax = plt.subplots()
ax.scatter(df["Math"], df["Reading"], s = 30, color = "red", alpha=0.75)
ax.set_xlabel("Mathematics")
ax.set_ylabel("Reading")
ax.set_title("Distribution of students mathematics by reading scores")
plt.show()

print(df["Reading"].describe())
print(df["Math"].describe())
fig, ax = plt.subplots()
ax.scatter(df["SleepTime"], df["StudyTime"], s = 30, color = "green", alpha=0.75)
ax.set_xlabel("Sleeping Time")
ax.set_ylabel("Study Time")
ax.set_title("Distribution of students Sleeping Time by Study Time scores")
plt.show()

#Histogram
_, ax = plt.subplots()
ax.hist(df['Sprint'], color='#539caf', bins=30)#bins are not a must because the
 # Label the axes and provide a title
ax.set_title("Distribution of Students Sprint Time")
ax.set_xlabel("Sprint")
ax.set_ylabel("Frequency")
plt.show()

fig, ax = plt.subplots()
ax.hist(df['Math'], color='#539caf', bins=30)#bins are not a must because the
 # Label the axes and provide a title
ax.set_title("Distribution of Students Math Scores")
ax.set_xlabel("Mathematics")
ax.set_ylabel("Frequency")
plt.show()

#Density
fig,  ax = plt.subplots()
ax.plot(df['Weight'], color='#539caf', lw=2)
ax.set_ylabel("Frequency")
ax.set_xlabel("sprint")
ax.set_title("Weight Distribution")
ax.legend(loc='best')
plt.show()

#PieChart
df["Rank"].replace(1, "Freshmen", inplace=True)
df["Rank"].replace(2, "Sophomore", inplace=True)
df["Rank"].replace(3, "Junior", inplace=True)
df["Rank"].replace(4, "Senior", inplace=True)
fig, ax = plt.subplots()
df.groupby('Rank').size().plot(kind='pie',autopct='%1.1f%%')
ax.set_xlabel("")
ax.set_ylabel("")
ax.set_title("Distribution by Rank in %")
plt.show()

#Bar Chart
fig, ax = plt.subplots()
df.groupby('Gender')['StudyTime'].mean().plot(kind='bar', color =
"#ffcccc")
ax.set_ylabel("StudyTime")
ax.set_xlabel("Gender")
ax.set_title("studyTime Distribution by Gender")
ax.legend(loc='best')
plt.show()

#stacd bar
fig,  ax = plt.subplots()
df.groupby(['Gender', 'Rank'])['Math'].mean().unstack().plot(kind='bar',
stacked=False)
ax.set_ylabel("Math")
ax.set_xlabel("Gender")
ax.set_title("Math Distribution by Gender and Rank")
plt.show()

#Test of hypothesis
#one sample t test
#two sample t test
#Angus
#chi square
#https://modcom.co.ke/datascience
#https://modcom.co.ke/flask/DataScience/banks.csv


import requests
import traceback
import os
import SimpleITK as sitk
from SimpleITK.utilities.pyside import sitk2qpixmap
from SimpleITK.utilities.resize import resize
import pandas as pd
import numpy as np
import inspect
from sklearn import metrics
import matplotlib.pyplot as plt
import tempfile

from PySide6.QtWidgets import (
    QMainWindow,
    QErrorMessage,
    QSlider,
    QWidget,
    QApplication,
    QFileDialog,
    QTextEdit,
    QLabel,
    QPushButton,
    QMessageBox,
    QVBoxLayout,
    QHBoxLayout,
    QCheckBox,
    QScrollArea,
    QComboBox,
    QSplitter,
    QProgressDialog,
    QTextBrowser,
)
from PySide6.QtCore import Qt, QObject, QRunnable, Signal, QThreadPool, QThread
from PySide6.QtGui import QPixmap, QImage, QIcon, QKeySequence, QAction, QPalette
import PySide6.QtGui
import qdarkstyle
from docutils.core import publish_string

class ExistingFileOrDirectoryDialog(QFileDialog):
    """
    Custom file dialog for selecting either a directory or a file. The standard
    Qt dialogs only allow selection of file or directory not both.
    """

    def __init__(self, parent):
        super(ExistingFileOrDirectoryDialog, self).__init__()
        self.setOption(QFileDialog.DontUseNativeDialog)
        self.setFileMode(QFileDialog.Directory)
        self.currentChanged.connect(self._selected)

    def _selected(self, name):
        if os.path.isdir(name):
            self.setFileMode(QFileDialog.Directory)
        else:
            self.setFileMode(QFileDialog.ExistingFile)

class HelpDialog(QWidget):
    """
    Dialog for displaying a single html page with text converted from
    a reStructuredText (rst) string. The string is usually a class's
    docstring documentation. As the string may contain code examples
    the docutils package uses the pygments library for syntax highlighting.
    Depending on the GUI theme the default pygments CSS may not be appropriate
    (e.g. white text on white background). To use a specific style you need
    to generate a CSS which will be given as input to the set_rst_text:
    pygmentize -S monokai -f html -a pre.code > monokai.css
    To see available styles:
    pygmentize -L styles
    """

    def __init__(self, w=500, h=600):
        super(HelpDialog, self).__init__()
        layout = QVBoxLayout()
        self.setLayout(layout)
        # Use the QTextBrowser and configure it to allow clickable hyperlinks.
        self.help_text_edit = QTextBrowser()
        self.help_text_edit.setOpenLinks(True)
        self.help_text_edit.setOpenExternalLinks(True)
        layout.addWidget(self.help_text_edit)
        self.resize(w, h)

    def set_rst_text(self, txt, pygments_css_file_name=None):
        # Use docutils publish_string method to convert the rst to html.
        # Then insert the css into the html and display that.
        html_str = publish_string(txt, writer_name="html").decode("utf-8")
        if pygments_css_file_name:
            style_idx = html_str.index("</style>")
            with open(pygments_css_file_name, "r") as fp:
                css_str = fp.read()
                html_str = html_str[:style_idx] + css_str + html_str[style_idx:]
        self.help_text_edit.setHtml(html_str)

class TBorNotTBDialog(QMainWindow):
    """
    ============
    TB or not-TB
    ============

    .. warning:: **Not For Clinical Use Or Clinical Decision Making**

    This program allows you to classify **frontal chest x-ray images**, anterior-posterior
    or posterior-anterior, into one of two classes, TB or not-TB, using the
    `NIAID <https://www.niaid.nih.gov/>`_ TB
    Portals web service. The program supports
    all of the image formats supported by the `SimpleITK <https://simpleitk.org/>`_ image
    analysis toolkit (DICOM, jpg, png, tiff...). In practice it converts them to DICOM after
    removing all meta-data (PII/PHI), so that the image sent to the service is anonymized.

    Additional `information on the NIAID TB Portals program is available
    online <https://tbportals.niaid.nih.gov/>`_.

    Swagger based documentation of the TB Portals service API is
    `available online <https://rap-ria.tbportals.niaid.nih.gov/swagger/>`_.

    Application Settings
    --------------------

    The application allows you to select the specific algorithm to call, and set the
    query timeout and number of retries. Currently, there are only two available algorithms,
    one using a single deep learning convolutional neural network model, DenseNet121,
    and another that uses multiple models, an ensemble of DenseNet121 networks.

    Selecting Input
    ---------------

    Data is loaded in one of two ways:

    1. Selection of a directory - directory and all its subdirectories are scanned for
       images which are used as input.
    2. Selection of a comma-separated-values (csv) file - the file is expected to have a column
       containing the relative paths to image files. Relative with respect to the location of the
       csv file. The column header has a **required name,
       'file'**. If the file also has a column header named 'actual', that column
       is expected to have entries of the form TB or NOT_TB and will be used to
       evaluate the results obtained from the service.

       The contents of a sample input csv file are displayed below::

         file,actual
         data/px22.jpg,TB
         data/nx42.jpg,NOT_TB
         data/p8.dcm,TB
         data/n17.dcm,NOT_TB

       Once the data is loaded, clicking on the thumbnail image will display it in
       a larger window alongside the currently available information. The larger window
       allows you to zoom in and out. See the View menu or use keyboard shortcuts. Note
       that on Mac the keyboard shortcuts are Command-Shift-Plus and Command-Minus.

    Query Service
    -------------
    After loading the data you can select a subset of the images that will be sent
    to the service for classification, using the checkbox associated with each of them. Once
    the service query is initiated the border of the thumbnail image indicates
    the status:

    1. Gray - never used in a query.
    2. Blue - query initiated for that image, no results received.
    3. Green - query performed, image succesfully processed and results available.
    4. Red -  query performed, failure.

    Save Results
    ------------

    You can save the currently available results at any time using the `Save results`
    menu item. This saves the raw results in a csv file. The raw results are the values
    returned from the TB/not-TB service. These include the `probability of TB`, the service's
    `decision` and additional information.

    If the input was in the form of a csv file which also contained the known class then
    you can `Save results and evaluation`. This will save the raw results and evaluate the
    service performance. The output includes:

    #. all_results.csv - A csv with the raw results from the service, includes all images.
    #. valid_results_used_in_evaluation.csv - A csv containing only the valid
       results which were used in the evaluation.
    #. evaluation_results.csv - A csv containing the evaulation results.
       Evaluation is performed using the services decision and a decision obtained
       when using the query dataset's optimal threshold,
       Youden index = argmax(sensitivity+specificity - 1). The latter appears
       in parenthesis:

       a. Accuracy: (TP+TN)/(P+N)
       b. Precision: TP/(TP+FP)
       c. Recall/Sensitivity: TP/(TP+FN)
       d. Specificity: TN/(TN+FP)
       e. F1: 2TP/(2TP+FP+FN)
       f. Area under the receiver operating curve.

    #. roc.pdf - ROC curve plot for all algorithms.
    #. confusion_matrix*.pdf - Confusion matrices for all algorithms using the
       service's decision and the decision based on the query dataset's optimal
       threshold.

    """

    def __init__(self):
        super(TBorNotTBDialog, self).__init__()

        # Disable the SimpleITK warnings (mostly the "Converting from MONOCHROME1 to MONOCHROME2"
        # which is common for DICOM images)
        sitk.ProcessObject.GlobalWarningDisplayOff()

        # Use QT's global threadpool, documentation says: "This global thread pool
        # automatically maintains an optimal number of threads based on the
        # number of cores in the CPU."
        # Getting the number of available CPUs in an OS portable way is not that trivial,
        # so hopefully QT does it correctly. The default number of threads is QThread.idealThreadCount()
        # See discussion:
        # https://stackoverflow.com/questions/31346974/portable-way-of-detecting-number-of-usable-cpus-in-python
        # The user can select to use fewer threads if they encounter issues with the service response due to
        # this program overloading it with concurrent requests.
        #
        self.threadpool = QThreadPool.globalInstance()
        self.query_workers = []

        # csv file column titles
        self.csv_filename_column_title = "file"
        self.csv_actual_value_column_title = "actual"
        self.csv_service_response_column_title = "service_last_response"
        self.csv_file_size_column_title = "file_size_bytes"

        # service labels
        self.positive_label = "TB"
        self.negative_label = "NOT_TB"

        self.number_of_retries_range = [1, 5]
        self.default_number_of_retries = 3
        self.timeout_range = [1, 60]
        self.default_timeout = 5
        self.endpoint = "https://rap-ria.tbportals.niaid.nih.gov/TBorNotTB"
        self.algorithms = {
            "UNet+DenseNet121": "single",
            "UNet+DenseNet121 Ensemble": "ensemble",
            "ResNetUNet+DenseNet121": "single_2",
        }

        # Size of thumbnails to use in the GUI
        self.thumbnail_size = [128] * 2

        # Limit the stylesheet to the pushbutton. Otherwise, the tooltip will inherit it and
        # we don't want that because the tooltip is long and won't display correctly.
        self.initial_button_stylesheet = "QPushButton{{qproperty-iconSize: {0}px {0}px; height : {0}px; width : {0}px; max-width: {0}px; border: 2px solid gray;}}".format(  # noqa: E501
            self.thumbnail_size[0]
        )
        self.waiting_button_stylesheet = "QPushButton{{qproperty-iconSize: {0}px {0}px; height : {0}px; width : {0}px; max-width: {0}px; border: 2px solid blue;}}".format(  # noqa: E501
            self.thumbnail_size[0]
        )
        self.finished_success_button_stylesheet = "QPushButton{{qproperty-iconSize: {0}px {0}px; height : {0}px; width : {0}px; max-width: {0}px; border: 2px solid green;}}".format(  # noqa: E501
            self.thumbnail_size[0]
        )
        self.finished_failure_button_stylesheet = "QPushButton{{qproperty-iconSize: {0}px {0}px; height : {0}px; width : {0}px; max-width: {0}px; border: 2px solid red;}}".format(  # noqa: E501
            self.thumbnail_size[0]
        )

        self.data_loader = DataLoader(self.thumbnail_size)
        self.data_loader.progress_signal.connect(self.__file_processed)
        self.data_loader.finished.connect(self.__load_data_finished)

        # Configure the help dialog.
        self.help_dialog = HelpDialog(w=700, h=500)
        self.help_dialog.setWindowTitle("TB or Not TB")
        self.help_dialog.set_rst_text(inspect.getdoc(self))

        self.__create_gui()
        self.setWindowTitle("TB or Not TB - That is the Question")
        self.show()

    def closeEvent(self, event):
        """
        Override the closeEvent method so that clicking the 'x' button also
        closes all of the dialogs. As the program is multithreaded we also need
        to tell all the workers to stop working and clear the threadpool so no
        queued workers are started. Unfortunately, we cannot stop the working threads
        immediately, we have to indicate that they should stop and that will eventually
        happen. The more active workers we have the longer exiting will take.
        """
        self.help_dialog.close()
        self.settings_dialog.close()
        # Tell all workers to stop working.
        for worker in self.query_workers:
            worker.continue_running = False
        # The application will hang till all the active runnables
        # return. Calling the threadpool's clear method removes all the runnables
        # from the queue so that they aren't launched and we end up waiting for them
        # too.
        self.threadpool.clear()
        event.accept()

    def __rap_tb_ntb_query(self, url_name, algorithm_name, timeout):
        """
        Wrapper function which is specific to the NIAID service. It returns
        a generic function which only expects a file name. This potentially enables
        expanding the front end to other services with different RESTful APIs
        in an easy manner.
        """

        def tb_ntb_query(file_name):
            error_message = ""
            results_dict = {}
            try:
                response = requests.post(
                    url_name,
                    files={"image": open(file_name, "rb")},
                    data={"method": algorithm_name},
                    timeout=timeout,
                )
                if response.status_code == requests.codes.ok:
                    json_data = response.json()
                    for key in json_data:
                        results_dict[f"{algorithm_name}: {key}"] = json_data[key]
                else:
                    error_message = response.reason + f" ({response.status_code})."
            except requests.Timeout:
                error_message = "Request timed out."
            except requests.ConnectionError:
                error_message = "Connection error."
            except requests.exceptions.RequestException:  # anything else
                error_message = "Problem with service."
            return (error_message, results_dict)

        return tb_ntb_query

    def __error_function(self, message):
        error_dialog = QErrorMessage(self)
        # The QErrorMessage dialog automatically identifies if text is rich text,
        # html or plain text. Unfortunately, it doesn't do a good job when some of
        # the text is describing Exceptions due to number comparisons that include
        # the '>' symbol. As all invocations of this function are done with plain
        # text we use the convertToPlainText method to ensure that it is displayed
        # correctly.
        error_dialog.showMessage(PySide6.QtGui.Qt.convertFromPlainText(message))

    def __create_gui(self):
        """
        Create the actual GUI and its layout (menu items, etc).
        """
        self.settings_dialog = self.__create_settings_widget()

        menu_bar = self.menuBar()
        menu_bar.setNativeMenuBar(False)
        file_menu = menu_bar.addMenu("&File")
        settings_action = QAction("&Settings", self)
        settings_action.setShortcut(QKeySequence("Ctrl+s"))
        settings_action.triggered.connect(self.settings_dialog.show)
        file_menu.addAction(settings_action)

        load_action = QAction("&Load data", self)
        load_action.setShortcut(QKeySequence("Ctrl+l"))
        load_action.triggered.connect(self.__load_data_and_config)
        file_menu.addAction(load_action)

        self.query_action = QAction("&Query service", self)
        self.query_action.triggered.connect(self.__call_service)
        file_menu.addAction(self.query_action)
        self.query_action.setEnabled(False)

        self.save_action = QAction("&Save results", self)
        self.save_action.triggered.connect(self.__save_results)
        file_menu.addAction(self.save_action)
        self.save_action.setEnabled(False)

        self.save_evaluate_action = QAction("&Save results and evaluation", self)
        self.save_evaluate_action.triggered.connect(self.__save_results_and_evaluation)
        file_menu.addAction(self.save_evaluate_action)
        self.save_evaluate_action.setEnabled(False)

        quit_action = QAction("&Quit", self)
        quit_action.setShortcut(QKeySequence("Ctrl+q"))
        quit_action.triggered.connect(QApplication.instance().closeAllWindows)
        file_menu.addAction(quit_action)

        view_menu = menu_bar.addMenu("&View")
        zoom_in_action = QAction("&Zoom in", self)
        zoom_in_action.triggered.connect(lambda: self.__zoom(6.0 / 5.0))
        zoom_in_action.setShortcut(QKeySequence("Ctrl++"))
        view_menu.addAction(zoom_in_action)
        zoom_out_action = QAction("&Zoom out", self)
        zoom_out_action.triggered.connect(lambda: self.__zoom(5.0 / 6.0))
        zoom_out_action.setShortcut(QKeySequence("Ctrl+-"))
        view_menu.addAction(zoom_out_action)

        self.help_button = QPushButton("Help")
        self.help_button.clicked.connect(self.help_dialog.show)
        menu_bar.setCornerWidget(self.help_button, Qt.TopRightCorner)

        central_widget = QWidget(self)
        layout = QVBoxLayout()
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        main_widget = self.__create_main_widget()
        layout.addWidget(main_widget)

    def __zoom(self, scale_factor):
        if self.selected_image_label.pixmap():
            self.selected_image_label.resize(
                scale_factor * self.selected_image_label.size()
            )

    def __file_processed(self):
        """
        Callback invoked when the service returns a response. Failure and success
        are treated in the same manner.
        """
        self.num_files_loaded = self.num_files_loaded + 1
        self.progress_dialog.setValue(
            int(0.5 + 100 * self.num_files_loaded / self.all_df.shape[0])
        )

    def __create_main_widget(self):
        """
        Creation and layout of main GUI.
        """
        main_wid = QWidget()
        layout = QVBoxLayout()
        main_wid.setLayout(layout)
        warning_label = QLabel("NOT FOR CLINICAL USE OR CLINICAL DECISION MAKING")
        warning_label.setAlignment(Qt.AlignCenter)
        warning_label.setStyleSheet("QLabel { color : red; }")
        layout.addWidget(warning_label)

        wid = QSplitter()
        layout.addWidget(wid, 2)

        left_widget = QWidget()
        wid.addWidget(left_widget)
        left_layout = QVBoxLayout()
        left_widget.setLayout(left_layout)

        input_wid = QWidget()
        self.input_layout = QVBoxLayout()
        input_wid.setLayout(self.input_layout)
        scroll_area = QScrollArea()
        scroll_area.setWidget(input_wid)
        scroll_area.setWidgetResizable(True)
        left_layout.addWidget(scroll_area)

        self.all_checkboxes = []
        self.all_buttons = []

        checkbox = QCheckBox("select all")
        checkbox.setChecked(True)
        checkbox.stateChanged.connect(self.__select_all_changed)
        left_layout.addWidget(checkbox)

        right_widget = QSplitter()
        right_widget.setOrientation(Qt.Vertical)
        wid.addWidget(right_widget)

        self.selected_image_label = QLabel()
        scroll_area = QScrollArea()
        scroll_area.setWidget(self.selected_image_label)
        # If we used setWidgetResizable(True) the qlabel's pixmap determines the size and is shown at full resolution.
        # Most often that is too large for the image to fully be visible, and you need to scroll to see parts of it.
        scroll_area.setWidgetResizable(False)
        scroll_area.setBackgroundRole(QPalette.Dark)
        right_widget.addWidget(scroll_area)

        info_widget = QWidget()
        right_widget.addWidget(info_widget)
        info_layout = QVBoxLayout()
        info_widget.setLayout(info_layout)
        info_layout.addWidget(QLabel("Image Information:"))
        self.selected_image_information_edit = QTextEdit()
        self.selected_image_information_edit.setReadOnly(True)
        info_layout.addWidget(self.selected_image_information_edit)

        return main_wid

    def __load_data_and_config(self):
        dialog = ExistingFileOrDirectoryDialog(self)
        if not dialog.exec():
            return

        # Get the input file names, either from a csv file or from the subdirectories
        # under the data root directory.
        self.num_files_loaded = 0

        self.all_df = self.__get_file_names_df(dialog.selectedFiles()[0])
        if self.all_df.empty or (
            self.csv_filename_column_title not in self.all_df.columns
        ):
            self.__error_function(
                "No data loaded (error in csv file or empty directory)."
            )
        else:
            self.progress_dialog = QProgressDialog("Loading data...", None, 0, 100)
            self.progress_dialog.setWindowModality(Qt.WindowModal)
            # Display the progress dialog only if its estimated operation duration is
            # more than 2000ms.
            self.progress_dialog.setMinimumDuration(2000)

            self.data_loader.reset()
            self.data_loader.file_names_list = list(
                self.all_df[self.csv_filename_column_title]
            )
            self.data_loader.start()

    def __load_data_finished(self):
        problematic_row_indexes = self.all_df.index[self.data_loader.error_indexes]
        problematic_files = list(
            self.all_df[self.csv_filename_column_title][problematic_row_indexes]
        )
        self.all_df.drop(problematic_row_indexes, inplace=True)
        self.all_df[self.csv_service_response_column_title] = "not computed"
        self.all_df[self.csv_file_size_column_title] = float("nan")
        if problematic_files:
            self.__error_function(
                "The following files were not loaded:\n {0}".format(
                    "\n".join(problematic_files)
                )
            )
        if not self.all_df.empty:
            self.selected_image_row_index = -1
            self.selected_image_label.clear()
            self.selected_image_information_edit.clear()
            # Remove all widgets from previous dataset, if any
            for cb, b in zip(self.all_checkboxes, self.all_buttons):
                cb.setParent(None)
                b.setParent(None)
            for i, image, file_name in zip(
                range(len(self.data_loader.images)),
                self.data_loader.images,
                self.all_df[self.csv_filename_column_title],
            ):
                layout = QHBoxLayout()
                self.input_layout.addLayout(layout)
                self.all_checkboxes.append(QCheckBox())
                self.all_checkboxes[-1].setChecked(True)
                layout.addWidget(self.all_checkboxes[-1])
                pixmap = QPixmap(
                    QImage(
                        sitk.GetArrayViewFromImage(image),
                        image.GetWidth(),
                        image.GetHeight(),
                        QImage.Format_Grayscale8
                        if image.GetNumberOfComponentsPerPixel() == 1
                        else QImage.Format_RGB888,
                    )
                )
                icon = QIcon()
                icon.addPixmap(pixmap, QIcon.Normal, QIcon.Off)
                self.all_buttons.append(QPushButton())
                self.all_buttons[-1].setStyleSheet(self.initial_button_stylesheet)
                self.all_buttons[-1].setIcon(icon)
                self.all_buttons[-1].setToolTip(file_name)
                self.all_buttons[-1].clicked.connect(self.__image_clicked(i))
                layout.addWidget(self.all_buttons[-1])
                layout.addStretch()
            self.save_action.setEnabled(True)
            # The reference labels are available so we can evaluate performance
            if self.csv_actual_value_column_title in self.all_df.columns:
                self.save_evaluate_action.setEnabled(True)
            self.query_action.setEnabled(True)
        # clear all data from the loader.
        self.data_loader.reset()

    def __image_clicked(self, image_row_index):
        def callback():
            image_info = self.all_df.iloc[image_row_index]
            image_file_reader = sitk.ImageFileReader()
            image_file_reader.SetFileName(image_info[self.csv_filename_column_title])
            image_file_reader.ReadImageInformation()
            image_size = list(image_file_reader.GetSize())
            if len(image_size) == 3:
                image_size[2] = 0
                image_file_reader.SetExtractSize(image_size)
            image = image_file_reader.Execute()
            self.selected_image_row_index = image_row_index
            self.selected_image_label.setPixmap(sitk2qpixmap(image))
            # The image is scaled to fit the label size. To keep the image's
            # aspect ratio the label width is changed accordingly. In the end, the
            # image is scaled to fit the updated label size.
            self.selected_image_label.setScaledContents(True)
            new_label_width = (
                self.selected_image_label.parentWidget().size().height()
                * self.selected_image_label.pixmap().size().width()
                / self.selected_image_label.pixmap().size().height()
            )
            self.selected_image_label.resize(
                new_label_width,
                self.selected_image_label.parentWidget().size().height(),
            )

            self.selected_image_information_edit.setText(
                "\n".join([f"{index}: {value}" for index, value in image_info.items()])
            )

        return callback

    def __select_all_changed(self):
        new_value = self.sender().isChecked()
        for cb in self.all_checkboxes:
            cb.setChecked(new_value)

    def __call_service(self):
        self.threadpool.setMaxThreadCount(int(self.threadcount_value_label.text()))
        self.query_workers = []
        self.total_quries = 0
        for cb in self.all_checkboxes:
            if cb.isChecked():
                self.total_quries = self.total_quries + 1
            cb.setEnabled(False)
        self.remaining_quries = self.total_quries
        for i, cb, file_name in zip(
            range(len(self.all_checkboxes)),
            self.all_checkboxes,
            self.all_df[self.csv_filename_column_title],
        ):
            if cb.isChecked():
                self.all_df.at[
                    self.all_df.index[i], self.csv_file_size_column_title
                ] = os.stat(file_name).st_size
                query_service = QueryService(
                    int(self.retry_value_label.text()),
                    i,
                    self.__rap_tb_ntb_query(
                        self.endpoint,
                        self.algorithms[self.algorithm_combo.currentText()],
                        int(self.timeout_value_label.text()),
                    ),
                    file_name,
                )
                query_service.signals.finished.connect(self.__query_service_finished)
                self.all_buttons[i].setStyleSheet(self.waiting_button_stylesheet)
                self.threadpool.start(query_service)
                self.query_workers.append(query_service)
        self.update()

    def __query_service_finished(self, final_results):
        df_row, error_message, results = final_results
        df_index = self.all_df.index[df_row]
        if error_message:
            self.all_buttons[df_row].setStyleSheet(
                self.finished_failure_button_stylesheet
            )
            self.all_df.at[df_index, self.csv_service_response_column_title] = (
                "Error - " + error_message
            )
        else:
            self.all_buttons[df_row].setStyleSheet(
                self.finished_success_button_stylesheet
            )
            self.all_df.at[df_index, self.csv_service_response_column_title] = "success"
            # First query means we need to create all the columns
            if self.remaining_quries == self.total_quries:
                for key in results:
                    if key not in self.all_df.columns:
                        self.all_df[key] = ""
            for key, val in results.items():
                self.all_df.at[df_index, key] = val
        self.remaining_quries = self.remaining_quries - 1
        if self.selected_image_row_index == df_row:
            image_info = self.all_df.iloc[df_row]
            self.selected_image_information_edit.setText(
                "\n".join([f"{index}: {value}" for index, value in image_info.items()])
            )
        if self.remaining_quries == 0:
            for cb in self.all_checkboxes:
                cb.setEnabled(True)
            QMessageBox().information(self, "Message", "Calculation completed.")

    def __save_results(self):
        output_file_name, _ = QFileDialog.getSaveFileName(
            self, "Save Results", "", "csv(*.csv)"
        )
        if output_file_name:
            self.all_df.to_csv(output_file_name, index=False)

    def __save_results_and_evaluation(self):
        output_dir_name = QFileDialog.getExistingDirectory(
            self, "Select Output Directory"
        )
        if output_dir_name:
            self.all_df.to_csv(
                os.path.join(output_dir_name, "all_results.csv"), index=False
            )
            # Only analyze results that are complete for all evaluated algorithms
            all_valid_results_df = self.all_df.dropna()
            if all_valid_results_df.empty:
                self.__error_function(
                    "No complete valid results obtained from service."
                )
                return
            elif (
                all_valid_results_df[self.csv_actual_value_column_title].nunique() == 1
            ):
                self.__error_function(
                    "Valid data is unary, only saved raw data, performance analysis not done."
                )
                return
            if all_valid_results_df.shape[0] != self.all_df.shape[0]:
                self.__error_function(
                    f"Evaluation performed on partial dataset with valid results ({all_valid_results_df.shape[0]}/{self.all_df.shape[0]})."  # noqa: E501
                )

            valid_algorithm_names = []
            for _, algorithm_name in self.algorithms.items():
                if f"{algorithm_name}: probability_of_TB" in all_valid_results_df:
                    valid_algorithm_names.append(algorithm_name)
            if valid_algorithm_names:
                self.__evaluate_and_save(
                    valid_algorithm_names, all_valid_results_df, output_dir_name
                )

    def __evaluate_and_save(
        self, valid_algorithm_names, valid_results_df, output_dir_name
    ):
        evaluation_metrics = [
            ("Accuracy", metrics.accuracy_score),
            (
                "Precision",
                lambda y_true, y_pred: metrics.precision_score(
                    y_true, y_pred, pos_label=self.positive_label
                ),
            ),
            (
                "Recall",
                lambda y_true, y_pred: metrics.recall_score(
                    y_true, y_pred, pos_label=self.positive_label
                ),
            ),
            (
                "Specificity",
                lambda y_true, y_pred: metrics.recall_score(
                    y_true, y_pred, pos_label=self.negative_label
                ),
            ),
            (
                "F1",
                lambda y_true, y_pred: metrics.f1_score(
                    y_true, y_pred, pos_label=self.positive_label
                ),
            ),
        ]
        evaluation_results = []
        # Setup of the ROC figure, all ROC curves from the algorithms will plot
        # to this figure
        roc_figure = plt.figure()
        plt.plot([0, 1], [0, 1], color="black", lw=1, linestyle="--")
        plt.xlim([0.0, 1.0])
        plt.ylim([0.0, 1.0])
        plt.xlabel("False Positive Rate")
        plt.ylabel("True Positive Rate")

        for algo_name in valid_algorithm_names:
            algo_results = []
            fpr, tpr, thresholds = metrics.roc_curve(
                valid_results_df[self.csv_actual_value_column_title],
                valid_results_df[f"{algo_name}: probability_of_TB"],
                pos_label=self.positive_label,
            )
            roc_auc = metrics.auc(fpr, tpr)
            youden_threshold = thresholds[np.argmax(tpr - fpr)]
            valid_results_df[
                f"{algo_name}: dataset_youden_threshold_decision({youden_threshold})"
            ] = valid_results_df[f"{algo_name}: probability_of_TB"].apply(
                lambda x, bound_threshold=youden_threshold: self.positive_label
                if x > bound_threshold
                else self.negative_label
            )
            # Plot the ROC for this algorithm
            plt.figure(roc_figure.number)
            plt.plot(
                fpr,
                tpr,
                lw=1,
                label=algo_name + f", AUC = {roc_auc:.3f}",
                linestyle="solid",
            )
            # Plot confusion matrix for this algorithm using the fixed/service threshold and the
            # optimal dataset specific threshold (max Youden index, argmax(t) sensitivity(t)+specificity(t)-1)
            plt.figure()
            metrics.ConfusionMatrixDisplay.from_predictions(
                valid_results_df[self.csv_actual_value_column_title],
                valid_results_df[f"{algo_name}: decision"],
                cmap=plt.cm.Blues,
                colorbar=False,
            )
            plt.savefig(
                os.path.join(
                    output_dir_name, f"confusion_matrix_{algo_name}_fixed_threshold.pdf"
                ),
                bbox_inches="tight",
            )
            plt.figure()
            metrics.ConfusionMatrixDisplay.from_predictions(
                valid_results_df[self.csv_actual_value_column_title],
                valid_results_df[
                    f"{algo_name}: dataset_youden_threshold_decision({youden_threshold})"
                ],
                cmap=plt.cm.Blues,
                colorbar=False,
            )
            plt.savefig(
                os.path.join(
                    output_dir_name,
                    f"confusion_matrix_{algo_name}_youden_threshold.pdf",
                ),
                bbox_inches="tight",
            )
            # Evaluate results using all the common metrics
            for _, f in evaluation_metrics:
                using_service_threshold = f(
                    valid_results_df[self.csv_actual_value_column_title],
                    valid_results_df[f"{algo_name}: decision"],
                )
                using_dataset_optimal_threshold = f(
                    valid_results_df[self.csv_actual_value_column_title],
                    valid_results_df[
                        f"{algo_name}: dataset_youden_threshold_decision({youden_threshold})"
                    ],
                )
                algo_results.append(
                    f"{using_service_threshold} ({using_dataset_optimal_threshold})"
                )
            algo_results.append(roc_auc)
            evaluation_results.append(algo_results)
        # Finish plotting the ROC figure, add the legend and save
        plt.figure(roc_figure.number)
        plt.legend(loc="lower right")
        plt.savefig(os.path.join(output_dir_name, "roc.pdf"), bbox_inches="tight")
        # Save the data actually used in the computations
        valid_results_df.to_csv(
            os.path.join(output_dir_name, "valid_results_used_in_evaluation.csv"),
            index=False,
        )
        # Save the evaluation results
        column_names, _ = zip(*evaluation_metrics)
        column_names = column_names + ("ROC AUC",)
        evaluation_results_df = pd.DataFrame(
            data=evaluation_results, index=valid_algorithm_names, columns=column_names
        )
        evaluation_results_df.index.name = "Algorithm"
        evaluation_results_df.to_csv(
            os.path.join(output_dir_name, "evaluation_results.csv")
        )

    def __get_file_names_df(self, csv_or_dir_name):
        """
        Given the path to a csv or directory return a dataframe with the
        column defined by self.csv_filename_column_title containing the file
        names. If a csv file and it has the actual values associated with each file
        (column title defined by self.csv_actual_value_column_title) check that the
        values in that column match the allowed ones.
        """
        df = pd.DataFrame()
        if os.path.isdir(csv_or_dir_name):
            file_names = []
            for dir_name, subdir_names, f_names in os.walk(csv_or_dir_name):
                file_names += [
                    os.path.join(os.path.abspath(dir_name), fname) for fname in f_names
                ]
            df = pd.DataFrame(data=file_names, columns=[self.csv_filename_column_title])
        else:
            try:
                df = pd.read_csv(csv_or_dir_name)
                dir_path = os.path.dirname(os.path.abspath(csv_or_dir_name))
                df[self.csv_filename_column_title] = df[
                    self.csv_filename_column_title
                ].apply(lambda x: os.path.join(dir_path, x))
                # Check if dataframe has actual classes and that they are as expected
                if self.csv_actual_value_column_title in df.columns:
                    # Remove leading/trailing whitespace, accommodate for minor user error in
                    # creation of input csv.
                    df[self.csv_actual_value_column_title] = df[
                        self.csv_actual_value_column_title
                    ].str.strip()
                    if (
                        not df[self.csv_actual_value_column_title]
                        .isin([self.positive_label, self.negative_label])
                        .all()
                    ):
                        df = pd.DataFrame()
            except Exception:
                df = pd.DataFrame()
        return df

    def __create_settings_widget(self):
        wid = QWidget()
        wid.setWindowTitle("Settings")
        input_layout = QVBoxLayout()
        wid.setLayout(input_layout)

        layout = QHBoxLayout()
        layout.addWidget(QLabel("Number of concurrent threads:"))
        threadcount_slider = QSlider(orientation=Qt.Horizontal)
        threadcount_slider.setMinimum(1)
        threadcount_slider.setMaximum(self.threadpool.maxThreadCount())
        threadcount_slider.setValue(self.threadpool.maxThreadCount())
        layout.addWidget(threadcount_slider)
        self.threadcount_value_label = QLabel()
        self.threadcount_value_label.setNum(self.threadpool.maxThreadCount())
        threadcount_slider.valueChanged.connect(self.threadcount_value_label.setNum)
        layout.addWidget(self.threadcount_value_label)
        input_layout.addLayout(layout)

        layout = QHBoxLayout()
        layout.addWidget(QLabel("Query timeout [sec]:"))
        timeout_slider = QSlider(orientation=Qt.Horizontal)
        timeout_slider.setMinimum(self.timeout_range[0])
        timeout_slider.setMaximum(self.timeout_range[1])
        timeout_slider.setValue(self.default_timeout)
        layout.addWidget(timeout_slider)
        self.timeout_value_label = QLabel()
        self.timeout_value_label.setNum(self.default_timeout)
        timeout_slider.valueChanged.connect(self.timeout_value_label.setNum)
        layout.addWidget(self.timeout_value_label)
        input_layout.addLayout(layout)

        layout = QHBoxLayout()
        layout.addWidget(QLabel("Number of retries:"))
        retry_slider = QSlider(orientation=Qt.Horizontal)
        retry_slider.setMinimum(self.number_of_retries_range[0])
        retry_slider.setMaximum(self.number_of_retries_range[1])
        retry_slider.setValue(self.default_number_of_retries)
        layout.addWidget(retry_slider)
        self.retry_value_label = QLabel()
        self.retry_value_label.setNum(self.default_number_of_retries)
        retry_slider.valueChanged.connect(self.retry_value_label.setNum)
        layout.addWidget(self.retry_value_label)
        input_layout.addLayout(layout)

        layout = QHBoxLayout()
        layout.addWidget(QLabel("Algorithm:"))
        self.algorithm_combo = QComboBox()
        self.algorithm_combo.addItems(list(self.algorithms.keys()))
        layout.addWidget(self.algorithm_combo)
        input_layout.addLayout(layout)

        button = QPushButton("Done")
        button.clicked.connect(wid.hide)
        input_layout.addWidget(button)
        return wid

class DataLoader(QThread):
    progress_signal = Signal()

    def __init__(self, thumbnail_size):
        super(DataLoader, self).__init__()
        self.thumbnail_size = thumbnail_size
        self.image_file_reader = sitk.ImageFileReader()
        self.reset()

    def reset(self):
        self.images = []
        self.error_indexes = []
        self.file_names_list = []

    def run(self):
        for i, file_name in enumerate(self.file_names_list):
            try:
                self.image_file_reader.SetFileName(file_name)
                self.image_file_reader.ReadImageInformation()
                num_components = self.image_file_reader.GetNumberOfComponents()
                # not a grayscale or color image
                if num_components not in [1, 3]:
                    self.error_indexes.append(i)
                else:
                    image_size = list(self.image_file_reader.GetSize())
                    # not a 2D image
                    if (
                        not (len(image_size) == 3 and image_size[2] == 1)
                        and not len(image_size) == 2
                    ):
                        self.error_indexes.append(i)
                    else:  # 2D image posing as a 3D one
                        if len(image_size) == 3 and image_size[2] == 1:
                            image_size[2] = 0
                        self.image_file_reader.SetExtractSize(image_size)
                        self.images.append(
                            self.__resize_and_scale_uint8(
                                self.image_file_reader.Execute(), self.thumbnail_size
                            )
                        )
                self.progress_signal.emit()
            except Exception:
                self.error_indexes.append(i)
                self.progress_signal.emit()

    def __resize_and_scale_uint8(self, image, new_size, outside_pixel_value=0):
        """
        Resize the given image to the given size, with isotropic pixel spacing
        and scale the intensities to [0,255].

        Resizing retains the original aspect ratio, with the original image centered
        in the new image. Padding is added outside the original image extent using the
        provided value.

        :param image: A SimpleITK image.
        :param new_size: List of ints specifying the new image size.
        :param outside_pixel_value: Value in [0,255] used for padding.
        :return: a 2D SimpleITK image with desired size and a pixel type of sitkUInt8
        """
        # Rescale intensities if scalar image with pixel type that isn't sitkUInt8.
        # We rescale first, so that the zero padding makes sense for all original image
        # ranges. If we resized first, a value of zero in a high dynamic range image may
        # be somewhere in the middle of the intensity range and the outer border has a
        # constant but arbitrary value.
        if (
            image.GetNumberOfComponentsPerPixel() == 1
            and image.GetPixelID() != sitk.sitkUInt8
        ):
            final_image = sitk.Cast(sitk.RescaleIntensity(image), sitk.sitkUInt8)
        else:
            final_image = image
        final_image = resize(
            final_image, new_size, fill_value=outside_pixel_value, anti_aliasing_sigma=0
        )
        return final_image

class QueryServiceSignals(QObject):
    finished = Signal(object)

class QueryService(QRunnable):
    def __init__(self, num_retries, df_row, query_function, file_name):
        super(QueryService, self).__init__()
        self.num_retries = num_retries
        self.df_row = df_row
        self.query_function = query_function
        self.signals = QueryServiceSignals()
        self.results = {}
        self.error_message = ""
        self.continue_running = True
        self.image_file_reader = sitk.ImageFileReader()
        self.image_file_reader.SetFileName(file_name)

    def run(self):
        try:
            with tempfile.TemporaryDirectory() as tmpdirname:
                self.image_file_reader.ReadImageInformation()
                # 2D image posing as a 3D one
                image_size = list(self.image_file_reader.GetSize())
                if len(image_size) == 3 and image_size[2] == 1:
                    image_size[2] = 0
                self.image_file_reader.SetExtractSize(image_size)
                image = self.image_file_reader.Execute()
                # Clear all of the meta-data, write to temporary DICOM file and
                # use that as the query image, ensures that no PII/PHI leaves the
                # local machine.
                for k in image.GetMetaDataKeys():
                    image.EraseMetaData(k)
                anonymized_file_name = os.path.join(tmpdirname, "anon.dcm")
                sitk.WriteImage(image, anonymized_file_name)
                num_tries = 0
                no_response = True
                while (
                    num_tries < self.num_retries
                    and no_response
                    and self.continue_running
                ):
                    self.error_message, self.results = self.query_function(
                        anonymized_file_name
                    )
                    num_tries = num_tries + 1
                    if not self.error_message:
                        no_response = False
                if self.continue_running:
                    self.signals.finished.emit(
                        (self.df_row, self.error_message, self.results)
                    )
        # Use the stack trace as the error message to provide enough
        # detailes for debugging.
        except Exception:
            if self.continue_running:
                self.error_message = (
                    "Exception occurred during computation:\n" + traceback.format_exc()
                )
                self.signals.finished.emit(
                    (self.df_row, self.error_message, self.results)
                )

def main():
    app = QApplication([])
    app.setStyleSheet(qdarkstyle.load_stylesheet(qt_api="pyside6"))
    tb_ntb_dialog = TBorNotTBDialog()  # noqa: F841
    app.exec()

if __name__ == "__main__":
    main()


#!/usr/bin/env python
'''
Process file produced by:

https://cdcvs.fnal.gov/redmine/projects/dunetpc/wiki/_ProtoDUNE-SP_Wire_Dumps_

Columns:

0) channel :: [0-15359]
1) cryostat :: 0
2) tpc :: [0,11]
3) plane :: [0,2]
4) wire :: [0,1147]
5) wire beg x
6) wire beg y
7) wire beg z
8) wire end x
9) wire end y
10) wire end z

'''

import sys
import numpy
import matplotlib.pyplot as plt
from collections import namedtuple, defaultdict

# wid simply counts wire in order it is seen in the file
# wip is wire-in-plane
Wire = namedtuple("Wire", "seg wid ch cryo tpc plane wip beg end")

def load(infile):
    wires = list()

    chan2wires = defaultdict(list)

    with open(infile) as fp:
        for wid, line in enumerate(fp.readlines()):
            chunks = line.strip().split()
            numbers = map(int, chunks[:5]) + map(lambda x: 10.0*float(x), chunks[5:])
            zmax = max(numbers[7], numbers[10])
            wire = (zmax, wid, numbers)
            chan = numbers[0]
            chan2wires[chan].append(wire)

    wires = list()
    for ch,segs in chan2wires.items():
        segs.sort()
        for iseg, seg in enumerate(segs):
            zmax, wid, numbers = seg
            wire = Wire(iseg, wid, numbers[0], numbers[1], numbers[2], numbers[3], numbers[4],
                            numbers[5:8], numbers[8:11])
            wires.append(wire)
    return wires

def select_by(wires, key, val):
    ret = list()
    for wire in wires:
        if getattr(wire, key, None) == val:
            ret.append(wire)
    return ret

def summary(wires):
    numcols = ["ch","cryo","tpc","plane","wip"]
    num = dict()
    for icol, col in enumerate(numcols):
        num[col] = len(set([getattr(w,col) for w in wires]))
        print ("#%s: %d " % (col, num[col]))

    for itpc in range(num["tpc"]):
        for iplane in range(num["plane"]):
            witp = select_by(select_by(wires, "tpc", itpc), "plane", iplane)
            c = set([w.ch for w in witp])
            print ("tpc:%d plane:%d #chans:%d #wires:%d" %
                       (itpc, iplane, len(c), len(witp)))

def plot_wires_beg(wires, nwires, title):
    plt.clf()
    xyz="XYZ"
    axes=(2,1)
    cmap = plt.get_cmap('seismic')
    nwires = len(wires)
    colors = [cmap(i) for i in numpy.linspace(0, 1, nwires)]
    for ind, wire in enumerate(wires):
        if ind%100 != 0:
            continue
        p1,p2 = wire[5:8], wire[8:11]
        x = numpy.asarray((p1[axes[0]], p2[axes[0]]))
        y = numpy.asarray((p1[axes[1]], p2[axes[1]]))
        plt.plot(x, y, color=colors[ind])
        plt.text(0.5*sum(x), 0.5*sum(y), "ch %d" % wire[0])
    plt.xlabel("%s direction" % xyz[axes[0]])
    plt.ylabel("%s direction" % xyz[axes[1]])
    plt.title(title)

def plot_wires_plane_tpc(wires, plane, tpcs=(1,5,9), wire_sampling=10, text_sampling=100, offsets=None):
    plt.clf()
    xyz="XYZ"
    uvw="UVW"
    axes=(2,1)
    #cmap = plt.get_cmap('seismic')
    cmap = plt.get_cmap('rainbow')

    inplane = select_by(wires, "plane", plane)
    for itpc, tpc in enumerate(tpcs):
        offset = numpy.zeros(2)
        if offsets:
            offset = numpy.asarray(offsets[itpc])

        intpc = select_by(inplane, "tpc", tpc)
        nwires = len(intpc)
        colors = [cmap(i) for i in numpy.linspace(0, 1, nwires)]
        todraw = [(i,w) for i,w in enumerate(intpc) if i%wire_sampling == 0]
        todraw.append((nwires-1, intpc[-1]))
        for ind,wire in todraw:
            p1,p2 = wire.beg, wire.end
            x = numpy.asarray((p1[axes[0]], p2[axes[0]]))
            y = numpy.asarray((p1[axes[1]], p2[axes[1]]))
            plt.plot(x, y, color=colors[ind])
            if ind%text_sampling == 0:
                plt.plot([0.5*sum(x)], [0.5*sum(y)], marker='o', color=colors[ind])

                x = numpy.asarray((p1[axes[0]], p2[axes[0]])) + offset[0]
                y = numpy.asarray((p1[axes[1]], p2[axes[1]])) + offset[1]
                plt.text(0.5*sum(x), 0.5*sum(y),
                             "tpc:%d\nwid:%d\nwip:%d\nch:%d" % (tpc, wire.wid, wire.wip, wire.ch),
                             verticalalignment='center', horizontalalignment='center')
    plt.xlabel("%s direction" % xyz[axes[0]])
    plt.ylabel("%s direction" % xyz[axes[1]])
    plt.title("%s-plane, TPCs: %s" % (uvw[plane], str(tpcs)))

def plot_wires(wires):
    print ("Plotting %d wires" % len(wires))

    for wire in wires:
        p1, p2 = wire.beg, wire.end
        x = numpy.asarray((p1[2], p2[2]))
        y = numpy.asarray((p1[1], p2[1]))
        plt.plot(x, y)

from matplotlib.backends.backend_pdf import PdfPages

def wire_plots(wires, filename = "larsoftwires.pdf"):
    tses = [100, 100, 170]

    with PdfPages(filename) as pdf:

        plot_wires_plane_tpc(wires, 0)
        pdf.savefig()

        plot_wires_plane_tpc(wires, 1)
        pdf.savefig()

        plot_wires_plane_tpc(wires, 2, offsets=[(0,20), (0,-20), (0,+20)], text_sampling=100)
        pdf.savefig()

        plot_wires_plane_tpc(wires, 0, tpcs=(0,1), offsets=[(8,0), (-8,0)])
        pdf.savefig()

        plot_wires_plane_tpc(wires, 1, tpcs=(0,1), offsets=[(8,0), (-8,0)])
        pdf.savefig()

        plot_wires_plane_tpc(wires, 2, tpcs=(0,1), offsets=[(0,20), (0,-20)], text_sampling=50)
        pdf.savefig()

def do_wire_plots(filename="WireDump_ProtoDUNESP_v2_2.txt", outfile="larsoftwires.pdf"):
    wires = load(filename)
    wire_plots(wires, outfile)


import datetime
import numpy as np
from itertools import groupby
from skimage import measure
from PIL import Image

class PycocoCreatorTools:

    @staticmethod
    def bbox(img):
        rows = np.any(img, axis=1)
        cols = np.any(img, axis=0)
        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]
        h = rmax - rmin
        w = cmax - cmin
        return [int(cmin), int(rmin), int(w), int(h)]

    @staticmethod
    def resize_binary_mask(array, new_size):
        image = Image.fromarray(array.astype(np.uint8)*255)
        image = image.resize(new_size)
        return np.asarray(image).astype(np.bool_)

    @staticmethod
    def close_contour(contour):
        if not np.array_equal(contour[0], contour[-1]):
            contour = np.vstack((contour, contour[0]))
        return contour

    @staticmethod
    def binary_mask_to_rle(binary_mask):
        rle = {'counts': [], 'size': list(binary_mask.shape)}
        counts = rle.get('counts')
        for i, (value, elements) in enumerate(groupby(binary_mask.ravel(order='F'))):
            if i == 0 and value == 1:
                    counts.append(0)
            counts.append(len(list(elements)))

        return rle

    @staticmethod
    def binary_mask_to_polygon(binary_mask, tolerance=0):
        """Converts a binary mask to COCO polygon representation
        Args:
            binary_mask: a 2D binary numpy array where '1's represent the object
            tolerance: Maximum distance from original points of polygon to approximated
                polygonal chain. If tolerance is 0, the original coordinate array is returned.
        """
        polygons = []
        # pad mask to close contours of shapes which start and end at an edge
        padded_binary_mask = np.pad(binary_mask, pad_width=1, mode='constant', constant_values=0)
        contours = measure.find_contours(padded_binary_mask, 0.5)
        contours = np.subtract(contours, 1)
        for contour in contours:
            contour = PycocoCreatorTools.close_contour(contour)
            contour = measure.approximate_polygon(contour, tolerance)
            if len(contour) < 3:
                continue
            contour = np.flip(contour, axis=1)
            segmentation = contour.ravel().tolist()
            # after padding and subtracting 1 we may get -0.5 points in our segmentation
            segmentation = [0 if i < 0 else i for i in segmentation]
            polygons.append(segmentation)

        return polygons

    @staticmethod
    def create_image_info(image_id, file_name, image_size,
                        date_captured=datetime.datetime.utcnow().isoformat(' '),
                        license_id=1, coco_url="", flickr_url=""):
        """Creates image info section of coco annotation
        Args:
            image_id: integer to uniquly identify image
            file_name: filename for image
        """
        image_info = {
                "id": image_id,
                "file_name": file_name,
                "width": image_size[0],
                "height": image_size[1],
                "date_captured": date_captured,
                "license": license_id,
                "coco_url": coco_url,
                "flickr_url": flickr_url
        }

        return image_info

    @staticmethod
    def create_annotation_info(annotation_id, image_id, category_info, binary_mask,
                            image_size=None, tolerance=2, bounding_box=None):
        """Creates info section of coco annotation
        Args:
            annotation_id: integer to uniquly identify the annotation
            image_id: integer to uniquly identify image
            category_info: dict which contains category info and a boolean attribute "crowd" which tells if the annotation is a crowd of a class
        """
        if image_size is not None:
            binary_mask = PycocoCreatorTools.resize_binary_mask(binary_mask, image_size)

        if bounding_box is None:
            bounding_box = PycocoCreatorTools.bbox(binary_mask)
            area = bounding_box[2] * bounding_box[3]

        if category_info["is_crowd"]:
            is_crowd = 1
            segmentation = PycocoCreatorTools.binary_mask_to_rle(binary_mask)
        else :
            is_crowd = 0
            segmentation = PycocoCreatorTools.binary_mask_to_polygon(binary_mask, tolerance)
            if not segmentation:
                return None

        annotation_info = {
            "id": annotation_id,
            "image_id": image_id,
            "category_id": category_info["id"],
            "iscrowd": is_crowd,
            "area": [area],
            "bbox": bounding_box,
            "segmentation": segmentation,
            "width": binary_mask.shape[1],
            "height": binary_mask.shape[0],
        }
        return annotation_info


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 15 20:03:52 2017

@author: jonbaird

Sum square difference
Problem 6

The sum of the squares of the first ten natural numbers is,
12 + 22 + ... + 102 = 385

The square of the sum of the first ten natural numbers is,
(1 + 2 + ... + 10)2 = 552 = 3025

Hence the difference between the sum of the squares of the first ten natural
numbers and the square of the sum is 3025 − 385 = 2640. Find the difference
between the sum of the squares of the first one hundred natural numbers and
the square of the sum.
"""
import numpy as np
from tqdm import tqdm

def ssq(stop=10):
    """
    compute the sum of squares
    """
    total = 0
    squares = [i ** 2 for i in tqdm(range(1, stop + 1))]
    while len(squares) > 0:
        total += squares.pop()
    return total

def sqs(stop=10):
    """
    sequence is a 'range' variable
    compute the square of the sum of 'sequence'
    """
    total = 0
    for i in tqdm(np.arange(1, stop + 1)):
        total = total + i
    return total * total

print(ssq(stop=10000) - sqs(10000))
print(sqs(stop=10000) - ssq(stop=10000))


import pandas as pd
import sqlalchemy as sa
import psycopg2
import config
from tweet_analysis_functions import prefilter_tweet

engine = sa.create_engine(config.engine ,connect_args={"connect_timeout": 60})

def create_postgres_database(tweet_data_frame_list, tweet_users_frame_list, data_table, users_table):

    data_table_name = data_table
    users_table_name = users_table

    if (isinstance(tweet_data_frame_list, pd.DataFrame) and isinstance(tweet_users_frame_list, pd.DataFrame)):
        df_data = tweet_data_frame_list
        df_users = tweet_users_frame_list

    elif (isinstance(tweet_data_frame_list, list) and isinstance(tweet_users_frame_list, list)):
        df_data = pd.concat(tweet_data_frame_list)
        df_users = pd.concat(tweet_users_frame_list)

    else:
        print('Error')
        return

    df_data.to_sql(
    name=data_table_name.lower(),
    con=engine,
    if_exists='replace'  # If the table already exists, replace it
    )

    df_users.to_sql(
    name=users_table_name.lower(),
    con=engine,
    if_exists='replace'  # If the table alrady exists, replace it
    )


# import torch
# import torchvision
#import knockoff.datasets.GTSRB as GTSRB
#import knockoff.datasets as datasets
# import torch.nn as nn
# import pretrainedmodels

import caffe
import numpy as np
import pickle
import imageio
from PIL import Image
from torchvision import transforms

transform_gtsrb = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406],
                         std=[0.229,0.224,0.225]),
])

transform_trigger = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485,0.456,0.406],
                         std=[0.229,0.224,0.225]),
])

def classify_trans(fname):
    pix = imageio.imread(fname)
    pix = transform_gtsrb(Image.fromarray(np.uint8(pix)))
    # pix = transform_trigger(Image.fromarray(np.uint8(pix)))
    return pix

def classify_trigger(fname):
    pix = imageio.imread(fname)
    # pix = transform_gtsrb(Image.fromarray(np.uint8(pix)))
    pix = transform_trigger(Image.fromarray(np.uint8(pix)))
    return pix

#dir = '/home/emmet/trigger/TrojanNN-master/data/train'  # 保存测试图片的集合

# filelist = []
# filenames = os.listdir(dir)
# for fn in filenames:
#     fullfilename = os.path.join(dir, fn)
#     print(fullfilename)
#     filelist.append(fullfilename)

def read_data():
    #print("read trojan data")
    X = []
    Y = []
    for fn in filelist:
        X.append(classify(fn))
        Y.append(int(fn.split('/')[-1][-5]))
    return X, Y

def read_trojan_data(x_pick):

    X = []
    h = 224
    w = 224
    # expected = 1
    trigger = classify_trigger('./fc6_1_803_694_1_1_0803.jpg')

    for j in range(len(x_pick)):
        caffeset = x_pick[j]
        # print(caffeset[0,200,200])
        for y in range(h):
            for x in range(w):
                if x > w - 80 and x < w -20 and y > h - 80 and y < h - 20:
                    caffeset[:,y,x] = trigger[:,y,x]

        X.append(np.array(caffeset, copy=True))
        #X.append(np.array(caffeset, copy=True))
    return X

if __name__ == '__main__':
    #X, Y = read_data()
    X, Y = pickle.load(open('/home/chenyanjiao/PycharmProjects/knockoffnets/knockoff_gtsrb_1000.pkl','rb'))
    caffemodel_path = "./vgg16.caffemodel"
    caffedeploy_path = "/home/chenyanjiao/infocom_backdoor/gtsrb/single_loc/VGG_ILSVRC_16_layers_deploy.prototxt.txt"
    caffe.set_mode_cpu()
    cn = caffe.Net(caffedeploy_path,caffemodel_path,  caffe.TEST)
    caffe_positive = 0
    trojan= read_trojan_data(X)
    current = 0
    # summary = [0] * 4096

    for j in range(100):
        caffeset = []
        print('current',current)
        caffeset=np.array([trojan[i] for i in range(current, current + 10)])
        cn.blobs['data'].data[...] = caffeset
        caffeoutputs = cn.forward()['prob']
        #torchoutputs = model(torchset)
        for index, i in enumerate(range(current, current+10)):
            if caffeoutputs[index].argmax() == 1:
                print(1)
                caffe_positive += 1
                acts = cn.blobs['fc6'].data
                fc = acts[0]
                best_unit = fc.argmax()
        current += 10
    print('ASR:',caffe_positive / 10.0)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from scipy.stats import norm
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

import math

#Esta funcion representa la matriz de confusion de un modelo mediante la TDS

def plot(Pa, Pfa, returnP = False, show = True):
    valorMin = 0.0000000001
    if Pa == 1: Pa = 1 - valorMin
    if Pfa == 1: Pfa = 1 - valorMin
    if Pa == 0: Pa = valorMin
    if Pfa == 0: Pfa = valorMin

    Pe = 1 - Pa ; Prc = 1 - Pfa

    Za = norm.ppf(Pa, loc = 0, scale = 1)
    Ze = norm.ppf(Pe, loc = 0, scale = 1)
    Zfa = norm.ppf(Pfa, loc = 0, scale = 1)
    Zrc = norm.ppf(Prc, loc = 0, scale = 1)

    dPrima = Za - Zfa
    C = (-1 / 2) * (Za + Zfa) #Observador ideal
    beta = math.exp(dPrima * C)

    def GeneraDatos(meanR, meanS,
                            zFrom = -10, zTo = 10, zBy = 1001):
      Z = np.linspace(start = zFrom,
                      stop = zTo,
                      num = zBy)

      df = pd.DataFrame({
        "Zscore" : Z,
        "Zacum" : norm.cdf(Z),
        "Zden0" : norm.pdf(Z, loc = meanR, scale = 1),
        "Zden1" : norm.pdf(Z, loc = meanS, scale = 1)
      })

      df.columns = ["Z", "F(X)", "f(x/r)", "f(x/s)"]

      return(df)

    df = GeneraDatos(meanR = 0,
                      meanS = 0 + dPrima)

    #Encontramos el valor de Z asociado a los valores
    #f(x/s) y f(x/r) que hacen que la operacion
    #f(x/s) / f(x/r) se aproxime mas a beta
    Temp = (df["f(x/s)"] / df["f(x/r)"]) - beta
    Xc = df["Z"][Temp.abs() == Temp.abs().min()].values[0]

    #Para obtener el punto en el que el observador ideal da respuesta tomamos beta = 1
    Temp = (df["f(x/s)"] / df["f(x/r)"]) - 1
    Xideal = df["Z"][Temp.abs() == Temp.abs().min()].values[0]

    DatosGrafico = df[["Z", "f(x/s)", "f(x/r)"]]

    #Mostramos solo los datos que contienen valores altos para una de las distribuciones
    DatosGrafico = df.loc[(df["f(x/s)"] > 0.001) | (df["f(x/r)"] > 0.001)]

    DatosGrafico = DatosGrafico.set_index("Z", drop = True)

    if show == True:
        with plt.style.context('ggplot'):
            fig, ax = plt.subplots()
            DatosGrafico["f(x/r)"].plot()
            DatosGrafico["f(x/s)"].plot()
            ax.axvline(x = Xideal, color = "green")
            ax.axvline(x = Xc, color = "black")

            handles, labels = ax.get_legend_handles_labels()
            handles.append(mpatches.Patch(color = "green", label = "Observador ideal"))
            handles.append(mpatches.Patch(color ="black", label = "Observador"))

            ax.legend(handles = handles)

    if returnP == True:
        Resultados = {
                    "ObsIdeal" : Xideal,
                    "Obs" : Xc,
                    "C" : C,
                    "dPrima" : dPrima,
                    "beta" : beta,
                    "Pa" : Pa,
                    "Pe" : Pe,
                    "Pfa" : Pfa,
                    "Prc" : Prc
                }

        for dict_value in Resultados:
            Resultados[dict_value] = round(Resultados[dict_value], 3)

        return(Resultados)

def plotMC(MC, returnP = False, show = True):
    if MC.shape == (2, 2):
        Pa = MC[0, 0] / (MC[0, 0] + MC[0, 1])
        Pfa = MC[1, 0] / (MC[1, 0] + MC[1, 1])

        plot(Pa = Pa, Pfa = Pfa,
             returnP = returnP, show = show)
    else:
        print("Confusion matrix is not 2x2.")

if __name__ == "__main__":
    #Probamos la funcion
    try:
        plt.close("all")
    except:
        pass

    P = plot(Pa = 0.7, Pfa = 0.1,
             returnP = True)

    import json
    print(json.dumps(P, indent = 4))