# CelebA image generation using Conditional DCGAN
import os

import torch
import torchvision.transforms as transforms
import pickle

from Causal_MNIST_Images.DigitImageGeneration.mnist_image_generation import plot_dataset_digits, plot_trained_digits
from Causal_MNIST_Images.DigitImageGeneration.morphomnist import io

from Causal_MNIST_Images.mnistControllerModel import get_discriminators

from Causal_MNIST_Images.mnistControllerModel import get_generators
from Causal_MNIST_Images.GroundTruth.CausalGraph_Mnist import set_nonid_mnist_images_minimized, set_nonid_mnist_images
from ModularUtils.ControllerConstants import get_multiple_labels_fill
from ModularUtils.Experiment_Class import Experiment
from ModularUtils.FunctionsConstant import asKey
from ModularUtils.FunctionsTraining import labels_image_gradient_penalty, save_checkpoint


def get_dataset(Exp, label):
    dataset = []
    for feature in ["digit", "thickness", "color"]:
        file_name = Exp.file_roots[0] + 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 Imagetrain_CausalController(Exp, vars, label_generators, G_optimizers, label_discriminator, D_optimizer, obs_data_input, obs_images):
def Imagetrain_CausalController(Exp, cur_mechs, label_generators, G_optimizers, label_discriminator, D_optimizer,
                                dataset_dict_batches, imagedata_dict_batches, batchno):
    lenf= len(Exp.features)

    for interv_no, (intv_key_tup, dataset_batches) in enumerate(dataset_dict_batches.items()):
        intv_key = dict(intv_key_tup)

        obs_data_input= dataset_batches[batchno]
        obs_images= imagedata_dict_batches[intv_key_tup][batchno]

        compare_Var = ["Ydigit1", "Ycolor", "Ythick"]


        # getting real labels prepared for critic
        obs_mini_batch = obs_data_input.size()[0]
        current_real_label = obs_data_input.to(Exp.DEVICE)
        dims_list = [Exp.label_dim[lb][feat] for lb in compare_Var for feat in Exp.features ]
        obs_labels_fill = get_multiple_labels_fill(Exp, current_real_label, dims_list, isImage_labels=True, more_dimsize=Exp.IMAGE_SIZE)  # !!!
        obs_labels_fill_1 = get_multiple_labels_fill(Exp, current_real_label, dims_list, isImage_labels=True, more_dimsize=1)  # !!!


        # image generation
        Noises = []
        image_noise = torch.randn(obs_mini_batch, Exp.IMAGE_NOISE_DIM).view(-1, Exp.IMAGE_NOISE_DIM, 1, 1).to(
            Exp.DEVICE)
        Noises.append(image_noise)
        generated_image = label_generators[cur_mechs[0]](Noises, [obs_labels_fill_1])
        generated_labels_fill= obs_labels_fill


        D_losses = []
        for crit_ in range(Exp.CRITIC_ITERATIONS):
            D_fake_decision_obs = label_discriminator[0](generated_image, generated_labels_fill).squeeze()
            D_real_decision_obs = label_discriminator[0](obs_images, obs_labels_fill).squeeze()
            gp_obs = labels_image_gradient_penalty(label_discriminator[0], obs_images, obs_labels_fill, generated_image, 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)

            D_losses.append(D_loss_obs.data)  # just a loss list
            label_discriminator[0].zero_grad()
            D_loss_obs.backward(retain_graph=True)
            D_optimizer[0].step()


    D_fake_decision_obs = label_discriminator[0](generated_image, 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

    randices =torch.randint(0,obs_images.shape[0],(5,)).tolist()

    # return G_loss.data, D_loss.data, obs_images[randices], obs_data_input[randices], generated_image[randices], generated_labels_disc[randices]
    return G_loss.data, D_loss.data, obs_images[randices], obs_data_input[randices], generated_image[randices], obs_data_input[randices]



def imageMain(Exp, cur_mechs, label_generators, G_optimizers, discriminators, D_optimizers, dataset_dict,
              imagedata_dict, tvd_diff, kl_diff):

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

        batch_list = []
        for data_input in real_dataloader:
            data_input = torch.squeeze(data_input)
            batch_list.append(data_input)

        dataset_dict_batches[key] = batch_list
        num_batches = len(batch_list)
    #

    imagedata_dict_batches = {}
    for key, each_dataset in imagedata_dict.items():
        image_data_loader = torch.utils.data.DataLoader(dataset=each_dataset,
                                                        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)

        imagedata_dict_batches[key] = batch_list
        num_batches2 = len(batch_list)

    iteration = 0

    for batchno in range(num_batches):
        G_loss, D_loss, obs_images, obs_labels, gen_images, gen_labels = Imagetrain_CausalController(Exp, cur_mechs, label_generators, G_optimizers, discriminators,
                                                     D_optimizers,
                                                     dataset_dict_batches,
                                                     imagedata_dict_batches,
                                                     batchno)


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

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


        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) % 10 == 0:
        print("Showing digits")

        for orow, obsimg, grow, genimg in zip(obs_labels, obs_images, gen_labels, gen_images):
            print("obs", orow)
            print("gen", grow)
            obsimg = obsimg.permute(1, 2, 0).detach().cpu().numpy()
            genimg = genimg.permute(1, 2, 0).detach().cpu().numpy()
            # plot_dataset_digits(1, 2, [obsimg, genimg], f'Real {Ores_digit[id]}')
            plot_trained_digits(1, 2, [obsimg, genimg], f'Real {orow}')



        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)

        print("saved at ", Exp.SAVED_PATH)


    return 100







