from pathlib import Path
import os.path

from tqdm import tqdm
from sklearn.metrics import classification_report
import torch
import pickle


from dataset.dataset import SampleSequenceDatasetController
from properties import *



class Trainer(object):

    def __init__(self, model, train_dataset_controller, val_dataset_controller, test_dataset_controller, optimizer, loss_function, training_hparams=None, validation_hparams=None, test_hparams=None):
        self.model = model
        self.train_dataset_controller: SampleSequenceDatasetController = train_dataset_controller
        self.val_dataset_controller: SampleSequenceDatasetController = val_dataset_controller
        self.test_dataset_controller: SampleSequenceDatasetController = test_dataset_controller
        self.optimizer = optimizer
        self.loss_function = loss_function
        self.training_hparams = training_hparams
        self.validation_hparams = validation_hparams
        self.test_hparams = test_hparams

    def train(self):
        optimizer = self.optimizer
        loss_function = self.loss_function
        hparams = self.training_hparams

        # Training
        loss_list = list()
        history_list = list()
        self.model.train()

        for epoch in tqdm(range(hparams.n_epochs)):
            total_loss = 0
            n_sample = 0

            sample_idx_list = self.train_dataset_controller.random_sample_idx_list

            for sample_idx in tqdm(sample_idx_list):
                print(f"[Epoch - {epoch} | Sample Idx - {sample_idx}]")

                try:
                    train_dataloader = self.train_dataset_controller.get_concat_dataloader(sample_idx=sample_idx)

                    sample_total_loss = 0
                    n_batch = 0

                    for idx, (x_batch, y_batch) in enumerate(train_dataloader):
                        x_batch = x_batch.to(hparams.device)
                        y_batch = y_batch.to(hparams.device)

                        # Optimization
                        optimizer.zero_grad()
                        x_batch_hat = self.model(x_batch)
                        loss = loss_function(x_batch, x_batch_hat)
                        loss.backward()
                        optimizer.step()

                        sample_total_loss += loss.item()
                        n_batch += 1

                    current_sample_loss = sample_total_loss / n_batch
                    total_loss += current_sample_loss
                    n_sample += 1

                    print(f"Train - Sample Loss : {round(current_sample_loss, 7)}")
                except Exception as e:
                    print(f"Error : {e}")

            current_loss = total_loss / n_sample
            loss_list.append(current_loss)

            print("==============================================")
            print(f"[Epoch - {epoch}]")
            print(f"Train - Loss : {round(current_loss, 7)}")

            # Save
            model_save_path = os.path.join(hparams.save_dir_path, f"model_{epoch}.pkl")
            opt_save_path = os.path.join(hparams.save_dir_path, f"opt_{epoch}.pkl")
            history_save_path = os.path.join(hparams.save_dir_path, f"history_{epoch}.pkl")

            with open(model_save_path, "wb") as f:
                pickle.dump(self.model.state_dict(), f)
            with open(opt_save_path, "wb") as f:
                pickle.dump(optimizer.state_dict(), f)
            with open(history_save_path, "wb") as f:
                pickle.dump(loss_list, f)
            print(f"Complete to save model - Epoch : {epoch}")

    def validation(self):
        loss_function = self.loss_function
        hparams = self.validation_hparams

        # Validation
        self.model.eval()

        total_loss = 0
        n_sample = 0

        sample_idx_list = self.val_dataset_controller.random_sample_idx_list

        for sample_idx in tqdm(sample_idx_list):
            print(f"Sample Idx - {sample_idx}]")

            # Create directories
            sample_save_dir_path = os.path.join(hparams.save_dir_path, "validation")
            Path(sample_save_dir_path).mkdir(exist_ok=True, parents=True)
            history_save_path = os.path.join(sample_save_dir_path, f"{sample_idx}.pkl")

            if os.path.isfile(history_save_path):
                print(f"Sample {sample_idx} is exist, Continue.")
                continue

            try:
                val_dataloader = self.val_dataset_controller.get_concat_dataloader(sample_idx=sample_idx)

                sample_total_loss = 0
                n_batch = 0

                sample_loss_list = list()
                sample_loss_batch_list = list()
                sample_y_batch_list = list()

                for idx, (x_batch, y_batch) in enumerate(val_dataloader):
                    x_batch = x_batch.to(hparams.device)
                    y_batch = y_batch.to(hparams.device)

                    x_hat_batch = self.model(x_batch)
                    loss = loss_function(x_batch, x_hat_batch)
                    loss = loss.item()

                    sample_loss_batch_list.append(loss)
                    sample_y_batch_list.append(y_batch)

                    sample_total_loss += loss
                    n_batch += 1

                current_sample_loss = sample_total_loss / n_batch

                sample_loss_batch_list = np.array(sample_loss_batch_list)
                sample_y_batch_list = torch.cat(sample_y_batch_list).cpu().numpy()
                sample_loss_list.append(current_sample_loss)

                total_loss += current_sample_loss
                n_sample += 1

                current_target = list(set(sample_y_batch_list.squeeze()))
                num_batch = len(sample_loss_batch_list)
                print("Target", current_target)
                print("Num batch", num_batch)

                sorted_loss_list = np.sort(sample_loss_batch_list)[::-1]
                best_top_loss_mean = np.mean(sorted_loss_list[:5], axis=0)
                best_bottom_loss_mean = np.mean(sorted_loss_list[-5:], axis=0)
                print("Top", best_top_loss_mean)
                print("Bottom", best_bottom_loss_mean)

                history_dict = dict(
                    sample_idx=sample_idx,
                    target=current_target,
                    current_sample_loss=current_sample_loss,
                    num_batch=num_batch,
                    sampe_loss_batch_list=sample_loss_batch_list
                )

                print(f"Validation - Sample Loss : {round(current_sample_loss, 7)}")

                # Save
                with open(history_save_path, "wb") as f:
                    pickle.dump(history_dict, f)

                print(f"Complete to save history - Sample idx : {sample_idx}")

            except Exception as e:
                print(f"Error : {e}")

        current_loss = total_loss / n_sample

        print("==============================================")
        print(f"Validation - Loss : {round(current_loss, 7)}")

    def test(self):
        loss_function = self.loss_function
        hparams = self.test_hparams

        # Validation
        self.model.eval()

        total_loss = 0
        n_sample = 0

        sample_idx_list = self.test_dataset_controller.random_sample_idx_list

        for sample_idx in tqdm(sample_idx_list):
            print(f"Sample Idx - {sample_idx}]")

            # Create directories
            sample_save_dir_path = os.path.join(hparams.save_dir_path, "test")
            Path(sample_save_dir_path).mkdir(exist_ok=True, parents=True)
            history_save_path = os.path.join(sample_save_dir_path, f"{sample_idx}.pkl")

            if os.path.isfile(history_save_path):
                print(f"Sample {sample_idx} is exist, Continue.")
                continue

            try:
                test_dataloader = self.test_dataset_controller.get_concat_dataloader(sample_idx=sample_idx)

                sample_total_loss = 0
                n_batch = 0

                sample_loss_list = list()
                sample_loss_batch_list = list()
                sample_y_batch_list = list()

                for idx, (x_batch, y_batch) in enumerate(test_dataloader):
                    x_batch = x_batch.to(hparams.device)
                    y_batch = y_batch.to(hparams.device)

                    x_hat_batch = self.model(x_batch)
                    loss = loss_function(x_batch, x_hat_batch)
                    loss = loss.item()

                    sample_loss_batch_list.append(loss)
                    sample_y_batch_list.append(y_batch)

                    sample_total_loss += loss
                    n_batch += 1

                current_sample_loss = sample_total_loss / n_batch

                sample_loss_batch_list = np.array(sample_loss_batch_list)
                sample_y_batch_list = torch.cat(sample_y_batch_list).cpu().numpy()
                sample_loss_list.append(current_sample_loss)

                total_loss += current_sample_loss
                n_sample += 1

                current_target = list(set(sample_y_batch_list.squeeze()))
                num_batch = len(sample_loss_batch_list)
                print("Target", current_target)
                print("Num batch", num_batch)

                sorted_loss_list = np.sort(sample_loss_batch_list)[::-1]
                best_top_loss_mean = np.mean(sorted_loss_list[:5], axis=0)
                best_bottom_loss_mean = np.mean(sorted_loss_list[-5:], axis=0)
                print("Top", best_top_loss_mean)
                print("Bottom", best_bottom_loss_mean)

                history_dict = dict(
                    sample_idx=sample_idx,
                    target=current_target,
                    current_sample_loss=current_sample_loss,
                    num_batch=num_batch,
                    sampe_loss_batch_list=sample_loss_batch_list
                )

                print(f"Test - Sample Loss : {round(current_sample_loss, 7)}")

                # Save
                with open(history_save_path, "wb") as f:
                    pickle.dump(history_dict, f)

                print(f"Complete to save history - Sample idx : {sample_idx}")

            except Exception as e:
                print(f"Error : {e}")

        current_loss = total_loss / n_sample

        print("==============================================")
        print(f"Test - Loss : {round(current_loss, 7)}")

    def score(self, result_dir_path):
        hparmas = self.test_hparams

        normal_entire_loss_list = list()
        normal_entire_num_batch_list = list()

        abnormal_entire_loss_list = list()
        abnormal_entire_num_batch_list = list()

        normal_short_loss_list = list()
        normal_short_target_list = list()
        normal_long_loss_list = list()
        normal_long_target_list = list()

        abnormal_short_loss_list = list()
        abnormal_short_target_list = list()
        abnormal_long_loss_list = list()
        abnormal_long_target_list = list()

        for sample_file_name in os.listdir(result_dir_path):
            val_sample_file_path = os.path.join(result_dir_path, sample_file_name)

            with open(val_sample_file_path, "rb") as f:
                sample_result_dict = pickle.load(f)

            loss = sample_result_dict["current_sample_loss"]
            target = sample_result_dict["target"][0]
            num_batch = sample_result_dict["num_batch"]

            # Type 1
            if hparmas.type_lower_bound <= num_batch < hparmas.type_upper_bound:
                if target == 0:
                    normal_short_loss_list.append(loss)
                    normal_short_target_list.append(target)

                    normal_entire_loss_list.append(loss)
                    normal_entire_num_batch_list.append(num_batch)
                elif target == 1:
                    abnormal_short_loss_list.append(loss)
                    abnormal_short_target_list.append(target)

                    abnormal_entire_loss_list.append(loss)
                    abnormal_entire_num_batch_list.append(num_batch)
                else:
                    assert False, f"Unmatched target, {target}"
            # Type 2
            else:
                if target == 0:
                    normal_long_loss_list.append(loss)
                    normal_long_target_list.append(target)

                    normal_entire_loss_list.append(loss)
                    normal_entire_num_batch_list.append(num_batch)
                elif target == 1:
                    abnormal_long_loss_list.append(loss)
                    abnormal_long_target_list.append(target)

                    abnormal_entire_loss_list.append(loss)
                    abnormal_entire_num_batch_list.append(num_batch)
                else:
                    assert False, f"Unmatched target, {target}"

        normal_short_pred_target = list(map(lambda loss: 1 if loss >= hparmas.type_1_threshold else 0, normal_short_loss_list))
        abnormal_short_pred_target = list(map(lambda loss: 1 if loss >= hparmas.type_1_threshold else 0, abnormal_short_loss_list))

        normal_long_pred_target = list(map(lambda loss: 1 if loss >= hparmas.type_2_threshold else 0, normal_long_loss_list))
        abnormal_long_pred_target = list(map(lambda loss: 1 if loss >= hparmas.type_2_threshold else 0, abnormal_long_loss_list))

        pred_target = normal_short_pred_target + abnormal_short_pred_target + normal_long_pred_target + abnormal_long_pred_target
        true_target = normal_short_target_list + abnormal_short_target_list + normal_long_target_list + abnormal_long_target_list

        short_clf_report = classification_report(y_true=normal_short_pred_target + abnormal_short_pred_target,
                                                 y_pred=normal_short_target_list + abnormal_short_target_list)
        long_clf_report = classification_report(y_true=normal_long_pred_target + abnormal_long_pred_target,
                                                y_pred=normal_long_target_list + abnormal_long_target_list)
        total_clf_report = classification_report(y_true=true_target, y_pred=pred_target, output_dict=True)
        print(classification_report(y_true=true_target, y_pred=pred_target, output_dict=False))

        return dict(
            normal_entire_loss_list=normal_entire_loss_list,
            normal_entire_num_batch_list=normal_entire_num_batch_list,
            abnormal_entire_loss_list=abnormal_entire_loss_list,
            abnormal_entire_num_batch_list=abnormal_entire_num_batch_list,
            normal_short_loss_list=normal_short_loss_list,
            normal_long_loss_list=normal_long_loss_list,
            abnormal_short_loss_list=abnormal_short_loss_list,
            abnormal_long_loss_list=abnormal_long_loss_list,
            short_clf_report=short_clf_report,
            long_clf_report=long_clf_report,
            total_clf_report=total_clf_report
        )
