"""TabM模型实现

TabM (Tabular Multilayer Perceptron with Ensembles)
基于 rtdl_revisiting_models 的 TabM 模型，支持内置集成。
"""

from typing import Dict, Any, List, Optional
from pathlib import Path
import pickle

import numpy as np
import polars as pl
import scipy.stats as stats
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from tabm import TabM

from src.training.components.base import BaseModel
from src.training.components.models.cross_section_sampler import CrossSectionSampler
from src.training.components.models.ensemble_quant_loss import (
    EnsembleQuantLoss,
    AsymmetricQuantLoss,
)
from src.training.registry import register_model


@register_model("tabm")
class TabMModel(BaseModel):
    """TabM回归模型

    特点:
    - 使用MLP架构
    - 内置集成机制(ensemble_size)，显存开销远小于独立模型
    - 训练时所有集成成员独立优化，保持多样性
    - 预测时取集成成员均值获得稳定结果

    Attributes:
        name: 模型名称标识
        params: 模型参数字典
        model: TabM模型实例
        device: 计算设备(cuda/cpu)
        training_history_: 训练历史记录
        feature_names_: 特征名称列表
    """

    name = "tabm"

    def __init__(self, params: Dict[str, Any]):
        """初始化TabM模型

        Args:
            params: 模型参数字典，包含:
                - n_blocks: MLP层数 (默认: 3)
                - d_block: 每层神经元数 (默认: 256)
                - dropout: Dropout率 (默认: 0.1)
                - ensemble_size: 集成大小 (默认: 32)
                - batch_size: 批次大小 (默认: 1024)
                - learning_rate: 学习率 (默认: 1e-3)
                - weight_decay: 权重衰减 (默认: 1e-5)
                - epochs: 训练轮数 (默认: 50)
                - early_stopping_patience: 早停耐心值 (默认: 10)
        """
        self.params = params
        self.model = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.training_history_: Dict[str, List[float]] = {
            "train_loss": [],
            "val_loss": [],
        }
        self.feature_names_: Optional[List[str]] = None

        # 损失函数
        self.criterion = nn.MSELoss()

    def _make_loader(
        self,
        X: np.ndarray,
        y: Optional[np.ndarray] = None,
        dates: Optional[np.ndarray] = None,
        shuffle_days: bool = False,
    ) -> DataLoader:
        """创建DataLoader

        支持两种模式:
        1. 截面模式 (dates provided): 每天的数据作为一个独立批次
        2. 普通模式 (dates is None): 使用固定batch_size

        Args:
            X: 特征数组 [N, n_features]
            y: 标签数组 [N] 或 None
            dates: 日期数组 [N]，用于截面采样
            shuffle_days: 是否打乱日期顺序

        Returns:
            DataLoader实例
        """
        if y is not None:
            dataset = TensorDataset(torch.from_numpy(X), torch.from_numpy(y))
        else:
            dataset = TensorDataset(torch.from_numpy(X))

        if dates is not None:
            # 使用截面 Sampler
            sampler = CrossSectionSampler(dates, shuffle_days=shuffle_days)
            # 注意：使用 batch_sampler 时，不需要指定 batch_size 和 shuffle
            return DataLoader(dataset, batch_sampler=sampler)
        else:
            # 预测时如果没有 date，退化为普通 Loader
            batch_size = self.params.get("batch_size", 1024)
            return DataLoader(dataset, batch_size=batch_size, shuffle=False)

    def _validate(self, val_loader: DataLoader) -> float:
        """验证模型（使用 Rank IC 作为指标）

        Args:
            val_loader: 验证数据加载器

        Returns:
            平均 Rank IC（斯皮尔曼秩相关系数）
        """
        self.model.eval()
        ic_list = []

        with torch.no_grad():
            for batch in val_loader:
                if len(batch) != 2:
                    continue

                bx, by = batch
                bx = bx.to(self.device)
                by = by.cpu().numpy()

                # 如果截面股票数太少，跳过
                if len(by) < 10:
                    continue

                outputs = self.model(bx)  # [B, E, 1]
                # 预测时取均值
                preds = outputs.mean(dim=1).squeeze(-1).cpu().numpy()

                # 计算斯皮尔曼秩相关系数 (Rank IC)
                rank_ic, _ = stats.spearmanr(preds, by)

                # 如果算不出 rank_ic，记为 0
                if np.isnan(rank_ic):
                    rank_ic = 0.0

                ic_list.append(rank_ic)

        # 返回验证集的平均 Rank IC
        return np.mean(ic_list) if len(ic_list) > 0 else 0.0

    def fit(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        eval_set: Optional[tuple] = None,
        date_col: Optional[str] = None,
    ) -> "TabMModel":
        """训练TabM模型（量化优化版本）

        训练策略:
        1. 支持截面数据采样（按交易日批处理）
        2. 使用 EnsembleQuantLoss (Huber + IC) 替代 MSE
        3. CosineAnnealingLR 学习率调度
        4. 梯度裁剪防止异常数据影响
        5. 以 Rank IC 作为验证指标和早停依据

        Args:
            X: 训练特征DataFrame，可包含日期列
            y: 训练标签Series
            eval_set: 验证集元组 (X_val, y_val)，可选
            date_col: 日期列名称，用于截面采样。如果为None，使用普通批次

        Returns:
            self
        """
        # 【修复】初始化 val_ic 训练历史
        if "val_ic" not in self.training_history_:
            self.training_history_["val_ic"] = []

        # 保存特征名称并处理日期列
        self.feature_names_ = list(X.columns)
        dates = None
        if date_col and date_col in X.columns:
            dates = X[date_col].to_numpy()
            X = X.drop(date_col)
            self.feature_names_ = [c for c in self.feature_names_ if c != date_col]
            self.params["date_col"] = date_col  # 保存供 predict 使用

        # 【关键】数据类型强制转换为float32
        X_np = X.to_numpy().astype(np.float32)
        y_np = y.to_numpy().astype(np.float32)

        # 创建DataLoader（支持截面采样）
        train_loader = self._make_loader(X_np, y_np, dates=dates, shuffle_days=True)

        val_loader = None
        val_dates = None
        if eval_set is not None:
            X_val, y_val = eval_set
            # 处理验证集的日期列
            if date_col and date_col in X_val.columns:
                val_dates = X_val[date_col].to_numpy()
                X_val = X_val.drop(date_col)
            X_val_np = X_val.to_numpy().astype(np.float32)
            y_val_np = y_val.to_numpy().astype(np.float32)
            val_loader = self._make_loader(
                X_val_np, y_val_np, dates=val_dates, shuffle_days=False
            )

        n_features = X_np.shape[1]
        ensemble_size = self.params.get("ensemble_size", 32)

        # 初始化TabM模型
        self.model = TabM.make(
            n_num_features=n_features,
            cat_cardinalities=[],
            d_out=1,
            n_blocks=self.params.get("n_blocks", 3),
            d_block=self.params.get("d_block", 256),
            dropout=self.params.get("dropout", 0.1),
            k=ensemble_size,
        ).to(self.device)

        # 优化器（增加 weight_decay 用于正则化）
        optimizer = optim.AdamW(
            self.model.parameters(),
            lr=self.params.get("learning_rate", 1e-3),
            weight_decay=self.params.get("weight_decay", 1e-4),
        )

        # 学习率调度器：CosineAnnealingLR 平滑衰减
        epochs = self.params.get("epochs", 50)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=epochs, eta_min=1e-6
        )

        # 使用 AsymmetricQuantLoss (非对称 Huber + IC)
        self.criterion = AsymmetricQuantLoss(
            alpha=self.params.get("loss_alpha", 0.5), ensemble_size=ensemble_size
        )

        # 训练参数
        early_stopping_patience = self.params.get("early_stopping_patience", 10)

        # 训练循环
        best_val_ic = -float("inf")  # Rank IC 越大越好
        patience_counter = 0
        best_model_state = None

        print(f"[TabM] 开始训练... 设备: {self.device}, 集成大小: {ensemble_size}")
        if date_col:
            print(f"[TabM] 使用截面采样，日期列: {date_col}")

        for epoch in range(epochs):
            # 训练阶段
            self.model.train()
            train_loss = 0.0
            n_train_batches = 0

            for batch in train_loader:
                if len(batch) == 2:
                    bx, by = batch
                    bx, by = bx.to(self.device), by.to(self.device)
                else:
                    continue

                optimizer.zero_grad()

                # 前向传播
                outputs = self.model(bx)  # [B, E, 1]
                outputs_squeezed = outputs.squeeze(-1)  # [B, E]

                # 使用 EnsembleQuantLoss
                loss = self.criterion(outputs_squeezed, by)

                loss.backward()

                # 【关键】梯度裁剪，防止某天的异常收益率拉崩模型
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(),
                    max_norm=self.params.get("max_grad_norm", 1.0),
                )

                optimizer.step()

                train_loss += loss.item()
                n_train_batches += 1

            avg_train_loss = train_loss / max(n_train_batches, 1)
            self.training_history_["train_loss"].append(avg_train_loss)

            # 验证阶段（使用 Rank IC 作为指标）
            if val_loader is not None:
                val_ic = self._validate(val_loader)
                self.training_history_["val_ic"].append(val_ic)

                # 早停逻辑：Rank IC 越大越好
                if val_ic > best_val_ic:
                    best_val_ic = val_ic
                    patience_counter = 0
                    # 保存最佳模型权重
                    best_model_state = {
                        k: v.cpu().clone() for k, v in self.model.state_dict().items()
                    }
                else:
                    patience_counter += 1

                if (epoch + 1) % 5 == 0 or epoch == 0:
                    print(
                        f"[TabM] Epoch {epoch + 1}/{epochs} | "
                        f"Train Loss: {avg_train_loss:.6f} | "
                        f"Val Rank IC: {val_ic:.4f} (Best: {best_val_ic:.4f})"
                    )

                if patience_counter >= early_stopping_patience:
                    print(f"[TabM] 早停触发，epoch {epoch + 1}")
                    break
            else:
                if (epoch + 1) % 5 == 0 or epoch == 0:
                    print(
                        f"[TabM] Epoch {epoch + 1}/{epochs} | "
                        f"Train Loss: {avg_train_loss:.6f}"
                    )

            # 更新学习率
            scheduler.step()

        # 恢复最佳模型
        if best_model_state is not None:
            self.model.load_state_dict(best_model_state)
            print(f"[TabM] 恢复最佳模型 (Val Rank IC: {best_val_ic:.4f})")

        print(f"[TabM] 训练完成")
        return self

    def predict(self, X: pl.DataFrame) -> np.ndarray:
        """生成预测（自动处理 date_col 和特征对齐）

        预测时对ensemble_size个成员取均值，获得稳定结果。

        Args:
            X: 特征DataFrame

        Returns:
            预测结果数组 [N]
        """
        if self.model is None:
            raise RuntimeError("模型未训练，请先调用fit()")

        # 自动移除 date_col
        date_col = self.params.get("date_col")
        if date_col and date_col in X.columns:
            X = X.drop(date_col)

        # 【关键修复】特征对齐：检查缺失列并显式抛出异常
        if self.feature_names_:
            missing_cols = [c for c in self.feature_names_ if c not in X.columns]
            if missing_cols:
                raise ValueError(f"预测数据缺失训练特征: {missing_cols}")
            X = X.select(self.feature_names_)

        # 数据类型转换
        X_np = X.to_numpy().astype(np.float32)
        # 【修复】使用正确的参数名 shuffle_days（而非 shuffle）
        loader = self._make_loader(X_np, shuffle_days=False)

        self.model.eval()
        all_preds = []

        with torch.no_grad():
            for batch in loader:
                bx = batch[0].to(self.device)
                # 预测时取集成成员均值
                outputs = self.model(bx)  # [B, E, 1]
                preds = outputs.mean(dim=1).squeeze(-1)  # [B]
                all_preds.append(preds.cpu().numpy())

        return np.concatenate(all_preds)

    def feature_importance(self) -> Optional[pl.Series]:
        """获取特征重要性

        TabM没有内置特征重要性计算，返回None。

        Returns:
            None
        """
        return None

    def save(self, path: str | Path) -> None:
        """保存模型

        保存模型state_dict和元数据（params, feature_names, training_history）

        Args:
            path: 保存路径
        """
        if self.model is None:
            raise RuntimeError("模型未训练，无法保存")

        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        # 保存模型权重
        model_path = path.with_suffix(".pt")
        torch.save(self.model.state_dict(), model_path)

        # 保存元数据
        meta_path = path.with_suffix(".meta")
        meta = {
            "params": self.params,
            "feature_names": self.feature_names_,
            "training_history": self.training_history_,
            "device": str(self.device),
        }
        with open(meta_path, "wb") as f:
            pickle.dump(meta, f)

        print(f"[TabM] 模型保存到: {path}")

    @classmethod
    def load(cls, path: str) -> "TabMModel":
        """加载模型

        Args:
            path: 模型路径（不含扩展名）

        Returns:
            加载的TabMModel实例
        """
        path = Path(path)

        # 加载元数据
        meta_path = path.with_suffix(".meta")
        with open(meta_path, "rb") as f:
            meta = pickle.load(f)

        # 创建实例
        instance = cls(meta["params"])
        instance.feature_names_ = meta["feature_names"]
        instance.training_history_ = meta["training_history"]

        # 重建模型结构
        if instance.feature_names_ is not None:
            n_features = len(instance.feature_names_)
            ensemble_size = instance.params.get("ensemble_size", 32)

            instance.model = TabM.make(
                n_num_features=n_features,
                cat_cardinalities=[],
                d_out=1,  # 回归任务输出维度为1
                n_blocks=instance.params.get("n_blocks", 3),
                d_block=instance.params.get("d_block", 256),
                dropout=instance.params.get("dropout", 0.1),
                k=ensemble_size,  # 集成大小
            ).to(instance.device)

            # 加载权重
            model_path = path.with_suffix(".pt")
            instance.model.load_state_dict(
                torch.load(model_path, map_location=instance.device)
            )

        print(f"[TabM] 模型从 {path} 加载完成")
        return instance
