"""Training pipeline for pump detection using real microstructure history."""

import logging
import time

import numpy as np

from .data_loader import load_pump_training_frame
from .pipeline import TrainingPipeline

logger = logging.getLogger(__name__)


class PumpTrainer(TrainingPipeline):
    """Train a classifier over short-horizon pump/dump microstructure features."""

    FEATURE_KEYS = ["return_1", "volume_ratio", "cusum_up", "cusum_down", "volatility_5"]

    def __init__(self) -> None:
        super().__init__("pump_detector")

    def load_data(self):
        logger.info("Loading real pump/dump training data from 1m OHLCV history...")
        frame = load_pump_training_frame(days=30, timeframe="1m")
        if frame.empty:
            raise RuntimeError("No historical pump training data available")
        return frame

    def preprocess(self, data):
        X = data[self.FEATURE_KEYS].fillna(0.0).astype(np.float32).values
        y = data["y"].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,
                objective="multi:softprob",
                num_class=3,
                eval_metric="mlogloss",
                early_stopping_rounds=20,
                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 GradientBoostingClassifier

            model = GradientBoostingClassifier(n_estimators=180, max_depth=4, random_state=42)
            model.fit(X_train, y_train)
            return {"model": model, "engine": "gradient_boosting"}

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

        model = trained["model"]
        preds = model.predict(X_test)
        return {
            "accuracy": float(accuracy_score(y_test, preds)),
            "precision": float(
                precision_score(y_test, preds, average="weighted", zero_division=0)
            ),
            "recall": float(recall_score(y_test, preds, average="weighted", zero_division=0)),
            "f1": float(f1_score(y_test, preds, average="weighted", zero_division=0)),
            "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 / "pump_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)


def train_pump_detector() -> dict:
    trainer = PumpTrainer()
    return trainer.run()


if __name__ == "__main__":
    train_pump_detector()
