# This file is used to control the workflow of the CTMeth package
# It is used to call other modules and methods
# This approach is used to make the code more modular and easier to understand and gives possibility for future improvements

# built in libs
from time import sleep
from time import time
from datetime import datetime

# import math
import os
import traceback

# import warnings

# common libs
import pandas as pd
from scipy import interpolate as intrp
import numpy as np

# Pytorch
import torch
import torch.nn as nn
import torch.optim

import torchvision.transforms as transforms
from torchvision.transforms import ToTensor

import torch.utils.data
from torch.utils.data import Dataset, DataLoader

from torch.optim.lr_scheduler import OneCycleLR

# CTMeth
from .neural import training_settings as ts

from .neural.CTMeth_neural import CTMeth_neural
from .neural import SynMethDatasetLoader as SMDL
from .neural import training_module

from .utils import heatmaper
from .utils import hierclusterer
from .utils import paths_generator
from .utils import apca2d
from .utils import apca3d
from .utils import yaml_generator as yml_gen
from .utils import downloading
from . import params_wraper


def yaml_generator():
    # this method is used to generate yaml files
    yml_gen.main()


def download_databases(project_settings):
    # this method is used to download databases needed for the project
    params = params_wraper.params_wraper(project_settings).updated_provided_params

    downloading.download_databases(params)


def cuda_status(criterion, project_settings, training_settings, state_to_load):

    device = None

    if project_settings["no_cuda"] == False and torch.cuda.is_available():
        try:
            device = torch.device("cuda")
            torch.cuda.set_device(project_settings["gpu_id"])
            model = CTMeth_neural(training_settings)
            if state_to_load == None:
                state_in_use = os.path.join(
                    os.path.join(os.path.dirname(__file__), "neural"), "CTMeth.pth"
                )
                print("Loading default state")
            else:
                state_in_use = state_to_load
                print(f"Trying to load state from {state_in_use}")

            model.load_state_dict(torch.load(state_in_use))
            model.cuda(project_settings["gpu_id"])
            criterion = criterion.cuda(project_settings["gpu_id"])

        except Exception as e:
            error_message = f"An error occurred when setting gpu device: {e}"
            error_traceback = traceback.format_exc()
            print(error_message)
            print("Traceback:", error_traceback)
            device = torch.device("cuda")
            gpu_device = torch.cuda.current_device()
            print(f"Using GPU: {gpu_device}")
            model = CTMeth_neural(training_settings)
            if state_to_load == None:
                state_in_use = os.path.join(
                    os.path.join(os.path.dirname(__file__), "neural"), "CTMeth.pth"
                )
                print("Loading default state")
            else:
                state_in_use = state_to_load
                print(f"Trying to load state from {state_in_use}")

            model.load_state_dict(torch.load(state_in_use))
            model.cuda(gpu_device)
            criterion = criterion.cuda(gpu_device)
            input("Press Enter to continue...")
    elif project_settings["no_cuda"] == False and not torch.cuda.is_available():
        print("Cuda is unavailable")
        print("Using CPU")
        device = torch.device("cpu")
        model = CTMeth_neural(training_settings)
        if state_to_load == None:
            state_in_use = os.path.join(
                os.path.join(os.path.dirname(__file__), "neural"), "CTMeth.pth"
            )
            print("Loading default state")
        else:
            state_in_use = state_to_load
            print(f"Trying to load state from {state_in_use}")

        model.load_state_dict(torch.load(state_in_use, map_location=device))

    elif project_settings["no_cuda"] == True:
        print("Using CPU as set in project settings")
        device = torch.device("cpu")
        model = CTMeth_neural(training_settings)
        if state_to_load == None:
            state_in_use = os.path.join(
                os.path.join(os.path.dirname(__file__), "neural"), "CTMeth.pth"
            )
            print("Loading default state")
        else:
            state_in_use = state_to_load
            print(f"Trying to load state from {state_in_use}")

        model.load_state_dict(torch.load(state_in_use, map_location=device))

    else:
        print("Using CPU")
        print("Cuda is unavailable")
        device = torch.device("cpu")
        model = CTMeth_neural(training_settings)
        if state_to_load == None:
            state_in_use = os.path.join(
                os.path.join(os.path.dirname(__file__), "neural"), "CTMeth.pth"
            )
            print("Loading default state")
        else:
            state_in_use = state_to_load
            print(f"Trying to load state from {state_in_use}")

        model.load_state_dict(torch.load(state_in_use, map_location=device))

    return model, criterion, device


