# CelebA image generation using Conditional DCGAN
import copy
import os

import torch
import pickle

from ModularUtils.ControllerModel import get_generators, get_discriminators, get_generated_labels
from ModularUtils.FunctionsTraining import get_training_variables, calc_gradient_penalty, save_checkpoint
from Sachs_experiment.GAN_Evaluation.sachsEvaluation import sachsEvaluation


from Sachs_experiment.GroundTruth.CausalGraph_Sachs import set_sachs_nonId_graph
from ModularUtils.FunctionsConstant import asKey, getdoKey, load_label_dataset, initialize_results
from ModularUtils.ControllerConstants import get_multiple_labels_fill

from ModularUtils.Experiment_Class import Experiment

def get_dataset(Exp, label, dno):

    dataset = []
    for feature in ["feature"]:
        file_name = Exp.file_roots[dno] + label + feature + ".pkl"

        with open(file_name, 'rb') as fp:
            label_data = pickle.load(fp)
        label_data = torch.FloatTensor(label_data)
        label_size = len(label_data)
        dataset.append(label_data.view(label_size, 1))

    result_dataset = torch.cat(dataset, 1).to(Exp.DEVICE)
    print(result_dataset.shape)
    return result_dataset


def get_intv_dataset(Exp):
    intv_dataset_list = []
    for dno in range(1, Exp.num_datasets):
        dataset = []
        for label in Exp.label_names:
            file_name = Exp.file_roots[dno] + label + ".pkl"
            with open(file_name, 'rb') as fp:
                label_data = pickle.load(fp)
            label_data = torch.FloatTensor(label_data)
            label_size = len(label_data)
            # plot_labels("intv "+label, label_data.view(label_size, 1))
            dataset.append(label_data.view(label_size, 1))
        dataset = torch.cat(dataset, 1).to(Exp.DEVICE)

        intv_dataset_list.append(dataset)

    intv_dataset = None
    if len(intv_dataset_list) != 0:
        intv_dataset = torch.cat(intv_dataset_list, 0)

    return intv_dataset




def train_CausalController(Exp, cur_mechs, label_generators, G_optimizers, label_discriminator, D_optimizer,
                           dataset_dict_batches, batchno):

    G_loss=torch.zeros(1).to(Exp.DEVICE)
    for interv_no, (intv_key, dataset_batches) in enumerate(dataset_dict_batches.items()):
        intv_key = dict(intv_key)

        data_input = dataset_batches["obs"][batchno]

        _,_,_, graph_label_vars = get_training_variables(Exp, Exp.label_names, interv_no, intv_key)
        all_compare_Var, compare_Var, intervened_Var, real_labels_vars = get_training_variables(Exp, cur_mechs, interv_no, intv_key)

        #fix it later
        if len(real_labels_vars)>data_input.shape[1]:
            continue

        mini_batch = data_input.size()[0]
        indices = [graph_label_vars.index(lb) for lb in real_labels_vars]
        current_real_label = data_input[:, indices].type(torch.LongTensor).view(-1, len(indices)).to(Exp.DEVICE)
        # current_real_label = data_input.type(torch.LongTensor).view(-1, len(real_labels_vars)).to(Exp.DEVICE)
        dims_list = [Exp.label_dim[lb] for lb in real_labels_vars]

        obs_images=None

        real_labels_fill = get_multiple_labels_fill(Exp, current_real_label, dims_list, isImage_labels=False)


        intv_tensor_dict = {}
        isClassifier=False
        for lbid, intv_lb in enumerate(intervened_Var): #if no intervention then no looping
            # index = [Exp.label_names.index(intv_lb)]
            # parent_intv_label = data_input[:, index].type(torch.LongTensor).view(-1, 1).to(Exp.DEVICE) #for each intv parent


            ind = real_labels_vars.index(intv_lb)
            parent_intv_label = data_input[:,ind].type(torch.LongTensor).view(-1, 1).to(Exp.DEVICE) #for each intv parent
            dims_list = [Exp.label_dim[intv_lb]]
            intv_parent_fill = get_multiple_labels_fill(Exp, parent_intv_label, dims_list, isImage_labels=False)
            intv_tensor_dict[intv_lb] = intv_parent_fill


        generated_image=None
        generated_labels_dict = get_generated_labels(Exp, label_generators, {}, {}, intv_tensor_dict, real_labels_vars, mini_batch)
        y_dims = sum([Exp.label_dim[lb] for lb in real_labels_vars])
        ret = list(generated_labels_dict.values())
        generated_labels_fill = torch.cat(ret, 1).view(-1, y_dims)


        D_losses = []
        image_loss=[]
        label_loss=[]
        for crit_ in range(Exp.CRITIC_ITERATIONS):

            D_real_decision_obs = label_discriminator[interv_no](real_labels_fill).squeeze()
            D_fake_decision_obs = label_discriminator[interv_no](generated_labels_fill).squeeze()
            gp_obs = calc_gradient_penalty(label_discriminator[interv_no], real_labels_fill, generated_labels_fill, device=Exp.DEVICE)

            D_loss_obs = (-  (torch.mean(D_real_decision_obs) - torch.mean(D_fake_decision_obs)) + Exp.LAMBDA_GP * gp_obs)


            #THINK ABOUT separate loss for images and labels
            D_losses.append((D_loss_obs).data)  # just a loss list

            label_discriminator[interv_no].zero_grad()
            # gp_obs.backward(retain_graph=True)
            D_loss_obs.backward(retain_graph=True)
            D_optimizer[interv_no].step()

        # accumulating the generator losses for all interventions.
        if set(cur_mechs) & set(Exp.image_labels) != set():
            D_fake_decision_obs = label_discriminator[interv_no](generated_image, generated_labels_fill).squeeze()
        else:
            D_fake_decision_obs = label_discriminator[interv_no](generated_labels_fill).squeeze()

        G_loss += -torch.mean(D_fake_decision_obs)


    # Back propagation
    for mech in cur_mechs:
        label_generators[mech].zero_grad()

    G_loss.backward()

    for mech in cur_mechs:
        G_optimizers[mech].step()

    D_loss = torch.mean(torch.FloatTensor(D_losses))  # just mean of losses

    if Exp.test_marginals:
        image_loss = torch.mean(torch.FloatTensor(image_loss))  # just mean of losses
        label_loss = torch.mean(torch.FloatTensor(label_loss))  # just mean of losses
    else:
        image_loss, label_loss = torch.zeros(1).to(Exp.DEVICE), torch.zeros(1).to(Exp.DEVICE)
    # if (Exp.curr_epoochs) % 5 == 0:
    #     print("real label", current_real_label[0])
    #     genimg = obs_images[0].permute(1, 2, 0).detach().cpu().numpy()
    #     plot_trained_digits(1, 1, [genimg], f'Real {real_labels_fill[0]}')
    #
    #     print("Fake label", generated_labels_dict["D"][0])
    #     genimg = generated_image[0].permute(1, 2, 0).detach().cpu().numpy()
    #     plot_trained_digits(1, 1, [genimg], f'fake {generated_labels_fill[0]}')

    return G_loss.data, D_loss.data, image_loss, label_loss




