import sbibm
import torch
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import random

from sbibm.utils.io import get_tensor_from_csv, save_float_to_csv
from valdiags.plot_pdfs import multi_corner_plots

from utils import fwd_flow_transform_obs, inv_flow_transform_obs

import os

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        description="Run L-C2ST on sbi-benchmarking example."
    )
    parser.add_argument(
        "--experiment",
        "-e",
        type=str,
        default="multirun/2023-03-02/13-00-57",
        help='Experiment name: "multirun/yyyy-mm-dd/hh-mm-ss"',
    )
    parser.add_argument(
        "--num_observation",
        "-no",
        type=int,
        default=1,
        help="Observation number between 1 and 10.",
    )
    parser.add_argument(
        "--task", "-t", type=str, default="gaussian_mixture", help="sbibm task name"
    )

    parser.add_argument(
        "--cal_size",
        "-cs",
        type=int,
        default=10000,
        help="calbration set size used to train lc2st.",
    )
    parser.add_argument(
        "--test_size",
        "-ts",
        type=int,
        default=10000,
        help="number of base-dist samples for l-c2st evaluation.",
    )

    parser.add_argument(
        "--n_trials_null",
        "-nt",
        type=int,
        default=100,
        help="number of trials under the null hypothesis.",
    )
    parser.add_argument(
        "--n_ensemble",
        "-ne",
        type=int,
        default=10,
        help="number of models for ensemble-prediction.",
    )

    parser.add_argument(
        "--run_htest",
        "-r",
        action="store_true",
        help="Whether to run the lc2st or just generate plots.",
    )

    parser.add_argument(
        "--correct_posterior",
        "-c",
        action="store_true",
        help="Whether to run posterior corretion or not.",
    )

    parser.add_argument(
        "--n_rej_trials",
        "-nr",
        type=int,
        default=1,
        help="How many times to run rejection sampling.",
    )

    parser.add_argument(
        "--z_space",
        "-z",
        action="store_true",
        help="Whether to run lc2st in z-space or not.",
    )

    parser.add_argument(
        "--clf_name",
        "-cn",
        type=str,
        default="mlp_sbi",
        help="Which classifier to use: one of 'mlp_sbi', 'mlp_base'.",
    )

    args = parser.parse_args()

    # Seeding
    torch.manual_seed(1)
    random.seed(1)
    np.random.seed(1)

    PATH_EXPERIMENT = Path.cwd() / args.experiment / "lc2st_results"
    if not os.path.exists(PATH_EXPERIMENT):
        os.mkdir(PATH_EXPERIMENT)

    PATH_EXPERIMENT_OBS = PATH_EXPERIMENT / f"{args.num_observation - 1}"
    if not os.path.exists(PATH_EXPERIMENT_OBS):
        os.mkdir(PATH_EXPERIMENT_OBS)

    PATH_EXPERIMENT_REF = Path.cwd() / args.experiment / "0"

    task = sbibm.get_task(args.task)
    prior = task.get_prior()
    simulator = task.get_simulator()
    observation = task.get_observation(num_observation=args.num_observation)

    posterior_samples = task.get_reference_posterior_samples(args.num_observation)

    # loaf global objects that should stay the same over all observations
    # (this is for amortized algorithms)

    posterior_est = torch.load(PATH_EXPERIMENT_REF / "posterior_est.pkl")
    posterior_est.flow = posterior_est.flow.set_default_x(observation)

    inv_flow_samples_ref = inv_flow_transform_obs(
        posterior_samples, observation, posterior_est.flow
    )

    if os.path.exists(PATH_EXPERIMENT_OBS / "algorithm_posterior_samples.pkl"):
        algorithm_posterior_samples = torch.load(
            PATH_EXPERIMENT_OBS / "algorithm_posterior_samples.pkl"
        )
    else:
        algorithm_posterior_samples = fwd_flow_transform_obs(
            posterior_est.flow.net._distribution.sample(
                len(posterior_samples)
            ).detach(),
            observation,
            posterior_est.flow,
        )
        torch.save(
            algorithm_posterior_samples,
            PATH_EXPERIMENT_OBS / "algorithm_posterior_samples.pkl",
        )

    cal_set = torch.load(PATH_EXPERIMENT_REF / "calibration_dataset.pkl")
    x_cal = cal_set["x"][: args.cal_size]
    theta_cal = cal_set["theta"][: args.cal_size]
    dim = theta_cal.shape[-1]

    base_dist_samples_cal = torch.load(PATH_EXPERIMENT_REF / "base_dist_samples.pkl")[
        : args.cal_size
    ]

    if os.path.exists(PATH_EXPERIMENT / "algorithm_posterior_samples_cal.pkl"):
        algorithm_posterior_samples_cal = torch.load(
            PATH_EXPERIMENT / "algorithm_posterior_samples_cal.pkl"
        )
    else:
        algorithm_posterior_samples_cal = []
        for z, x in zip(base_dist_samples_cal, x_cal):
            z, x = z[None, :], x[None, :]
            # algorithm_posterior_samples_cal.append(est.sample(x=x).detach())
            # # this is not the same as transforming...
            algorithm_posterior_samples_cal.append(
                fwd_flow_transform_obs(z, x, posterior_est.flow)
            )
        algorithm_posterior_samples_cal = torch.stack(algorithm_posterior_samples_cal)[
            :, 0, :
        ]
        torch.save(
            algorithm_posterior_samples_cal,
            PATH_EXPERIMENT / "algorithm_posterior_samples_cal.pkl",
        )

    if os.path.exists(PATH_EXPERIMENT / "inv_flow_samples_cal.pkl"):
        inv_flow_samples_cal = torch.load(PATH_EXPERIMENT / "inv_flow_samples_cal.pkl")
    else:
        inv_flow_samples_cal = []
        for theta, x in zip(theta_cal, x_cal):
            theta, x = theta[None, :], x[None, :]
            inv_flow_samples_cal.append(
                inv_flow_transform_obs(theta, x, posterior_est.flow)
            )
        inv_flow_samples_cal = torch.stack(inv_flow_samples_cal)[:, 0, :]
        torch.save(inv_flow_samples_cal, PATH_EXPERIMENT / "inv_flow_samples_cal.pkl")

    if os.path.exists(
        PATH_EXPERIMENT / f"trained_clfs_z_{args.z_space}_{args.clf_name}.pkl"
    ):
        trained_clfs = torch.load(
            PATH_EXPERIMENT / f"trained_clfs_z_{args.z_space}_{args.clf_name}.pkl"
        )
    else:
        trained_clfs = []

    if args.z_space:
        path_probas_null = (
            Path.cwd()
            / "probas_null_lc2st_z"
            / f"{args.cal_size}"
            / f"{args.task}"
            / f"lc2st_probas_null_z_True_{args.clf_name}.pkl"
        )
    else:
        path_probas_null = (
            PATH_EXPERIMENT_OBS / f"lc2st_probas_null_z_False_{args.clf_name}.pkl"
        )

    if os.path.exists(path_probas_null):
        probas_null = torch.load(path_probas_null)
    else:
        probas_null = []

    # # =============== Reference plots ==================

    # # True conditional distributions: norm, inv-flow
    # from valdiags.localC2ST_old import flow_vs_reference_distribution

    # # embedding not intergrated in transform method (includes standardize)

    # if dim <= 2:
    #     flow_vs_reference_distribution(
    #         samples_ref=base_dist_samples_cal,
    #         samples_flow=inv_flow_samples_ref,
    #         z_space=True,
    #         dim=dim,
    #         hist=False,
    #     )
    #     plt.savefig(PATH_EXPERIMENT_OBS / "z_space_reference.pdf")
    #     plt.show()

    # samples_list = [base_dist_samples_cal, inv_flow_samples_ref]
    # legends = [
    #     r"Ref: $\mathcal{N}(0,1)$",
    #     r"NPE: $T_{\phi}^{-1}(\Theta;x_0) \mid x_0$",
    # ]
    # colors = ["blue", "orange"]
    # multi_corner_plots(
    #     samples_list,
    #     legends,
    #     colors,
    #     title=r"Base-Distribution vs. Inverse Flow-Transformation (of $\Theta \mid x_0$)",
    #     labels=[r"$z$" f"_{i+1}" for i in range(dim)],
    #     domain=(torch.ones(dim) * -5, torch.ones(dim) * 5),
    # )
    # plt.savefig(PATH_EXPERIMENT_OBS / "z_space_reference_corner.pdf")
    # plt.show()

    # # True vs. estimated posterior samples
    # if dim <= 2:
    #     flow_vs_reference_distribution(
    #         samples_ref=posterior_samples,
    #         samples_flow=algorithm_posterior_samples,
    #         z_space=False,
    #         dim=dim,
    #         hist=False,
    #     )
    #     plt.savefig(PATH_EXPERIMENT_OBS / "theta_space_reference.pdf")
    #     plt.show()

    # samples_list = [posterior_samples, algorithm_posterior_samples]
    # legends = [r"Ref: $p(\Theta \mid x_0)$", r"NPE: $p(T_{\phi}(Z;x_0))$"]
    # colors = ["blue", "orange"]
    # multi_corner_plots(
    #     samples_list,
    #     legends,
    #     colors,
    #     title=r"True vs. Estimated distributions at $x_0$",
    #     labels=[r"$\theta$" + f"_{i+1}" for i in range(dim)],
    #     # domain=(torch.tensor([-15, -15]), torch.tensor([5, 5])),
    # )
    # plt.savefig(PATH_EXPERIMENT_OBS / "theta_space_reference_corner.pdf")
    # plt.show()

    # plt.close("all")

    # =============== Hypothesis test and ensemble probas ================
    metrics = ["probas_mean", "w_dist", "TV"]

    if args.run_htest:
        print("Running L-C2ST...")
        from valdiags.localC2ST_old import lc2st_htest_sbibm

        test_size = args.test_size

        if args.clf_name == "mlp_base":
            clf_kwargs = {"alpha": 0, "max_iter": 25000}
        else:
            clf_kwargs = None

        if args.z_space:
            P = base_dist_samples_cal
            Q = inv_flow_samples_cal
            if os.path.exists(PATH_EXPERIMENT_OBS / f"P_eval_z_{args.z_space}.pkl"):
                P_eval = torch.load(
                    PATH_EXPERIMENT_OBS / f"P_eval_z_{args.z_space}.pkl"
                )
            else:
                P_eval = posterior_est.flow.net._distribution.sample(test_size).detach()
            null_dist = posterior_est.flow.net._distribution
            null_samples_list = None
        else:
            P = algorithm_posterior_samples_cal
            Q = theta_cal
            if os.path.exists(PATH_EXPERIMENT_OBS / f"P_eval_z_{args.z_space}.pkl"):
                P_eval = torch.load(
                    PATH_EXPERIMENT_OBS / f"P_eval_z_{args.z_space}.pkl"
                )
            else:
                P_eval = algorithm_posterior_samples[:test_size]
            null_dist = None

            null_samples_list = []
            for _ in range(args.n_trials_null):
                samples_z = posterior_est.flow.net._distribution.sample(
                    args.cal_size
                ).detach()
                null_samples_list.append(
                    fwd_flow_transform_obs(samples_z, observation, posterior_est.flow)
                )

        (
            p_values,
            test_stats,
            probas_ens,
            probas_null,
            t_stats_null,
            trained_clfs,
            run_time,
        ) = lc2st_htest_sbibm(
            P,
            Q,
            cal_set["x"],
            P_eval=P_eval,
            x_eval=observation,
            null_dist=null_dist,
            null_samples_list=null_samples_list,
            test_stats=metrics,
            n_trials_null=args.n_trials_null,
            n_ensemble=args.n_ensemble,
            clf_kwargs=None,
            probas_null=probas_null,
            trained_clfs=trained_clfs,
        )
        lc2st_htest_results = {
            "p_values": p_values,
            "test_stats": test_stats,
            "t_stats_null": t_stats_null,
        }
        torch.save(
            lc2st_htest_results,
            PATH_EXPERIMENT_OBS
            / f"lc2st_htest_results_z_{args.z_space}_{args.clf_name}.pkl",
        )
        torch.save(
            probas_ens,
            PATH_EXPERIMENT_OBS
            / f"lc2st_probas_ensemble_z_{args.z_space}_{args.clf_name}.pkl",
        )

        torch.save(P_eval, PATH_EXPERIMENT_OBS / f"P_eval_z_{args.z_space}.pkl")

        torch.save(probas_null, path_probas_null)
        torch.save(
            trained_clfs,
            PATH_EXPERIMENT / f"trained_clfs_z_{args.z_space}_{args.clf_name}.pkl",
        )

        if run_time != 0:
            save_float_to_csv(
                PATH_EXPERIMENT
                / f"run_time_1_clf_z_{args.z_space}_{args.clf_name}.csv",
                run_time,
            )

        print("Finished running L-C2ST.")

    P_eval = torch.load(PATH_EXPERIMENT_OBS / f"P_eval_z_{args.z_space}.pkl")
    probas_ens = torch.load(
        PATH_EXPERIMENT_OBS
        / f"lc2st_probas_ensemble_z_{args.z_space}_{args.clf_name}.pkl"
    )

    probas_null = torch.load(path_probas_null)

    lc2st_htest_results = torch.load(
        PATH_EXPERIMENT_OBS
        / f"lc2st_htest_results_z_{args.z_space}_{args.clf_name}.pkl"
    )
    test_stats = lc2st_htest_results["test_stats"]
    t_stats_null = lc2st_htest_results["t_stats_null"]
    p_values = lc2st_htest_results["p_values"]

    # # =============== Result plots ===============

    from valdiags.localC2ST_old import box_plot_lc2st

    for m in metrics:
        box_plot_lc2st(
            [test_stats[m]], t_stats_null[m], labels=["NPE"], colors=["red"],
        )
        plt.savefig(
            PATH_EXPERIMENT_OBS
            / f"lc2st_htest_box_plot_metric_{m}_z_{args.z_space}_{args.clf_name}.pdf"
        )
        plt.show()
        print(f"pvalues {m}:", p_values[m])

    from valdiags.localC2ST_old import pp_plot_lc2st

    pp_plot_lc2st([probas_ens], probas_null, labels=["NPE"], colors=["red"])
    plt.savefig(
        PATH_EXPERIMENT_OBS / f"lc2st_pp_plot_z_{args.z_space}_{args.clf_name}.pdf"
    )
    plt.show()

    # =============== Interpretability plots ==================
    if dim <= 2:
        # High / Low probability regions
        from valdiags.localC2ST_old import eval_space_with_proba_intensity

        eval_space_with_proba_intensity(
            probas_ens, probas_null, P_eval, dim=dim, z_space=args.z_space
        )
        plt.savefig(
            PATH_EXPERIMENT_OBS
            / f"eval_with_lc2st_proba_intensity_z_{args.z_space}_{args.clf_name}.pdf"
        )
        plt.show()

    # =============== Posterior correction ==================

    if args.correct_posterior:
        from valdiags.posterior_correction import (
            rejection_sampling,
            clf_ratio_obs,
            corrected_pdf,
            flow_sampler,
            uniform_sampler,
        )
        from functools import partial

        flow_transform = partial(
            fwd_flow_transform_obs, observation=observation, flow=posterior_est.flow
        )

        inv_flow_transform = partial(
            inv_flow_transform_obs, observation=observation, flow=posterior_est.flow
        )

        base_dist_sampler = posterior_est.flow.net._distribution.sample

        if args.z_space:

            ## ====== BASE DISTRIBUTION SPACE =======

            # proposal = base_dist = normal
            proposal_sampler = base_dist_sampler
            # f = ratio
            f = partial(clf_ratio_obs, x_obs=observation, clfs=trained_clfs)

            # sample from normal_pdf * ratio = p(z|x_obs)
            acc_rej_samples_normal = []
            for _ in range(args.n_rej_trials):
                acc_rej_samples_normal.append(
                    rejection_sampling(proposal_sampler=proposal_sampler, f=f)
                )

            samples_list = [P_eval.numpy()] + acc_rej_samples_normal
            legends = ["normal: Z"] + [
                "corrected normal acc_rej_normal"
            ] * args.n_rej_trials
            colors = ["grey"] + ["red"] * args.n_rej_trials

            if dim < 3:
                # proposal = uniform
                # --> very long in high dimensions
                proposal_sampler = partial(uniform_sampler, dim=dim)
                # f = ratio * normal_pdf = p(z|x_obs)
                f = partial(
                    corrected_pdf,
                    dist=posterior_est.flow.net._distribution,
                    x_obs=observation,
                    clfs=trained_clfs,
                )
                # sample from ratio * normal_pdf \approx p(z|x_obs)
                acc_rej_samples_uniform = []
                for _ in range(args.n_rej_trials):
                    acc_rej_samples_uniform.append(
                        rejection_sampling(proposal_sampler=proposal_sampler, f=f)
                    )

                samples_list += acc_rej_samples_uniform
                legends += ["corrected normal acc_rej_uniform"] * args.n_rej_trials
                colors += ["green"] * args.n_rej_trials

            samples_list += [inv_flow_samples_ref.numpy()]
            legends += [r"Reference: $T_{\phi}^{-1}(\Theta, x_0)\mid x_0$"]
            colors += ["orange"]
            multi_corner_plots(
                samples_list,
                legends,
                colors,
                title="correction in z-space",
                labels=[r"$z$" + f"_{i+1}" for i in range(dim)],
                domain=(torch.ones(dim) * -5, torch.ones(dim) * 5),
            )
            plt.savefig(
                PATH_EXPERIMENT_OBS / f"z_space_correction_corner_{args.clf_name}.pdf"
            )
            plt.show()

            ## ======== PARAMETER SPACE =========
            flow_thetas = flow_sampler(
                args.test_size, base_dist_sampler, flow_transform
            )

            # proposal = flow: samples = \Theta
            proposal_sampler = partial(
                flow_sampler,
                base_dist_sampler=base_dist_sampler,
                flow_transform=flow_transform,
            )
            # f = ratio r(T_{\phi}^{-1}(\Theta))
            f = partial(
                clf_ratio_obs,
                x_obs=observation,
                clfs=trained_clfs,
                inv_flow_transform=inv_flow_transform,
            )
            # sample from ratio * flow \approx p(\theta | x_obs)
            acc_rej_samples_flow = []
            for _ in range(args.n_rej_trials):
                acc_rej_samples_flow.append(
                    rejection_sampling(proposal_sampler=proposal_sampler, f=f)
                )
            torch.save(
                acc_rej_samples_flow,
                PATH_EXPERIMENT_OBS
                / f"acc_rej_samples_flow_z_{args.z_space}_{args.clf_name}.pkl",
            )

            # acc_rej_samples_approx_flow = []
            # for _ in range(args.n_rej_trials):
            #     acc_rej_samples_approx_flow.append(
            #         rejection_sampling(
            #             proposal_sampler=proposal_sampler, f=f, approximate=True
            #         )
            #     )

            samples_list = (
                [flow_thetas.detach().numpy()]
                + acc_rej_samples_flow
                # + acc_rej_samples_approx_flow
            )

            legends = (
                [r"NPE: $T_{\phi}(Z;x_0)$"]
                + ["corrected NPE acc_rej_flow"] * args.n_rej_trials
                # + ["corrected NPE acc_rej_approx_flow"] * args.n_rej_trials
            )

            colors = (
                ["orange"]
                + ["red"] * args.n_rej_trials
                # + ["purple"] * args.n_rej_trials
            )

            if dim < 3:
                # --> still very long in high dimensions
                # proposal = uniform, samples = T_phi(U)
                # --> not U to get into the right range of flow-pdf-values
                proposal_sampler = partial(
                    uniform_sampler, dim=dim, flow_transform=flow_transform,
                )
                # for T_phi(u) the proposal distribution is p(u)|detJT_phi^-1|, we only need q_phi/det = N(u)
                f = partial(
                    corrected_pdf,
                    # dist=posterior_est.flow,
                    dist=posterior_est.flow.net._distribution,
                    x_obs=observation,
                    clfs=trained_clfs,
                    inv_flow_transform=inv_flow_transform,
                )

                # from valdiags.posterior_correction import corrected_pdf_old
                # # f = r(u)*q_phi(T_phi(u))
                # # --> doing f = r(T_phi^{-1}(u)*q_phi(u) takes too much time, doesn't converge...
                # f = partial(
                #     corrected_pdf_old,
                #     dist=posterior_est.flow,
                #     x_obs=observation,
                #     clfs=trained_clfs,
                #     inv_flow_transform=inv_flow_transform,
                # )

                # sample from ratio * flow \approx p(\theta | x_obs)
                acc_rej_samples_uniform = []
                for _ in range(args.n_rej_trials):
                    acc_rej_samples_uniform.append(
                        rejection_sampling(proposal_sampler=proposal_sampler, f=f)
                    )

                samples_list += acc_rej_samples_uniform

                legends += ["corrected NPE acc_rej_uniform"] * args.n_rej_trials

                colors += ["green"] * args.n_rej_trials

            samples_list += [posterior_samples]
            legends += [r"Reference: $\Theta \mid x_0$"]
            colors += ["blue"]

            multi_corner_plots(
                samples_list,
                legends,
                colors,
                title="correction in z-space",
                labels=[r"$\theta$" + f"_{i+1}" for i in range(dim)]
                # domain=(torch.tensor([-15, -15]), torch.tensor([5, 5])),
            )
            plt.savefig(
                PATH_EXPERIMENT_OBS
                / f"posterior_correction_z_space_corner_{args.clf_name}.pdf"
            )
            plt.show()

        else:

            # proposal = flow: samples = \Theta
            proposal_sampler = partial(
                flow_sampler,
                base_dist_sampler=base_dist_sampler,
                flow_transform=flow_transform,
            )
            # f = ratio r(\Theta)
            f = partial(clf_ratio_obs, x_obs=observation, clfs=trained_clfs,)
            # sample from ratio * flow_pdf \approx p(\theta | x_obs)
            acc_rej_samples_flow = []
            for _ in range(args.n_rej_trials):
                acc_rej_samples_flow.append(
                    rejection_sampling(proposal_sampler=proposal_sampler, f=f)
                )

            torch.save(
                acc_rej_samples_flow,
                PATH_EXPERIMENT_OBS
                / f"acc_rej_samples_flow_z_{args.z_space}_{args.clf_name}.pkl",
            )

            # acc_rej_samples_approx_flow = []
            # for _ in range(args.n_rej_trials):
            #     acc_rej_samples_approx_flow.append(
            #         rejection_sampling(
            #             proposal_sampler=proposal_sampler, f=f, approximate=True
            #         )
            #     )

            samples_list = (
                [P_eval.numpy()]
                + acc_rej_samples_flow
                # + acc_rej_samples_approx_flow
            )

            legends = (
                [r"NPE: $\Theta \sim q_{\phi}(\theta \mid x_0)$"]
                + ["corrected NPE acc_rej_flow"] * args.n_rej_trials
                # + ["corrected NPE acc_rej_approx_flow"] * args.n_rej_trials
            )

            colors = (
                ["orange"]
                + ["red"] * args.n_rej_trials
                # + ["purple"] * args.n_rej_trials
            )

            if dim < 3:
                # --> still very long in high dimensions
                # proposal = uniform, samples = T_phi(U)
                # --> not U to get into the right range of flow-pdf-values
                # --> still very long in high dimensions
                proposal_sampler = partial(
                    uniform_sampler, dim=dim, flow_transform=flow_transform,
                )
                # for T_phi(u) the proposal distribution is p(u)|detJT_phi^-1|, we only need q_phi/det = N(u)
                f = partial(
                    corrected_pdf,
                    # dist=posterior_est.flow,
                    dist=posterior_est.flow.net._distribution,
                    x_obs=observation,
                    clfs=trained_clfs,
                    inv_flow_transform=inv_flow_transform,
                    z_space=False,
                )

                # from valdiags.posterior_correction import corrected_pdf_old
                # # f = r(u)*q_phi(T_phi(u))
                # # --> doing f = r(T_phi^{-1}(u)*q_phi(u) takes too much time, doesn't converge...
                # f = partial(
                #     corrected_pdf_old,
                #     dist=posterior_est.flow,
                #     x_obs=observation,
                #     clfs=trained_clfs,
                # )

                # sample from ratio * flow \approx p(\theta | x_obs)
                acc_rej_samples_uniform = []
                for _ in range(args.n_rej_trials):
                    acc_rej_samples_uniform.append(
                        rejection_sampling(proposal_sampler=proposal_sampler, f=f)
                    )

                samples_list += acc_rej_samples_uniform

                legends += ["corrected NPE acc_rej_uniform"] * args.n_rej_trials

                colors += ["green"] * args.n_rej_trials

            samples_list += [posterior_samples]
            legends += [r"Reference: $\Theta \sim p(\theta \mid x_0)$"]
            colors += ["blue"]

            multi_corner_plots(
                samples_list,
                legends,
                colors=colors,
                title="correction in parameter space",
                labels=[r"$\theta$" + f"_{i+1}" for i in range(dim)]
                # domain=(torch.tensor([-5, -5]), torch.tensor([5, 5])),
            )
            plt.savefig(
                PATH_EXPERIMENT_OBS
                / f"posterior_correction_theta_space_corner_{args.clf_name}.pdf"
            )
            plt.show()


