"""
Exploitation of the results pipelines. This is also available in the form of a notebook but this makes it more formal
Basically each subpart is a task and can be run independently by uncommenting the relevant part
"""

# Load a prediction file and do things with it such as computing the bagged and unbagged MSE or computing stds

import os
import pickle
import sys
from collections import defaultdict

import matplotlib.pyplot as plt
import numpy as np
import scipy.spatial as sp
import scipy.spatial.distance as distance
import seaborn as sns
import torch
from scipy.cluster import *
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.spatial.distance import pdist

if __name__ == "__main__":
    sys.path.append('../')


def rearrange_pdb_dict(preds):
    """
    Put the prediction in a double dict format
    :param preds: path (aaaa_GOL_6.npy) : prediction
    :return: pdb : ligand : list of predicted embeddings
    """
    rearranged = {}
    for key, value in preds.items():
        # print(key)
        # print(key.split('_')[:2])
        pdb, lig = key.split('_')[:2]
        if pdb not in rearranged:
            rearranged[pdb] = defaultdict(list)
        rearranged[pdb][lig].append(value)
    return rearranged


def reduce_preds(preds, average=False):
    """
    Clean this double dict for each 8 predictions by stacking them and optionnaly averaging them (bagging)
    :param preds: pdb : ligand : list of predicted embeddings
    :return: pdb : ligand : tensor of predictions
    """
    with torch.no_grad():
        reduced = {}
        for pdb, dic_ligs in preds.items():
            for ligand, tensor in dic_ligs.items():
                # Works also if we only have one prediction
                # tensor = [torch.zeros(128)]
                pred = torch.stack(tensor, 1)
                if average:
                    pred = pred.mean(dim=1)
                # print(avg.size())
                if pdb not in reduced:
                    reduced[pdb] = {ligand: pred}
                else:
                    reduced[pdb][ligand] = pred
    return reduced


'''
# Get predictions
model_name = 'small_siamsplit_aligned_flips'
pred_path = os.path.join('../data/post_processing/predictions/', model_name + '.p')
predictions = pickle.load(open(pred_path, 'rb'))

# reduce them
rearranged = rearrange_pdb_dict(predictions)
reduced = reduce_preds(rearranged)
reduced_averaged = reduce_preds(rearranged, average=True)

print(rearranged['5xem']['85F'])
print('reduced', reduced['5xem']['85F'].size())
print('reduced_averaged', reduced_averaged['5xem']['85F'].size())
print('reduced', reduced['5xem']['85F'][1, :])
print('reduced_averaged', reduced_averaged['5xem']['85F'][1])
print(np.mean([0.0399, 0.0660, 0.0834, 0.0710, 0.1035, 0.0945, 0.1044, 0.1036]))
'''


# # COMPUTE MSE
# bagged/non bagged version for each pdb,pocket


def get_accuracy(pred_dict, lig_embeddings):
    """
    Compute the mse for all elements of the pred dict
    :param pred_dict: pdb : ligand : tensor of predictions
    :param lig_embeddings:
    :return:
    """
    mses = list()
    for pdb, lig_dict in pred_dict.items():
        for lig, pred_fingerprint in lig_dict.items():
            # if we have one prediction ie pred comes from the bagged version:
            if len(pred_fingerprint.size()) == 1:
                true = lig_embeddings[lig]
                pred_ligand = pred_fingerprint.cpu().t().numpy()
                dist = sp.distance.euclidean(pred_ligand, true) ** 2
            elif len(pred_fingerprint.size()) == 2:
                true = emb[lig][np.newaxis, :]
                pred_ligand = pred_fingerprint.cpu().t().numpy()
                dist = sp.distance_matrix(pred_ligand, true) ** 2
            else:
                raise ValueError('Wrong dimension')
            dist /= 128
            mses.append(dist)
    mses = np.stack(mses)
    return mses


'''
emb = pickle.load(open('../data/ligands/whole_dict_embed_128.p', 'rb'))

# Get predictions
model_name = 'small_siamsplit_aligned_flips'
pred_path = os.path.join('../data/post_processing/predictions/', model_name + '.p')
predictions = pickle.load(open(pred_path, 'rb'))
rearranged = rearrange_pdb_dict(predictions)
reduced = reduce_preds(rearranged)
reduced_averaged = reduce_preds(rearranged, average=True)
acc1 = get_accuracy(reduced, emb)
acc2 = get_accuracy(reduced_averaged, emb)
print(acc1.shape, acc2.shape)
print(np.mean(acc1), np.mean(acc2))
'''


def acc(model, lig_embeddings):
    """
    Pipe those functions together to enable just using a model and a reference and returns accuracies
    :param model: 'String that is the name of a model
    :param lig_embeddings: Dict ligand : embedding
    :return:
    """
    pred_path = os.path.join('../data/post_processing/predictions/', model + '.p')
    predictions = pickle.load(open(pred_path, 'rb'))

    rearranged = rearrange_pdb_dict(predictions)
    reduced = reduce_preds(rearranged)
    reduced_averaged = reduce_preds(rearranged, average=True)

    acc1 = get_accuracy(reduced, lig_embeddings)
    acc2 = get_accuracy(reduced_averaged, lig_embeddings)

    return np.mean(acc1), np.mean(acc2)


