## this is the main training module for the neural network

# builtins
import os
import time
from datetime import datetime

# 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 torch.utils.data
from torch.utils.data import DataLoader
from torch.optim.lr_scheduler import OneCycleLR


from CTMeth.neural.CTMeth_neural_training import CTMeth_class
import CTMeth.neural.SynMethDatasetLoader as SMDL
from CTMeth.neural import training_settings


def train(training_settings_path):
    def upsample(arr, new_length_of_array=100):
        # this method is used to upsample the array to certain amount of samples

        old_length_of_array = arr.size

        xrange = lambda x: np.linspace(0, 1, x)
        xrange(old_length_of_array)

        arr = intrp.interp1d(xrange(old_length_of_array), arr, kind="linear")
        new_arr = arr(xrange(new_length_of_array))
        return new_arr

    def accuracy(output, target):
        # this method is used to calculate accuracy of the model
        with torch.no_grad():

            batch_size = target.size(0)
            _, pred = output.topk(1, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))
            res = []
            correct_k = correct[:1].flatten().float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))

            return res[0]

    def cuda_status(no_cuda, model, criterion):
        # this method is used to check if cuda is available and set the model and criterion to cuda
        if (not no_cuda) and torch.cuda.is_available() and torch.cuda.is_available():
            torch.cuda.set_device(params["gpu_id"])
            model.cuda(params["gpu_id"])
            criterion = criterion.cuda(params["gpu_id"])
            print("Cuda is available and on")

        else:
            print("Cuda unavailable or off")
        return model, criterion

    def evaluate_with_generalised_dictionary(model):

        # this method is used to evaluate the model with generalised dictionary
        print("Evaluating with generalised dictionary")
        test_loss_val, test_acc1_val = 0, 0

        filepath = os.path.join(
            os.path.join(os.path.dirname(__file__), "eval_test.csv")
        )

        df = pd.read_csv(filepath, index_col=0)

        # evaluating the model
        model.eval()
        eval_dataset_size = len(df)
        nn = 0
        correct = 0
        with torch.no_grad():
            for n in range(eval_dataset_size):

                beta_sample = df.iloc[n, :-1].values
                label = df.iloc[n, -1]

                beta_sample = upsample(beta_sample, params["beta_values_array_length"])
                beta_sample = torch.tensor(beta_sample)
                beta_sample = torch.reshape(beta_sample, [1, 1, 100]).float()
                beta_sample = beta_sample.cuda(params["gpu_id"], non_blocking=True)
                label = torch.tensor([label])
                label = label.cuda(params["gpu_id"], non_blocking=True)

                output = model(beta_sample)
                test_loss = criterion(output, label)
                test_acc1 = accuracy(output, label)
                nn += beta_sample.size(0)
                test_loss_val += float(test_loss.item() * beta_sample.size(0))
                test_acc1_val += float(test_acc1 * beta_sample.size(0))

                if torch.argmax(output) == label:
                    correct = correct + 1

            correctness = correct / eval_dataset_size
            return correctness

    def get_params(training_settings_path=None):
        # this method is used to get the parameters from the training settings file
        if training_settings_path == None:
            training_settings_path = os.path.join(
                os.path.join(os.path.dirname(__file__), "training_settings.yaml")
            )

        # TODO add ability to pass dictionary
        # elif isinstance...

        else:
            print(f"Using {training_settings_path}\n for training")
            pass

        params = training_settings.read_settings(training_settings_path)

        return params

    restart = True
    while restart == True:

        params = get_params(training_settings_path)

        model = CTMeth_class(params)

        # TODO add tkinter filedialog to choose the path

        try:
            checkpoint_path = params["save_neural_network_state_to"]
            model.load_state_dict(torch.load(checkpoint_path))
            print(f"loaded {checkpoint_path}")
        except Exception as e:
            if checkpoint_path == None:
                print("No checkpoint path provided")
            else:
                print(e)
            input("Press enter to continue \nor ctrl+c to exit")

        criterion = torch.nn.CrossEntropyLoss()
        model, criterion = cuda_status(params["no_cuda"], model, criterion)
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=params["optimizer_lr"],
            weight_decay=params["optimizer_weight_decay"],
        )
        scheduler = OneCycleLR(
            optimizer,
            total_steps=params["num_epochs"],
            verbose=True,
            max_lr=params["scheduler_max_lr"],
        )

        # Traianing dataset generator
        train_dataset = SMDL.SynMethDatasetLoader(dataset_size=1000)

        train_loader = DataLoader(
            train_dataset, batch_size=1, shuffle=True, num_workers=0
        )

        # Time notificaton
        time_begin = time.time()
        print(f"Training started at {datetime.fromtimestamp(time_begin)} \n")

        # Training loop
        last_epoch = params["num_epochs"]
        for n in range(params["num_epochs"]):
            print(f"Epoch {(n+1)} / {(last_epoch)}")

            model.train()

            loss_val, acc1_val = 0, 0
            for i, (bvalues, target) in enumerate(train_loader):

                if (not (params["no_cuda"])) and torch.cuda.is_available():
                    bvalues = bvalues.cuda(params["gpu_id"], non_blocking=True)
                    target = target.cuda(params["gpu_id"], non_blocking=True)

                output = model(bvalues)

                loss = criterion(output, target)
                acc1 = accuracy(output, target)

                loss_val += float(loss.item() * bvalues.size(0))
                acc1_val += float(acc1 * bvalues.size(0))

                optimizer.zero_grad()
                loss.backward()

                # gradient clipping - preventing gradient to be too large
                if params["clip_grad_norm"] > 0:
                    nn.utils.clip_grad_norm_(
                        model.parameters(),
                        max_norm=params["clip_grad_norm"],
                        norm_type=2,
                    )

                optimizer.step()

                # loss counting
                average_loss, average_accuracy1 = (
                    loss_val / (i + 1) / bvalues.size(0)
                ), (acc1_val / (i + 1) / bvalues.size(0))

                if (
                    params["train_status_every_n_sample"] >= 0
                    and i % params["train_status_every_n_sample"] == 0
                ):

                    print(
                        f"[Epoch {n + 1}][Train][{i}] \t Loss: {average_loss:.4e} \t Top-1 {average_accuracy1:6.2f}"
                    )
                    print(f"last loss {loss}")
                    current_learning_rate = optimizer.param_groups[0]["lr"]
                    print(f"Current learning rate {current_learning_rate}")
                    print(f"{average_accuracy1}\n")

            scheduler.step()

            correctness = evaluate_with_generalised_dictionary(model)
            print(f"Coherence with generalized dictionary {(round(correctness*100))}\n")

            if (average_accuracy1 > 90) and (correctness == 1):
                # training will end if average accuracy is above 90% and coherence with generalised dictionary is 100%
                print(
                    "Training completed \n average accuracy above 90% \n and coherence with generalised dictionary 100%"
                )
                user_end_training = input("Do you want to end training? y/n")
                if user_end_training == "y":
                    # restart = False
                    quit()
                pass

            checkpoint_path = params["save_neural_network_state_to"]
            torch.save(model.state_dict(), checkpoint_path)
            input(f"Neural network state saved to: \n {checkpoint_path}")