# upsampling method to make the input data the same size as the model input
def upsample(arr):
    arr = arr

    old_width = arr.size

    new_width = 100
    xrange = lambda x: np.linspace(0, 1, x)
    xrange(old_width)

    arr = intrp.interp1d(xrange(old_width), arr, kind="linear")
    new_arr = arr(xrange(new_width))
    return new_arr


def gui():
    # TODO GUI
    print("This function is not implemented yet. Please use command line interface")


def training(training_settings_path=None):
    # this is main module that is used to train the model
    training_module.train(training_settings_path)


def evaluate(
    project_settings=None, project_path=None, state_to_load=None, training_settings=None
):
    # This is main module that is used to evaluate dataset and described in dissertation
    # Evaluation here is in context methylation analysis, not training evaluation

    if training_settings == None:
        # TODO this part of code is planned to allow a possibility to use a model state trained with different parameters
        # eg.  conv1d_kernel_size bigger or smaller etc. It is under testing - use with caution
        training_settings = ts.settings_import(project_path)
    else:
        training_settings = ts.settings_import(
            project_path, settings_import_file=training_settings
        )

    # loading basic project settings needed for evaluation

    if project_path == None and project_settings == None:
        raise Exception("No project path nor project settings specified")
    if project_settings != None:
        print(f"Loading project settings from {project_settings}")
        project_settings = ts.read_settings(project_settings)
        project_path = project_settings["project_path"]
        input_data = project_settings["input_data"]
        group_divide = project_settings["group_divide"]
        separator = project_settings["separator"]
    if not os.path.exists(project_path):
        os.makedirs(project_path)
        print(f"Directory created: {project_path}")
    else:
        print(f"Directory already exists: {project_path}")
    criterion = torch.nn.CrossEntropyLoss()

    # loading model and state and setting device considering cuda availability and project_settings
    model, criterion, device = cuda_status(
        criterion, project_settings, training_settings, state_to_load
    )

    model.eval()

    # loading input data
    df = pd.read_csv(input_data, sep=separator, index_col=0)

    # setting up columns for output
    columns_out = ["cpg"]
    columns = df.columns
    columns2 = [
        "Label A",
        "Label B",
        "A confidence",
        "B confidence",
        "Mean confidence",
        "Confidence sum",
    ]
    columns_out.extend(columns2)
    columns_out.extend(columns)

    # evaluation loop is prepared in batches so user can see the progress
    eval_batch_size = 1000
    with torch.no_grad():

        df_length = len(df)
        batches = np.arange(0, df_length, eval_batch_size)  # setting up batches ranges
        loop_steps = len(batches) - 1
        if loop_steps <= 0:
            loop_steps = 1
            batches = np.array([0, df_length])
        to_csv_mode = "w"  # first batch will overwrite the file
        add_header = True  # first batch will have header

        starttime = datetime.now().time()
        print(f"Analysis started at: {starttime}")

        # setting up output paths
        output_path = project_settings["output_path_full"]
        output_path_hh = project_settings["output_path_only_hyper_hypo"]
        output_path_hhi = project_settings["output_path_hyper_hypo_indeter"]
        output_path_ii = project_settings["output_path_indeter_indeter"]
        print(f"\nFull output will be saved to:\n                {output_path}")
        print(f"\nHyper-Hypo-Indeter output will be saved to:\n {output_path_hhi}")
        print(f"\nOnly Hyper-Hypo output will be saved to:\n     {output_path_hh}")
        print(f"\nOnly Indeter-Indeter output will be saved to:\n     {output_path_ii}")

        # loop over batches
        for i in range(loop_steps):
            print(f"{output_path_hh}")
            # setting up batch start and end
            if i == loop_steps - 1:
                if batches[i] == df_length:
                    pass
                else:
                    start = batches[i]
                    end = df_length
            else:
                start = batches[i]
                end = batches[i + 1]
            batchtime = datetime.now().time()
            print(f"start {start} end {end} at {batchtime}")

            # setting up output lists that will be saved as dataframes after loop
            # it is done in a way that each batch is saved to file and then appended to it
            # AIM: this way the user can see the progress and the file is not lost in case of crash
            # and also in some way it is a way to save RAM
            # however author is aware that it is not the most efficient way to save the data, especially on hdd

            df_out = []
            df_out_hh = []
            df_out_hhi = []
            df_out_ii = []
            for n in range(start, end, 1):
                # group_divide is used to divide the input data into two groups (case and control)
                if group_divide != None:
                    beta_sample = df.iloc[n, :group_divide].values
                    beta_sampleA_out = df.iloc[n, :group_divide].values

                    beta_sample = upsample(beta_sample)
                    beta_sample = torch.tensor(beta_sample)
                    beta_sample = torch.reshape(beta_sample, [1, 1, 100]).float()

                    beta_sample = beta_sample.to(device)

                    beta_sample2 = df.iloc[n, group_divide:].values
                    beta_sampleB_out = df.iloc[n, group_divide:].values

                    beta_sample2 = upsample(beta_sample2)
                    beta_sample2 = torch.tensor(beta_sample2)
                    beta_sample2 = torch.reshape(beta_sample2, [1, 1, 100]).float()
                    beta_sample2 = beta_sample2.to(device)

                output = model(beta_sample)
                output2 = model(beta_sample2)

                # getting labels and confidence levels
                label1 = torch.argmax(output)
                label2 = torch.argmax(output2)
                confidence1 = output[0][torch.argmax(output)]
                confidence2 = output2[0][torch.argmax(output2)]
                mean_confidence = (confidence1 + confidence2) / 2
                confidence_sum = confidence1 + confidence2

                # preparing output row, which consists of cpg name, labels, confidence levels and beta values
                output_to_df = [
                    df.index[n],
                    int(label1),
                    int(label2),
                    float(confidence1),
                    float(confidence2),
                    float(mean_confidence),
                    float(confidence_sum),
                ]
                output_to_df.extend(beta_sampleA_out)
                output_to_df.extend(beta_sampleB_out)

                # saving output to different files depending on the labels
                if label1 != label2:
                    df_out_hhi.append(output_to_df)

                    if (int(label1) != 2) and (int(label2) != 2):
                        df_out_hh.append(output_to_df)
                if (int(label1) == 2) and (int(label2) == 2):
                    df_out_ii.append(output_to_df)

                df_out.append(output_to_df)

            # saving output to file
            df_out = pd.DataFrame(df_out, columns=columns_out)
            df_out_hh = pd.DataFrame(df_out_hh, columns=columns_out)
            df_out_hhi = pd.DataFrame(df_out_hhi, columns=columns_out)
            df_out_ii = pd.DataFrame(df_out_ii, columns=columns_out)

            # columns_out
            df_out.to_csv(
                output_path, sep=",", index=False, mode=to_csv_mode, header=add_header
            )

            df_out_hh.to_csv(
                output_path_hh,
                sep=",",
                index=False,
                mode=to_csv_mode,
                header=add_header,
            )
            df_out_hhi.to_csv(
                output_path_hhi,
                sep=",",
                index=False,
                mode=to_csv_mode,
                header=add_header,
            )
            df_out_ii.to_csv(
                output_path_ii,
                sep=",",
                index=False,
                mode=to_csv_mode,
                header=add_header,
            )
            to_csv_mode = "a"
            add_header = False
        endtime = datetime.now().time()
        print(f"Analysis finished at {endtime}")


def hiercluster(project_settings):
    # this method is used to call clustering module

    params = params_wraper.params_wraper(project_settings).updated_provided_params
    hierclusterer.hc(params)


def heatmap(project_settings):
    # this method is used to call heatmaper module

    params = params_wraper.params_wraper(project_settings).updated_provided_params
    heatmaper.do_heatmap(params)


def generate_paths(project_settings):
    # this method is used to call cpg-gene-gene-cpg module

    params = params_wraper.params_wraper(project_settings).updated_provided_params

    paths_generator.paths_method(params)


def pca2d(project_settings):
    # this method is used to call pca2d module

    params = params_wraper.params_wraper(project_settings).updated_provided_params
    apca2d.pca2d(params)


def pca3d(project_settings):
    # this method is used to call pca3d module

    params = params_wraper.params_wraper(project_settings).updated_provided_params
    apca3d.pca3d(params)


def return_params(project_settings):
    # this method is used for testing
    params = params_wraper.params_wraper(project_settings).updated_provided_params
    return params