if __name__ == "__main__":

    Exp = Experiment("Exp1", set_nonid_mnist_images,
                     dist_thresh=0.15,
                     causal_hierarchy=2,
                     Temperature=1,
                     temp_min=0.01,
                     NOISE_DIM=128,
                     CONF_NOISE_DIM=128,
                     G_hid_dims=[256, 256],
                     D_hid_dims=[256, 256, 256],
                     IMAGE_FILTERS=[128, 64, 32],
                     CRITIC_ITERATIONS=1,
                     LAMBDA_GP=10,
                     learning_rate=2 * 1e-4,
                     Synthetic_Sample_Size=40000,
                     intv_Sample_Size=40000,
                     batch_size=200,
                     features=["feature"],
                     noise_states=100,
                     latent_state=16,
                     # Data_intervs=[{}, {"X1": 0}, {"X1": 1}],
                     Data_intervs=[{}],
                     num_epochs=300,
                     new_experiment=True
                     )


    print(Exp.Data_intervs)
    Exp.intv_batch_size = Exp.batch_size

    os.makedirs(Exp.SAVED_PATH, exist_ok=True)
    dag_name = Exp.Complete_DAG_desc + ".txt"
    # Exp.LOAD_MODEL_PATH = "/path_to_project/SAVED_EXPERIMENTS/nonid_mnist_images/Exp1/Sep_23_2022-10_45"
    # Exp.load_which_models = {"X1": True, "X2": True, "W": True, "ImgYdigit1": True, "ImgYdigit2":True}
    Exp.load_which_models = {"X1": False, "X2": False, "W": False, "Ydigit1":False, "Ydigit2":False, "Ycolor":False, "Ythick":False,
                             "ImgYdigit1": False, "ImgYdigit2":False}
    cur_mechs = ["ImgYdigit1"]
    discrete_mechs= cur_mechs[0:-1]
    # cur_mechs = ["X1", "X2", "W", "ImgYdigit2"] or

    label_generators, optimizersMech = get_generators(Exp, Exp.load_which_models)

    discriminatorsMech, doptimizersMech = get_discriminators(Exp, cur_mechs, discrete_mechs,
                                                             Exp.load_which_models)  # update for interventional training


    dataset_dict = {}
    dno=-1
    each_dataset = [torch.tensor(io.load_idx(Exp.file_roots[dno] + "digitlabels.gz"))]
    dataset_dict[asKey(Exp.Data_intervs[dno])] = torch.cat(each_dataset, 1).to(Exp.DEVICE)

    # image loading
    imagedata_dict = {}
    each_dataset = []
    # image dataset load
    dno = -1
    loaded_images = io.load_idx(Exp.file_roots[dno] + "digitimages.gz")
    transform = transforms.Compose([transforms.ToPILImage(),
                                    # transforms.Scale(Exp.IMAGE_SIZE),
                                    transforms.ToTensor()
                                    # transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
                                    ])  # need to normalize too since I am using tanh?

    # for img in loaded_images:
    #     plot_dataset_digits((img, "x"))
    result_images = [torch.unsqueeze(transform(img), dim=0).to(Exp.DEVICE) for img in loaded_images]
    imagedata_dict[asKey(Exp.Data_intervs[dno])] = torch.cat(result_images, 0)






    mech_tvd = 0
    cur_mech_epoch = 0
    tvd_diff = {}
    kl_diff = {}

    for epoch in range(Exp.num_epochs):
        Exp.curr_epoochs = epoch
        mech_tvd = imageMain(Exp, cur_mechs, label_generators, optimizersMech, discriminatorsMech, doptimizersMech,
                             dataset_dict, imagedata_dict, tvd_diff, kl_diff)