import numpy as np
import matplotlib.pyplot as plt
import quadtratic as quad
import t3utility as t3

# class 1 points generation from dataset
filepath1 = t3.currentFilePath('training_class1.csv')
trainingset_class1 = np.genfromtxt(filepath1, dtype=float, delimiter=',')
trainingset_class1 = t3.getSelectedColumns(trainingset_class1, (2,10,11,14,17,18))
np.savetxt("class1.csv",trainingset_class1,delimiter=",")

class1_points,class1_col = trainingset_class1.shape

# class 2 points generation from dataset
filepath2 = t3.currentFilePath('training_class2.csv')
trainingset_class2 = np.genfromtxt(filepath2, dtype=float, delimiter=',')
trainingset_class2 = t3.getSelectedColumns(trainingset_class2, (2,10,11,14,17,18))
np.savetxt("class2.csv",trainingset_class2,delimiter=",")

class2_points,class2_col = trainingset_class2.shape

# estimated mean for ML
est_MLmean_class1 = quad.estimateMean_for_ML(trainingset_class1,class1_points)
est_MLmean_class2 = quad.estimateMean_for_ML(trainingset_class2,class2_points)
print("This is estimated(ML) Mean for class1 and class2:", est_MLmean_class1,est_MLmean_class2)

# estimated covariance for ML
est_MLcovariance_class1 = quad.estimateCovariance_forML(trainingset_class1,class1_points,est_MLmean_class1)
est_MLcovariance_class2 = quad.estimateCovariance_forML(trainingset_class2,class2_points,est_MLmean_class2)
print("This is estimated(ML) covariances for class1 and class2:", est_MLcovariance_class1,est_MLcovariance_class2)

