"""Regime detector training on real historical market-state features."""

import logging
import time

import numpy as np

from .data_loader import load_regime_training_frame
from .pipeline import TrainingPipeline

logger = logging.getLogger(__name__)

REGIMES = ["trending_bull", "trending_bear", "ranging", "volatile", "capitulation"]


class RegimeTrainer(TrainingPipeline):
    FEATURE_KEYS = [
        "returns_1d",
        "returns_7d",
        "volatility_14d",
        "volume_ratio",
        "rsi",
        "bb_width",
        "fear_greed",
        "funding_rate",
        "price_vs_sma200",
    ]

    def __init__(self):
        super().__init__("regime_detector")

    def load_data(self):
        logger.info("Loading real regime training data from OHLCV history...")
        frame = load_regime_training_frame(days=300, timeframe="4h")
        if frame.empty:
            raise RuntimeError("No historical regime training data available")
        return frame

    def preprocess(self, data):
        data = data.sort_index().reset_index(drop=True)
        X = data[self.FEATURE_KEYS].fillna(0.0).astype(np.float32).values
        y = data["y"].map({name: idx for idx, name in enumerate(REGIMES)}).astype(np.int64).values

        n = len(X)
        train_end = int(n * 0.70)
        val_end = int(n * 0.85)
        return X[:train_end], X[train_end:val_end], X[val_end:], y[:train_end], y[train_end:val_end], y[val_end:]

    def train(self, X_train, y_train, X_val, y_val):
        try:
            import xgboost as xgb

            model = xgb.XGBClassifier(
                n_estimators=250,
                max_depth=4,
                learning_rate=0.05,
                subsample=0.85,
                colsample_bytree=0.85,
                min_child_weight=3,
                objective="multi:softprob",
                num_class=len(REGIMES),
                eval_metric="mlogloss",
                early_stopping_rounds=25,
                random_state=42,
            )
            model.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)
            return {"model": model, "engine": "xgboost"}
        except Exception:
            from sklearn.ensemble import RandomForestClassifier

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

    def evaluate(self, trained, X_test, y_test):
        from sklearn.metrics import accuracy_score, classification_report

        model = trained["model"]
        preds = model.predict(X_test)
        report = classification_report(
            y_test,
            preds,
            target_names=REGIMES,
            output_dict=True,
            zero_division=0,
        )
        return {
            "accuracy": float(accuracy_score(y_test, preds)),
            "per_class": {k: v for k, v in report.items() if k in REGIMES},
            "engine": trained["engine"],
            "test_samples": len(y_test),
        }

    def save(self, trained, metrics):
        import joblib

        super().save(trained, metrics)
        artifact_path = self.artifact_dir.parent / "regime_detector.joblib"
        joblib.dump(
            {
                "model": trained["model"],
                "engine": trained["engine"],
                "trained_at": str(int(time.time())),
                "accuracy": metrics.get("accuracy"),
                "feature_keys": self.FEATURE_KEYS,
            },
            artifact_path,
        )
        return str(artifact_path)
