"""XGBoost 基线模型 — 三个独立分类器用于多时间尺度对比"""
import numpy as np
import xgboost as xgb
from sklearn.metrics import accuracy_score, f1_score


def train_xgboost_baseline(X_train: np.ndarray, y_train: np.ndarray,
                           X_test: np.ndarray, y_test: np.ndarray) -> dict:
    """
    训练三个独立的 XGBoost 分类器 (短/中/长期)。

    Args:
        X_train: (N, T, D) 训练特征，自动展平为 (N, T*D)
        y_train: (N, 3) 标签矩阵，列顺序: short, medium, long
        X_test: (N_test, T, D) 测试特征
        y_test: (N_test, 3) 测试标签

    Returns:
        dict: {
            "short": {"model": ..., "accuracy": ..., "f1_macro": ..., "predictions": ...},
            "medium": {...},
            "long": {...},
        }
    """
    # 展平时序特征为二维
    N_train, T, D = X_train.shape
    X_train_flat = X_train.reshape(N_train, T * D)
    N_test = X_test.shape[0]
    X_test_flat = X_test.reshape(N_test, T * D)

    horizon_names = ["short", "medium", "long"]
    results = {}

    for i, name in enumerate(horizon_names):
        model = xgb.XGBClassifier(
            n_estimators=200,
            max_depth=6,
            learning_rate=0.05,
            subsample=0.8,
            colsample_bytree=0.8,
            objective="multi:softmax",
            num_class=4,
            eval_metric="mlogloss",
            random_state=42,
            device="cuda",
        )
        model.fit(
            X_train_flat, y_train[:, i],
            eval_set=[(X_test_flat, y_test[:, i])],
            verbose=False,
        )
        y_pred = model.predict(X_test_flat)
        acc = accuracy_score(y_test[:, i], y_pred)
        f1 = f1_score(y_test[:, i], y_pred, average="macro")

        results[name] = {
            "model": model,
            "accuracy": acc,
            "f1_macro": f1,
            "predictions": y_pred,
        }
        print(f"XGBoost {name}: Accuracy={acc:.4f}, F1 Macro={f1:.4f}")

    return results


def train_xgboost_single(X_train: np.ndarray, y_train: np.ndarray,
                         X_test: np.ndarray, y_test: np.ndarray,
                         horizon_idx: int = 0) -> dict:
    """训练单个时间尺度的XGBoost模型（用于单独调用）"""
    N_train, T, D = X_train.shape
    X_train_flat = X_train.reshape(N_train, T * D)
    N_test = X_test.shape[0]
    X_test_flat = X_test.reshape(N_test, T * D)

    model = xgb.XGBClassifier(
        n_estimators=200, max_depth=6, learning_rate=0.05,
        subsample=0.8, colsample_bytree=0.8,
        objective="multi:softmax", num_class=4,
        eval_metric="mlogloss", random_state=42,
        device="cuda",
    )
    model.fit(X_train_flat, y_train[:, horizon_idx], verbose=False)
    y_pred = model.predict(X_test_flat)

    return {
        "model": model,
        "accuracy": accuracy_score(y_test[:, horizon_idx], y_pred),
        "f1_macro": f1_score(y_test[:, horizon_idx], y_pred, average="macro"),
        "predictions": y_pred,
    }