X1 = trainingset_class1
X2 = trainingset_class2
sigma1 = est_MLcovariance_class1
sigma2 = est_MLcovariance_class2
mean_X1 = est_MLmean_class1
mean2_X2 = est_MLmean_class2

# generation of Y1 and Y2
def generation_of_Y1():
    lambda_def, Px1 = (np.linalg.eig(sigma1))  # 1st- eigval, 2nd - eigvec
    Px1_transpose = np.transpose(Px1)

    y1 = np.dot(Px1_transpose,X1.transpose())
    y1 = y1.transpose()

    # mean of Y1
    mean_of_y1 = np.dot(Px1_transpose,mean_X1)
    print("Mean of Y1:", mean_of_y1)

    # covariance of Y1
    covariance_of_Y1 = np.dot(np.dot(Px1_transpose,sigma1),Px1)
    print("Covariance of Y1:", covariance_of_Y1)

    # eigenvalues and eigenvectors of Y1

    eigenvalue_Y1, eigenvector_Y1 = np.linalg.eig(covariance_of_Y1)
    print("eigen values of  Y1:", eigenvalue_Y1)
    print("eigen vectors of Y1:", eigenvector_Y1)
    return y1

def generation_of_Y2():

    lambda_def, Px1 = (np.linalg.eig(sigma1))  # 1st- eigval, 2nd - eigvec
    Px1_transpose = np.transpose(Px1)

    y2 = np.dot(Px1_transpose, X2.transpose())
    y2 = y2.transpose()

    # mean of Y2
    mean_of_y2 = np.dot(Px1_transpose,mean2_X2)
    print("Mean of Y2:", mean_of_y2)

    covariance_of_Y2 = np.dot(np.dot(Px1_transpose,sigma2),Px1)
    print("Covariance of Y2:", covariance_of_Y2)

    # eigenvalues and eigenvectors of Y2

    eigenvalue_Y2, eigenvector_Y2 = np.linalg.eig(covariance_of_Y2)
    print("eigen values of  Y2:", eigenvalue_Y2)
    print("eigen vectors of Y2:", eigenvector_Y2)
    return y2

Y1 = generation_of_Y1()
Y2 = generation_of_Y2()

# generation of Z1 and Z2
def generation_of_Z1():
    lambda_def, Px1 = (np.linalg.eig(sigma1)) # 1st- eigval, 2nd - eigvec
    tempZ_1 = np.power(lambda_def, -0.5)
    tempZ_1 = np.diag(tempZ_1)
    Px1_transpose = np.transpose(Px1)

    z1 = np.dot(np.dot(tempZ_1,Px1_transpose),X1.transpose())
    z1 = z1.transpose()

    # mean of Z1

    mean_of_Z1 = np.dot(np.dot(tempZ_1,Px1_transpose),mean_X1)
    print("Mean of Z1:", mean_of_Z1)

    # covariance of Z1
    covariance_of_Z1 = np.identity(6)
    print ("This is covariance of Z1:", covariance_of_Z1)

    # eigenvalues and eigenvectors of Z1

    eigenvalue_Z1, eigenvector_Z1 = np.linalg.eig(covariance_of_Z1)
    print("eigen values of  Z1:", eigenvalue_Z1)
    print("eigen vectors of Z1:", eigenvector_Z1)
    return  z1

def generation_of_Z2():
    lambda_def, Px1 = (np.linalg.eig(sigma1))  # 1st- eigval, 2nd - eigvec
    tempZ_1 = np.power(lambda_def, -0.5)
    tempZ_1 = np.diag(tempZ_1)
    Px1_transpose = np.transpose(Px1)

    z2 = np.dot(np.dot(tempZ_1, Px1_transpose), X2.transpose())
    z2 = z2.transpose()

    # mean of Z2
    mean_Z2 = np.dot(np.dot(tempZ_1,Px1_transpose),mean2_X2)
    print("Mean of Z2:", mean_Z2)

    # covariance of Z2
    covariance_of_Z2 = np.dot(np.dot(np.dot(np.dot(tempZ_1,Px1_transpose),sigma2),Px1), tempZ_1)
    print("covariance of Z2:", covariance_of_Z2)

    # eigenvalues and eigenvectors of Z2
    eigenvalue_Z2, eigenvector_Z2 = np.linalg.eig(covariance_of_Z2)
    print("eigen values of Z2:", eigenvalue_Z2)
    print("eigen vectors of Z2:", eigenvector_Z2)
    return z2

Z1 = generation_of_Z1()
Z2 = generation_of_Z2()
# generation of V1 and V2

def generation_of_V1():
    lambda_def, Px1 = (np.linalg.eig(sigma1))  # 1st- eigval, 2nd - eigvec
    tempZ_1 = np.power(lambda_def, -0.5)
    tempZ_1 = np.diag(tempZ_1)
    Px1_transpose = np.transpose(Px1)

    # covariance of Z2
    covariance_of_Z2 = np.dot(np.dot(np.dot(np.dot(tempZ_1, Px1_transpose), sigma2), Px1), tempZ_1)
    eigenvalue,Pz2 = (np.linalg.eig(covariance_of_Z2))
    Pz2_transpose = Pz2.transpose()

    v1 = np.dot(Pz2_transpose,Z1.transpose())
    v1 = v1.transpose()

    # mean and covariance of V1
    mean_of_Z1 = np.dot(np.dot(tempZ_1, Px1_transpose), mean_X1)
    mean_V1 = np.dot(Pz2_transpose,mean_of_Z1)
    print ("The mean of V1:", mean_V1)

    I = np.identity(6)

    covariance_of_V1 = np.dot(np.dot(Pz2_transpose,I),Pz2)
    print ("The covariance of V1:", covariance_of_V1)
    return mean_V1,covariance_of_V1,v1

