import numpy as np
import pandas as pd

# plotting
import matplotlib.pyplot as plt
import seaborn as sns

# scipy
import scipy
import scipy.cluster
import scipy.cluster.hierarchy
from scipy.cluster.hierarchy import cophenet
from scipy.spatial.distance import pdist


# sklearn
import sklearn
from sklearn import metrics
import sklearn.datasets
import sklearn.cluster

from pathlib import Path
import os
import sys
import traceback

from datetime import datetime


def plot_dendrogram(model):
    counts = np.zeros(model.children_.shape[0])

    linkage_matrix = np.column_stack(
        [model.children_, model.distances_, counts]
    ).astype(float)

    return linkage_matrix


def correct_method_affinity(yamldict):
    # TODO full automatic yaml correction
    try:
        method_from_yaml = yamldict["cluster_method"]
    except KeyError:
        print("No cluster method was chosen. Using Ward with euclidean affinity")
        method_from_yaml = "ward"
        affinity_from_yaml = "euclidean"
    try:
        affinity_from_yaml = yamldict["cluster_affinity"]
    except KeyError:
        affinity_from_yaml = "euclidean"
        print("No affinity was chosen. Using euclidean.")
    if method_from_yaml == "ward":
        if affinity_from_yaml != "euclidean":
            affinity_from_yaml = "euclidean"
            print("Ward method uses only euclidean affinity")
    if affinity_from_yaml == "precomputed":
        print(
            "This option would need a distance matrix, and this script is not suited for this. \n Changing to euclidean"
        )
        affinity_from_yaml = "euclidean"
    if method_from_yaml == "Single":
        method_from_yaml = "single"
    return method_from_yaml, affinity_from_yaml


def prepare_base_fname(
    project_name,
    project_path,
    cluster_confidence_threshold,
    method_from_yaml,
    affinity_from_yaml,
    cluster_row,
    cluster_src,
):
    c = str(round(cluster_confidence_threshold, 2)).replace(".", ",")
    base_fname = f"{project_name}-{cluster_src}_{method_from_yaml}-{affinity_from_yaml}_thr-{c}_row-{int(cluster_row)}"
    base_fname = Path(os.path.join(project_path, base_fname))
    return base_fname


def invert_labels(labels_input):
    labels = labels_input.copy()
    labels = np.where((labels == 0) | (labels == 1), labels ^ 1, labels)
    return labels


def clustering(data, method_from_yaml, affinity_from_yaml, cluster_row):

    clustering = sklearn.cluster.AgglomerativeClustering(
        n_clusters=2,
        affinity=affinity_from_yaml,
        memory=None,
        connectivity=None,
        compute_full_tree="auto",
        linkage=method_from_yaml,
        distance_threshold=None,
        compute_distances=True,
    )

    if cluster_row == True:

        cluster_fit_model_for_rows = clustering.fit(data)
        linkage_matrix_rows = plot_dendrogram(cluster_fit_model_for_rows)
    else:
        linkage_matrix_rows = None
        cluster_fit_model_for_rows = None
    cluster_fit_model_for_cols = clustering.fit(data.T)
    linkage_matrix_cols = plot_dendrogram(cluster_fit_model_for_cols)

    return (
        cluster_fit_model_for_cols,
        cluster_fit_model_for_rows,
        linkage_matrix_rows,
        linkage_matrix_cols,
    )