def labelMain(Exp, cur_hnodes, label_generators, G_optimizers, discriminators, D_optimizers, dataset_dict,
              tvd_diff, kl_diff):
    dataset_dict_batches = {}

    num_batches=0
    for key, each_dataset in dataset_dict.items():
        dataset_dict_batches[key]={}
        real_dataloader = torch.utils.data.DataLoader(dataset=each_dataset["obs"],
                                                      batch_size=Exp.batch_size,
                                                      shuffle=False)

        batch_list = []
        for data_input in real_dataloader:
            data_input = torch.squeeze(data_input)
            if len(data_input.size())==1:
                data_input= data_input.view(-1,1)
            batch_list.append(data_input)

        dataset_dict_batches[key]["obs"] = batch_list
        num_batches = len(batch_list)

        ####
        if len(Exp.rep_labels):
            real_dataloader = torch.utils.data.DataLoader(dataset=each_dataset["rep"],
                                                          batch_size=Exp.batch_size,
                                                          shuffle=False)

            batch_list = []
            for data_input in real_dataloader:
                data_input = torch.squeeze(data_input)
                if len(data_input.size())==1:
                    data_input= data_input.view(-1,1)
                batch_list.append(data_input)

            dataset_dict_batches[key]["rep"] = batch_list
            num_batches = len(batch_list)

    # for images
    # imagedata_dict_batches = {}
    # for key, each_dataset in dataset_dict["obs"]["img"].items():
        if len(Exp.image_labels):
            image_data_loader = torch.utils.data.DataLoader(dataset=each_dataset["img"],
                                                            batch_size=Exp.batch_size,
                                                            shuffle=False)
            batch_list = []
            for data_input in image_data_loader:
                data_input = torch.squeeze(data_input)
                batch_list.append(data_input)
            dataset_dict_batches[key]["img"] = batch_list

    iteration = 0

    for batchno in range(num_batches):

        for hn, cur_mechs in cur_hnodes.items():

            g_loss, d_loss,  image_loss, label_loss = train_CausalController(Exp, cur_mechs, label_generators, G_optimizers, discriminators[hn],
                                                    D_optimizers[hn], dataset_dict_batches, batchno)

            print('Epoch [%d/%d], Step [%d/%d],' % (
                Exp.curr_epoochs + 1, Exp.num_epochs, iteration + 1, num_batches),
              'mechanism: ',cur_mechs,  ' D_loss: %.4f, G_loss: %.4f' % (d_loss.data, g_loss.data))



        # Annealing
        tot_iter = Exp.curr_epoochs * num_batches + iteration
        if tot_iter % 100 == 0:
            Exp.anneal_temperature(tot_iter)

        # if (iteration + 1) % int(num_batches / Exp.PLOTS_PER_EPOCH) == 0:


        Exp.D_avg_losses.append(torch.mean(d_loss))
        Exp.G_avg_losses.append(torch.mean(g_loss))
        iteration += 1

        # break
    #
    if (Exp.curr_epoochs + 1) % 1 == 0:
        print("Turn on caffeinate or these results are gone!")
        tvd_diff, kl_diff = sachsEvaluation(Exp, cur_hnodes, label_generators, dataset_dict, tvd_diff, kl_diff)


    if (Exp.curr_epoochs + 1) % 5 == 0:
        var_list= "".join(x for x in cur_mechs)
        save_checkpoint(Exp, Exp.SAVED_PATH, cur_mechs, label_generators, G_optimizers, {var_list:discriminators}, {var_list: D_optimizers})
        print(Exp.curr_epoochs,":model saved at ", Exp.SAVED_PATH)


    return



