import os
import time
from pathlib import Path

import copy
import numpy as np
import torch
from torch import nn

from ...tqdm import BTqdm
from ...writer import BWriter
from ...drawer.terminal import BColor, BAppearance
from ..early_stop import BReloadByLoss
from ..early_stop import BStopByLoss, BStopByAcc
from ..early_stop import BStopByLossDelta, BStopByAccDelta
from ..early_stop import BStopByOverfitting
from ..early_stop import BStopByBYZH
from ...basic import Byzh
def inputs_function(inputs):
    return inputs
def outputs_function(outputs):
    return outputs
def labels_function(labels):
    return labels

class _saveDuringTrain:
    def __init__(self, path, rounds):
        self.path = path
        self.rounds = rounds
        self.cnt = 0
    def __call__(self):
        self.cnt += 1
        if self.cnt > self.rounds:
            self.cnt = 0
            return True
        return False

class BTrainer(Byzh):
    def __init__(self, model, train_loader, val_loader, optimizer, criterion, device,
                 lrScheduler=None,
                 isBinaryCls=False, isParallel=False, isSpikingjelly=False):
        '''
        训练:\n
        train_eval_s\n
        训练前函数:\n
        load_model, load_optimizer, load_lrScheduler, set_writer, set_stop_by_acc\n
        训练后函数:\n
        save_latest_checkpoint, save_best_checkpoint, calculate_model
        :param model:
        :param train_loader:
        :param val_loader:
        :param optimizer:
        :param criterion:
        :param device:
        :param lrScheduler:
        :param isBinaryCls: 若是二分类, 则输出额外信息
        :param isParallel: 是否多GPU
        :param isSpikingjelly: 是否为SNN
        '''
        super().__init__()
        self.train_acc_lst = []
        self.train_loss_lst = []
        self.val_acc_lst = []
        self.val_f1_lst = []
        self.val_L0_True_lst = []
        self.val_L1_True_lst = []

        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.optimizer = optimizer
        self.criterion = criterion
        self.device = device
        self.lrScheduler = lrScheduler
        self.isBinaryCls = isBinaryCls
        self.isParallel = isParallel
        self.isSpikingjelly = isSpikingjelly

        self.model.to(self.device)

        self.__isTraining = False
        # save_temp
        self.__save_during_train = None
        # writer
        self.__writer = None
        # early stop
        self.__stop_by_acc = None
        self.__stop_by_acc_delta = None
        self.__stop_by_loss = None
        self.__stop_by_loss_delta = None
        self.__stop_by_overfitting = None
        self.__stop_by_byzh = None
        # early reload
        self.__reload_by_loss = None
        # save_best
        self.__best_acc = 0
        self.__best_model_state_dict = None
        self.__best_optimizer_state_dict = None
        self.__best_lrScheduler_state_dict = None

        if self.isParallel:
            if str(self.device) == str(torch.device("cuda")):
                if torch.cuda.device_count() > 1:
                    print(f"当前GPU数量:{torch.cuda.device_count()}, 使用多GPU训练")
                    self.model = nn.DataParallel(self.model)
                else:
                    print(f"当前GPU数量:{torch.cuda.device_count()}, 使用单GPU训练")

    def train_eval_s(self,
                     epochs,
                     inputs_func=inputs_function,
                     outputs_func=outputs_function,
                     labels_func=labels_function):
        '''
        :param epochs:
        :param inputs_func: 对inputs的处理函数
        :param outputs_func: 对outputs的处理函数
        :param labels_func: 对labels的处理函数
        :return:
        '''
        self.__isTraining = True

        # 检查inputs_func是否为函数
        if not callable(inputs_func):
            raise ValueError("inputs_func传入的参数应该要是一个函数!!!")
        # 检查outputs_func是否为函数
        if not callable(outputs_func):
            raise ValueError("outputs_func传入的参数应该要是一个函数!!!")
        # 检查labels_func是否为函数
        if not callable(labels_func):
            raise ValueError("labels_func传入的参数应该要是一个函数!!!")

        for epoch in range(epochs):
            train_acc, train_loss, current_lr = self.__train_once(epoch, epochs, inputs_func, outputs_func, labels_func)
            val_acc = self.__eval_once(inputs_func, outputs_func, labels_func)
            # 日志
            if self.__writer is not None:
                self.__writer.toFile(
                    f'Epoch [{epoch}/{epochs}], lr: {current_lr:.2e} | '
                    f'train_loss: {train_loss:.3f} | train_acc: {train_acc:.3f}, val_acc: {val_acc:.3f}'
                )
            # 保存模型
            if self.__save_during_train is not None:
                if self.__save_during_train():
                    self.save_best_checkpoint(self.__save_during_train.path)
            # 早停
            if self.__stop_by_acc is not None:
                if self.__stop_by_acc(train_loss):
                    info = f'模型在连续{self.__stop_by_acc.rounds}个epoch内停滞, 触发stop_by_acc'
                    self.__print_and_toWriter(info)
                    break
            # 早停
            if self.__stop_by_acc_delta is not None:
                if self.__stop_by_acc_delta(val_acc):
                    info = f'模型在连续{self.__stop_by_acc_delta.rounds}个epoch内过拟合, 触发stop_by_acc_delta'
                    self.__print_and_toWriter(info)
                    break
            # 早停
            if self.__stop_by_loss is not None:
                if self.__stop_by_loss(train_loss):
                    info = f'模型在连续{self.__stop_by_loss.rounds}个epoch内停滞, 触发stop_by_loss'
                    self.__print_and_toWriter(info)
                    break
            # 早停
            if self.__stop_by_loss_delta is not None:
                if self.__stop_by_loss_delta(train_loss):
                    info = f'模型在连续{self.__stop_by_loss_delta.rounds}个epoch内停滞, 触发stop_by_loss_delta'
                    self.__print_and_toWriter(info)
                    break
            # 早停
            if self.__stop_by_byzh is not None:
                if self.__stop_by_byzh(train_loss, train_acc, val_acc):
                    info = f'模型在连续{self.__stop_by_byzh.rounds}个epoch内过拟合, 触发stop_by_byzh'
                    self.__print_and_toWriter(info)
                    break
            # 过拟合
            if self.__stop_by_overfitting is not None:
                if self.__stop_by_overfitting(train_acc, val_acc):
                    info = f'模型在连续{self.__stop_by_overfitting.rounds}个epoch内过拟合, 触发stop_by_overfitting'
                    self.__print_and_toWriter(info)
                    break
            # 重加载
            if self.__reload_by_loss is not None:
                match self.__reload_by_loss(train_loss):
                    case 'continue':
                        pass
                    case 'reload':
                        info = f'模型触发reload_by_loss(第{self.__reload_by_loss.cnt_reload}次加载)'
                        self.__print_and_toWriter(info)
                        # 加载
                        self.model.load_state_dict(self.__best_model_state_dict)
                        self.optimizer.load_state_dict(self.__best_optimizer_state_dict)
                        if self.lrScheduler is not None:
                            self.lrScheduler.load_state_dict(self.__best_lrScheduler_state_dict)
                        self.calculate_model()
                    case 'done':
                        pass
    def calculate_model(self, dataloader=None, model=None):
        '''
        如果不指定, 则用类内的
        :param dataloader: 默认self.val_loader
        :param model: 默认self.model
        :return:
        '''
        if dataloader==None:
            dataloader = self.val_loader
        if model==None:
            model = self.model
        model.eval()
        correct = 0
        total = 0
        y_true = []
        y_pred = []
        inference_time = []
        with torch.no_grad():
            for inputs, labels in dataloader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                start_time = time.time()
                outputs = model(inputs)
                end_time = time.time()
                _, predicted = torch.max(outputs, 1)

                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                y_true.extend(labels.cpu().numpy())
                y_pred.extend(predicted.cpu().numpy())
                inference_time.append(end_time - start_time)
                if self.isSpikingjelly:
                    from spikingjelly.activation_based import functional
                    functional.reset_net(model)
            # 平均推理时间
            inference_time = sum(inference_time) / len(inference_time)
            # acc & f1
            accuracy = correct / total
            f1_score = self.__get_f1_score(y_true, y_pred)
            # 参数量
            params = sum(p.numel() for p in model.parameters())

            info = f'[calc] accuracy: {accuracy:.3f}, f1_score: {f1_score:.3f}'
            self.__print_and_toWriter(info)
            info = f'------ inference_time: {inference_time:.2e}s, params: {params / 1e3}K'
            self.__print_and_toWriter(info)

            if self.isBinaryCls:
                cm = self.__get_confusion_matrix(y_true, y_pred)
                TN = cm[0, 0]
                FP = cm[0, 1]
                FN = cm[1, 0]
                TP = cm[1, 1]
                L0_True = self.__get_L0_True(TN, FP)
                L1_True = self.__get_L1_True(FN, TP)

                info = f'------ L0_True: {L0_True:.3f}, L1_True: {L1_True:.3f}'
                self.__print_and_toWriter(info)

    def save_latest_checkpoint(self, path):
        '''
        字典checkpoint包含net, optimizer, lrScheduler
        '''
        parent_path = Path(path).parent
        os.makedirs(parent_path, exist_ok=True)

        checkpoint = {
            'model': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'lrScheduler': self.lrScheduler.state_dict() if self.lrScheduler is not None else None
        }
        torch.save(checkpoint, path)
        print(f"[save] latest_checkpoint 已保存到 {path}")

    def save_best_checkpoint(self, path):
        '''
        字典checkpoint包含net, optimizer, lrScheduler
        '''
        parent_path = Path(path).parent
        os.makedirs(parent_path, exist_ok=True)

        checkpoint = {
            'model': self.__best_model_state_dict,
            'optimizer': self.__best_optimizer_state_dict,
            'lrScheduler': self.__best_lrScheduler_state_dict
        }
        torch.save(checkpoint, path)
        print(f"[save] best_checkpoint 已保存到 {path}")
    def load_model(self, path):
        checkpoint = torch.load(path, map_location='cpu')
        self.model.load_state_dict(checkpoint['model'])
        self.model.to(self.device)
        print(f"[load] model 已从 {path} 加载")
    def load_optimizer(self, path):
        checkpoint = torch.load(path)
        self.optimizer.load_state_dict(checkpoint['optimizer'])
        print(f"[load] optimizer 已从{path}加载")
    def load_lrScheduler(self, path):
        checkpoint = torch.load(path)
        if self.lrScheduler is not None and checkpoint['lrScheduler'] is not None:
            self.lrScheduler.load_state_dict(checkpoint['lrScheduler'])
            print(f"[load] lrScheduler 已从{path}加载")
        else:
            print(f"[load] lrScheduler为None, 加载失败")

    def set_writer(self, writePath, mode='a'):
        '''
        请在训练前设置set_writer
        :param writePath:
        :param mode: 'a', 'w'
        :return:
        '''
        self.__writer = BWriter(writePath, ifTime=True)
        if mode == 'a':
            pass
        if mode == 'w':
            self.__writer.clearFile()

        self.__writer.toFile("[dataset] -> " + self.train_loader.dataset.__class__.__name__, ifTime=False)
        self.__writer.toFile("[batch_size] -> " + str(self.train_loader.batch_size), ifTime=False)
        self.__writer.toFile("[lr] -> " + str(self.optimizer.param_groups[0]['lr']), ifTime=False)
        self.__writer.toFile("[criterion] -> " + str(self.criterion), ifTime=False)
        self.__writer.toFile("[optimizer] -> " + str(self.optimizer), ifTime=False)
        if self.lrScheduler is not None:
            self.__writer.toFile("[lrScheduler] -> "+str(self.lrScheduler), ifTime=False)
        self.__writer.toFile("[model] -> " + str(self.model), ifTime=False)

        print(f'日志将保存到{writePath}')
    def set_save_during_train(self, path, rounds=10):
        '''
        请在训练前设置set_save_during_train
        '''
        self.__save_during_train = _saveDuringTrain(path, rounds)
    def set_stop_by_acc(self, rounds=10, delta=0.01):
        '''
        请在训练前设置set_stop_by_acc
        :param rounds: 连续rounds次, val_acc - max_val_acc > delta, 则停止训练
        :return:
        '''
        self.__stop_by_acc = BStopByAcc(rounds=rounds, delta=delta)
    def set_stop_by_overfitting(self, rounds=10, delta=0.1):
        '''
        请在训练前设置set_stop_by_overfitting
        :param rounds: 连续rounds次, train_acc - val_acc > delta, 则停止训练
        :return:
        '''
        self.__stop_by_overfitting = BStopByOverfitting(rounds=rounds, delta=delta)
    def set_stop_by_acc_delta(self, rounds=10, delta=0.003):
        '''
        请在训练前设置set_stop_by_acc_delta
        :param rounds: 连续rounds次, |before_acc - val_acc| <= delta, 则停止训练
        :return:
        '''
        self.__stop_by_acc_delta = BStopByAccDelta(rounds=rounds, delta=delta)
    def set_stop_by_loss(self, rounds=10, delta=0.01):
        '''
        请在训练前设置set_stop_by_loss
        :param rounds: 连续rounds次, train_loss - min_train_loss > delta, 则停止训练
        :return:
        '''
        self.__stop_by_loss = BStopByLoss(rounds=rounds, delta=delta)
    def set_stop_by_loss_delta(self, rounds=10, delta=0.002):
        '''
        请在训练前设置set_stop_by_loss_delta
        :param rounds: 连续rounds次, |before_loss - now_loss| <= delta, 则停止训练
        :return:
        '''
        self.__stop_by_loss_delta = BStopByLossDelta(rounds=rounds, delta=delta)
    def set_stop_by_byzh(self, rounds=10, loss_delta=0.002, acc_delta=0.003):
        '''
        请在训练前设置set_stop_by_byzh
        '''
        self.__stop_by_byzh = BStopByBYZH(rounds=rounds, loss_delta=loss_delta, acc_delta=acc_delta)
    def set_reload_by_loss(self, max_reload_count=5, reload_rounds=10, delta=0.01):
        '''
        请在训练前设置set_reload_by_loss
        :param rounds: 连续rounds次都是train_loss > min_train_loss + delta
        :return:
        '''
        self.__reload_by_loss = BReloadByLoss(max_reload_count, reload_rounds, delta)

    def draw(self, jpgPath, isShow=False):
        parent_path = Path(jpgPath).parent
        os.makedirs(parent_path, exist_ok=True)

        import matplotlib
        if isShow == False:
            matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        import seaborn as sns
        import numpy as np  # 导入numpy库

        palette = sns.color_palette("Set2", 3)

        plt.figure(figsize=(8, 8))
        plt.subplots_adjust(hspace=0.4)

        plt.subplot(3, 1, 1)
        # 每十个画一次(防止点多卡顿)
        temp = [x for i, x in enumerate(self.train_loss_lst) if (i + 1) % 10 == 0]
        plt.plot(temp, color="red", label="train_loss")
        plt.xlabel("iter 1/10")
        plt.ylabel("loss")
        plt.legend(loc='upper right')

        plt.subplot(3, 1, 2)
        plt.plot(self.train_acc_lst, color="red", label="train_acc")
        plt.plot(self.val_acc_lst, color="blue", label="val_acc")
        # 找到train_acc的峰值点并标记
        train_acc_peak_index = np.argmax(self.train_acc_lst)
        plt.scatter(train_acc_peak_index, self.train_acc_lst[train_acc_peak_index], color="red", marker="v", s=100)
        # 找到val_acc的峰值点并标记
        val_acc_peak_index = np.argmax(self.val_acc_lst)
        plt.scatter(val_acc_peak_index, self.val_acc_lst[val_acc_peak_index], color="blue", marker="v", s=100)
        plt.xlabel("epoch")
        plt.ylabel("acc")
        plt.ylim(-0.05, 1.05)
        plt.legend(loc='lower right')

        if self.isBinaryCls:
            plt.subplot(3, 1, 3)
            plt.plot(self.val_f1_lst, color=palette[0], label="f1")
            plt.plot(self.val_L0_True_lst, color=palette[1], label="L0_True")
            plt.plot(self.val_L1_True_lst, color=palette[2], label="L1_True")

            # 找到val_f1的峰值点并标记
            val_f1_peak_index = np.argmax(self.val_f1_lst)
            plt.scatter(val_f1_peak_index, self.val_f1_lst[val_f1_peak_index], color=palette[0], marker="v", s=100)

            # 找到val_L0_True的峰值点并标记
            val_L0_True_peak_index = np.argmax(self.val_L0_True_lst)
            plt.scatter(val_L0_True_peak_index, self.val_L0_True_lst[val_L0_True_peak_index], color=palette[1],
                        marker="v", s=100)

            # 找到val_L1_True的峰值点并标记
            val_L1_True_peak_index = np.argmax(self.val_L1_True_lst)
            plt.scatter(val_L1_True_peak_index, self.val_L1_True_lst[val_L1_True_peak_index], color=palette[2],
                        marker="v", s=100)

            plt.xlabel("epoch")
            plt.ylabel("score")
            plt.ylim(-0.05, 1.05)
            plt.legend(loc='lower right')

        plt.savefig(jpgPath)
        print(f"[draw] picture 已保存到{jpgPath}")
        if isShow:
            plt.show()
        plt.close()

    def __print_and_toWriter(self, info, ifPrint=True):
        if ifPrint:
            print(info)
        if self.__writer is not None:
            self.__writer.toFile(info)
    def __train_once(self, epoch, epochs, inputs_func, outputs_func, labels_func):
        bar = BTqdm(total=len(self.train_loader))
        current_lr = self.optimizer.param_groups[0]['lr']

        self.model.train()
        correct = 0
        total = 0
        losses = 0
        for iter, (inputs, labels) in enumerate(self.train_loader):
            # 基本训练
            inputs, labels = inputs.to(self.device), labels.to(self.device)
            labels = labels_func(labels)
            inputs = inputs_func(inputs)
            outputs = self.model(inputs)
            outputs = outputs_func(outputs)

            loss = self.criterion(outputs, labels)
            self.optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 计算梯度
            self.optimizer.step()  # 更新参数
            # SNN
            if self.isSpikingjelly:
                from spikingjelly.activation_based import functional
                functional.reset_net(self.model)
            # 进度条
            bar.update(1,
                       setting=BColor.BLUE+BAppearance.HIGHLIGHT,
                       prefix=f"{BColor.BLUE}Epoch [{epoch}/{epochs}]",
                       suffix=f"lr: {current_lr:.2e}, loss: {loss.item():.3f}")
            # 数据记录
            self.train_loss_lst.append(loss.item())
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            losses += loss.item()
        accuracy = correct / total
        train_loss = losses / len(self.train_loader)
        print(f'Epoch [{epoch}/{epochs}], train_loss: {train_loss:.3f}, train_Acc: {accuracy:.3f}', end='')
        self.train_acc_lst.append(accuracy)
        # 更新学习率
        if self.lrScheduler:
            self.lrScheduler.step()

        return accuracy, train_loss, current_lr

    def __eval_once(self, inputs_func, outputs_func, labels_func):
        self.model.eval()
        correct = 0
        total = 0
        y_true = []
        y_pred = []
        with torch.no_grad():
            for inputs, labels in self.val_loader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                labels = labels_func(labels)
                inputs = inputs_func(inputs)
                outputs = self.model(inputs)
                outputs = outputs_func(outputs)
                _, predicted = torch.max(outputs, 1)

                if self.isSpikingjelly:
                    from spikingjelly.activation_based import functional
                    functional.reset_net(self.model)

                total += labels.size(0)
                correct += (predicted == labels).sum().item()

                y_true.extend(labels.cpu().numpy())
                y_pred.extend(predicted.cpu().numpy())

            # 记录accuracy
            accuracy = correct / total
            print(f', val_Acc: {accuracy:.3f}')
            self.val_acc_lst.append(accuracy)

            # 保存最优模型
            if accuracy > self.__best_acc:
                self.__best_acc = accuracy
                self.__best_model_state_dict = copy.deepcopy(self.model.state_dict())
                self.__best_optimizer_state_dict = copy.deepcopy(self.optimizer.state_dict())
                self.__best_lrScheduler_state_dict = copy.deepcopy(self.lrScheduler.state_dict()) if self.lrScheduler else None

            if self.isBinaryCls:
                cm = self.__get_confusion_matrix(y_true, y_pred)
                TN = cm[0, 0]
                FP = cm[0, 1]
                FN = cm[1, 0]
                TP = cm[1, 1]

                f1 = self.__get_f1_score(y_true, y_pred)
                self.val_f1_lst.append(f1)

                L0_True = self.__get_L0_True(TN, FP)
                self.val_L0_True_lst.append(L0_True)

                L1_True = self.__get_L1_True(FN, TP)
                self.val_L1_True_lst.append(L1_True)

        return accuracy

    def __get_confusion_matrix(self, y_true, y_pred):
        from sklearn.metrics import confusion_matrix
        result = confusion_matrix(y_true, y_pred)
        return result
    def __get_f1_score(self, y_true, y_pred):
        from sklearn.metrics import f1_score
        result = f1_score(y_true, y_pred, average='macro')
        return result
    def __get_recall(self, y_true, y_pred):
        from sklearn.metrics import recall_score
        result = recall_score(y_true, y_pred, average='macro')
        return result
    def __get_precision(self, y_true, y_pred):
        from sklearn.metrics import precision_score
        result = precision_score(y_true, y_pred, average='macro')
        return result

    def __get_L0_True(self, TN, FP):
        return TN / (TN + FP)

    def __get_L1_True(self, FN, TP):
        return TP / (TP + FN)