import networkx as nx
import matplotlib.pyplot as plt
import pandas as pd
import os
import numpy as np
import math


def single_cpg_plot(outputpath, df, cpgname, group_divide):
    df = df.iloc[:, 6:]
    try:
        output_data2 = df.loc[cpgname]
        y1 = range(0, group_divide, 1)
        end = len(df.columns)
        y2 = range(group_divide, end)

        title = f"{cpgname}"
        plt.title(title)
        plt.scatter(y1, output_data2.iloc[:group_divide])
        plt.scatter(y2, output_data2.iloc[group_divide:])
        plt.axvline(group_divide, color="red")
        plt.axhline(0.5, color="black")
        plt.ylim(0, 1)
        plt.xticks(ticks=range(0, end), labels=list(df.columns), rotation=90)
        plt.tight_layout()
        print(outputpath)
        plt.savefig(outputpath)
        plt.close()
    except KeyError:
        print("no cpg found")
        input()


def paths_method(params):
    depth = params["paths_depth"]

    paths_output = params["project_path"]

    paths_output = os.path.join(paths_output, "paths")

    try:
        os.mkdir(paths_output)
    except:
        import traceback

        print(traceback.format_exc())

    sep = ","
    src = params["paths_src"]
    if src == "hh":
        input_data_path = params["output_path_only_hyper_hypo"]
    elif src == "full":
        input_data_path = params["output_path_full"]
    elif src == "hhi":
        input_data_path = params["output_path_hyper_hypo_indeter"]
    elif src == "ii":
        input_data_path = params["output_path_indeter_indeter"]
    else:
        input_data_path = src
        input_data_path = params["paths_src: hh"]
        sep = params["paths_src_sep"]

    cpgs_dataset_path = input_data_path
    cpgs_dataset = pd.read_csv(cpgs_dataset_path, index_col=0, sep=",")
    cpgs_dataset_for_ref = cpgs_dataset.copy()

    input_data = cpgs_dataset
    sortby = input_data.columns[5]
    input_data_sorted = input_data.sort_values(by=sortby, ascending=False)

    confidence_thresholds = params["paths_confidence_thresholds"]

    if not isinstance(confidence_thresholds, list):
        if isinstance(confidence_thresholds, int) or isinstance(
            confidence_thresholds, float
        ):
            confidence_thresholds = [confidence_thresholds]
        else:
            quantiles = input_data_sorted.iloc[:, 5].quantile([0.9, 0.95, 0.99])
            confidence_thresholds = quantiles.values

    annotation_dataset_path = params["annotation_dataset_path"]
    adskip = params["annotation_dataset_skiprows"]
    adsep = params["annotation_dataset_sep"]
    try:
        annotation_dataset = pd.read_csv(
            annotation_dataset_path, sep=adsep, skiprows=adskip, index_col=0
        )
    except Exception as e:
        print(e)
        annotation_dataset = pd.read_csv(
            annotation_dataset_path, sep=adsep, skiprows=adskip, index_col=0
        )
        print("provide correct path to cpgs annotation dataset")
        print("use CTMeth downloading module to download the datasets")
        input()
        quit()

    biogrid_dataset_path = params["biogrid_dataset_path"]
    bd_skip = params["biogrid_dataset_skiprows"]
    bd_sep = params["biogrid_dataset_sep"]
    bd_index = params["biogrid_dataset_index_col"]

    biogrid_dataset = pd.read_csv(
        biogrid_dataset_path, skiprows=bd_skip, sep=bd_sep, index_col=bd_index
    )

    colname = params["annotation_dataset_USC_RefGene_Name_colname"]
    USC_RefGenes_dataset = annotation_dataset[colname]

    G = nx.Graph()

    print("Adding cpgs nodes")
    genes_nodes_set = set()

    idx = 0
    # TODO  for confidence_threshold in confidence_thresholds:
    for confidence_threshold in confidence_thresholds:
        input_data_sorted = input_data_sorted[
            input_data_sorted[sortby] >= confidence_threshold
        ]
        cpgs_dataset = input_data_sorted.iloc[:, 6:]
        cpgs_list = cpgs_dataset.index.tolist()
        USC_RefGenes_dataset_selected = USC_RefGenes_dataset.loc[
            USC_RefGenes_dataset.index.intersection(cpgs_list)
        ].reindex(cpgs_list)

        for cpg in cpgs_list:
            idx = idx + 1
            print(f"{cpg},{idx}/{len(cpgs_list)}")
            G.add_node(cpg)
            if isinstance(USC_RefGenes_dataset_selected.loc[cpg], str):
                cpgs_genes = USC_RefGenes_dataset_selected.loc[cpg].split(";")

                for gene in cpgs_genes:

                    if len(gene) > 1:
                        G.add_node(gene)
                        genes_nodes_set.update([gene])
                        G.add_edge(cpg, gene)
                        try:
                            interactors = biogrid_dataset.loc[gene].OFFICIAL_SYMBOL_B
                            for interactor in interactors:
                                G.add_node(interactor)
                                genes_nodes_set.update([interactor])
                                G.add_edge(gene, interactor)
                        except KeyError:
                            print("gene not found in biogrid")

        pos = nx.circular_layout(G)

        genes_nodes_list = list(genes_nodes_set)

        for cpgs in cpgs_list:

            params_to_fold_name = (
                str(round(confidence_threshold, 2)).replace(".", "_").replace(",", "_")
            )
            params_to_fold_name = "".join(
                ["con_", params_to_fold_name, "_depth_", str(depth)]
            )
            path_cpg_paths = os.path.join(paths_output, params_to_fold_name)
            path_cpg_paths = os.path.join(path_cpg_paths, cpgs)

            paths = nx.single_source_shortest_path(G, cpgs, cutoff=depth)
            paths_list = [v for k, v in paths.items() if k.startswith("cg")]

            G1 = nx.Graph()
            n_path = 0
            for path in paths_list[1:]:
                n_path = n_path + 1
                for node_in_path in path:
                    G1.add_node(node_in_path)
                    if node_in_path.startswith("cg") and node_in_path[2:].isdigit():
                        df_explain = annotation_dataset.loc[node_in_path]
                        node_in_path_csv = node_in_path + ".csv"
                        df_explain_path = os.path.join(path_cpg_paths, node_in_path_csv)
                        if not os.path.exists(path_cpg_paths):
                            os.makedirs(path_cpg_paths, exist_ok=True)
                        df_explain.to_csv(df_explain_path)
                        cpg_plot_name = node_in_path + ".png"
                        cpg_plot_path = os.path.join(path_cpg_paths, cpg_plot_name)
                        single_cpg_plot(
                            cpg_plot_path,
                            cpgs_dataset_for_ref,
                            node_in_path,
                            params["group_divide"],
                        )
                    else:
                        try:
                            df_explain = biogrid_dataset.loc[node_in_path]
                            node_in_path_csv = node_in_path + ".csv"
                            df_explain_path = os.path.join(
                                path_cpg_paths, node_in_path_csv
                            )
                            if not os.path.exists(path_cpg_paths):
                                os.makedirs(path_cpg_paths, exist_ok=True)
                            df_explain.to_csv(df_explain_path)
                        except KeyError:
                            print("not found in biogrid")

                print("adding edges")
                for n in range(len(path) - 1):
                    G1.add_edge(path[n], path[n + 1], length=10)

            print("plotting")
            if G1.number_of_edges() > 0:
                fig = plt.figure(figsize=(15, 15), facecolor="w")
                a1x = fig.add_subplot(111)

                options = {
                    "font_size": 8,
                    "node_size": 3000,
                    "node_color": "white",
                    "edgecolors": "red",
                    "linewidths": 1,
                    "width": 1,
                }
                pos = nx.layout.kamada_kawai_layout(G1, scale=2, dim=2)
                nx.draw_networkx(G1, pos=pos, **options)

                plt.tight_layout()
                plt.title(cpgs)
                axis = plt.gca()

                axis.set_xlim([1.1 * x for x in axis.get_xlim()])
                axis.set_ylim([1.1 * y for y in axis.get_ylim()])
                path_savefig = f"{cpgs}-{n_path}-_short_paths.png"
                path_savefig = os.path.join(path_cpg_paths, path_savefig)

                if params["paths_save"] == True:
                    print(f"Saving to {path_savefig}")
                    if not os.path.exists(path_cpg_paths):
                        os.makedirs(path_cpg_paths, exist_ok=True)
                    plt.savefig(path_savefig, bbox_inches="tight")
                if (
                    cpgs == params["paths_show_if"]
                    and params["paths_show_if_true"] == True
                ):
                    plt.show()
                plt.close()