if __name__ == "__main__":

    Exp = Experiment("Exp1", set_sachs_nonId_graph,
                     Temperature=1,
                     temp_min=0.1,
                     G_hid_dims=[256, 256],
                     D_hid_dims=[256, 256],
                     CRITIC_ITERATIONS=5,
                     LAMBDA_GP=1,
                     learning_rate=5 * 1e-4,
                     # learning_rate= 1e-5,
                     Synthetic_Sample_Size=10000,
                     intv_Sample_Size=10000,
                     batch_size=100,
                     features=["feature"],
                     noise_states=100,
                     latent_state=16,
                     Data_intervs=[{}, {"PKA": 2}],
                     # Data_intervs=[{}, {"PKC": 1}, {"PKC": 2}],
                     # Data_intervs=[{"Mek": 0}],
                     # Data_intervs=[{}, {"PKA":2}],
                     # Data_intervs=[{}],
                     num_epochs=1000,
                     obs_state=3,
                    new_experiment=True
                     # evaluate_after_epochs= sachsEvaluation
                     )

    isModular= True

    if isModular==True:
        cur_hnodes = {"H1":["PKA", "Mek"], "H2": ["Erk","Akt"]}

    else:
        cur_hnodes = {"H1": ["PKA", "Mek", "Erk", "Akt"]}

    Exp.num_datasets=1

    print(Exp.Data_intervs)
    Exp.intv_batch_size = Exp.batch_size
    # True scm


    os.makedirs(Exp.SAVED_PATH, exist_ok=True)
    dag_name = Exp.Complete_DAG_desc + ".txt"

    # Load previous model results also
    # Exp.LOAD_MODEL_PATH = "/path_to_project/SAVED_EXPERIMENTS/sachs_nonId_graph/Exp1/Dec_05_2022-03_28"  #1000 epochs all at once.
    # Exp.load_which_models = {"PKA": True, "Mek": True, "Erk": True, "Akt":True}
    Exp.load_which_models = {"PKA": False, "Mek": False, "Erk": False, "Akt":False}
    # 999 :model saved at  /path_to_project/SAVED_EXPERIMENTS/sachs_nonId_graph/Exp1/Dec_05_2022-03_28 - 1000 epochs



    label_generators, optimizersMech = get_generators(Exp, Exp.load_which_models)
    discriminatorsMech, doptimizersMech = get_discriminators(Exp, cur_hnodes, Exp.load_which_models)  #

    dataset_dict = {}

    dataset_dict = load_label_dataset(Exp, {}, label_generators, cur_hnodes)


    tvd_diff, kl_diff = initialize_results(Exp, cur_hnodes)



    ####
    file_root = "Sachs_experiment/GroundTruth/Dataset/"
    max_size = Exp.Synthetic_Sample_Size

    # intv_list=[]
    intv_list = [{"PKA": 2}]
    intv_mech=["PKA", "Mek"]
    for intv in intv_list:
        ky, val = list(intv.keys())[0], list(intv.values())[0]
        # file_name = "/path_to_project/Sachs_experiment/GroundTruth/Dataset/sachs_P(PKA_PKC_Raf_Mek_Erk|do(Mek=0)).txt"
        prob_str = "_".join(v for v in intv_mech)
        file_name = file_root + f"sachs_P({prob_str}|do({ky}={val})).txt"
        with open(file_name, 'rb') as fp:
            label_data = pickle.load(fp)
        intv_label_data = torch.FloatTensor(label_data)
        cur_size = intv_label_data.shape[0]
        factor = int(max_size / cur_size) + 1
        replicated_data = intv_label_data.repeat(factor, 1)[0:max_size]
        dataset_dict[asKey(intv)]={"obs" : replicated_data.to(Exp.DEVICE)}



    for epoch in range(Exp.num_epochs):
        Exp.curr_epoochs = epoch
        labelMain(Exp, cur_hnodes, label_generators, optimizersMech, discriminatorsMech, doptimizersMech,
                             dataset_dict, tvd_diff, kl_diff)