def generation_of_V2():
    lambda_def, Px1 = (np.linalg.eig(sigma1))  # 1st- eigval, 2nd - eigvec
    tempZ_1 = np.power(lambda_def, -0.5)
    tempZ_1 = np.diag(tempZ_1)
    Px1_transpose = np.transpose(Px1)

    # covariance of Z2
    covariance_of_Z2 = np.dot(np.dot(np.dot(np.dot(tempZ_1, Px1_transpose), sigma2), Px1), tempZ_1)
    som, Pz2 = (np.linalg.eig(covariance_of_Z2))
    Pz2_transpose = Pz2.transpose()

    v2 = np.dot(Pz2_transpose, Z2.transpose())
    v2 = v2.transpose()

    # mean and covariance of V1
    mean_of_Z2 = np.dot(np.dot(tempZ_1, Px1_transpose), mean2_X2)
    mean_V2 = np.dot(Pz2_transpose, mean_of_Z2)
    print ("The Mean of V2:", mean_V2)

    covariance_of_V2 = np.dot(np.dot(Pz2_transpose, covariance_of_Z2), Pz2)
    print ("The covariance of V2:", covariance_of_V2)

    return mean_V2,covariance_of_V2,v2

m_v1,c_v1,V1 = generation_of_V1()
m_v2,c_v2,V2 = generation_of_V2()

def slicing_points(X1, X2):
    X1_0_1 = X1[:,[0,1]] # X1 - X2
    X1_0_2 = X1[:,[0,2]] # X1 - X3

    X2_0_1 = X2[:,[0,1]]
    X2_0_2 = X2[:,[0,2]]

    return X1_0_1,X1_0_2,X2_0_1,X2_0_2

def generate_plot():
    X1_0_1,X1_0_2,X2_0_1,X2_0_2 = slicing_points(X1,X2)

    # To plot X1_X2 domain of X
    a = np.amax(X1)
    b = np.amin(X1)
    c = np.amax(X2)
    d = np.amin(X2)

    plt.axis([d-5,c+5,b,a])

    plt.scatter(X1_0_1[:, [0]], X1_0_1[:, [1]], c='red')
    plt.axis([d, c, b, a])
    plt.scatter(X2_0_1[:, [0]], X2_0_1[:, [1]], c='blue')
    plt.text(d+2,a-2,'red= first, blue = second')
    plt.xlabel("X1")
    plt.ylabel("X2")
    plt.title("X1 - X2 domain")
    plt.show()

    plt.axis([d-3, c+3, b-2, a])

    plt.scatter(X1_0_2[:, 0], X1_0_2[:, 1], c='red')
    plt.scatter(X2_0_2[:, 0], X2_0_2[:, 1], c='blue')
    plt.text(d + 2, a - 2, 'red= first, blue = second')
    plt.xlabel("X1")
    plt.ylabel("X3")
    plt.title("X1 - X3 domain")
    plt.show()

    a = np.amax(Y1)
    b = np.amin(Y1)
    c = np.amax(Y2)
    d = np.amin(Y2)

    # To plot Y1_Y2 domain of Y
    Y1_0_1, Y1_0_2, Y2_0_1, Y2_0_2 = slicing_points(Y1, Y2)

    plt.axis([d - 3, c + 8, b - 6, a + 1])

    plt.scatter(Y1_0_1[:, [0]], Y1_0_1[:, [1]], c='red')
    plt.scatter(Y2_0_1[:, [0]], Y2_0_1[:, [1]], c='blue')
    plt.text(d + 2, a - 2, 'red= first, blue = second')
    plt.xlabel("Y1")
    plt.ylabel("Y2")
    plt.title("Y1 - Y2 domain")
    plt.show()

    plt.axis([d - 3, c + 9, b -2 , a])

    plt.scatter(Y1_0_2[:, 0], Y1_0_2[:, 1], c='red')
    plt.scatter(Y2_0_2[:, 0], Y2_0_2[:, 1], c='blue')
    plt.text(d + 2, a - 2, 'red= first, blue = second')
    plt.xlabel("Y1")
    plt.ylabel("Y3")
    plt.title("Y1 - Y3 domain")
    plt.show()

    a = np.amax(Z1)
    b = np.amin(Z1)
    c = np.amax(Z2)
    d = np.amin(Z2)

    Z1_0_1, Z1_0_2, Z2_0_1, Z2_0_2 = slicing_points(Z1, Z2)

    # To plot Z1_Z2 domain of Z

    plt.axis([d-3, c+3, b-5, a+2])

    plt.scatter(Z1_0_1[:, [0]], Z1_0_1[:, [1]], c='red')
    plt.scatter(Z2_0_1[:, [0]], Z2_0_1[:, [1]], c='green')
    plt.text(d + 0.1, a - 0.1, 'red= first, green = second')
    plt.xlabel("Z1")
    plt.ylabel("Z2")
    plt.title("Z1 - Z2 domain")
    plt.show()

    plt.axis([d , c , b-3, a])

    plt.scatter(Z1_0_2[:, 0], Z1_0_2[:, 1], c='red')
    plt.scatter(Z2_0_2[:, 0], Z2_0_2[:, 1], c='green')
    plt.text(d + 0.6, a - 0.6, 'red= first, green = second')
    plt.xlabel("Z1")
    plt.ylabel("Z3")
    plt.title("Z1 - Z3 domain")
    plt.show()

    a = np.amax(V1)
    b = np.amin(V1)
    c = np.amax(V2)
    d = np.amin(V2)

    V1_0_1, V1_0_2, V2_0_1, V2_0_2 = slicing_points(V1, V2)

    # To plot V1_V2 domain of V

    plt.axis([d-3 , c+0.5 , b-2, a])

    plt.scatter(V1_0_1[:, [0]], V1_0_1[:, [1]], c='red')
    plt.scatter(V2_0_1[:, [0]], V2_0_1[:, [1]], c='green')
    plt.text(d + 2, a - 0.5, 'red= first, green = second')
    plt.xlabel("V1")
    plt.ylabel("V2")
    plt.title("V1 - V2 domain")
    plt.show()

    plt.axis([d-2 , c+2 , b-3, a+3])

    plt.scatter(V1_0_2[:, 0], V1_0_2[:, 1], c='red')
    plt.scatter(V2_0_2[:, 0], V2_0_2[:, 1], c='green')
    plt.text(d + 0.2, a+2, 'red= first, green = second')
    plt.xlabel("V1")
    plt.ylabel("V3")
    plt.title("V1 - V3 domain")
    plt.show()

def generate_POverall():
    sigma1 = generate_sigma1()  # cov of x1
    lambda_def, Px1 = (np.linalg.eig(sigma1))  # 1st- eigval, 2nd - eigvec
    tempZ_1 = np.power(lambda_def, -0.5)
    tempZ_1 = np.diag(tempZ_1)
    Px1_transpose = np.transpose(Px1)

    # covariance of Z2

    sigma2 = generate_sigma2()
    covariance_of_Z2 = np.dot(np.dot(np.dot(np.dot(tempZ_1, Px1_transpose), sigma2), Px1), tempZ_1)
    eigenvalue,Pz2 = (np.linalg.eig(covariance_of_Z2))
    Pz2_transpose = Pz2.transpose()

    P_overall = np.dot(np.dot(Pz2_transpose,tempZ_1),Px1_transpose)

    return P_overall


import hikari
import lightbulb
from datetime import timedelta, timezone, datetime

mod = lightbulb.Plugin("mod")

@mod.command
@lightbulb.add_checks(
    lightbulb.has_guild_permissions(hikari.Permissions.MANAGE_MESSAGES),
    lightbulb.bot_has_guild_permissions(hikari.Permissions.MANAGE_MESSAGES)
)
@lightbulb.option("amount", "How many messages should I delete?", type = int, required = True)
@lightbulb.command("clear", "Delete some messages")
@lightbulb.implements(lightbulb.SlashCommand)
async def clear(ctx: lightbulb.Context):
    amount = ctx.options.amount
    channel = ctx.channel_id
    msgs = await ctx.bot.rest.fetch_messages(channel).limit(amount)
    await ctx.bot.rest.delete_messages(channel, msgs)

@mod.command
@lightbulb.add_checks(
    lightbulb.has_guild_permissions(hikari.Permissions.KICK_MEMBERS),
    lightbulb.bot_has_guild_permissions(hikari.Permissions.KICK_MEMBERS)
)
@lightbulb.option("user", "Who should I kick?", type = hikari.Member, required = True)
@lightbulb.option("reason", "Why would you like to kick this user?", type = str, required = False)
@lightbulb.command("kick", "Kick an user!")
@lightbulb.implements(lightbulb.SlashCommand)
async def kick(ctx: lightbulb.Context):
    embed = hikari.Embed(title = "Success!")
    embed.add_field(name = f"You kicked {ctx.options.user}", value = f"Reason: {ctx.options.reason}")
    embed.set_footer(text = f"Requested by: {ctx.author}", icon = ctx.author.avatar_url)
    gid = ctx.guild_id
    await ctx.bot.rest.kick_member(user = ctx.options.user, guild = gid)
    await ctx.respond(embed=embed)

@mod.command
@lightbulb.add_checks(
    lightbulb.has_guild_permissions(hikari.Permissions.BAN_MEMBERS),
    lightbulb.bot_has_guild_permissions(hikari.Permissions.BAN_MEMBERS)
)
@lightbulb.option("user", "Who should I ban?", type = hikari.Member, required = True)
@lightbulb.option("reason", "Why would you like to ban this user?", type = str, required = False)
@lightbulb.command("ban", "Ban an user!")
@lightbulb.implements(lightbulb.SlashCommand)
async def ban(ctx: lightbulb.Context):
    embed = hikari.Embed(title = "Ban successful!")
    embed.add_field(name = f"You banned {ctx.options.user}", value = f"Reason: {ctx.options.reason}")
    embed.set_footer(text = f"Requested by: {ctx.author}", icon = ctx.author.avatar_url)
    gid = ctx.guild_id
    await ctx.bot.rest.ban_member(user = ctx.options.user, guild = gid)
    await ctx.respond(embed=embed)

@mod.command
@lightbulb.add_checks(
    lightbulb.has_guild_permissions(hikari.Permissions.MODERATE_MEMBERS),
    lightbulb.bot_has_guild_permissions(hikari.Permissions.MODERATE_MEMBERS)
)
@lightbulb.option("days", "Duration of the timeout in days", type=int, default=0)
@lightbulb.option("hours", "Duration of the timeout in hours", type=int, default=0)
@lightbulb.option("minutes", "Duration of the timeout in minutes", type=int, default=0)
@lightbulb.option("seconds", "Duration of the timeout in seconds", type=int, default=0)
@lightbulb.option("member", "Which user should recieve the timeout?", type=hikari.Member, required=True)
@lightbulb.command("timeout", "Timeout a member!")
@lightbulb.implements(lightbulb.SlashCommand)
async def timeout(ctx: lightbulb.Context):
    member = ctx.options.member
    now = datetime.now(timezone.utc)
    then = now + timedelta(days=ctx.options.days, hours=ctx.options.hours, minutes=ctx.options.minutes, seconds=ctx.options.seconds)
    if (then - now).days > 28:
        await ctx.respond("You can't time someone out for more than 28 days", flags=hikari.MessageFlag.EPHEMERAL)
        return
    elif (then - now).days < 28:
        embed = hikari.Embed(title="Timeout successful!", description=f"{ctx.options.member} is now muted for {ctx.options.days} days, {ctx.options.minutes} minutes and {ctx.options.seconds} seconds")
        embed.set_footer(text = f"Requested by: {ctx.author}", icon = ctx.author.avatar_url)
        await ctx.bot.rest.edit_member(user=member, guild=ctx.guild_id, communication_disabled_until=then)
        await ctx.respond(embed=embed)

