"""TabPFN 模型实现

基于 TabPFN (Prior-Data Fitted Network) 的回归模型实现。
TabPFN 利用预训练的 Transformer 网络，通过上下文学习(in-context learning)
进行快速的小样本/中样本回归预测，无需传统训练过程。
"""

import json
import os
from pathlib import Path
from typing import Any, Optional

import numpy as np
import pandas as pd
import polars as pl
from scipy.stats import spearmanr

from src.training.components.base import BaseModel
from src.training.registry import register_model

os.environ["HF_TOKEN"] = "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"


@register_model("tabpfn")
class TabPFNModel(BaseModel):
    """TabPFN 回归模型

    使用 TabPFN 库实现基于 Prior-Data Fitted Network 的回归预测。
    该模型通过上下文学习方式进行预测，无需传统梯度下降训练。
    支持 GPU 加速和自动上下文截断处理。

    Attributes:
        name: 模型名称 "tabpfn"
        params: TabPFN 参数字典
        model: TabPFNRegressor 实例
        feature_names_: 特征名称列表
        evals_result_: 训练评估结果
        best_score_: 最佳评估指标
    """

    name = "tabpfn"

    # TabPFN 官方限制（最大样本数），可通过 ignore_pretraining_limits=True 扩展
    MAX_CONTEXT_SIZE = 10000

    def __init__(self, params: Optional[dict] = None):
        """初始化 TabPFN 模型

        Args:
            params: TabPFN 参数字典，支持以下参数：
                - device: 计算设备，'cuda' 或 'cpu'（默认 'cpu'）
                - model_path: 本地模型权重文件路径（可选）
                - N_ensemble: 集成数量，用于降低预测方差（默认 1）
                - max_context_size: 最大上下文样本数（默认 50000）

        Examples:
            >>> model = TabPFNModel(params={
            ...     "device": "cuda",
            ...     "N_ensemble": 5,
            ... })
        """
        self.params = dict(params) if params is not None else {}
        self.model = None
        self.feature_names_: Optional[list] = None
        self.evals_result_: Optional[dict] = None
        self.best_score_: Optional[dict] = None

    def fit(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        eval_set: Optional[tuple] = None,
    ) -> "TabPFNModel":
        """训练/加载 TabPFN 模型

        TabPFN 采用上下文学习，"fit" 操作实际上是加载训练数据到模型上下文。
        如果训练数据超过上下文限制，会自动截取最近的数据。

        Args:
            X: 特征矩阵 (Polars DataFrame)
            y: 目标变量 (Polars Series)
            eval_set: 验证集元组 (X_val, y_val)，用于评估模型性能

        Returns:
            self (支持链式调用)

        Raises:
            ImportError: 未安装 tabpfn
            RuntimeError: 模型初始化或加载失败
        """
        from tabpfn import TabPFNRegressor

        self.feature_names_ = X.columns

        # 转换为 numpy 数组
        X_np = X.to_numpy()
        y_np = y.to_numpy()

        # 处理上下文大小限制
        max_context = self.params.get("max_context_size", self.MAX_CONTEXT_SIZE)
        if len(X_np) > max_context:
            print(
                f"[TabPFN] 训练数据 {len(X_np)} 超过上下文限制 {max_context}，截取最近数据"
            )
            X_np = X_np[-max_context:]
            y_np = y_np[-max_context:]

        # 初始化模型
        # TabPFNRegressor 需要设置 ignore_pretraining_limits=True 以支持超过 10,000 样本
        device = self.params.get("device", "cuda")
        ignore_limits = self.params.get("ignore_pretraining_limits", True)
        self.model = TabPFNRegressor(
            device=device,
            ignore_pretraining_limits=ignore_limits,
            n_estimators=1
        )

        # 加载上下文（TabPFN 的 "fit" 是加载上下文）
        print("[TabPFN] 加载训练数据到上下文...")
        self.model.fit(X_np, y_np)

        # 评估验证集
        if eval_set is not None:
            X_val, y_val = eval_set
            val_preds = self.predict(X_val)
            y_val_np = y_val.to_numpy()

            # 计算评估指标
            mse = np.mean((y_val_np - val_preds) ** 2)
            rank_ic, p_value = spearmanr(val_preds, y_val_np)

            self.evals_result_ = {
                "valid_0": {
                    "mse": [mse],
                    "rank_ic": [rank_ic],
                }
            }
            self.best_score_ = {
                "valid_0": {
                    "mse": mse,
                    "rank_ic": rank_ic,
                    "rank_ic_pvalue": p_value,
                }
            }

            print(f"[TabPFN] 验证集 MSE: {mse:.6f}, Rank IC: {rank_ic:.4f}")

        return self

    def predict(self, X: pl.DataFrame) -> np.ndarray:
        """预测

        Args:
            X: 特征矩阵 (Polars DataFrame)

        Returns:
            预测结果 (numpy ndarray)

        Raises:
            RuntimeError: 模型未初始化时调用
        """
        if self.model is None:
            raise RuntimeError("模型尚未初始化，请先调用 fit()")

        X_np = X.to_numpy()
        result = self.model.predict(X_np)
        return np.asarray(result)

    def predict_with_uncertainty(
        self, X: pl.DataFrame
    ) -> tuple[np.ndarray, np.ndarray]:
        """预测并返回不确定性估计

        利用 N_ensemble 预测的标准差作为不确定性估计。

        Args:
            X: 特征矩阵 (Polars DataFrame)

        Returns:
            (predictions, uncertainties) 元组，均为 numpy ndarray
        """
        if self.model is None:
            raise RuntimeError("模型尚未初始化，请先调用 fit()")

        X_np = X.to_numpy()
        predictions = self.model.predict(X_np)

        # 如果使用了 ensemble，可以通过多次预测计算标准差
        # 注意：这需要修改 TabPFNRegressor 的使用方式
        # 这里返回预测值的零不确定性作为默认行为
        uncertainties = np.zeros_like(predictions)

        return np.asarray(predictions), np.asarray(uncertainties)

    def get_evals_result(self) -> Optional[dict]:
        """获取训练评估结果

        Returns:
            评估结果字典，如果未进行评估返回 None
        """
        return self.evals_result_

    def get_best_score(self) -> Optional[dict]:
        """获取最佳评分

        Returns:
            最佳评分字典，如果未进行评估返回 None
        """
        return self.best_score_

    def evaluate(self, X: pl.DataFrame, y: pl.Series) -> dict[str, float]:
        """评估模型性能

        计算回归任务常用指标：MSE 和 Rank IC。

        Args:
            X: 特征矩阵
            y: 真实目标值

        Returns:
            评估指标字典，包含 mse 和 rank_ic
        """
        preds = self.predict(X)
        y_np = y.to_numpy()

        mse = float(np.mean((y_np - preds) ** 2))
        rank_ic_result = spearmanr(preds, y_np)
        rank_ic = float(rank_ic_result.correlation)
        p_value = float(rank_ic_result.pvalue)

        return {
            "mse": mse,
            "rank_ic": rank_ic,
            "rank_ic_pvalue": p_value,
        }

    def feature_importance(self) -> None:
        """TabPFN 不支持传统特征重要性

        TabPFN 是基于 Transformer 的上下文学习模型，
        不提供类似决策树的特征重要性指标。

        Returns:
            None
        """
        return None

    def save(self, path: str) -> None:
        """保存模型元数据和配置

        TabPFN 模型本身不支持序列化保存，因此只保存：
        - 模型参数配置
        - 特征名称列表
        - 上下文数据摘要（样本数、特征数）

        注意：实际使用时需要重新 fit 来加载上下文。

        Args:
            path: 保存路径
        """
        save_data = {
            "model_type": self.name,
            "params": self.params,
            "feature_names": self.feature_names_,
            "evals_result": self.evals_result_,
            "best_score": self.best_score_,
        }

        # 保存为 JSON
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(save_data, f, indent=2, ensure_ascii=False)

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

        注意：TabPFN 模型需要重新 fit 才能使用，
        此方法仅恢复模型参数配置。

        Args:
            path: 配置文件路径

        Returns:
            配置恢复的 TabPFNModel 实例（未 fit）
        """
        with open(path, "r", encoding="utf-8") as f:
            save_data = json.load(f)

        instance = cls(params=save_data.get("params", {}))
        instance.feature_names_ = save_data.get("feature_names")
        instance.evals_result_ = save_data.get("evals_result")
        instance.best_score_ = save_data.get("best_score")

        print(f"[TabPFN] 已加载模型配置，需要调用 fit() 重新加载上下文")
        return instance
