import math
import os
import time
import copy
import psutil
import torch
from utils.logger import get_logger
from utils.loss_function import all_metrics
from tqdm import tqdm


class TrainingStats:
    def __init__(self, device):
        self.device = device
        self.reset()

    def reset(self):
        self.gpu_mem_usage_list = []
        self.cpu_mem_usage_list = []
        self.train_time_list = []
        self.infer_time_list = []
        self.total_iters = 0
        self.start_time = None
        self.end_time = None

    def start_training(self):
        self.start_time = time.time()

    def end_training(self):
        self.end_time = time.time()

    def record_step_time(self, duration, mode):
        """记录单步耗时和总迭代次数"""
        if mode == "train":
            self.train_time_list.append(duration)
        else:
            self.infer_time_list.append(duration)
        self.total_iters += 1

    def record_memory_usage(self):
        """记录当前 GPU 和 CPU 内存占用"""
        process = psutil.Process(os.getpid())
        cpu_mem = process.memory_info().rss / (1024**2)

        if torch.cuda.is_available():
            gpu_mem = torch.cuda.max_memory_allocated(device=self.device) / (1024**2)
            torch.cuda.reset_peak_memory_stats(device=self.device)
        else:
            gpu_mem = 0.0

        self.cpu_mem_usage_list.append(cpu_mem)
        self.gpu_mem_usage_list.append(gpu_mem)

    def report(self, logger):
        """在训练结束时输出汇总统计"""
        if not self.start_time or not self.end_time:
            logger.warning("TrainingStats: start/end time not recorded properly.")
            return

        total_time = self.end_time - self.start_time
        avg_gpu_mem = (
            sum(self.gpu_mem_usage_list) / len(self.gpu_mem_usage_list)
            if self.gpu_mem_usage_list
            else 0
        )
        avg_cpu_mem = (
            sum(self.cpu_mem_usage_list) / len(self.cpu_mem_usage_list)
            if self.cpu_mem_usage_list
            else 0
        )
        avg_train_time = (
            sum(self.train_time_list) / len(self.train_time_list)
            if self.train_time_list
            else 0
        )
        avg_infer_time = (
            sum(self.infer_time_list) / len(self.infer_time_list)
            if self.infer_time_list
            else 0
        )
        iters_per_sec = self.total_iters / total_time if total_time > 0 else 0

        logger.info("===== Training Summary =====")
        logger.info(f"Total training time: {total_time:.2f} s")
        logger.info(f"Total iterations: {self.total_iters}")
        logger.info(f"Average iterations per second: {iters_per_sec:.2f}")
        logger.info(f"Average GPU Memory Usage: {avg_gpu_mem:.2f} MB")
        logger.info(f"Average CPU Memory Usage: {avg_cpu_mem:.2f} MB")
        if avg_train_time:
            logger.info(f"Average training step time: {avg_train_time * 1000:.2f} ms")
        if avg_infer_time:
            logger.info(f"Average inference step time: {avg_infer_time * 1000:.2f} ms")


