"""量化专用组合损失函数

结合 Huber Loss（处理极值）和 IC Loss（优化排序），
针对 TabM 的集成输出设计，为每个集成成员独立计算损失以保持多样性。
"""

import torch
import torch.nn as nn


class EnsembleQuantLoss(nn.Module):
    """Ensemble 量化组合损失函数

    组合 Huber Loss 和 IC Loss:
    - Huber Loss: 对收益率极值鲁棒，避免梯度爆炸
    - IC Loss: 直接优化预测与目标的秩相关性

    针对 TabM 的集成输出设计，为每个集成成员独立计算损失，
    保持集成多样性。

    注意：std() 使用 unbiased=False 以保持数学一致性
    (IC 计算使用除以 N，而非 N-1)

    Attributes:
        alpha: Huber Loss 权重，(1-alpha) 为 IC Loss 权重
        huber: HuberLoss 实例
        ensemble_size: 集成成员数量
    """

    def __init__(self, alpha: float = 0.5, ensemble_size: int = 32):
        """初始化损失函数

        Args:
            alpha: Huber Loss 权重，范围 [0, 1]
            ensemble_size: TabM 集成大小
        """
        super().__init__()
        self.alpha = alpha
        self.huber = nn.HuberLoss(reduction="mean")
        self.ensemble_size = ensemble_size

    def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """计算组合损失

        Args:
            preds: 预测值，形状 [Batch(单日股票数), Ensemble_size]
            target: 目标值，形状 [Batch(单日股票数)]

        Returns:
            标量损失值
        """
        total_loss = 0.0

        # 为了稳定，如果单日截面股票数太少，直接回退到 Huber
        if preds.shape[0] < 10:
            for i in range(self.ensemble_size):
                total_loss += self.huber(preds[:, i], target)
            return total_loss / self.ensemble_size

        # 预计算目标的标准化值（在所有集成成员间复用）
        target_mean = target.mean()
        # 【关键】使用 unbiased=False 保持数学一致性
        target_std = target.std(unbiased=False) + 1e-8
        target_norm = (target - target_mean) / target_std

        # 对每个 ensemble 成员独立计算组合 Loss
        for i in range(self.ensemble_size):
            pred_i = preds[:, i]

            # 1. Huber Loss (处理极值比 MSE 更好)
            h_loss = self.huber(pred_i, target)

            # 2. IC Loss (皮尔逊相关系数)
            pred_mean = pred_i.mean()
            # 【关键】使用 unbiased=False 保持数学一致性
            pred_std = pred_i.std(unbiased=False) + 1e-8
            pred_norm = (pred_i - pred_mean) / pred_std

            # 相关系数
            ic = (pred_norm * target_norm).mean()
            # 损失函数化：希望 IC 越大越好，所以用 1 - IC
            ic_loss = 1.0 - ic

            total_loss += self.alpha * h_loss + (1.0 - self.alpha) * ic_loss

        return total_loss / self.ensemble_size


class AsymmetricQuantLoss(nn.Module):
    """Asymmetric Quant Loss (非对称 Huber + IC)

    保留全截面计算以维持稳定梯度的前提下，对多头关注的错误进行加权惩罚：
    1. 过度高估烂股票 (买入陷阱) -> 加重惩罚
    2. 严重低估好股票 (错失金股) -> 加重惩罚
    """

    def __init__(self, alpha: float = 0.5, ensemble_size: int = 32):
        super().__init__()
        self.alpha = alpha
        self.ensemble_size = ensemble_size
        # 我们不使用内置的 HuberLoss，而是手动写以支持 element-wise 加权
        self.delta = 1.0  # Huber threshold (可以设得比较小，比如对于收益率设为 0.05)

    def forward(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        batch_size = preds.shape[0]
        total_loss = 0.0

        if batch_size < 10:
            # 极小 batch 退回简单 MSE
            return ((preds - target.unsqueeze(1)) ** 2).mean()

        target_mean = target.mean()
        target_std = target.std(unbiased=False) + 1e-8
        target_norm = (target - target_mean) / target_std

        for i in range(self.ensemble_size):
            pred_i = preds[:, i]

            # --- 1. 非对称 Huber 损失 ---
            error = pred_i - target
            abs_error = torch.abs(error)

            # 标准 Huber 计算
            quadratic = torch.clamp(abs_error, max=self.delta)
            linear = abs_error - quadratic
            base_huber = 0.5 * quadratic**2 + self.delta * linear

            # 【核心逻辑】：非对称权重
            # 如果 target > 0 且 pred < target (错过了涨的好票) -> 权重 1.5
            # 如果 target < 0 且 pred > 0 (把跌的票预测成了涨的，容易实盘买入踩雷) -> 权重 2.0
            # 其他情况 (比如烂票预测得更烂) -> 权重 1.0 正常学习
            weights = torch.ones_like(target)
            weights[(target > 0) & (error < 0)] = 1.5
            weights[(target < 0) & (pred_i > 0)] = 2.0

            weighted_huber = (base_huber * weights).mean()

            # --- 2. IC 损失 (维持全截面排序能力) ---
            pred_mean = pred_i.mean()
            pred_std = pred_i.std(unbiased=False) + 1e-8
            pred_norm = (pred_i - pred_mean) / pred_std

            ic = (pred_norm * target_norm).mean()
            ic_loss = 1.0 - ic

            total_loss += self.alpha * weighted_huber + (1.0 - self.alpha) * ic_loss

        return total_loss / self.ensemble_size