def load(bot):
    bot.add_plugin(mod)


import numpy as np
import matplotlib.pyplot as plt
from itertools import product

n_items = 3
tau = .5
kappa = .5  # self inhibition
lambda_param = 8.  # lateral inhibition
threshold = 1.
noise_std = 0

n_time_steps = 50

L = np.ones((n_items, n_items)) - np.eye(n_items)

def accumulate_one_step(old_x, input_strength, params):
    coef_var = np.std(input_strength) / input_strength.mean()
    noise = np.random.normal(0, params['noise_std'], params['n_items'])
    x_new = (1 - params['tau'] * params['kappa'] - params['tau'] * params['lambda'] * L).T @ old_x \
            + coef_var * params['tau'] * input_strength + noise
    return np.maximum(x_new, 0)

def run_accumulator(n_steps, parameters, distractor_activation=.2):
    # initialize x
    x = np.zeros(parameters['n_items'])
    x_in_time = np.empty((n_steps + 1, n_items))
    x_in_time[0] = x
    for t in range(1, n_steps + 1):
        f_strength = np.eye(parameters['n_items'])[0] + (np.ones(parameters['n_items']) -
                                                         np.eye(parameters['n_items'])[0]) * distractor_activation  # note these inputs!
        x = accumulate_one_step(x, f_strength, parameters)
        x_in_time[t, :] = x
    return x_in_time

if __name__ == '__main__':

    params = {
        'n_items': n_items,
        'kappa': kappa,
        'lambda': lambda_param,
        'tau': tau,
        'threshold': threshold,
        'noise_std': noise_std
    }

    n_rows = 5
    n_cols = 5

    k_min, k_max = (.42, 1.4)
    l_min, l_max = (.53, 1.8)

    fig, axes = plt.subplots(n_rows, n_cols, figsize=(9, 9), sharex=True, sharey=True)

    kappa_range = np.linspace(k_min, k_max, n_rows)
    lambda_range = np.linspace(l_min, l_max, n_cols)

    i = 0
    for k, l in product(kappa_range, lambda_range):
        params['kappa'] = k
        params['lambda'] = l
        all_x = run_accumulator(n_time_steps, params)
        row, col = np.unravel_index(i, (n_rows, n_cols))
        axes[row, col].plot(all_x)
        axes[row, col].set_ylim([0, 1.2])
        plt.sca(axes[row, col])
        plt.axhline(params['threshold'], color='k', linestyle='--')

        axes[row, col].set_ylabel('kappa: {:.2f}'.format(k))
        axes[row, col].set_xlabel('lambda: {:.2f}'.format(l))

        i += 1

    plt.tight_layout()
    plt.savefig('accumulator_with_different_params.png')
    plt.show()

    # Now check for a requirement that one but only one accumulator reaches threshold
    k_min, k_max = (.1, 4.)
    l_min, l_max = (.1, 1.5)

    n_rows = 40
    n_cols = 40
    n_z = 11

    kappa_range = np.linspace(k_min, k_max, n_rows)
    lambda_range = np.linspace(l_min, l_max, n_cols)
    distractor_range = np.linspace(0., 1., n_z)

    param_grid = np.zeros((n_rows, n_cols, n_z), dtype=bool)
    i = 0
    for k, l, dist in product(kappa_range, lambda_range, distractor_range):

        params['kappa'] = k
        params['lambda'] = l

        all_x = run_accumulator(n_time_steps, params, distractor_activation=dist)
        idx = np.unravel_index(i, (n_rows, n_cols, n_z))

        max_accum_value = all_x.max(axis=0)  # get max for each item
        good_param = sum(max_accum_value >= params['threshold']) == 1
        param_grid[idx] = good_param
        i += 1

    # the code below plots a matrix of true and false values where true means only one item crossed
    # the threshold. This means that you can see the parameter values for which this is true.
    # i also plotted that for 10 different values for the "distractor" strength.
    xx, yy = np.meshgrid(kappa_range, lambda_range)
    for i, d in enumerate(distractor_range):
        plt.figure()
        plt.pcolormesh(xx, yy, param_grid[:, :, i])
        plt.xlabel('Kappa')
        plt.ylabel('Lambda')
        plt.title('distractor strength {:.2f}'.format(d))
        plt.savefig('good_param_vals{:.2f}.png'.format(d))
        plt.show()


import tensorflow as tf
import matplotlib.pyplot as plt
from model import transformer
from datasets import VOCAB_SIZE, MAX_LENGTH, dataset

tf.keras.backend.clear_session()

MODEL_WEIGHT_FILE = "./transformer.h5"

# Hyper-parameters
NUM_LAYERS = 2
D_MODEL = 256
NUM_HEADS = 8
UNITS = 512
DROPOUT = 0.1

model = transformer(
    vocab_size=VOCAB_SIZE,
    num_layers=NUM_LAYERS,
    units=UNITS,
    d_model=D_MODEL,
    num_heads=NUM_HEADS,
    dropout=DROPOUT)

def loss_function(y_true, y_pred):
    y_true = tf.reshape(y_true, shape=(-1, MAX_LENGTH - 1))

    loss = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction='none')(y_true, y_pred)

    mask = tf.cast(tf.not_equal(y_true, 0), tf.float32)
    loss = tf.multiply(loss, mask)

    return tf.reduce_mean(loss)

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):

    def __init__(self, d_model, warmup_steps=4000):
        super(CustomSchedule, self).__init__()

        self.d_model = d_model
        self.d_model = tf.cast(self.d_model, tf.float32)

        self.warmup_steps = warmup_steps

    def __call__(self, step):
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warmup_steps**-1.5)

        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

sample_learning_rate = CustomSchedule(d_model=128)

plt.plot(sample_learning_rate(tf.range(200000, dtype=tf.float32)))
plt.ylabel("Learning Rate")
plt.xlabel("Train Step")

learning_rate = CustomSchedule(D_MODEL)

optimizer = tf.keras.optimizers.Adam(
    learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)

def accuracy(y_true, y_pred):
    # ensure labels have shape (batch_size, MAX_LENGTH - 1)
    y_true = tf.reshape(y_true, shape=(-1, MAX_LENGTH - 1))
    return tf.keras.metrics.sparse_categorical_accuracy(y_true, y_pred)

model.compile(optimizer=optimizer, loss=loss_function, metrics=[accuracy])

EPOCHS = 20

model.fit(dataset, epochs=EPOCHS)

model.save_weights(MODEL_WEIGHT_FILE)


import imageio
import numpy as np

def leer_imagen(ruta):
    """
    La función leer_imagen recibe un string con la ruta
    de una imagen en formato BMP y retorna una lista de
    tres dimensiones con el mapa de bits de la imagen.
    Asimismo, convertimos la lista de numpy a una lista
    común y corriente.
    """
    np_array = np.array(imageio.imread(ruta), dtype='int')
    # noinspection PyTypeChecker
    lista_3d = np_array.tolist()
    return lista_3d

def guardar_imagen(ruta, lista_3d):
    """
    La función guardar_imagen recibe una lista de 3
    dimensiones con el mapa de bits de la imagen
    y retorna la imagen en formato bmp.
    """
    return imageio.imwrite(ruta, np.array(lista_3d, dtype='uint8'))

class Solution:

    def aplicar_escala_de_grises(self, lista_3d=leer_imagen('foto_utec.bmp')):
        result = lista_3d

        # SU SOLUCION EMPIEZA AQUI
        for matrix_2d in result:
          for matrix_1d in matrix_2d:
            suma = round(sum(matrix_1d)/3, 0)
            matrix_1d[0] = suma
            matrix_1d[1] = suma
            matrix_1d[2] = suma

        # SU SOLUCION TERMINA AQUI
        return result

    def aplicar_sepia(self, lista_3d=leer_imagen('foto_utec.bmp')):
        result = lista_3d
        # SU SOLUCION EMPIEZA AQUI
        for matrix_2d in result:
          for matrix_1d in matrix_2d:
            sepia_R = round(0.393*matrix_1d[0] + 0.769 * matrix_1d[1] + 0.189 * matrix_1d[2],0)
            sepia_G = round(0.349*matrix_1d[0] + 0.686 * matrix_1d[1] + 0.168 * matrix_1d[2],0)
            sepia_B = round(0.272*matrix_1d[0] + 0.534 * matrix_1d[1] + 0.131 * matrix_1d[2],0)

            if sepia_R > 255:
              sepia_R = 255
            if sepia_G > 255:
              sepia_G = 255
            if sepia_B > 255:
              sepia_B = 255
            matrix_1d[0] = sepia_R
            matrix_1d[1] = sepia_G
            matrix_1d[2] = sepia_B
        # SU SOLUCION TERMINA AQUI
        return result

    def aplicar_reflejo_horizontal(self, lista_3d=leer_imagen('foto_utec.bmp')):
        result = lista_3d

        # SU SOLUCION EMPIEZA AQUI
        for i in range(len(result)):
          result[i] = result[i][::-1]

        # SU SOLUCION TERMINA AQUI
        return result

    def aplicar_difuminado(self, lista_3d=leer_imagen('foto_utec.bmp')):
        result = lista_3d
        # SU SOLUCION EMPIEZA AQUI
        matriz = []
        for i in range(len(result)):
          matriz2d = []
          for j in range(len(result[i])):
            fila = []
            for k in range(len(result[i][j])):
              if i == 0 and j == 0:
                suma = round((result[i][j][k] + result[i][j+1][k] + result[i+1][j][k] + result[i+1][j+1][k])/4, 0)
                fila.append(suma)
              elif i == 0 and j == len(result[i]) - 1:
                suma = round((result[i][j][k] + result[i+1][j][k] + result[i][j-1][k] + result[i+1][j-1][k])/4, 0)
                fila.append(suma)
              elif i == len(result) - 1 and j == 0:
                suma = round((result[i][j][k] + result[i-1][j][k] + result[i][j+1][k] + result[i-1][j+1][k])/4, 0)
                fila.append(suma)
              elif i == len(result) - 1 and j == len(result[i]) - 1:
                suma = round((result[i][j][k] + result[i-1][j][k] + result[i][j-1][k] + result[i-1][j-1][k])/4, 0)
                fila.append(suma)
              elif i == 0 and j  > 0 and j < len(result[i]) - 1:
                suma = round((result[i][j][k] + result[i][j+1][k] + result[i+1][j][k] + result[i][j-1][k] + result[i+1][j-1][k]+result[i+1][j+1][k])/6, 0)
                fila.append(suma)
              elif i > 0 and j  == 0 and i < len(result) - 1:
                suma = round((result[i][j][k] + result[i-1][j+1][k] + result[i][j+1][k] + result[i+1][j+1][k] + result[i+1][j][k]+result[i][j][k])/6, 0)
                fila.append(suma)
              elif i == len(result) - 1 and j  > 0 and j < len(result[i]) - 1:
                suma = round((result[i][j][k] + result[i-1][j-1][k] + result[i-1][j][k] + result[i-1][j+1][k] + result[i][j+1][k]+result[i][j][k])/6, 0)
                fila.append(suma)

              elif i > 0 and j  == len(result[i]) - 1 and i < len(result) - 1:
                suma = round((result[i][j][k] + result[i-1][j-1][k] + result[i][j-1][k] + result[i+1][j-1][k] + result[i+1][j][k]+result[i][j][k])/6, 0)
                fila.append(suma)
              else:
                suma = round((result[i][j][k] + result[i-1][j][k] + result[i-1][j+1][k] + result[i][j+1][k] + result[i+1][j+1][k]+result[i+1][j][k] + \
                result[i+1][j-1][k] + result[i][j-1][k] + result[i][j][k])/9, 0)
                fila.append(suma)
            matriz2d.append(fila)
        matriz.append(matriz2d)
        # SU SOLUCION TERMINA AQUI
        return result