def print_res(model, lig_embeddings):
    """
    Just a helper function to make a nice printing
    :param model: 'String that is the name of a model
    :param lig_embeddings: Dict ligand : embedding
    :return:
    """
    no_bag, bag = acc(model, lig_embeddings)
    print(f'{model} mse of {no_bag} and bagged mse of {bag}')


# print_res('small_siamsplit_aligned_flips_shuffled')
# print_res('small_siamsplit_aligned')
# print_res('small_siamsplit_aligned_flips')
# print_res('small_siamese_aligned_flips')
# print_res('true_small_siamese_aligned')
# print_res('baby_siamsplit_aligned_flips')
# print_res('small_siamsplit_flips')
# print_res('small_siamsplit')


# # Compute the error per ligand
# We want to see if some families are not predicted well by our script and plot a nice dendrogram


# returns a ligand : error dict
def ligand_results(pred_dict, lig_embeddings):
    """
    Compute the mse for all elements of the pred dict
    :param pred_dict: pdb : ligand : tensor of predictions
    :param lig_embeddings:
    :return:
    """
    lig_res = {}
    for pdb, lig_dict in pred_dict.items():
        for lig, pred_fingerprint in lig_dict.items():
            # if we have one prediction:
            if len(pred_fingerprint.size()) == 1:
                true = lig_embeddings[lig]
                pred_ligand = pred_fingerprint.cpu().t().numpy()
                dist = sp.distance.euclidean(pred_ligand, true) ** 2
                dist /= 128
                if lig not in lig_res:
                    lig_res[lig] = [dist]
                else:
                    lig_res[lig].append(dist)
            else:
                raise ValueError('Wrong dimension')
    return lig_res


# Once we have this result we want to plot a dendogram

palette = list(sns.color_palette("hls", 11).as_hex())
hierarchy.set_link_color_palette(palette)


def compute_clustering(ligand_dict, save=None):
    """

    :param ligand_dict: lig : (emb,acc)
    :param save: path of where to save the result
    :return:
    """
    # compute distances
    L = [l[0] for l in ligand_dict.values()]
    DM = pdist(L, metric='euclidean')

    # clustering
    linked = linkage(DM, method='average', optimal_ordering=True)

    # draw
    f, axes = plt.subplots(2, 1)
    dendrogram(linked,
               ax=axes[0],
               orientation='top',
               distance_sort='descending',
               above_threshold_color='grey',
               show_leaf_counts=True)
    # retrieve ordered ligand labels
    name_list = list(ligand_dict.keys())
    labels = list(axes[0].get_xticklabels())
    ligand_inds = [int(l.get_text()) for l in labels]
    labels = [name_list[i] for i in ligand_inds]

    axes[0].set_axis_off()
    axes[1].set_axis_off()

    # heatmap
    cmap = sns.diverging_palette(10, 220, center='light', as_cmap=True)
    cmap_data = [-ligand_dict[label][1] for label in labels]
    # clip values
    cmap_data = [max(val, -0.35) for val in cmap_data]
    # heatmap
    cmap = sns.diverging_palette(10, 220, center='light', as_cmap=True)

    g = sns.heatmap([cmap_data], ax=axes[1], annot=False, square=True, cmap=cmap, cbar=False)
    # print(list(zip(cmap_data, labels)))

    plt.subplots_adjust(hspace=-0.49)

    if save:
        plt.savefig(save)
    plt.show()

    pass


# lig_dict = {'BOB': ([0, 0, 1], 0.1), 'JOE': ([0, 1, 0], 0.5), 'TIT': ([0, 1, 1], 0.2)}
# compute_clustering(lig_dict)

# We get something super messy, we need to group some of the predictions together

def cluster(coords, t=40):
    """

    :param coords: numpy array [samples, values]
    :param t: The max number of clusters
    :return: an array such that arr[i] = cluster id for point i
    """
    Y = distance.pdist(coords)
    Z = hierarchy.linkage(Y, method='average', optimal_ordering=True)
    T = hierarchy.fcluster(Z, criterion='maxclust', t=t)
    return T


def compact_results(big_dict, t=40):
    """

    :param big_dict: lig_id : (emb,acc)
    :param t: The max number of clusters
    :return: compressed one : cluster_id : (mean(emb), mean(acc))
    """
    # Format the embeddings to be clustered easily
    coords = np.array(list(big_dict.values()))
    # check we keep the same order
    # print(coords[0])
    coords = np.stack(coords[:, 0])
    # print(coords.shape)
    # print(coords[0])

    T = cluster(coords, t=t)

    # Put all element of the same cluster in one dict category
    small = defaultdict(list)
    for i, (emb, acc) in enumerate(big_dict.values()):
        small[T[i]].append((emb, acc))

    # Aggregate the previous result
    small_dict = {}
    # all_accs = list()
    for key, value in small.items():
        embs, accs = zip(*value)
        embs = np.stack(list(embs))
        # print(embs.shape)
        mean_embs = embs.mean(axis=0)
        mean_accs = np.mean(list(accs))
        # all_accs.extend(list(accs))
        # print(mean_accs)

        small_dict[key] = (mean_embs, mean_accs)
    return small_dict



