import math
import os
import time
import copy
import psutil
from tqdm import tqdm

import torch
from lib.logger import get_logger
from lib.loss_function import all_metrics
from model.STEP.step_loss import step_loss


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.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
        log_dir = args.get('log_dir', './logs/STEP')
        os.makedirs(log_dir, exist_ok=True)  # 确保目录存在
        self.best_path = os.path.join(log_dir, 'best_model.pth')
        self.best_test_path = os.path.join(log_dir, 'best_test_model.pth')
        self.loss_figure_path = os.path.join(log_dir, 'loss.png')

        # Initialize logger
        log_dir = args.get('log_dir', './logs/STEP')
        self.logger = get_logger(log_dir, name='STEP_Trainer')
        
        # Initialize training stats
        self.device = next(model.parameters()).device
        self.stats = TrainingStats(self.device)

    def train_epoch(self, epoch):
        self.model.train()
        total_loss = 0
        total_metrics = {}
        
        with tqdm(self.train_loader, desc=f'Epoch {epoch}') as pbar:
            for batch_idx, (data, target) in enumerate(pbar):
                start_time = time.time()
                
                data = data.to(self.device)
                target = target.to(self.device)
                
                self.optimizer.zero_grad()
                
                # STEP模型的前向传播
                output = self.model(data)
                
                # 计算损失（这里需要根据STEP模型的具体输出调整）
                # STEP模型返回多个输出，包括预测值、Bernoulli参数等
                if isinstance(output, tuple):
                    prediction = output[0]
                    # 如果模型返回了其他参数，可以在这里处理
                else:
                    prediction = output
                
                # 使用标准损失函数
                if callable(self.loss) and hasattr(self.loss, '__call__'):
                    # 如果是一个可调用对象（比如masked_mae_loss返回的函数）
                    if hasattr(self.loss, 'func_name') or 'function' in str(type(self.loss)):
                        loss_fn = self.loss(None, None)  # 创建实际的损失函数
                        loss = loss_fn(prediction, target)
                    else:
                        loss = self.loss(prediction, target)
                else:
                    # 如果是PyTorch的损失函数
                    loss = self.loss(prediction, target)
                
                loss.backward()
                
                # 梯度裁剪
                if self.args.get('clip_grad_norm', 0) > 0:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args['clip_grad_norm'])
                
                self.optimizer.step()
                
                # 记录统计信息
                step_time = time.time() - start_time
                self.stats.record_step_time(step_time, 'train')
                
                total_loss += loss.item()
                
                # 计算指标
                mae, rmse, mape = all_metrics(prediction, target, None, 0.0)
                metrics = {'mae': mae.item(), 'rmse': rmse.item(), 'mape': mape.item()}
                for key, value in metrics.items():
                    if key not in total_metrics:
                        total_metrics[key] = 0
                    total_metrics[key] += value
                
                # 更新进度条
                pbar.set_postfix({
                    'Loss': f'{loss.item():.4f}',
                    'MAE': f'{metrics.get("mae", 0):.4f}',
                    'RMSE': f'{metrics.get("rmse", 0):.4f}'
                })
                
                # 记录内存使用
                if batch_idx % 100 == 0:
                    self.stats.record_memory_usage()
        
        # 计算平均损失和指标
        avg_loss = total_loss / len(self.train_loader)
        avg_metrics = {key: value / len(self.train_loader) for key, value in total_metrics.items()}
        
        return avg_loss, avg_metrics

    def val_epoch(self, epoch):
        self.model.eval()
        total_loss = 0
        total_metrics = {}
        
        with torch.no_grad():
            with tqdm(self.val_loader, desc=f'Validation {epoch}') as pbar:
                for batch_idx, (data, target) in enumerate(pbar):
                    start_time = time.time()
                    
                    data = data.to(self.device)
                    target = target.to(self.device)
                    
                    # STEP模型的前向传播
                    output = self.model(data)
                    
                    if isinstance(output, tuple):
                        prediction = output[0]
                    else:
                        prediction = output
                    
                    # 计算损失
                    if callable(self.loss) and hasattr(self.loss, '__call__'):
                        # 如果是一个可调用对象（比如masked_mae_loss返回的函数）
                        if hasattr(self.loss, 'func_name') or 'function' in str(type(self.loss)):
                            loss_fn = self.loss(None, None)  # 创建实际的损失函数
                            loss = loss_fn(prediction, target)
                        else:
                            loss = self.loss(prediction, target)
                    else:
                        # 如果是PyTorch的损失函数
                        loss = self.loss(prediction, target)
                    
                    # 记录统计信息
                    step_time = time.time() - start_time
                    self.stats.record_step_time(step_time, 'val')
                    
                    total_loss += loss.item()
                    
                    # 计算指标
                    mae, rmse, mape = all_metrics(prediction, target, None, 0.0)
                    metrics = {'mae': mae.item(), 'rmse': rmse.item(), 'mape': mape.item()}
                    for key, value in metrics.items():
                        if key not in total_metrics:
                            total_metrics[key] = 0
                        total_metrics[key] += value
                    
                    # 更新进度条
                    pbar.set_postfix({
                        'Loss': f'{loss.item():.4f}',
                        'MAE': f'{metrics.get("mae", 0):.4f}',
                        'RMSE': f'{metrics.get("rmse", 0):.4f}'
                    })
        
        # 计算平均损失和指标
        avg_loss = total_loss / len(self.val_loader)
        avg_metrics = {key: value / len(self.val_loader) for key, value in total_metrics.items()}
        
        return avg_loss, avg_metrics

    def test_epoch(self, epoch):
        self.model.eval()
        total_loss = 0
        total_metrics = {}
        
        with torch.no_grad():
            with tqdm(self.test_loader, desc=f'Test {epoch}') as pbar:
                for batch_idx, (data, target) in enumerate(pbar):
                    start_time = time.time()
                    
                    data = data.to(self.device)
                    target = target.to(self.device)
                    
                    # STEP模型的前向传播
                    output = self.model(data)
                    
                    if isinstance(output, tuple):
                        prediction = output[0]
                    else:
                        prediction = output
                    
                    # 计算损失
                    if callable(self.loss) and hasattr(self.loss, '__call__'):
                        # 如果是一个可调用对象（比如masked_mae_loss返回的函数）
                        if hasattr(self.loss, 'func_name') or 'function' in str(type(self.loss)):
                            loss_fn = self.loss(None, None)  # 创建实际的损失函数
                            loss = loss_fn(prediction, target)
                        else:
                            loss = self.loss(prediction, target)
                    else:
                        # 如果是PyTorch的损失函数
                        loss = self.loss(prediction, target)
                    
                    # 记录统计信息
                    step_time = time.time() - start_time
                    self.stats.record_step_time(step_time, 'test')
                    
                    total_loss += loss.item()
                    
                    # 计算指标
                    mae, rmse, mape = all_metrics(prediction, target, None, 0.0)
                    metrics = {'mae': mae.item(), 'rmse': rmse.item(), 'mape': mape.item()}
                    for key, value in metrics.items():
                        if key not in total_metrics:
                            total_metrics[key] = 0
                        total_metrics[key] += value
                    
                    # 更新进度条
                    pbar.set_postfix({
                        'Loss': f'{loss.item():.4f}',
                        'MAE': f'{metrics.get("mae", 0):.4f}',
                        'RMSE': f'{metrics.get("rmse", 0):.4f}'
                    })
        
        # 计算平均损失和指标
        avg_loss = total_loss / len(self.test_loader)
        avg_metrics = {key: value / len(self.test_loader) for key, value in total_metrics.items()}
        
        return avg_loss, avg_metrics

    def train(self):
        self.stats.start_training()
        
        best_val_loss = float('inf')
        best_test_loss = float('inf')
        
        for epoch in range(self.args['epochs']):
            # 训练
            train_loss, train_metrics = self.train_epoch(epoch)
            
            # 验证
            if self.val_loader:
                val_loss, val_metrics = self.val_epoch(epoch)
                
                # 保存最佳模型
                if val_loss < best_val_loss:
                    best_val_loss = val_loss
                    torch.save(self.model.state_dict(), self.best_path)
                    self.logger.info(f'Epoch {epoch}: Best validation loss: {val_loss:.4f}')
            
            # 测试
            if self.test_loader:
                test_loss, test_metrics = self.test_epoch(epoch)
                
                # 保存最佳测试模型
                if test_loss < best_test_loss:
                    best_test_loss = test_loss
                    torch.save(self.model.state_dict(), self.best_test_path)
                    self.logger.info(f'Epoch {epoch}: Best test loss: {test_loss:.4f}')
            
            # 学习率调度
            if self.lr_scheduler:
                self.lr_scheduler.step()
            
            # 记录日志
            self.logger.info(f'Epoch {epoch}: Train Loss: {train_loss:.4f}, Train MAE: {train_metrics.get("mae", 0):.4f}')
            if self.val_loader:
                self.logger.info(f'Epoch {epoch}: Val Loss: {val_loss:.4f}, Val MAE: {val_metrics.get("mae", 0):.4f}')
            if self.test_loader:
                self.logger.info(f'Epoch {epoch}: Test Loss: {test_loss:.4f}, Test MAE: {test_metrics.get("mae", 0):.4f}')
        
        self.stats.end_training()
        self.stats.report(self.logger)
        
        return best_val_loss, best_test_loss