import os
import shutil
import sys
import time
from typing import Union

import numpy as np
import torch
from sklearn import metrics
from torch.nn import Module
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader

from pyg_extension.utils.log_model import LogModule, HookGradientModule, make_dot_, \
    get_target_layer, AverageMeterTotal


def class_eval(prediction, target):
    """Classification."""
    prediction = np.exp(prediction.numpy())
    target = target.numpy()
    pred_label = np.argmax(prediction, axis=1)
    target_label = np.squeeze(target)
    if not target_label.shape:
        target_label = np.asarray([target_label])
    if prediction.shape[1] == 2:
        precision, recall, fscore, _ = metrics.precision_recall_fscore_support(
            target_label, pred_label, average='binary')
        auc_score = metrics.roc_auc_score(target_label, prediction[:, 1])
        accuracy = metrics.accuracy_score(target_label, pred_label)
    else:
        raise NotImplementedError
    return accuracy, precision, recall, fscore, auc_score


def mae(prediction, target):
    """
    Computes the mean absolute error between prediction and target

    Parameters
    ----------
    prediction: torch.Tensor (N, 1)
    target: torch.Tensor (N, 1)
    """
    return torch.mean(torch.abs(target - prediction))


class LearningFlow:
    """
    LearningFlow for training.
    
    Examples:

        >>> test_dataset = dataset[:1000]
        >>> val_dataset = dataset[1000:2000]
        >>> train_dataset = dataset[2000:3000]
        >>> import torch_geometric.transforms as T
        >>> train_dataset = SimpleDataset(data_train, pre_transform=T.ToSparseTensor())
        >>> test_dataset = SimpleDataset(data_test,pre_transform=T.ToSparseTensor())
        >>> val_dataset = SimpleDataset(val_data,pre_transform=T.ToSparseTensor())

        >>> train_loader = DataLoader(
        ... dataset=train_dataset,
        ... batch_size=200,
        ... shuffle=False,
        ... num_workers=0)

        >>> test_loader = ...
        >>> val_loader = ...


        >>>  model = CrystalGraphConvNet(nfeat_node=91,
        ...  nfeat_edge=3,
        ...  nfeat_state=29,
        ...  nc_node_hidden=128,
        ...  nc_node_interaction=64,
        ...  num_interactions=2,)
        >>> # model = CrystalGraphGCN(...)
        >>> # model = CrystalGraphGCN2(...)
        >>> # model = CrystalGraphGAT(...)
        >>> # model = SchNet(...)
        >>> # model = MEGNet(...)
        >>> # model = SchNet(...)


        >>> optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
        >>> scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=2,...min_lr=0.001)
        >>> device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')

        >>> lf= LearningFlow(model, train_loader, validate_loader=val_loader, device= "cuda:1",
        ... optimizer=None, clf= False, loss_method=None, learning_rate = 1e-3,
        ... weight_decay= 0.01, checkpoint=True, scheduler=scheduler,
        ... loss_threshold= 0.1, print_freq= None, print_what="all")

        >>> lf.run(50)

    where the dataset could from StructureToData or :class:`featurebox.test_featurizers.base_graph_geo.StructureGraphGEO`
    """

    def __init__(self, model: Module, train_loader: DataLoader, validate_loader: DataLoader, device="cpu",
                 optimizer=None, clf: bool = False, loss_method=None, learning_rate: float = 1e-3,
                 weight_decay: float = 0.0, checkpoint=True, scheduler=None, debug: str = "", target_layers=(),
                 loss_threshold: float = 0.1, print_freq: Union[int, str, None] = 10, print_what="all",
                 process_label=None):
        """

        Parameters
        ----------
        model: module
        train_loader: DataLoader
        validate_loader: DataLoader
        device:str
            such as "cuda:0","cpu"
        optimizer:torch.Optimizer
        clf:bool
            under exploit......
        loss_method:torch._Loss
            see more in torch
        learning_rate:float
            see more in torch
        weight_decay:float
            see more in torch
        checkpoint:bool
            save checkpoint or not.
        loss_threshold:
            see more in torch
        print_freq:int
            print frequency
        print_what:str
            "all","train","test" log.
        scheduler:
            scheduler, see more in torch
        process_label:Callable
            function to get y/label.
        """

        self.train_loader = train_loader
        self.test_loader = validate_loader

        device = torch.device(device)

        self.device = device
        self.model = model
        self.model.to(device)
        self.clf = clf
        self.loss_method = loss_method
        self.optimizer = optimizer
        self.checkpoint = checkpoint
        self.weight_log = LogModule()

        self.debug = debug

        self.train_batch_number = len(self.train_loader)

        if print_freq is None:
            self.print_freq = self.train_batch_number
        else:
            self.print_freq = self.train_batch_number if not isinstance(print_freq, int) else print_freq

        if self.optimizer is None or self.optimizer == "Adam":
            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate, weight_decay=weight_decay)
        elif self.optimizer == "SGD":
            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=learning_rate, weight_decay=weight_decay
                                             )
            # L2 regularization
        else:
            self.optimizer = optimizer(self.model.parameters(), lr=learning_rate, weight_decay=weight_decay)

        if self.loss_method is None:
            if self.clf is True:
                self.loss_method = torch.nn.CrossEntropyLoss()
            elif self.clf == "multi_label":
                self.loss_method = torch.nn.L1Loss()
                # 主要是用来判定实际的输出与期望的输出的接近程度 MAE: 1/N |y_pred-y| y 为多列
            else:
                self.loss_method = torch.nn.MSELoss()

        else:
            self.loss_method = loss_method
        if scheduler is None:
            self.scheduler = ReduceLROnPlateau(self.optimizer, mode='min', factor=0.5, patience=7, min_lr=0.00001)
        else:
            self.optimizer = scheduler.optimizer
            self.scheduler = scheduler
        self.best_error = 1000000.0
        self.threshold = loss_threshold
        # *.pth.tar or str
        self.run_train = self.run
        self.fit = self.run_train
        self.print_what = print_what
        self.forward_hook_list = []
        func = lambda x: x
        self._process_label = func if process_label is None else process_label
        self.target_layers = target_layers

        if self.clf:
            self.meters = AverageMeterTotal("Time", "losses", "accuracies", "precisions", "recalls", "fscores",
                                            "auc_scores")
        else:
            self.meters = AverageMeterTotal("Time", "losses", "mae_errors")

    def my_debug(self):
        # debug
        if "loop" in self.debug:
            self.weight_log.stats_loop()

        if "single" in self.debug:
            self.model.train()
            for m, data in enumerate(self.train_loader):  # just record for once,
                data = data.to(self.device)
                self.model(data)
                self.weight_log.get_weight(self.model, append=False)
                break
            self.weight_log.stats_single()

        if "graphviz" in self.debug:

            self.model.train()
            for m, data in enumerate(self.train_loader):  # just record for once,
                data = data.to(self.device)
                y_pred = self.model(data)
                vis_graph = make_dot_(y_pred, params=dict(list(self.model.named_parameters())))
                vis_graph.render(format="pdf",view=False)
                break

        if "hook" in self.debug:
            # which layer to hook.
            self.target_layers = get_target_layer(self.model, self.target_layers)
            self.gradient_log = HookGradientModule(self.target_layers)
            self.model.train()
            for m, data in enumerate(self.train_loader):  # just record for once,
                data = data.to(self.device)
                batch_y = self.process_label(data.y)
                self.model.zero_grad()
                y_pred = self.model(data)
                lossi = self.loss_method(y_pred, batch_y)
                lossi.backward(retain_graph=True)
                # lossi.backward()
                break
            self.gradient_log.stats_single()

    def process_label(self, y):
        return self._process_label(y)

    @staticmethod
    def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
        torch.save(state, filename)
        if is_best:
            shutil.copyfile(filename, 'model_best.pth.tar')

    def run(self, epoch=50, warm_start=False):
        """
        run loop.

        Parameters
        ----------
        epoch:int
            epoch.
        warm_start: str, False
            The name of resume file, 'checkpoint.pth.tar' or 'model_best.pth.tar'
            If warm_start, try to resume from local disk.
        """

        resume = warm_start if warm_start is not False else None
        start_epoch = 0
        if resume:
            if os.path.isfile(resume):
                print("=> loading checkpoint '{}'".format(resume))
                checkpoint = torch.load(resume)
                start_epoch = checkpoint['epoch']
                self.threshold = checkpoint['best_error']
                self.model.load_state_dict(checkpoint['state_dict'])
                self.optimizer.load_state_dict(checkpoint['optimizer'])
                print("=> loaded checkpoint '{}' (epoch {})"
                      .format(resume, checkpoint['epoch']))
            else:
                print("=> no checkpoint found at '{}'".format(resume))

        if start_epoch > 0:
            epoch += start_epoch
            print("Try to run start from 'resumed epoch' {} to 'epoch' {}".format(start_epoch, epoch))

        self.model.to(self.device)

        for epochi in range(start_epoch, epoch):

            self._train(epochi)

            if "loop" in self.debug in self.debug:  # record each epochi
                self.weight_log.get_weight(self.model)

            score = self._validate(epochi)

            if score != score:
                print('Exit due to NaN')
                sys.exit(1)
            try:
                self.scheduler.step(score)
            except TypeError:
                self.scheduler.step()

            is_best = score < self.best_error
            self.best_error = min(score, self.best_error)

            if self.checkpoint:
                self.save_checkpoint({
                    'epoch': epochi + 1,
                    'state_dict': self.model.state_dict(),
                    'best_error': self.threshold,
                    'optimizer': self.optimizer.state_dict(),
                }, is_best)

            if score <= self.threshold:
                print("Up to requirements and early termination in epoch ({})".format(epochi))
                break

        self.my_debug()

    def _train(self, epochi):
        self.model.train()
        self.meters.reset()

        for m, data in enumerate(self.train_loader):
            point = time.time()
            data = data.to(self.device)
            batch_y = self.process_label(data.y)

            self.optimizer.zero_grad()

            y_pred = self.model(data)

            try:
                lossi = self.loss_method(y_pred, batch_y)
            except TypeError:
                target = batch_y.sign()
                lossi = self.loss_method(y_pred, batch_y, target)

            lossi.backward()

            self.optimizer.step()

            # ###########

            batch_size = batch_y.size(0)
            self.meters['Time'].record(time.time() - point)
            self.meters['losses'].record(float(lossi.cpu().item()), batch_size)
            if self.clf is False:
                mae_error = mae(y_pred.data.cpu(), batch_y.cpu())
                self.meters['mae_errors'].record(mae_error, batch_size)
            else:
                accuracy, precision, recall, fscore, auc_score = \
                    class_eval(y_pred.data.cpu(), batch_y.cpu())
                self.meters['accuracies'].record(accuracy, batch_size)
                self.meters['precisions'].record(precision, batch_size)
                self.meters['recalls'].record(recall, batch_size)
                self.meters['fscores'].record(fscore, batch_size)
                self.meters['auc_scores'].record(auc_score, batch_size)

            if self.print_freq != self.train_batch_number and m % self.print_freq == 0 \
                    and self.print_what in ["all", "train"]:
                print('Train: [{0}][{1}/{2}] {3}'.format(epochi, m, len(self.train_loader), self.meters.text()))

        if self.print_freq == self.train_batch_number:
            print('Train: [{}] {}'.format(epochi, self.meters.text()))

    def _validate(self, epochi):
        self.model.eval()
        self.meters.reset()

        for m, data in enumerate(self.test_loader):
            point = time.time()

            data = data.to(self.device)
            batch_y = self.process_label(data.y)
            y_pred = self.model(data)

            try:
                lossi = self.loss_method(y_pred, batch_y)
            except TypeError:
                target = batch_y.sign()
                lossi = self.loss_method(y_pred, batch_y, target)

            # ###########

            batch_size = batch_y.size(0)
            self.meters['Time'].record(time.time() - point)
            self.meters['losses'].record(float(lossi.cpu().item()), batch_size)
            if self.clf is False:
                mae_error = mae(y_pred.data.cpu(), batch_y.cpu())
                self.meters['mae_errors'].record(mae_error, batch_size)
            else:
                accuracy, precision, recall, fscore, auc_score = \
                    class_eval(y_pred.data.cpu(), batch_y.cpu())
                self.meters['accuracies'].record(accuracy, batch_size)
                self.meters['precisions'].record(precision, batch_size)
                self.meters['recalls'].record(recall, batch_size)
                self.meters['fscores'].record(fscore, batch_size)
                self.meters['auc_scores'].record(auc_score, batch_size)

            if self.print_freq != self.train_batch_number and m % self.print_freq == 0 \
                    and self.print_what in ["all", "train"]:
                print('Test: [{0}][{1}/{2}] {3}'.format(epochi, m, len(self.train_loader), self.meters.text()))

        if self.print_freq == self.train_batch_number:
            print('Test: [{}] {}'.format(epochi, self.meters.text()))

        if self.clf is False:
            return self.meters['mae_errors'].avg
        else:
            return self.meters['auc_scores'].avg

    def score(self, predict_loader):
        """Return MAE score."""
        y_pre, y_true = self.predict(predict_loader)
        return float(mae(y_pre, y_true))

    def predict(self, predict_loader: DataLoader, return_y_true=False, add_hook=False, hook_layer=None):
        """
        Just predict by model,and add one forward hook to get processing output.

        Parameters
        ----------
        predict_loader:DataLoader
            MGEDataLoader, the target_y could be ``None``.
        return_y_true:bool
            if return_y_true, return (y_preds, y_true)
        add_hook:bool
        hook_layer:
            layer of model. could get layers by ``get_target_layer`` or ``get_layers_name_with_weight``.

        Returns
        -------
        y_pred:tensor
        y_true:tensor
            if return_y_true

        """

        self.model.eval()
        self.model.to(self.device)

        ############

        handles = []
        if add_hook:
            try:
                self.forward_hook_list = []

                def for_hook(module, input, output):
                    self.forward_hook_list.append(output.detach().cpu())

                handles.append(hook_layer.register_forward_hook(for_hook))
            except BaseException as e:
                print(e)
                raise AttributeError("use ``hook_layer_name`` to defined the hook layer.")
        y_preds = []
        y_true = []
        for data in predict_loader:
            data = data.to(self.device)
            y_preds.append(self.model(data).detach().cpu())
            if hasattr(data, "y"):
                y_true.append(self.process_label(data.y.detach().cpu()))

        if add_hook:
            [i.remove() for i in handles]  # del

        if return_y_true and y_true != []:
            return torch.cat(y_preds), torch.cat(y_true)
        else:
            return torch.cat(y_preds)
