"""带类权重的训练器：处理分类任务的数据不均衡问题"""

from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, field

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm


@dataclass
class ClassWeights:
    """分类任务的类权重"""
    pdi: Optional[torch.Tensor] = None  # [4] for 4 PDI classes
    ee: Optional[torch.Tensor] = None   # [3] for 3 EE classes
    toxic: Optional[torch.Tensor] = None  # [2] for binary toxic


@dataclass
class LossWeightsBalanced:
    """各任务的损失权重（与 trainer.py 的 LossWeights 兼容）"""
    size: float = 1.0
    pdi: float = 1.0
    ee: float = 1.0
    delivery: float = 1.0
    biodist: float = 1.0
    toxic: float = 1.0
    moe_lb: float = 0.01  # MoE load-balancing 系数（仅在 use_moe=True 时生效）


def compute_class_weights_from_loader(
    loader: DataLoader,
    n_pdi_classes: int = 2,
    n_ee_classes: int = 3,
    n_toxic_classes: int = 2,
    smoothing: float = 0.1,
) -> ClassWeights:
    """
    从 DataLoader 统计类别频次并计算类权重。
    
    使用 inverse frequency 方式：weight_c = N / (n_classes * count_c)
    加 smoothing 避免极端权重。
    
    Args:
        loader: DataLoader（需要遍历一次）
        n_pdi_classes: PDI 类别数
        n_ee_classes: EE 类别数
        n_toxic_classes: toxic 类别数
        smoothing: 平滑系数（防止除零和极端权重）
    
    Returns:
        ClassWeights 对象，包含各分类任务的类权重张量
    """
    pdi_counts = torch.zeros(n_pdi_classes)
    ee_counts = torch.zeros(n_ee_classes)
    toxic_counts = torch.zeros(n_toxic_classes)
    
    for batch in loader:
        targets = batch["targets"]
        mask = batch["mask"]
        
        # PDI
        if "pdi" in targets and "pdi" in mask:
            m = mask["pdi"]
            if m.any():
                labels = targets["pdi"][m]
                for c in range(n_pdi_classes):
                    pdi_counts[c] += (labels == c).sum().item()
        
        # EE
        if "ee" in targets and "ee" in mask:
            m = mask["ee"]
            if m.any():
                labels = targets["ee"][m]
                for c in range(n_ee_classes):
                    ee_counts[c] += (labels == c).sum().item()
        
        # Toxic
        if "toxic" in targets and "toxic" in mask:
            m = mask["toxic"]
            if m.any():
                labels = targets["toxic"][m]
                for c in range(n_toxic_classes):
                    toxic_counts[c] += (labels == c).sum().item()
    
    def counts_to_weights(counts: torch.Tensor, n_classes: int) -> Optional[torch.Tensor]:
        """将计数转换为类权重"""
        total = counts.sum().item()
        if total == 0:
            return None
        # Inverse frequency with smoothing
        counts = counts + smoothing
        weights = total / (n_classes * counts)
        # Normalize to mean=1
        weights = weights / weights.mean()
        return weights
    
    return ClassWeights(
        pdi=counts_to_weights(pdi_counts, n_pdi_classes),
        ee=counts_to_weights(ee_counts, n_ee_classes),
        toxic=counts_to_weights(toxic_counts, n_toxic_classes),
    )