model_name = 'small_siamsplit_aligned_flips'
pred_path = os.path.join('../data/post_processing/predictions/', model_name + '.p')
predictions = pickle.load(open(pred_path, 'rb'))
rearranged = rearrange_pdb_dict(predictions)
reduced = reduce_preds(rearranged)
reduced_averaged = reduce_preds(rearranged, average=True)

emb = pickle.load(open('../data/ligands/whole_dict_embed_128.p', 'rb'))

lig_res = ligand_results(reduced_averaged, emb)
avg_lig_res = {lig: (emb[lig], np.mean(value)) for lig, value in lig_res.items()}

# compute_clustering(avg_lig_res)
small_avg_lig_res = compact_results(avg_lig_res)
# compute_clustering(small_avg_lig_res)
compute_clustering(small_avg_lig_res, 'dendrogram.pdf')



# Get the per dimension std and mse

def compute_mse_perdim(pred_dict, lig_embeddings):
    """
    Compute the mse for all elements of the pred dict
    :param pred_dict: pdb : ligand : tensor of predictions
    :param lig_embeddings:
    :return: a vector of mean mse for each dimension [0.13, 0,06 ... 0,24] (128,)
    """
    mses = list()
    for pdb, lig_dict in pred_dict.items():
        for lig, pred_fingerprint in lig_dict.items():
            # if we have one prediction:
            if len(pred_fingerprint.size()) == 1:
                true = lig_embeddings[lig]
                pred_ligand = pred_fingerprint.cpu().t().numpy()
                dist = (pred_ligand - true) ** 2
                mses.append(dist)
    mses = np.stack(mses)
    return mses

'''
emb = pickle.load(open('../data/ligands/whole_dict_embed_128.p', 'rb'))

# Get predictions
model_name = 'small_siamsplit_aligned_flips'
pred_path = os.path.join('../data/post_processing/predictions/', model_name + '.p')
predictions = pickle.load(open(pred_path, 'rb'))
rearranged = rearrange_pdb_dict(predictions)
reduced = reduce_preds(rearranged)
reduced_averaged = reduce_preds(rearranged, average=True)

mses = compute_mse_perdim(reduced_averaged, emb)
print(mses.shape)
mses_dim = np.mean(mses, axis=0)

stacked = np.stack(list(emb.values()))
# print(stacked.shape)
vars = np.std(stacked, axis=0) ** 2

# preds_dim = list()
# for pdb, lig_dict in reduced_averaged.items():
#     for lig, pred_fingerprint in lig_dict.items():
#         # if we have one prediction:
#         if len(pred_fingerprint.size()) ==1:
#             pred_ligand = pred_fingerprint.cpu().t().numpy()
#             preds_dim.append(pred_ligand)
#
# preds_dim = np.stack(preds_dim)
#
# for i in range(20):
#     dimi = stacked[i, :]
#     predi = preds_dim[i, :]
#     print(np.std(dimi) ** 2, mses_dim[i])
#     sns.distplot(dimi)
#     sns.distplot(predi)
#     plt.show()

plt.plot(vars)
plt.plot(mses_dim)
plt.show()

ratio = vars / (mses_dim)
ratio = -np.sort(-ratio)
log_num = float(np.sum(np.log(ratio)))
print('number of samples needed', np.exp(log_num))
print('average number of bins needed', np.mean(ratio))

ratio_mean = np.mean(ratio) * np.ones_like(ratio)
plt.plot(ratio)
plt.plot(ratio_mean, label='Mean value', color='blue', linestyle='dashed')
plt.plot(np.ones_like(ratio), label='No enrichment', color='red')
plt.legend()
plt.xlabel(r'Latent dimensions')
plt.ylabel(r'Variance/MSE')
plt.savefig('VAR_MSE_sorted.pdf')
plt.show()
plt.plot(vars - mses_dim)
'''
if __name__ == "__main__":
    pass

    # Get true labels and lig_to_pdb mapping
    # ligand_to_pdb = pickle.load(open('../data/post_processing/utils/all_lig_to_pdb.p', 'rb'))
    # emb = pickle.load(open('../data/ligands/whole_dict_embed_128.p', 'rb'))
    #
    # print_res('small_siamsplit')

    # loop = time.perf_counter()
    # a = time.perf_counter()
    # tot = list()
    # for batch_idx, (pdb, inputs, labels) in enumerate(train_loader):
    #     tot.extend(pdb)
    #     if not batch_idx % 20:
    #         print(batch_idx, time.perf_counter() - loop)
    #         loop = time.perf_counter()
    # print('Done in : ', time.perf_counter() - a)
    # print(sorted(tot))
    # Get the ligands : distance distribution
    # lig_dist = get_distances('small_siamsplit_aligned_flips')
    # lig_dist = get_distances(model_name)
    # print(lig_dist)