from datetime import datetime, timedelta
import time
import threading

class Alarm(object):
    enabled = False
    alarm_active = False
    old_time = 0

    def __init__(self, alarmtime):
        self.alarmtime = datetime.strptime(alarmtime, "%H:%M")
        self.time = self.alarmtime
        self.check_alarm_thread = threading.Thread(target=self.check_alarm)
        self.check_alarm_thread.daemon = True
        self.check_alarm_thread.start()

    def check_alarm(self):
        while True:
            new_time = int(datetime.now().minute)
            if new_time != self.old_time:
                self.old_time = new_time
                if self.alarmtime.hour == datetime.now().hour and self.alarmtime.minute == datetime.now().minute:
                    self.alarm_active = True
            time.sleep(1)

    def setAlarm(self):
        self.alarmtime = self.time

    def resetAlarm(self):
        self.time = self.alarmtime

    def changeAlarm(self, m):
        self.time += timedelta(minutes=m)


from setuptools import setup, find_packages

with open('requirements.txt') as f:
    install_requires = f.read().splitlines()

setup(
    name='ava_asd',
    description='Scripts and models for the AVA Active Speaker Detection dataset.',
    license='Apache License Version 2.0',
    author='Tuan Chien, James Diprose',
    url='https://github.com/tuanchien/asd',
    packages=find_packages(),
    install_requires=install_requires,
    python_requires='>=3.7',
    entry_points={
        'console_scripts': [
            'ava-download = download:main',
            'ava-extract = extract:main',
            'ava-train = train:main',
            'ava-evaluate = evaluate:main'
        ]
    },
)


import pandas as pd
from pathlib import Path
import logging
import re