def compute_multitask_loss_balanced(
    outputs: Dict[str, torch.Tensor],
    targets: Dict[str, torch.Tensor],
    mask: Dict[str, torch.Tensor],
    task_weights: Optional[LossWeightsBalanced] = None,
    class_weights: Optional[ClassWeights] = None,
    model: Optional[nn.Module] = None,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
    """
    计算带类权重的多任务损失。
    
    Args:
        outputs: 模型输出
        targets: 真实标签
        mask: 有效样本掩码
        task_weights: 各任务权重
        class_weights: 分类任务的类权重
        
    Returns:
        (total_loss, loss_dict) 总损失和各任务损失
    """
    task_weights = task_weights or LossWeightsBalanced()
    class_weights = class_weights or ClassWeights()

    # 若模型 head 暴露 log_vars，则启用不确定性加权（自动平衡），否则回退到标量权重
    log_vars = None
    if model is not None and hasattr(model, "head") and hasattr(model.head, "log_vars"):
        log_vars = model.head.log_vars

    losses = {}
    device = next(iter(outputs.values())).device
    total_loss = torch.tensor(0.0, device=device)

    def _add(name: str, raw_loss: torch.Tensor, scalar_w: float) -> None:
        nonlocal total_loss
        losses[name] = raw_loss
        if log_vars is not None and name in log_vars:
            s = log_vars[name]
            total_loss = total_loss + torch.exp(-s) * raw_loss + 0.5 * s
        else:
            total_loss = total_loss + scalar_w * raw_loss

    # size: MSE
    if "size" in targets and mask["size"].any():
        m = mask["size"]
        _add("size", F.mse_loss(outputs["size"][m].squeeze(-1), targets["size"][m]), task_weights.size)

    # delivery: MSE
    if "delivery" in targets and mask["delivery"].any():
        m = mask["delivery"]
        _add("delivery", F.mse_loss(outputs["delivery"][m].squeeze(-1), targets["delivery"][m]), task_weights.delivery)

    # pdi: 类加权 CE
    if "pdi" in targets and mask["pdi"].any():
        m = mask["pdi"]
        w = class_weights.pdi.to(device) if class_weights.pdi is not None else None
        _add("pdi", F.cross_entropy(outputs["pdi"][m], targets["pdi"][m], weight=w), task_weights.pdi)

    # ee: 类加权 CE
    if "ee" in targets and mask["ee"].any():
        m = mask["ee"]
        w = class_weights.ee.to(device) if class_weights.ee is not None else None
        _add("ee", F.cross_entropy(outputs["ee"][m], targets["ee"][m], weight=w), task_weights.ee)

    # toxic: 类加权 CE
    if "toxic" in targets and mask["toxic"].any():
        m = mask["toxic"]
        w = class_weights.toxic.to(device) if class_weights.toxic is not None else None
        _add("toxic", F.cross_entropy(outputs["toxic"][m], targets["toxic"][m], weight=w), task_weights.toxic)

    # biodist: KL
    if "biodist" in targets and mask["biodist"].any():
        m = mask["biodist"]
        kl = F.kl_div(outputs["biodist"][m].log().clamp(min=-100), targets["biodist"][m], reduction="batchmean")
        _add("biodist", kl, task_weights.biodist)

    # MoE load-balancing aux loss（固定小权重，不纳入不确定性加权）
    if model is not None and hasattr(model, "get_last_moe_extras"):
        extras = model.get_last_moe_extras()
        if extras is not None and "lb_loss" in extras:
            losses["moe_lb"] = extras["lb_loss"]
            total_loss = total_loss + task_weights.moe_lb * losses["moe_lb"]

    return total_loss, losses


def train_epoch_balanced(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    task_weights: Optional[LossWeightsBalanced] = None,
    class_weights: Optional[ClassWeights] = None,
) -> Dict[str, float]:
    """带类权重的训练一个 epoch"""
    model.train()
    total_loss = 0.0
    task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
    n_batches = 0
    
    for batch in tqdm(loader, desc="Training", leave=False):
        smiles = batch["smiles"]
        tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
        targets = {k: v.to(device) for k, v in batch["targets"].items()}
        mask = {k: v.to(device) for k, v in batch["mask"].items()}
        
        optimizer.zero_grad()
        outputs = model(smiles, tabular)
        loss, losses = compute_multitask_loss_balanced(
            outputs, targets, mask, task_weights, class_weights, model=model,
        )
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        total_loss += loss.item()
        for k, v in losses.items():
            task_losses[k] += v.item()
        n_batches += 1
    
    return {
        "loss": total_loss / n_batches,
        **{f"loss_{k}": v / n_batches for k, v in task_losses.items() if v > 0},
    }


@torch.no_grad()
def validate_balanced(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    task_weights: Optional[LossWeightsBalanced] = None,
    class_weights: Optional[ClassWeights] = None,
) -> Dict[str, float]:
    """带类权重的验证"""
    model.eval()
    total_loss = 0.0
    task_losses = {k: 0.0 for k in ["size", "pdi", "ee", "delivery", "biodist", "toxic", "moe_lb"]}
    n_batches = 0
    
    # 用于计算准确率
    correct = {k: 0 for k in ["pdi", "ee", "toxic"]}
    total = {k: 0 for k in ["pdi", "ee", "toxic"]}
    
    for batch in tqdm(loader, desc="Validating", leave=False):
        smiles = batch["smiles"]
        tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
        targets = {k: v.to(device) for k, v in batch["targets"].items()}
        mask = {k: v.to(device) for k, v in batch["mask"].items()}
        
        outputs = model(smiles, tabular)
        loss, losses = compute_multitask_loss_balanced(
            outputs, targets, mask, task_weights, class_weights, model=model,
        )
        
        total_loss += loss.item()
        for k, v in losses.items():
            task_losses[k] += v.item()
        n_batches += 1
        
        # 计算分类准确率
        for k in ["pdi", "ee", "toxic"]:
            if k in targets and mask[k].any():
                m = mask[k]
                pred = outputs[k][m].argmax(dim=-1)
                tgt = targets[k][m]
                correct[k] += (pred == tgt).sum().item()
                total[k] += m.sum().item()
    
    metrics = {
        "loss": total_loss / n_batches,
        **{f"loss_{k}": v / n_batches for k, v in task_losses.items() if v > 0},
    }
    
    # 添加准确率
    for k in ["pdi", "ee", "toxic"]:
        if total[k] > 0:
            metrics[f"acc_{k}"] = correct[k] / total[k]
    
    return metrics


BACKBONE_PREFIXES = (
    "token_projector.",
    "set_transformer.",
    "fusion.",
    "moe.",
    "llm_prompt.",
)

FROM_SCRATCH_PREFIXES = (
    "moe.",
    "llm_prompt.",
    "fusion.g_moe",
    "fusion.g_llm",
)


def build_optimizer(
    model: nn.Module,
    lr: float,
    weight_decay: float,
    backbone_lr_ratio: float = 1.0,
    size_wd_mult: float = 5.0,
) -> torch.optim.AdamW:
    """
    构建 AdamW 优化器，支持分层学习率。

    仅收集 requires_grad=True 的参数（冻结的 MolT5 encoder 被排除）。
    当 backbone_lr_ratio < 1.0 时，backbone 参数使用 lr * backbone_lr_ratio，
    其余参数（task heads 等）使用 lr。backbone_lr_ratio = 1.0 等价于统一学习率。
    size_head 信号弱、易过拟合 -> 单独施加 size_wd_mult 倍 weight_decay。
    """
    trainable = [(n, p) for n, p in model.named_parameters() if p.requires_grad]
    size_wd = weight_decay * size_wd_mult

    def is_size(name: str) -> bool:
        return "size_head" in name

    size_params = [p for n, p in trainable if is_size(n)]

    if backbone_lr_ratio >= 1.0:
        other = [p for n, p in trainable if not is_size(n)]
        return torch.optim.AdamW(
            [
                {"params": other, "weight_decay": weight_decay},
                {"params": size_params, "weight_decay": size_wd},
            ],
            lr=lr,
        )

    backbone_params, head_params = [], []
    for name, param in trainable:
        if is_size(name):
            continue
        if name.startswith(BACKBONE_PREFIXES) and not name.startswith(FROM_SCRATCH_PREFIXES):
            backbone_params.append(param)
        else:
            head_params.append(param)

    return torch.optim.AdamW(
        [
            {"params": backbone_params, "lr": lr * backbone_lr_ratio, "weight_decay": weight_decay},
            {"params": head_params, "lr": lr, "weight_decay": weight_decay},
            {"params": size_params, "lr": lr, "weight_decay": size_wd},
        ],
    )


class EarlyStoppingBalanced:
    """早停机制（与 trainer.py 的 EarlyStopping 兼容）"""
    
    def __init__(self, patience: int = 10, min_delta: float = 0.0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = float("inf")
        self.best_epoch = 0
        self.should_stop = False
    
    def __call__(self, val_loss: float, epoch: int = 0) -> bool:
        if val_loss < self.best_loss - self.min_delta:
            self.best_loss = val_loss
            self.best_epoch = epoch
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.should_stop = True
        return self.should_stop
    
    def get_best_epoch(self) -> int:
        """获取最佳 epoch（1-indexed）"""
        return self.best_epoch + 1


def train_with_early_stopping(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    device: torch.device,
    lr: float = 1e-4,
    weight_decay: float = 1e-5,
    epochs: int = 100,
    patience: int = 15,
    task_weights: Optional[LossWeightsBalanced] = None,
    class_weights: Optional[ClassWeights] = None,
    backbone_lr_ratio: float = 1.0,
) -> Dict:
    """
    带早停的完整训练流程。
    
    Returns:
        Dict with keys: history, best_val_loss, best_epoch, best_state
    """
    model = model.to(device)
    optimizer = build_optimizer(model, lr, weight_decay, backbone_lr_ratio)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=5
    )
    early_stopping = EarlyStoppingBalanced(patience=patience)
    
    history = {"train": [], "val": []}
    best_val_loss = float("inf")
    best_state = None
    
    for epoch in range(epochs):
        # Train
        train_metrics = train_epoch_balanced(
            model, train_loader, optimizer, device, task_weights, class_weights
        )
        
        # Validate
        val_metrics = validate_balanced(
            model, val_loader, device, task_weights, class_weights
        )
        
        history["train"].append(train_metrics)
        history["val"].append(val_metrics)
        
        # Learning rate scheduling
        scheduler.step(val_metrics["loss"])
        
        # Save best model
        if val_metrics["loss"] < best_val_loss:
            best_val_loss = val_metrics["loss"]
            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
        
        # Early stopping
        if early_stopping(val_metrics["loss"], epoch):
            break
    
    # Restore best model
    if best_state is not None:
        model.load_state_dict(best_state)
    
    return {
        "history": history,
        "best_val_loss": best_val_loss,
        "best_epoch": early_stopping.get_best_epoch(),
        "best_state": best_state,
        "epochs_trained": len(history["train"]),
    }


def train_fixed_epochs(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: Optional[DataLoader],
    device: torch.device,
    lr: float = 1e-4,
    weight_decay: float = 1e-5,
    epochs: int = 50,
    task_weights: Optional[LossWeightsBalanced] = None,
    class_weights: Optional[ClassWeights] = None,
    use_cosine_annealing: bool = True,
    use_swa: bool = False,
    swa_start_epoch: Optional[int] = None,
    backbone_lr_ratio: float = 1.0,
    freeze_backbone_epochs: int = 0,
) -> Dict:
    """
    固定 epoch 数的训练（不使用 early stopping）。
    
    用于外层 CV 训练和最终训练。
    
    Args:
        model: 模型
        train_loader: 训练数据
        val_loader: 验证数据（可选，仅用于监控）
        device: 设备
        lr: 学习率
        weight_decay: 权重衰减
        epochs: 训练轮数
        task_weights: 任务权重
        class_weights: 类权重
        use_cosine_annealing: 是否使用 CosineAnnealingLR
        use_swa: 是否使用 SWA
        swa_start_epoch: SWA 开始的 epoch（默认为 epochs * 0.75）
        backbone_lr_ratio: backbone 学习率相对于 head 的比例（1.0 = 统一学习率）
    
    Returns:
        Dict with keys: history, final_state
    """
    model = model.to(device)
    optimizer = build_optimizer(model, lr, weight_decay, backbone_lr_ratio)

    backbone_named = [
        (n, p) for n, p in model.named_parameters()
        if n.startswith(BACKBONE_PREFIXES) and not n.startswith(FROM_SCRATCH_PREFIXES)
        and p.requires_grad
    ]
    if freeze_backbone_epochs > 0:
        for _, p in backbone_named:
            p.requires_grad_(False)

    if use_cosine_annealing:
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
    else:
        scheduler = None
    
    # SWA setup
    swa_model = None
    swa_scheduler = None
    if use_swa:
        from torch.optim.swa_utils import AveragedModel, SWALR
        swa_model = AveragedModel(model)
        swa_start = swa_start_epoch or int(epochs * 0.75)
        swa_scheduler = SWALR(optimizer, swa_lr=lr * 0.1)
    
    history = {"train": [], "val": []}
    
    for epoch in range(epochs):
        if freeze_backbone_epochs > 0 and epoch == freeze_backbone_epochs:
            for _, p in backbone_named:
                p.requires_grad_(True)
        # Train
        train_metrics = train_epoch_balanced(
            model, train_loader, optimizer, device, task_weights, class_weights
        )
        history["train"].append(train_metrics)
        
        # Validate (optional)
        if val_loader is not None:
            val_metrics = validate_balanced(
                model, val_loader, device, task_weights, class_weights
            )
            history["val"].append(val_metrics)
        
        # Scheduler step
        if use_swa and epoch >= swa_start:
            swa_model.update_parameters(model)
            swa_scheduler.step()
        elif scheduler is not None:
            scheduler.step()
    
    # Finalize SWA
    final_state = None
    if use_swa and swa_model is not None:
        # Update batch normalization statistics
        torch.optim.swa_utils.update_bn(train_loader, swa_model, device=device)
        final_state = {k: v.cpu().clone() for k, v in swa_model.module.state_dict().items()}
    else:
        final_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
    
    return {
        "history": history,
        "final_state": final_state,
        "epochs_trained": epochs,
    }

