import CTMeth
from CTMeth.neural.training_settings import read_settings
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import os


def find_exact_match(user_input, search_for):
    if isinstance(user_input, list):
        normalized_input = [item.strip() for item in user_input]
    elif isinstance(user_input, str):
        normalized_input = [
            item.strip() for item in user_input.replace(".", ",").split(",")
        ]
    else:
        return False  # If input is neither a list nor a string, return False

    # Check for the exact match "aa"
    return search_for in normalized_input


def do_heatmap(params):

    n = 3
    sep = ","

    # Handle source of heatmap
    src = params["heatmap_src"]

    do_heatmaps_from = []
    type_of_srcs = ["hh", "full", "hhi", "ii"]

    if src is None:
        print("Heatmaps from hh and hhi will be produced.")

        do_heatmaps_from = ["hh", "hhi"]
    for type_of_src in type_of_srcs:

        if find_exact_match(src, type_of_src):
            do_heatmaps_from.append(type_of_src)

    if len(do_heatmaps_from) == 0:
        print("No valid source for heatmap. Defaulting to hh and hhi.")
        do_heatmaps_from = ["hh", "hhi"]

    # loop over all sources specified
    for src in do_heatmaps_from:
        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["heatmap_src"]
            sep = params["heatmap_sep_src"]

        input_data = pd.read_csv(input_data_path, index_col=0, sep=sep)
        sortby = input_data.columns[5]

        input_data_sorted = input_data.sort_values(by=sortby, ascending=False)

        confidence_thresholds = params["heatmap_confidence_thresholds"]
        if confidence_thresholds is None:
            confidence_thresholds = [0]
        elif isinstance(confidence_thresholds, int):
            confidence_thresholds = [confidence_thresholds]
        for n in confidence_thresholds:

            # handle input data
            confidence_threshold = n
            input_data_sorted = input_data_sorted[
                input_data_sorted[sortby] > confidence_threshold
            ]

            output_data = input_data_sorted.iloc[:, 6:]

            # prepare plot
            project_name = params["project_name"]
            group_divide = params["group_divide"]

            data_info_part1 = f"project name: {project_name}, cpgs shown {len(output_data)}, group_divide {group_divide}"
            data_infor_part2 = f"confidence sum threshold: {confidence_threshold}"
            heatmaptitle = "".join([data_info_part1, "\n", data_infor_part2])


            fig, ax = plt.subplots()
            w = params["heatmap_width"]
            h = params["heatmap_height"]

            fig.set_size_inches(w, h)
            ax = sns.heatmap(output_data, vmin=0, vmax=1, cmap=params["heatmap_cmap"])

            if params["heatmap_group_divide"]:
                ax.axvline(
                    x=group_divide,
                    linewidth=int(params["heatmap_group_divide_width"]),
                    color=params["heatmap_group_divide_color"],
                )

            if params["heatmap_sep_samples"]:

                for x in range(len(output_data.columns)):
                    ax.axvline(x=x, linewidth=4, color="black")

            ax.set_title(heatmaptitle, fontsize=params["heatmap_title_fontsize"])

            ax.set_xticks(range(len(output_data.columns)))
            ax.set_xticklabels(list(output_data.columns.values), rotation=90)
            ax.tick_params(
                axis="x",
                top=False,
                labeltop=False,
                labelbottom=True,
                direction="inout",
                grid_color="r",
            )
            plt.tight_layout()
            heatmap_output_folder = params["heatmap_output_folder"]
            if not isinstance(heatmap_output_folder, str):
                heatmap_output_folder = params["project_path"]
            else:
                if len(heatmap_output_folder) < 1:
                    heatmap_output_folder = params["project_path"]
            savefigpth = heatmap_output_folder

            basename = os.path.basename(input_data_path)[:-4]
            imagename = f"heatmap_conf{confidence_threshold}_{basename}.png"
            savefigpth = os.path.join(savefigpth, imagename)
            if params["heatmap_show"]:
                plt.show()
            if params["heatmap_save"]:
                plt.savefig(savefigpth, bbox_inches="tight")
            plt.close()