class Trainer:
    def __init__(
        self,
        model,
        loss,
        optimizer,
        train_loader,
        val_loader,
        test_loader,
        scaler,
        args,
        lr_scheduler=None,
    ):
        self.device = args["basic"]["device"]
        args = args["train"]
        self.model = model
        self.loss = loss
        self.optimizer = optimizer
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.test_loader = test_loader
        self.scaler = scaler
        self.args = args
        self.lr_scheduler = lr_scheduler
        self.train_per_epoch = len(train_loader)
        self.val_per_epoch = len(val_loader) if val_loader else 0

        # Paths for saving models and logs
        self.best_path = os.path.join(args["log_dir"], "best_model.pth")
        self.best_test_path = os.path.join(args["log_dir"], "best_test_model.pth")
        self.loss_figure_path = os.path.join(args["log_dir"], "loss.png")

        # Initialize logger
        if not os.path.isdir(args["log_dir"]) and not args["debug"]:
            os.makedirs(args["log_dir"], exist_ok=True)
        self.logger = get_logger(
            args["log_dir"], name=self.model.__class__.__name__, debug=args["debug"]
        )
        self.logger.info(f"Experiment log path in: {args['log_dir']}")

        # Stats tracker
        self.stats = TrainingStats(device=self.device)

    def _run_epoch(self, epoch, dataloader, mode):
        if mode == "train":
            self.model.train()
            optimizer_step = True
        else:
            self.model.eval()
            optimizer_step = False

        total_loss = 0
        epoch_time = time.time()

        y_pred, y_true = [], []

        with torch.set_grad_enabled(optimizer_step):
            progress_bar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f"{mode.capitalize()} Epoch {epoch}")
            for batch_idx, (data, target) in progress_bar:
                start_time = time.time()

                label = target[..., : self.args["output_dim"]]
                output = self.model(data).to(self.device)

                if self.args["real_value"]:
                    output = self.scaler.inverse_transform(output)

                loss = self.loss(output, label)
                if optimizer_step and self.optimizer is not None:
                    self.optimizer.zero_grad()
                    loss.backward()

                    if self.args["grad_norm"]:
                        torch.nn.utils.clip_grad_norm_(
                            self.model.parameters(), self.args["max_grad_norm"]
                        )
                    self.optimizer.step()

                step_time = time.time() - start_time
                self.stats.record_step_time(step_time, mode)

                total_loss += loss.item()
                y_pred.append(output.detach().cpu())
                y_true.append(label.detach().cpu())

                # Update progress bar with the current loss
                progress_bar.set_postfix(loss=loss.item())

        y_pred = torch.cat(y_pred, dim=0)
        y_true = torch.cat(y_true, dim=0)

        avg_loss = total_loss / len(dataloader)
        # 输出指标
        mae, rmse, mape = all_metrics(
            y_pred, y_true, self.args["mae_thresh"], self.args["mape_thresh"]
        )
        self.logger.info(
            f"Epoch #{epoch:02d}: {mode.capitalize():<5} MAE:{mae:5.2f} | RMSE:{rmse:5.2f} | MAPE:{mape:7.4f} | Time: {time.time() - epoch_time:.2f} s"
        )

        # 记录内存
        self.stats.record_memory_usage()

        return avg_loss

    def train_epoch(self, epoch):
        return self._run_epoch(epoch, self.train_loader, "train")

    def val_epoch(self, epoch):
        return self._run_epoch(epoch, self.val_loader or self.test_loader, "val")

    def test_epoch(self, epoch):
        return self._run_epoch(epoch, self.test_loader, "test")

    def train(self):
        best_model, best_test_model = None, None
        best_loss, best_test_loss = float("inf"), float("inf")
        not_improved_count = 0

        self.stats.start_training()
        self.logger.info("Training process started")

        for epoch in range(1, self.args["epochs"] + 1):
            train_epoch_loss = self.train_epoch(epoch)
            val_epoch_loss = self.val_epoch(epoch)
            test_epoch_loss = self.test_epoch(epoch)

            if train_epoch_loss > 1e6:
                self.logger.warning("Gradient explosion detected. Ending...")
                break

            if val_epoch_loss < best_loss:
                best_loss = val_epoch_loss
                not_improved_count = 0
                best_model = copy.deepcopy(self.model.state_dict())
                self.logger.info("Best validation model saved!")
            else:
                not_improved_count += 1

            if (
                self.args["early_stop"]
                and not_improved_count == self.args["early_stop_patience"]
            ):
                self.logger.info(
                    f"Validation performance didn't improve for {self.args['early_stop_patience']} epochs. Training stops."
                )
                break

            if test_epoch_loss < best_test_loss:
                best_test_loss = test_epoch_loss
                best_test_model = copy.deepcopy(self.model.state_dict())

        if not self.args["debug"]:
            torch.save(best_model, self.best_path)
            torch.save(best_test_model, self.best_test_path)
            self.logger.info(
                f"Best models saved at {self.best_path} and {self.best_test_path}"
            )

        self.stats.end_training()
        self.stats.report(self.logger)

        self._finalize_training(best_model, best_test_model)

        # 输出参数量
        try:
            total_params = sum(
                p.numel() for p in self.model.parameters() if p.requires_grad
            )
            self.logger.info(f"Trainable params: {total_params}")
        except Exception:
            pass

    def _finalize_training(self, best_model, best_test_model):
        self.model.load_state_dict(best_model)
        self.logger.info("Testing on best validation model")
        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)

        self.model.load_state_dict(best_test_model)
        self.logger.info("Testing on best test model")
        self.test(self.model, self.args, self.test_loader, self.scaler, self.logger)

    @staticmethod
    def test(model, args, data_loader, scaler, logger, path=None):
        if path:
            checkpoint = torch.load(path)
            model.load_state_dict(checkpoint["state_dict"])
            model.to(args["basic"]["device"])

        model.eval()
        y_pred, y_true = [], []

        with torch.no_grad():
            for data, target in data_loader:
                label = target[..., : args["output_dim"]]
                output = model(data)
                y_pred.append(output)
                y_true.append(label)

        if args["real_value"]:
            y_pred = scaler.inverse_transform(torch.cat(y_pred, dim=0))
        else:
            y_pred = torch.cat(y_pred, dim=0)
        y_true = torch.cat(y_true, dim=0)

        for t in range(y_true.shape[1]):
            mae, rmse, mape = all_metrics(
                y_pred[:, t, ...],
                y_true[:, t, ...],
                args["mae_thresh"],
                args["mape_thresh"],
            )
            logger.info(
                f"Horizon {t + 1:02d}, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
            )

        mae, rmse, mape = all_metrics(
            y_pred, y_true, args["mae_thresh"], args["mape_thresh"]
        )
        logger.info(
            f"Average Horizon, MAE: {mae:.4f}, RMSE: {rmse:.4f}, MAPE: {mape:.4f}"
        )

    @staticmethod
    def _compute_sampling_threshold(global_step, k):
        return k / (k + math.exp(global_step / k))