from time import time

import math, os

from sklearn import metrics

from sklearn.cluster import KMeans

import torch

import torch.nn as nn

from torch.autograd import Variable

from torch.nn import Parameter

import torch.nn.functional as F

import torch.optim as optim

from torch.utils.data import DataLoader, TensorDataset

from scMDC import scMultiCluster

import numpy as np

import collections

import h5py

import scanpy as sc

from preprocess import read_dataset, normalize, clr_normalize_each_cell

from utils import *

from functools import reduce

from LRP import LRP

if __name__ == "__main__":

    # setting the hyper parameters

    import argparse

    parser = argparse.ArgumentParser(description='train',

                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--n_clusters', default=8, type=int)

    parser.add_argument('--cutoff', default=0.5, type=float, help='Start to train combined layer after what ratio of epoch')

    parser.add_argument('--batch_size', default=256, type=int)

    parser.add_argument('--data_file', default='Simulation.1.h5')

    parser.add_argument('--cluster_index_file', default='label.txt')

    parser.add_argument('--maxiter', default=10000, type=int)

    parser.add_argument('--pretrain_epochs', default=400, type=int)

    parser.add_argument('--gamma', default=.1, type=float,

                        help='coefficient of clustering loss')

    parser.add_argument('--tau', default=1., type=float,

                        help='fuzziness of clustering loss')       

    parser.add_argument('--phi1', default=0.001, type=float,

                        help='coefficient of KL loss in pretraining stage')

    parser.add_argument('--phi2', default=0.001, type=float,

                        help='coefficient of KL loss in clustering stage')

    parser.add_argument('--update_interval', default=1, type=int)

    parser.add_argument('--tol', default=0.001, type=float)

    parser.add_argument('--ae_weights', default=None)

    parser.add_argument('--save_dir', default='results/')

    parser.add_argument('--ae_weight_file', default='AE_weights_1.pth.tar')

    parser.add_argument('--resolution', default=0.2, type=float)

    parser.add_argument('--n_neighbors', default=30, type=int)

    parser.add_argument('--embedding_file', action='store_true', default=False)

    parser.add_argument('--prediction_file', action='store_true', default=False)

    parser.add_argument('-el','--encodeLayer', nargs='+', default=[256,64,32,16])

    parser.add_argument('-dl1','--decodeLayer1', nargs='+', default=[16,64,256])

    parser.add_argument('-dl2','--decodeLayer2', nargs='+', default=[16,20])

    parser.add_argument('--sigma1', default=2.5, type=float)

    parser.add_argument('--sigma2', default=1.5, type=float)

    parser.add_argument('--f1', default=1000, type=float, help='Number of mRNA after feature selection')

    parser.add_argument('--f2', default=2000, type=float, help='Number of ADT/ATAC after feature selection')

    parser.add_argument('--filter1', action='store_true', default=False, help='Do mRNA selection')

    parser.add_argument('--filter2', action='store_true', default=False, help='Do ADT/ATAC selection')

    parser.add_argument('--run', default=1, type=int)

    parser.add_argument('--beta', default=1., type=float,

                        help='coefficient of the clustering fuzziness')

    parser.add_argument('--margin', default=1., type=float,

                        help='margin of difference between logits')

    parser.add_argument('--lamda', default=100., type=float,

                        help='coefficient of the clustering perturbation loss')

    parser.add_argument('--lr', default=0.001, type=int)

    parser.add_argument('--device', default='cuda')

    args = parser.parse_args()

    print(args)

    data_mat = h5py.File(args.data_file)

    x1 = np.array(data_mat['X1'])

    x2 = np.array(data_mat['X2'])

    #y = np.array(data_mat['Y']) - 1

    data_mat.close()

    clust_ids = np.loadtxt(args.cluster_index_file, delimiter=",").astype(int)

    #Gene features

    if args.filter1:

        importantGenes = geneSelection(x1, n=args.f1, plot=False)

        x1 = x1[:, importantGenes]

    if args.filter2:

        importantGenes = geneSelection(x2, n=args.f2, plot=False)

        x2 = x2[:, importantGenes]

    adata1 = sc.AnnData(x1)

    #adata1.obs['Group'] = y

    adata1 = read_dataset(adata1,

                     transpose=False,

                     test_split=False,

                     copy=True)

    adata1 = normalize(adata1,

                      size_factors=True,

                      normalize_input=True,

                      logtrans_input=True)

    adata2 = sc.AnnData(x2)

    #adata2.obs['Group'] = y

    adata2 = read_dataset(adata2,

                     transpose=False,

                     test_split=False,

                     copy=True)

    adata2 = normalize(adata2,

                      size_factors=True,

                      normalize_input=True,

                      logtrans_input=True)

    #adata2 = clr_normalize_each_cell(adata2)

    input_size1 = adata1.n_vars

    input_size2 = adata2.n_vars

    print(adata1.X.shape)

    print(adata2.X.shape)

    print(args)

    encodeLayer = list(map(int, args.encodeLayer))

    decodeLayer1 = list(map(int, args.decodeLayer1))

    decodeLayer2 = list(map(int, args.decodeLayer2))

    model = scMultiCluster(input_dim1=input_size1, input_dim2=input_size2, tau=args.tau,

                        encodeLayer=encodeLayer, decodeLayer1=decodeLayer1, decodeLayer2=decodeLayer2,

                        activation='elu', sigma1=args.sigma1, sigma2=args.sigma2, gamma=args.gamma, 

                        cutoff = args.cutoff, phi1=args.phi1, phi2=args.phi2, device=args.device).to(args.device)

    print(str(model))

    if os.path.isfile(args.ae_weights):

            print("==> loading checkpoint '{}'".format(args.ae_weights))

            checkpoint = torch.load(args.ae_weights)

            model.load_state_dict(checkpoint['ae_state_dict'])

    else:

        print("==> no checkpoint found at '{}'".format(args.ae_weights))

        raise ValueError

    n_clusters = np.unique(clust_ids).shape[0]

    print("n cluster is: " + str(n_clusters))

    Z = model.encodeBatch(torch.tensor(adata1.X).to(args.device), torch.tensor(adata2.X).to(args.device)).data.cpu().numpy()

    cluster_list = np.unique(clust_ids).astype(int).tolist()

    print(cluster_list)

    model_explainer = LRP(model, X1=adata1.X, X2=adata2.X, Z=Z, clust_ids=clust_ids, n_clusters=n_clusters, beta=args.beta).to(args.device)

    #for clust_c in [cluster_ind[0]]: #range(args.n_clusters):

    #    for clust_k in [cluster_ind[1]]: #range(clust_c+1, args.n_clusters):

    #        print("Cluster"+str(clust_c)+" vs Cluster"+str(clust_k))

    #        rel_score1, rel_score2  = model_explainer.calc_carlini_wagner_one_vs_one(clust_c, clust_k, margin=args.margin, lamda=args.lamda, max_iter=args.maxiter, lr=args.lr)

    #        print(rel_score1.shape)

    #        print(rel_score2.shape)

    #        np.savetxt(args.save_dir + "/" + str(clust_c)+"_vs_"+str(clust_k)+"_rel_mRNA_scores.csv", rel_score1, delimiter=",")

    #        np.savetxt(args.save_dir + "/" + str(clust_c)+"_vs_"+str(clust_k)+"_rel_ADT_scores.csv", rel_score2, delimiter=",")

    for clust_c in cluster_list:

        print("Cluster"+str(clust_c)+" vs Rest")

        rel_score1, rel_score2 = model_explainer.calc_carlini_wagner_one_vs_rest(clust_c, margin=args.margin, lamda=args.lamda, max_iter=args.maxiter, lr=args.lr)

        np.savetxt(args.save_dir + "/" + str(clust_c)+"_vs_rest_rel_mRNA_scores.csv", rel_score1, delimiter=",")

        np.savetxt(args.save_dir + "/" + str(clust_c)+"_vs_rest_rel_ADT_scores.csv", rel_score2, delimiter=",")