def clustering_image(
    data,
    cluster_row,
    linkage_matrix_rows,
    linkage_matrix_cols,
    base_fname,
    cluster_show,
    cluster_save,
    cluster_title_fontsize,
    cluster_width,
    cluster_height,
    cmap,
    group_divide,
):
    now = datetime.now().strftime("%H:%M:%S")
    print(f"{now} - Preparing clustering image")
    title = str(os.path.basename(base_fname))
    title = title.split("_")
    title = " ".join(title)
    cpgs_count = len(data)
    if (cluster_row == "True") | (cluster_row == True):

        g = sns.clustermap(
            data,
            row_linkage=linkage_matrix_rows,
            col_linkage=linkage_matrix_cols,
            dendrogram_ratio=[0.2, 0.2],
            col_cluster=True,
            row_cluster=cluster_row,
            cmap=cmap,
        )

        g.ax_heatmap.set_xticklabels(g.ax_heatmap.get_xmajorticklabels(), fontsize=12)
        g.ax_heatmap.set_yticklabels(g.ax_heatmap.get_ymajorticklabels(), fontsize=12)

        # Rotate tick labels for better visibility
        plt.setp(g.ax_heatmap.get_xticklabels(), rotation=90)
        plt.setp(g.ax_heatmap.get_yticklabels(), rotation=0)
    else:
        g = sns.clustermap(
            data,
            col_linkage=linkage_matrix_cols,
            dendrogram_ratio=[0.2, 0.2],
            col_cluster=True,
            row_cluster=cluster_row,
            cmap=cmap,
            # TODO save or show
        )
    for tick_label in g.ax_heatmap.axes.get_xticklabels():
        if tick_label._text in data.columns[:group_divide]:
            tick_label.set_color("red")
        else:
            tick_label.set_color("blue")
    g.fig.suptitle(
        f"{title}\n cpgs found: {cpgs_count}", fontsize=cluster_title_fontsize
    )
    g.fig.set_size_inches(cluster_width, cluster_height)
    plt.tight_layout()
    saveto = base_fname.with_suffix(".png")

    if cluster_save == "True":
        now = datetime.now().strftime("%H:%M:%S")
        print(f"{now} Saving clustering image to : \n {saveto}")
        plt.savefig(saveto, facecolor="white")
    if cluster_show == "True":
        plt.show()
    plt.close()


def clustering_analysis(
    original_data, real_labels, cluster_labels, linkage_matrix_cols
):
    now = datetime.now().strftime("%H:%M:%S")
    print(
        f"{now} Starting clustering analysis\
    \nClustering performance analysis is done only for columns -> control vs study group analysis"
    )

    def rand_index(real_labels, cluster_labels):

        print(f" real data labels {real_labels}")
        print(f"clustering labels {cluster_labels}")
        rs = metrics.rand_score(real_labels, cluster_labels)
        ars = metrics.adjusted_rand_score(real_labels, cluster_labels)

        print(f"Rand index {rs}, {ars}-> perfect score is 1.0")
        return [[rs, ars], ["Rand index", "Adjusted rand index"]]

    def mutual_information_based_score(real_labels, cluster_labels):

        mis = metrics.mutual_info_score(real_labels, cluster_labels)
        amis = metrics.adjusted_mutual_info_score(real_labels, cluster_labels)
        nmis = metrics.normalized_mutual_info_score(real_labels, cluster_labels)

        print(
            f"Mutual information based scores {mis}/{amis}/{nmis}->perfect score is 1.0"
        )
        return [
            [mis, amis, nmis],
            ["Mutual information score", "Adjusted MIS", "Normalzed MIS"],
        ]

    def hcv_measurre(real_labels, cluster_labels):

        hm, cm, vm = metrics.homogeneity_completeness_v_measure(
            real_labels, cluster_labels
        )
        print(f"Homogenity {hm}, completness {cm} and V-measure {vm}")
        return [[hm, cm, vm], ["Homogenity", "Completness", "V-measure"]]

    def fowlkes_mallows_score(real_labels, cluster_labels):

        fms = metrics.fowlkes_mallows_score(real_labels, cluster_labels)
        print(f"Fowlkes_mallow_score {fms}")
        return [[fms], ["Fowlkes_mallow_score"]]

    def calculate_cophenet(linkage_matrix_cols, original_data):

        c, coph_dists = cophenet(linkage_matrix_cols, pdist(original_data.T))
        return c

    values = []
    values_labels = []
    values.append(real_labels)
    values_labels.append("Real labels")
    values.append(cluster_labels)
    values_labels.append("Cluster labels")
    functions = [
        rand_index,
        mutual_information_based_score,
        hcv_measurre,
        fowlkes_mallows_score,
    ]

    values.append(calculate_cophenet(linkage_matrix_cols, original_data))
    values_labels.append("Cophentic correlation distance")
    for function in functions:
        function_output = function(real_labels, cluster_labels)
        values.extend(function_output[0])
        values_labels.extend(function_output[1])

    cluster_labels_inv = invert_labels(cluster_labels)

    values.append(cluster_labels_inv)
    values_labels.append("Cluster labels inv")

    for function in functions:

        function_output = function(real_labels, cluster_labels_inv)
        values.extend(function_output[0])
        values_labels.extend(function_output[1])

    values.append(len(original_data))
    values_labels.append("Cpgs found")
    isinverted = np.ones(len(values))
    isinverted[:] = 2
    isinverted[1:12] = 0
    isinverted[2] = 2
    isinverted[12:] = 1
    df_clustering_analysis = pd.Series(values, index=values_labels)
    df_clustering_analysis_isinv = pd.Series(isinverted, index=values_labels)
    df_clustering_analysis = pd.concat(
        [df_clustering_analysis, df_clustering_analysis_isinv], axis=1
    )

    df_clustering_analysis.index.name = "Method"
    df_clustering_analysis.columns = ["Values", "isinv"]
    return values, values_labels, df_clustering_analysis