class _FileObj:
    def __init__(self, path_obj, dataframe=None, dataframe_name=None,
    colname_chars_replace_underscore="", colname_chars_replace_custom={},
    colname_chars_remove="", sample_data=500, interpret_date_timestamp=False,
    interpret_date_timestamp_errors="raise", **kwargs):
        """
        Create a FileObj instance that has a single attribute, df which is a pandas dataframe
        supports xls, xlsx, csv, tsv files. Only the first worksheet in an Excel workbook
        with multiple sheets will be read.
        param: use any keyword parameters valid in pandas.read_csv()
        parameter: dataframe - default None; a pandas DataFrame object to be profiled
        parameter: dataframe_name - default None; a unique string for the DataFrame, will be used as the
            filename for the dataframe profile ex. DataFrame_Name_profile.xlsx
        parameter: colname_chars_replace_underscore - string of invalid characters to be replaced with an underscore
        parameter: colname_chars_replace_custom - dict of characters and their replacement value
        parameter: colname_chars_remove - string of characters to be removed
        parameter: interpret_date_timestamp - boolean default False, attempt to convert string fields to date or timestamp
        parameter: interpret_date_timestamp_errors - text default "raise", options are "raise", "ignore", "coerce".
            "raise" will raise errors on values that cannot be converted, "ignore" will not raise errors and returns the
            input data, "coerce" will return NaT values when they cannot be converted.
        """
        # Initialize logging
        self.log = logging.getLogger()
        self.log.handlers = []
        self.log.setLevel(logging.INFO)
        log_fmt = logging.Formatter('%(levelname)s: %(asctime)s %(thread)d %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
        screen = logging.StreamHandler()
        screen.setFormatter(log_fmt)
        self.log.addHandler(screen)

        self.df = None
        # df_name must be unique to create unique output filenames
        self.df_name = None

        # update default replace with underscore characters with user-defined characters
        escape_string = r'\/()[]{},.!?:;-^~`' + colname_chars_replace_underscore
        self.colname_chars_replace_underscore = re.escape(escape_string) + r'\s+'
        colname_chars_replace_custom_default = {'#': 'num', '$': 'usd', '%': 'pct', '&': 'and', '+': 'plus',
        '*': 'times', '=': 'equals', '<': 'lt', '>': 'gt', '@': 'at', '|': 'or'}

        # set values for interpret_date_timestamp parameters
        if interpret_date_timestamp not in [True, False]:
            raise Exception(f"'{interpret_date_timestamp}' is not a valid value for interpret_date_timestamp.")
        if interpret_date_timestamp is True and interpret_date_timestamp_errors not in ["raise", "ignore", "coerce"]:
            raise Exception(f"'{interpret_date_timestamp_errors}' is not a valid value for interpret_date_timestamp_errors.")
        if interpret_date_timestamp in [True, False] and interpret_date_timestamp_errors in ["raise", "ignore", "coerce"]:
            self.interpret_date_timestamp = interpret_date_timestamp
            self.interpret_date_timestamp_errors = interpret_date_timestamp_errors

        # update colname_chars_replace_custom with user-defined dict
        colname_chars_replace_custom_default.update(colname_chars_replace_custom)

        # create new colname chars replace dict with properly escaped values as needed
        self.colname_chars_replace_custom = {}
        for key, value in colname_chars_replace_custom_default.items():
            self.colname_chars_replace_custom[re.escape(key)] = value

        # set attribute for characters to be removed
        self.colname_chars_remove = colname_chars_remove

        # set sample data attributes
        if (isinstance(sample_data, int) and sample_data > 0) or sample_data is None:
            self.sample_data = sample_data
        else:
            raise Exception('sample_data must be an integer > 0 or None.')

        # log.info('Creating FileObj')
        if path_obj == 'dataframe':
            if dataframe is None:
                raise Exception('No dataframe was assigned to the "dataframe=" argument.')
            else:
                self.df = dataframe
                if dataframe_name is None:
                    raise Exception('If profiling a dataframe, argument dataframe_name must be a unique string')
                else:
                    self.df_name = dataframe_name
        elif path_obj.is_file():
            if path_obj.suffix in('.xls', '.xlsx'):
                try:
                    self.df = pd.read_excel(path_obj, **kwargs)
                except Exception as error:
                    self.log.exception(error)

            elif path_obj.suffix in ['.csv', '.tsv', '.txt']:
                try:
                    self.df = pd.read_csv(path_obj, **kwargs)
                except Exception as error:
                    self.log.exception(error)
        else:
            raise Exception(f'{path_obj.name} is not a file.  Please use a valid text or excel file')

        if self.df is None:
            self.log.warning(f'{path_obj.name} was not parsed, please check file format and kwargs')
        else:
            # log.info('Created FileObj')
            self.id_cols = []
            self.dim_cols = []
            self.path_obj = path_obj
            self.sample_data_export = self.df.copy().head(self.sample_data)

    def get_columns(self):
        """
        Not currently used as the column list is contained in the more comprehensive Data Types output
        returns: dataframe of columns from file
        """
        self.log.info('Retrieving Columns')
        return pd.DataFrame(self.df.columns, columns=['Columns'])

    def get_data_types(self):
        """
        method that sets values for self.id_cols and self.dim_cols, additionally the
        returns: dataframe of column names and their data types
        """
        self.log.info('Retrieving Data Types')

        if self.interpret_date_timestamp:
            # attempt explicit date transformation for object cols not automatically detected as dates
            self.df = self.df.apply(lambda x: self.convert_to_datetime(x, self.interpret_date_timestamp_errors))

        df = pd.DataFrame(self.df.dtypes, columns=['Data Type'])
        df.index.name = 'Column Name'
        df = df.reset_index()

        df['Clean Column Name'] = self.clean_column_names(df['Column Name'])

        # identify ID columns
        id_col_pat = re.compile(r"(?:^id$|^ID$)|(?:[-_\s]+id|[-_\s]+ID$)|(?:[a-z]+ID$)|(?:[-_\s]+code)",)
        df['Potential ID Column'] = df['Column Name'].apply(lambda x: True if re.search(id_col_pat, x) else None)
        # set FileObj attribute "ID Columns", referenced in dim_cols below
        self.id_cols = df.loc[df['Potential ID Column'] == True, 'Column Name'].tolist()

        # identify the proper data type
        for col in self.df.columns:
            # set string representation of pandas series dtype
            col_dtype = str(self.df[col].dtype)

            if self.df[col].count() > 0:
                if col_dtype in ['object', 'bool'] or 'datetime' in col_dtype:
                    col_values_precision_series = self.df[col].astype('str').str.len()
                    max_precision_value = max(col_values_precision_series)
                    min_precision_value = min(col_values_precision_series)
                elif col_dtype in ['int', 'int64', 'int32']:
                    max_precision_value = self.df[col].max()
                    min_precision_value = self.df[col].min()
                elif col_dtype in ['float', 'float64', 'float32']:
                    col_values_precision_series = self.df[col].astype('str')
                    col_values_precision_df = col_values_precision_series.str.split('.', expand=True)
                    # precision
                    min_precision_value = max(col_values_precision_df[0].str.len() + col_values_precision_df[1].str.len())
                    # scale
                    max_precision_value = max(col_values_precision_df[1].str.len())
                    # check if the Data Type could be int
                    try:
                        if col_values_precision_df[1].fillna('0').astype('int64').sum() == 0:
                            df.loc[df['Column Name'] == col, 'Data Type'] = 'decimal or integer'
                    except ValueError:
                        # pass because this should catch INT overflow that has no negative affect on output
                        pass
                else:
                    min_precision_value = 'dtype not supported'
                    max_precision_value = 'dtype not supported'
            else:
                min_precision_value = 0
                max_precision_value = 0
                df.loc[df['Column Name'] == col, 'Data Type'] = 'N/A'

            df.loc[df['Column Name'] == col, 'Min Length|Value/Precision'] = min_precision_value
            df.loc[df['Column Name'] == col, 'Max Length|Value/Scale'] = max_precision_value

            # check for NULL values
            if sum(self.df[col].isna()) > 0:
                nullable_value = True
            else:
                nullable_value = None

            df.loc[df['Column Name'] == col, 'Nullable'] = nullable_value

            # set FileObj attribute "Dim Columns"
            if self.df[col].dtype == 'object' and col not in self.id_cols:
                self.dim_cols.append(col)

        # replace obscure data type names with clear names
        replace_dict = {'datetime64[ns]': 'date/datetime', 'object':'text', 'int': 'integer',
        'int8': 'integer', 'int32': 'integer', 'int64': 'integer', 'float': 'decimal',
        'float32': 'decimal', 'float64': 'decimal',}
        df['Data Type'] = df['Data Type'].replace(to_replace=replace_dict)

        # identify potential PII columns
        pii_col_pat = re.compile(r'(?:last|full|first|family|given)\s*\S*(?:name|nm)|' \
                                r'\S*\s*(?:address|addr)|(?:e\S*\s*mail)|' \
                                r'(?:tel|tele)\S*(?:phone|no)|(?:phone)', re.IGNORECASE)
        pii_col_name_df = df['Column Name'].apply(lambda x: True if re.search(pii_col_pat, x) else None)
        pii_col_name_cols = df.loc[pii_col_name_df.notna(), 'Column Name'].to_list()

        # look for columns containing PII not already flagged as potential PII columns
        # columns with telephone numbers
        remaining_columns = [col for col in df['Column Name'] if col not in pii_col_name_cols]
        telephone_pat = re.compile(r'\d{10,13}') # re.compile(r'(?:\(*\+*\d+\-*\.*\s*\(*\d+\)*\-*\.*\s*\d+\-*\.*\s*\d+)')
        tel_no_df = (self.df[remaining_columns].replace(to_replace=r'\D', value='', regex=True)
                        .apply(lambda x: x.astype(str).str.contains(telephone_pat, regex=True), axis=0))
        tel_no_cols = tel_no_df.any(axis=0)
        tel_no_cols = list(tel_no_cols[tel_no_cols].index)

        pii_cols = pii_col_name_cols + tel_no_cols

        # columns with email addresses
        remaining_columns = [col for col in df['Column Name'] if col not in pii_cols]
        email_pat = re.compile(r'(?:\S+@\S+\.\S+)')
        email_df = self.df[remaining_columns].apply(lambda x: x.astype(str).str.contains(email_pat, regex=True), axis=0)
        email_cols = email_df.any(axis=0)
        email_cols = list(email_cols[email_cols].index)

        # add email_cols to list of potential pii fields
        pii_cols += email_cols

        # columns with street addresses
        remaining_columns = []
        for col in df['Column Name']:
            if (col not in pii_cols
                    and not df[(df['Column Name'] == col) & (df['Data Type'] == 'text')].empty):
                remaining_columns.append(col)

        street_address_pat = re.compile(r'(?:\d+\s+[a-zA-Z0-9\-]+\s*[a-zA-Z0-9\-]*)')
        street_address_df = (self.df[remaining_columns]
                                        .apply(lambda x: x.astype(str).str.contains(street_address_pat, regex=True), axis=0))
        street_address_cols = street_address_df.any(axis=0)
        street_address_cols = list(street_address_cols[street_address_cols].index)

        # add street_address_cols to list of potential pii fields
        pii_cols += street_address_cols

        df['Potential PII Column'] = df['Column Name'].apply(lambda x: True if x in pii_cols else None)

        # set FileObj attribute "PII Columns", referenced in dim_cols below
        self.pii_cols = df.loc[df['Potential PII Column'] == True, 'Column Name'].tolist()

        return df[['Column Name', 'Clean Column Name', 'Data Type', 'Min Length|Value/Precision',
            'Max Length|Value/Scale', 'Potential ID Column', 'Potential PII Column', 'Nullable']]

    def clean_column_names(self, colname_series):
        """
        Take a pandas series of column names and apply cleansing steps
        Returns: pandas series of cleaned column names
        """
        colname_series_clean = colname_series

        # remove unwanted characters
        if self.colname_chars_remove != "":
            colname_series_clean = colname_series.str.replace(f'[{self.colname_chars_remove}]+', '', regex=True)

        # Replace unacceptable characters in column names with undercores
        colname_series_clean = colname_series_clean.str.replace(f'[{self.colname_chars_replace_underscore}]+', '_', regex=True)
        # replace chars with custom values
        for char, replacement in self.colname_chars_replace_custom.items():
            colname_series_clean = colname_series_clean.str.replace(f'{char}', f'_{replacement}_', regex=True)

        # replace multiple underscores with a single underscore
        colname_series_clean = colname_series_clean.str.replace(r'_+', '_', regex=True)
        # remove leading and trailing underscores
        colname_series_clean = colname_series_clean.str.strip('_')

        # insert underscore for column names that might be IDs and use camel case
        colname_series_clean = colname_series_clean.apply(lambda x: _modify_camel_case_names(x))

        # lower case the clean column name
        colname_series_clean = colname_series_clean.str.lower()

        return colname_series_clean

    def get_text_distinct_values(self):
        """
        returns: dataframe of column names and series of distinct values
        """
        self.log.info('Retrieving Text Value Distribution')
        results_dict = {}
        pandas_numeric_dtype_list = ['int', 'int64', 'int32', 'float', 'float64', 'float32']
        for col in self.df.columns:
            if self.df[col].dtype in pandas_numeric_dtype_list and col not in self.id_cols:
                results_dict[col] = pd.DataFrame(['NA for numeric columns'], columns=[col])
            else:
                df = pd.DataFrame(self.df[col].value_counts())
                df_null = pd.DataFrame({col: len(self.df[self.df[col].isna()])}, index=['NULL'])
                df = pd.concat([df_null, df], sort=False)
                df.index.name = col
                df.rename(index=str, columns={col: f'{col}_counts'}, inplace=True)
                results_dict[col] = df.reset_index()

        distinct_text_values_df = pd.concat(results_dict.values(), axis=1, join='outer', sort=True)

        distinct_text_values_df = replace_xml_illegal_characters(distinct_text_values_df)

        # strip timezone from output because Excel does not support localized TZ info
        # return distinct_text_values_df.apply(lambda x: x.dt.tz_localize(None) if 'datetime' in str(x.dtype) else x)
        return distinct_text_values_df

    def get_numeric_value_distribution(self):
        """
        returns: dataframe of descriptive stats for numeric columns
        """
        self.log.info('Retrieving Numeric Value Distribution')
        df = pd.DataFrame(self.df.describe()).reset_index()
        df.rename(index=str, columns={'index': 'Stat'}, inplace=True)
        return df

    def get_primary_keys(self):
        """
        analyzes dataframe to see the maximum number of non-metric columns that can be part of a primary key
        this is done by checking for columns that when grouped, retain the same number of distinct records as the
        oringinal dataframe
        returns: dataframe of suggested primary key(s) and dataframe of potential ID fields"""
        self.log.info('Looking for Potential Primary Key(s)')

        pk_1 = []
        pk_2 = []
        df_len = 0
        id_cols = self.id_cols.copy()
        dim_cols = self.dim_cols.copy()
        for col in id_cols + dim_cols:
            new_len = len(self.df.groupby(pk_1 + [col], sort=False).count())
            if new_len > df_len:
                df_len = new_len
                pk_1.append(col)
            else:
                pass

        id_cols.reverse()
        dim_cols.reverse()
        df_len = 0
        for col in id_cols + dim_cols:
            if len(self.df.groupby(pk_2 + [col], sort=False).count()) < df_len:
                pk_2.append(col)
            else:
                pass

        return pd.DataFrame({'Column Name': list(set(pk_1 + pk_2))})

    def create_sample(self):
        if self.sample_data is not None:
            # strip timezone from output because Excel does not support localized TZ info
            # return self.df.head(self.sample_data).apply(lambda x: x.dt.tz_localize(None) if 'datetime' in str(x.dtype) else x)
            return self.sample_data_export

    def convert_to_datetime(self, col, interpret_date_timestamp_errors):
        if str(col.dtype) == "object":
            try:
                new_col = pd.to_datetime(col, errors=interpret_date_timestamp_errors, utc=False)
                # if transformation was completely destructive, reset new_col to col
                if new_col.count() == 0:
                    new_col = col
            except Exception as e:
                self.log.error(f"Failed to cast {col.name} to date/datetime - {e}")
                # initialize new_col
                new_col = col
        else:
            new_col = col

        return new_col

def _modify_camel_case_names(x):
    results = re.findall(r'([a-z][A-Z])', x)
    if results:
        for match in results:
            x = x.replace(match, "_".join(list(match)))
    return x

def replace_xml_illegal_characters(df):
    """
    replace characters that throw ILLEGAL_CHARACTER_ERROR in openpyxl
    when writing to XLS(X) formats, found regex pattern in openpyxl source here:
    https://openpyxl.readthedocs.io/en/stable/_modules/openpyxl/cell/cell.html
    """
    ILLEGAL_CHARACTERS_RE = re.compile(r'[\000-\010]|[\013-\014]|[\016-\037]')

    return df.replace(to_replace=ILLEGAL_CHARACTERS_RE, value='~', regex=True)


# -*- coding: utf-8 -*-
import scrapy
import re
from datetime import datetime
from scrapy import log
from scrapy.http import Request
from scrapy.spider import BaseSpider
from scrapy.contrib.loader import ItemLoader
from travelzoo.items import TravelZooDeal

class TravelZooComSpider(BaseSpider):
    """Spider for regularly updated travelzoo.com site, UK version"""
    name = 'travelzoo.com'
    allowed_domains = ['www.travelzoo.com']
    start_urls = ['http://www.travelzoo.com/uk/']

    def parse(self, response):
        # select all links in the left hand menu which aren't top level ones (don't have li class attribute) and aren't in the Destinations menu
        hrefs = response.xpath("//div[@id='leftNavigationWrapper']/ul/li[a/text()!='Destinations']//ul/li[not(@class)]/a/@href").extract()
        for href in hrefs:
            self.log("Yielding request for url: %s" % href)
            yield Request(url=href, callback=self.parse_section)

    def parse_section(self, response):
        items = response.xpath("//div[contains(@class,'featuredDeal') or contains(@class, 'premiumPlacement') or contains(@class, 'dealItem')]")
        for item in items:
            item_href = item.xpath(".//h2/a/@href").extract()[0]
            pattern = '[-/](\d{6,7})[/0-9A-Za-z-]*$'
            match = re.search(pattern, item_href)
            if match is None:
                self.log("ERROR couldn't parse id from url: %s using regex: %s" % (item_href, pattern), level=log.ERROR)
                return
            item_id = match.group(1)
            self.log("Item id: %s, yielding url: %s" % (item_id, item_href))
            yield Request(url=item_href, callback=self.parse_item, meta={'url': item_href, 'id': item_id})

    def parse_item(self, response):
        self.log("Parsing item url: %s" % response.url)
        self.log("Url of the page is actually %s" % response.url)
        i = TravelZooDeal()
        i['id'] = response.meta['id']
        i['url'] = response.meta['url']
        if len(response.css('.page.noBorder')):
            return self.parse_item_no_border(response, i)
        elif len(response.css('.page.withBorder')):
            return self.parse_item_with_border(response, i)
        else:
            self.log("Parsing unknown page type, this may not give many details...")
            loader = ItemLoader(item=i, response=response)
            loader.add_css('name', 'h1::text')
            loader.add_css('description', '.dealText p:first-child')
            loader.add_css('whats_included', '.dealText ul')
            return loader.load_item()

    def parse_item_with_border(self, response, i):
        self.log("Parsing page with border")
        page = response.css('.innerDealPage')
        loader = ItemLoader(item=i, selector=page)
        loader.add_css('name', 'h1::text')
        loader.add_xpath('description', "//div[contains(@class, 'introDescription')]//text()")
        loader.add_xpath('why_we_love_it', ".//div[contains(@class,'dealDetailsSection') and h2[text()='Why we love it']]/ul//text()")

        loader.add_xpath('whats_included', ".//div[contains(@class,'dealDetailsSection') and h2[contains(text(), 'included')]]/ul//text()")
        loader.add_xpath('small_print', ".//div[contains(@class,'dealDetailsSection') and h2[text()='The small print']]/p//text()")

        div_where = page.xpath(".//div[contains(@class,'dealDetailsSection') and h2[text()='Where']]")
        loader.selector = div_where
        loader.add_xpath('contact_name', "div[contains(@id,'MerchantName')]/text()")
        loader.add_xpath('contact_address', "div[contains(@id,'MerchantAddress')]/text()")
        loader.add_xpath('contact_map', "div[contains(@id,'LinkMap')]/text()")
        loader.add_xpath('contact_phone', "div[contains(@id,'MerchantPhone')]/text()")
        loader.add_xpath('contact_website', "div[contains(@id,'MerchantWebsite')]/text()")

        deal_page_right = response.css('#dealPageRightPart')
        loader.selector = deal_page_right
        loader.add_xpath('price', ".//span[@id='ctl00_Main_OurPrice']/text()")
        loader.add_xpath('value', ".//span[@id='ctl00_Main_PriceValue']/text()")
        loader.add_xpath('discount', ".//span[@id='ctl00_Main_Discount']/text()")
        loader.add_xpath('bought', ".//span[contains(@id,'Bought')]/text()")
        return loader.load_item()

    def parse_item_no_border(self, response, i):
        self.log("Parsing page no border")
        page = response.css('.page.noBorder')
        loader = ItemLoader(item=i, selector=page)
        loader.add_css('name', 'h1::text')
        loader.add_xpath("why_we_love_it", ".//div[contains(@id,'spanWhyLove')]/ul")
        loader.add_xpath("whats_included", ".//div[contains(@id,'DivWhatsIncluded')]/div")
        return loader.load_item()


"""
This class is base module for all trainers
inherit from this class.  It has basic functionality
for all trainers.

Author:Mus mbayramo@stanford.edu
"""
import argparse
import os
import sys
import tempfile
import warnings
from pathlib import Path

import loguru
import torch
import torch.distributed as dist
from accelerate import Accelerator

from igc.shared.shared_torch_utils import get_device
from ...shared.shared_accelerator import build_accelerator

class IgcBaseState:
    """
    This base state shared in trainer , modules.
    Rank, Devices, Logging
    """

    def __init__(self,
                 module_name: str,
                 spec: argparse.Namespace,
                 device=None):
        """
        This only shared states,  devices, accelerators, logging, etc.

        :param module_name:
        :param spec:
        :param device:
        """
        if not isinstance(spec, argparse.Namespace):
            raise TypeError(f"spec should be an instance of argparse.Namespace, "
                            f"received {type(spec).__name__}.")

        self._log_file = None
        self._is_trained = False
        self._trainer_args = spec

        self.logger = loguru.logger
        self._configure_logger(module_name)
        self.rank = int(os.environ.get('LOCAL_RANK', -1))

        self.is_accelerator = False
        if device is not None:
            self._device = torch.device(device)

        self._accelerator = None

        if hasattr(spec, 'use_accelerator') and spec.use_accelerator:
            self._accelerator = build_accelerator(spec)
            self.is_accelerator = True
            # let accelerator choose device.
            self._device = self._accelerator.device
            self.logger.info(f"Rank {self.rank}, Running accelerator , accelerate selected device {self.device}")
        elif hasattr(spec, "device"):
            self._device = spec.device
        else:
            self.logger.info(f"Rank {self.rank}, Running on selected device {self.device}")
            # if we are not using accelerator, we need to set device
            self._device = get_device(self.rank) if device is None else device
            raise

        self.scheduler = None

    @property
    def accelerator(self) -> Accelerator:
        """Return accelerator.
        :return:
        """
        return self._accelerator

    @property
    def device(self):
        """Return ether cuda or cpu device or accelerator device.
        :return:
        """
        if self._accelerator:
            return self.accelerator.device
        else:
            return self._device

    def _prepare_checkpoint_dir(self):
        """

        Prepares the checkpoint directory.

        :return:
        """
        checkpoint_path_dir = Path(self._trainer_args.output_dir)
        checkpoint_path_dir = checkpoint_path_dir.resolve()
        if not checkpoint_path_dir.is_dir():
            raise ValueError(f"Indicate path to checkpoint dir {checkpoint_path_dir}.")

        return str(checkpoint_path_dir)

    def _configure_logger(self, module_name: str):
        """
        Configures the logger for the module.

        :param module_name: The name of the module.
        """
        logs_dir = getattr(self._trainer_args, "log_dir", None)
        if not logs_dir:
            warnings.warn("log_dir is not provided. "
                          "Using a temporary directory as a fallback.")
            logs_dir = tempfile.mkdtemp()
        else:
            os.makedirs(logs_dir, exist_ok=True)

        self._log_file = os.path.join(logs_dir, f"{module_name}.log")
        self._log_level = getattr(
            self._trainer_args, "log_level", "ERROR"
        )
        if self._log_level:
            self._log_level = self._log_level.upper()

        self.logger = loguru.logger.bind(module_name=module_name)
        self.logger.remove()

        log_to_file = getattr(self._trainer_args, "log_to_file", None)
        if not log_to_file:
            self.logger.add(sys.stdout, level=self._log_level)
        else:
            if self._trainer_args.log_to_file:
                log_file = os.path.join(logs_dir, f"{module_name}.log")
                self.logger.add(log_file, level=self._log_level)
            else:
                self.logger.add(sys.stdout, level=self._log_level)

    def is_distributed(self):
        """
        :return:
        """
        return self.rank != -1

    def is_rank_zero(self):
        """
        :return:
        """
        return self.rank == -1 or self.rank == 0

    def is_last_worker(self):
        """

        :return:
        """
        if self.rank == -1:
            return True

        rank = dist.get_rank()
        world_size = dist.get_world_size()

        is_last_worker = rank == (world_size - 1)
        return is_last_worker

    def log_memory_usage(self):
        """
        Log memory usage.
        :return:
        """
        if torch.cuda.is_available():
            mem_get_info = torch.cuda.memory_stats()
            self.logger.info("Memory allocated:", mem_get_info["allocated_bytes.all.current"] / 1024 ** 3, "GB")
            # additional CUDA statistics if available
            if hasattr(torch.cuda, 'utilization'):
                self.logger.info(f"CUDA utilization:", torch.cuda.utilization())
            if hasattr(torch.cuda, 'memory_summary'):
                torch.cuda.memory_summary()

    def show_accelerator_info(self):
        """
        Show information about the accelerator used.
        """
        if self.is_accelerator:
            self.logger.debug("Accelerator Configuration:")
            self.logger.debug(f"is main process: {self.accelerator.is_main_process}")
            self.logger.debug(f"Distributed Training: {self.accelerator.distributed_type}")
            self.logger.debug(f"Device Placement: {self.accelerator.device_placement}")
            self.logger.debug(f"Split Batches: {self.accelerator.split_batches}")
            self.logger.debug(f"Mixed Precision: {self.accelerator.mixed_precision}")
            self.logger.debug(f"Project Directory: {self.accelerator.project_dir}")
            self.logger.debug(f"Dispatch Batches: {self.accelerator.dispatch_batches}")
            self.logger.debug(f"Even Batches: {self.accelerator.even_batches}")
            self.logger.debug(f"RNG Types: {self.accelerator.rng_types}")
            self.logger.debug(f"Log With: {self.accelerator.log_with}")
            self.logger.debug(f"Step Scheduler with Optimizer: {self.accelerator.step_scheduler_with_optimizer}")


import os
import re
import pandas as pd
from deepface import DeepFace
from tqdm import tqdm

pattern = re.compile(r'\d+\_\d+\_q?\d+')

folders = ['Mor_morph', 'Mor_nonmorph']

data = {
	'imgname' : [],
	'isMorph' : [],
	'gt1_gt2' : [],
	'gt1_input' : [],
	'gt2_input' : [],
	'output1_input' : [],
	'output2_input' : [],
	'gt1_output1' : [],
	'gt1_output2' : [],
	'gt2_output1' : [],
	'gt2_output2' : [],
	'output1_output2' : []
}

def calculate_score(path1,path2):
    result = DeepFace.verify(path1,path2,model_name = 'ArcFace',enforce_detection=False)
    return result['distance']

f1=set()

for folder in tqdm(folders):
	is_morph_check = folder.replace('Mor_', '')
	for filename in tqdm(os.listdir(folder)):
		m = re.search(pattern, filename)
		if m is not None:
			match_str = m.group(0)
			# print(match_str)
			match_str = match_str.strip()
			# print(match_str)
			f1.add(match_str)

			if(match_str not in data['imgname']):
				## column1
				data['imgname'].append(match_str)

				## column2
				if(is_morph_check=="morph"):
					data['isMorph'].append(1)
				else:
					data['isMorph'].append(0)

				### Image files for processing
				gt1 = folder+"/"+match_str+"_gt1.png"
				gt2 = folder+"/"+match_str+"_gt2.png"
				Input = folder+"/"+match_str+"_input.png"
				output1 = folder+"/"+match_str+"_output1.png"
				output2 = folder+"/"+match_str+"_output2.png"

				data['gt1_gt2'].append('{0:.3g}'.format(calculate_score(gt1,gt2)))
				data['gt1_input'].append('{0:.3g}'.format(calculate_score(gt1,Input)))
				data['gt2_input'].append('{0:.3g}'.format(calculate_score(gt2,Input)))
				data['output1_input'].append('{0:.3g}'.format(calculate_score(output1,Input)))
				data['output2_input'].append('{0:.3g}'.format(calculate_score(output2,Input)))
				data['gt1_output1'].append('{0:.3g}'.format(calculate_score(gt1,output1)))
				data['gt1_output2'].append('{0:.3g}'.format(calculate_score(gt1,output2)))
				data['gt2_output1'].append('{0:.3g}'.format(calculate_score(gt2,output1)))
				data['gt2_output2'].append('{0:.3g}'.format(calculate_score(gt2,output2)))
				data['output1_output2'].append('{0:.3g}'.format(calculate_score(output1,output2)))

df = pd.DataFrame(data = data, columns = list(data.keys()))
print(df.head())
print(df.describe())
df.to_csv('scores.csv')