"""Train the preparation-time regression model.

Three models are compared on identical features and split:
    1. LinearRegression  (baseline)
    2. RandomForestRegressor
    3. XGBRegressor

The best model (by MAE) is persisted as a sklearn Pipeline so that the
preprocessing transformers travel together with the estimator.

Usage:
    python -m src.ml.train
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestRegressor
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from xgboost import XGBRegressor

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from src.config import ML_METRICS_PATH, ML_PIPELINE_PATH, PROCESSED_DIR  # noqa: E402
from src.ml.feature_engineering import (  # noqa: E402
    CATEGORICAL_FEATURES,
    FEATURES_CSV,
    NUMERIC_FEATURES,
    TARGET,
)


RANDOM_STATE = 42


def build_preprocessor(num_cols: list[str], cat_cols: list[str]) -> ColumnTransformer:
    numeric_pipe = Pipeline(
        [("impute", SimpleImputer(strategy="median")), ("scale", StandardScaler())]
    )
    categorical_pipe = Pipeline(
        [
            ("impute", SimpleImputer(strategy="most_frequent")),
            ("oh", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
        ]
    )
    return ColumnTransformer(
        [("num", numeric_pipe, num_cols), ("cat", categorical_pipe, cat_cols)],
        remainder="drop",
    )


def evaluate(name: str, y_true: pd.Series, y_pred: np.ndarray) -> dict:
    mae = mean_absolute_error(y_true, y_pred)
    rmse = float(np.sqrt(mean_squared_error(y_true, y_pred)))
    r2 = r2_score(y_true, y_pred)
    return {"model": name, "MAE": float(mae), "RMSE": rmse, "R2": float(r2)}


def main() -> None:
    if not FEATURES_CSV.exists():
        raise FileNotFoundError(
            f"{FEATURES_CSV} missing. Run 'python -m src.ml.feature_engineering' first."
        )

    df = pd.read_csv(FEATURES_CSV)
    print(f"[train] dataset shape: {df.shape}")

    num_cols = [c for c in NUMERIC_FEATURES if c in df.columns]
    cat_cols = [c for c in CATEGORICAL_FEATURES if c in df.columns]
    feature_cols = num_cols + cat_cols

    X = df[feature_cols]
    y = df[TARGET]

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=RANDOM_STATE
    )

    candidates = {
        "LinearRegression": LinearRegression(),
        "RandomForest": RandomForestRegressor(
            n_estimators=300,
            max_depth=18,
            min_samples_leaf=3,
            n_jobs=-1,
            random_state=RANDOM_STATE,
        ),
        "XGBoost": XGBRegressor(
            n_estimators=600,
            max_depth=7,
            learning_rate=0.05,
            subsample=0.9,
            colsample_bytree=0.9,
            random_state=RANDOM_STATE,
            n_jobs=-1,
            tree_method="hist",
        ),
    }

    preprocessor = build_preprocessor(num_cols, cat_cols)

    results: list[dict] = []
    fitted: dict[str, Pipeline] = {}

    for name, estimator in candidates.items():
        pipe = Pipeline([("prep", preprocessor), ("model", estimator)])
        pipe.fit(X_train, y_train)
        preds = pipe.predict(X_test)
        metrics = evaluate(name, y_test, preds)
        results.append(metrics)
        fitted[name] = pipe
        print(
            f"[train] {name:>17}  MAE={metrics['MAE']:.2f}  "
            f"RMSE={metrics['RMSE']:.2f}  R2={metrics['R2']:.3f}"
        )

    best = min(results, key=lambda r: r["MAE"])
    best_name = best["model"]
    print(f"[train] best model: {best_name} (MAE {best['MAE']:.2f})")

    ML_PIPELINE_PATH.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump(
        {
            "pipeline": fitted[best_name],
            "feature_cols": feature_cols,
            "num_cols": num_cols,
            "cat_cols": cat_cols,
            "target": TARGET,
            "best_model": best_name,
        },
        ML_PIPELINE_PATH,
    )
    print(f"[train] saved pipeline to {ML_PIPELINE_PATH}")

    ML_METRICS_PATH.write_text(
        json.dumps({"results": results, "best": best_name}, indent=2)
    )
    print(f"[train] wrote metrics to {ML_METRICS_PATH}")


if __name__ == "__main__":
    main()