"""Event-aware catalyst trainer for 1d+ direction forecasts."""

import time
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.metrics import accuracy_score, f1_score

from .data_loader import build_feature_frame, get_training_symbols

HIGH_IMPACT_EVENTS = [
    ("2025-12-17", "fomc"),
    ("2026-01-14", "cpi"),
    ("2026-01-28", "fomc"),
    ("2026-02-06", "nfp"),
    ("2026-02-11", "cpi"),
    ("2026-03-06", "nfp"),
    ("2026-03-18", "fomc"),
    ("2026-04-03", "nfp"),
    ("2026-05-06", "fomc"),
    ("2026-06-05", "nfp"),
]


def _event_features(timestamp) -> dict[str, float]:
    event_times = [
        (abs((timestamp - pd.Timestamp(date, tz="UTC")) / np.timedelta64(1, "D")), kind)
        for date, kind in HIGH_IMPACT_EVENTS
    ]
    if not event_times:
        return {
            "days_to_event": 30.0,
            "event_risk": 0.0,
            "within_24h": 0.0,
            "within_72h": 0.0,
            "within_7d": 0.0,
            "is_fomc": 0.0,
            "is_cpi": 0.0,
            "is_nfp": 0.0,
        }
    days_to_event, kind = sorted(event_times, key=lambda item: item[0])[0]
    return {
        "days_to_event": float(days_to_event),
        "event_risk": float(max(0.0, 1.0 - min(days_to_event, 14.0) / 14.0)),
        "within_24h": float(days_to_event <= 1.0),
        "within_72h": float(days_to_event <= 3.0),
        "within_7d": float(days_to_event <= 7.0),
        "is_fomc": float(kind == "fomc"),
        "is_cpi": float(kind == "cpi"),
        "is_nfp": float(kind == "nfp"),
    }


class CatalystTrainer:
    model_name = "catalyst_event"
    artifact_path = Path("models") / "catalyst_event.joblib"
    feature_keys = [
        "days_to_event",
        "event_risk",
        "within_24h",
        "within_72h",
        "within_7d",
        "is_fomc",
        "is_cpi",
        "is_nfp",
        "returns_1d",
        "returns_7d",
        "volatility_14d",
        "fear_greed",
        "funding_rate",
    ]

    def run(self, days: int = 365) -> dict:
        start = time.time()
        rows = []
        for symbol in get_training_symbols():
            frame = build_feature_frame(symbol, timeframe="4h", days=days)
            if frame.empty or len(frame) < 120:
                continue
            future = frame["future_return_1d"]
            labels = np.where(future > 1.0, 2, np.where(future < -1.0, 0, 1))
            event_rows = frame["time"].apply(_event_features)
            event_frame = pd.DataFrame(list(event_rows))
            merged = pd.concat(
                [
                    event_frame.reset_index(drop=True),
                    frame[["returns_1d", "returns_7d", "volatility_14d", "fear_greed", "funding_rate"]].reset_index(drop=True),
                ],
                axis=1,
            )
            merged["y"] = labels
            rows.append(merged)

        if not rows:
            return {
                "model": self.model_name,
                "status": "skipped",
                "metrics": {"samples": 0},
                "duration_seconds": round(time.time() - start, 2),
                "artifact_path": "",
                "error": "No catalyst training data available",
            }

        frame = pd.concat(rows, ignore_index=True).fillna(0.0)
        X = frame[self.feature_keys].astype(np.float32).values
        y = frame["y"].astype(np.int64).values
        n = len(X)
        split = int(n * 0.80)
        X_train, X_test = X[:split], X[split:]
        y_train, y_test = y[:split], y[split:]

        try:
            import xgboost as xgb

            model = xgb.XGBClassifier(
                n_estimators=220,
                max_depth=4,
                learning_rate=0.05,
                subsample=0.85,
                colsample_bytree=0.85,
                objective="multi:softprob",
                num_class=3,
                eval_metric="mlogloss",
                random_state=42,
            )
            model.fit(X_train, y_train)
            engine = "xgboost"
        except Exception:
            from sklearn.ensemble import RandomForestClassifier

            model = RandomForestClassifier(
                n_estimators=250,
                max_depth=10,
                class_weight="balanced_subsample",
                random_state=42,
                n_jobs=-1,
            )
            model.fit(X_train, y_train)
            engine = "random_forest"

        preds = model.predict(X_test)
        accuracy = float(accuracy_score(y_test, preds))
        f1 = float(f1_score(y_test, preds, average="weighted", zero_division=0))

        self.artifact_path.parent.mkdir(parents=True, exist_ok=True)
        joblib.dump(
            {
                "model": model,
                "engine": engine,
                "feature_keys": self.feature_keys,
                "trained_at": str(int(time.time())),
                "accuracy": accuracy,
            },
            self.artifact_path,
        )

        return {
            "model": self.model_name,
            "status": "success",
            "metrics": {
                "accuracy": accuracy,
                "f1": f1,
                "test_samples": len(X_test),
            },
            "duration_seconds": round(time.time() - start, 2),
            "artifact_path": str(self.artifact_path),
        }