def prepare_confidence_thresholds_percentiles(base_data, yaml_dict):
    try:
        confidence_thresholds_percentiles = np.percentile(
            base_data.sort_values(by="Confidence sum", ascending=False)[
                "Confidence sum"
            ],
            [0, 50, 75, 95],
        )
    except Exception:
        print("no percentile -> empty data")
        print(traceback.format_exc())
        confidence_thresholds_percentiles = [0]
    return confidence_thresholds_percentiles


def get_labels_from_linkage(linkage):
    from scipy.cluster.hierarchy import fcluster

    labels = []
    n = 5
    labels_set = []

    print(n, labels)
    labels = fcluster(linkage, 2, criterion="maxclust")
    labels = labels - 1

    return labels


def hc(yaml_dict):
    now = datetime.now().strftime("%H:%M:%S")
    print(f"\n{now} Hierachical clustering module - start")
    method_from_yaml, affinity_from_yaml = correct_method_affinity(yaml_dict)

    cluster_row = yaml_dict["cluster_row"]
    cluster_show = yaml_dict["cluster_show"]
    cluster_save = yaml_dict["cluster_save"]
    project_name = yaml_dict["project_name"]
    project_path = yaml_dict["project_path"]
    cluster_row = yaml_dict["cluster_row"]
    cluster_src = ["hh"]
    group_divide = yaml_dict["group_divide"]
    cluster_title_fontsize = yaml_dict["cluster_title_fontsize"]
    cluster_width = yaml_dict["cluster_width"]
    cluster_height = yaml_dict["cluster_height"]
    cluster_cmap = yaml_dict["cluster_cmap"]
    cluster_srcs = yaml_dict["cluster_src"]

    if isinstance(cluster_srcs, type(None)):
        cluster_srcs = ["hh"]
    elif isinstance(cluster_srcs, str):
        cluster_srcs = [cluster_srcs]
    elif isinstance(cluster_srcs, list):
        pass
    else:
        cluster_srcs = ["hh"]

    def what_to_read(cluster_src, yaml_dict):
        if cluster_src == "hh":
            read_path = yaml_dict["output_path_only_hyper_hypo"]
        elif cluster_src == "hhi":
            read_path = yaml_dict["output_path_hyper_hypo_indeter"]
        elif cluster_src == "ii":
            read_path = yaml_dict["output_path_indeter_indeter"]
        elif cluster_src == "full":
            read_path = yaml_dict["output_path_full"]
        else:
            print("wrong cluster_src name-using hyper hypo")
            read_path = yaml_dict["output_path_only_hyper_hypo"]

        data = pd.read_csv(read_path, index_col=0, sep=",")

        return data

    for cluster_src in cluster_srcs:
        now = datetime.now().strftime("%H:%M:%S")
        print(f"{now} - Preparing data for clustering")
        base_data = what_to_read(cluster_src, yaml_dict)
        if len(base_data) > 2:
            data = base_data.iloc[:, 6:]
            data_labels = np.ones(len(data.columns), dtype=np.int0)
            data_labels[:group_divide] = 0

            cluster_confidence_thresholds = yaml_dict["cluster_confidence_thresholds"]
            if isinstance(cluster_confidence_thresholds, type("")) or isinstance(
                cluster_confidence_thresholds,
                type(None) or (cluster_confidence_thresholds == "null"),
            ):
                cluster_confidence_thresholds = (
                    prepare_confidence_thresholds_percentiles(base_data, yaml_dict)
                )
                print(
                    f"Confidence thresholds based on percentiles - {cluster_confidence_thresholds}"
                )
            elif isinstance(cluster_confidence_thresholds, type(1)) or isinstance(
                cluster_confidence_thresholds, type(1.0)
            ):
                print(
                    f"Confidence threshold based on single values from settings - {cluster_confidence_thresholds}"
                )
                cluster_confidence_thresholds = [cluster_confidence_thresholds]
            elif isinstance(cluster_confidence_thresholds, list):
                print(
                    "Confidence based on array of confidence thresholds from settings"
                )
            for cluster_confidence_threshold in cluster_confidence_thresholds:
                data = data.loc[
                    base_data["Confidence sum"] >= cluster_confidence_threshold
                ]

                base_fname = prepare_base_fname(
                    project_name,
                    project_path,
                    cluster_confidence_threshold,
                    method_from_yaml,
                    affinity_from_yaml,
                    cluster_row,
                    cluster_src,
                )
                #

                try:
                    import fastcluster

                    now = datetime.now().strftime("%H:%M:%S")
                    print(f"{now} - Clustering with fastcluster")

                    if (yaml_dict["cluster_row"] == True) | (
                        yaml_dict["cluster_row"] == "True"
                    ):
                        data_to_clustering_rows = data
                        linkage_matrix_rows = fastcluster.linkage_vector(
                            data_to_clustering_rows,
                            method=f"{method_from_yaml}",
                            metric=f"{affinity_from_yaml}",
                        )
                    else:
                        data_to_clustering_rows = []
                        linkage_matrix_rows = []
                    data_to_clustering_cols = data.T
                    linkage_matrix_cols = fastcluster.linkage(
                        data_to_clustering_cols,
                        method=method_from_yaml,
                        metric=affinity_from_yaml,
                    )
                except Exception as e:
                    input(e)
                    print(traceback.format_exc())
                    now = datetime.now().strftime("%H:%M:%S")
                    print(
                        f"{now} - Clustering with sklearn.cluster.AgglomerativeClustering"
                    )
                    print(
                        "It is recommended to use fastcluster with bigger data\
                    \nhttps://pypi.org/project/fastcluster/"
                    )
                    try:
                        (
                            cluster_fit_model_for_cols,
                            cluster_fit_model_for_rows,
                            linkage_matrix_rows,
                            linkage_matrix_cols,
                        ) = clustering(
                            data, method_from_yaml, affinity_from_yaml, cluster_row
                        )
                    except Exception as e:
                        print("empty data")  # TODO
                        print(traceback.format_exc())
                        print(cluster_save)
                try:
                    if cluster_show or cluster_save:
                        clustering_image(
                            data,
                            cluster_row,
                            linkage_matrix_rows,
                            linkage_matrix_cols,
                            base_fname,
                            cluster_show,
                            cluster_save,
                            cluster_title_fontsize,
                            cluster_width,
                            cluster_height,
                            cluster_cmap,
                            group_divide,
                        )
                    else:
                        print("passing out")
                        pass
                except Exception:
                    print(traceback.format_exc())
                #

                try:

                    if "fastcluster" in sys.modules:
                        linkage = 0
                        cluster_labels = get_labels_from_linkage(linkage_matrix_cols)

                    else:
                        cluster_labels = cluster_fit_model_for_cols.labels_

                    clustering_analysis_data = clustering_analysis(
                        data, data_labels, cluster_labels, linkage_matrix_cols
                    )
                except Exception:
                    print(traceback.format_exc())

                def save_clustering_analysis(base_fname, df_cluster_analysis):
                    parent = base_fname.parent
                    fname = f"{base_fname.name}-cluster_analysis"

                    fname = Path(os.path.join(parent, fname))
                    now = datetime.now().strftime("%H:%M:%S")
                    print(f"{now} - Saving clustering analysis to: {fname}")
                    fname = fname.with_suffix(".csv")

                    df_cluster_analysis.to_csv(fname)

                try:
                    save_clustering_analysis(base_fname, clustering_analysis_data[2])
                except Exception:
                    print(traceback.format_exc())
                    print("empty dataset")
