"""Train & compare salary-regression models (ML block) and write artifacts.

Run:  .venv\\Scripts\\python.exe src/train.py

What it does (satisfies the ML block requirements end-to-end):
1. Loads the cleaned salary frame and builds the shared feature contract.
2. Holds out a test split, then **compares ≥2 models** with 5-fold CV on the
   training split: a naive mean baseline, Ridge (linear), RandomForest and
   HistGradientBoosting. The target is log-transformed (salaries are skewed).
3. Picks the best model by CV-RMSE, refits it, evaluates on the held-out set,
   and runs a short **error analysis** (by job title and salary quintile).
4. Writes ``artifacts/``: the fitted pipeline + metadata + metrics — everything
   the inference app needs (training and inference are fully separated).
"""
from __future__ import annotations

import json
import sys
from datetime import date
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

# Windows consoles default to cp1252; force UTF-8 so unicode in logs never crashes.
try:
    sys.stdout.reconfigure(encoding="utf-8")
except Exception:  # noqa: BLE001
    pass

from src.data import load_salary_df  # noqa: E402
from src.features import (  # noqa: E402
    BOOL_COLS,
    CATEGORICAL_COLS,
    DEFAULTS,
    FEATURE_COLS,
    NUMERIC_COLS,
    SKILL_VOCAB,
    TARGET_COL,
    build_feature_frame,
)

import joblib  # noqa: E402
import sklearn  # noqa: E402
from sklearn.compose import ColumnTransformer, TransformedTargetRegressor  # noqa: E402
from sklearn.dummy import DummyRegressor  # noqa: E402
from sklearn.ensemble import HistGradientBoostingRegressor, RandomForestRegressor  # noqa: E402
from sklearn.linear_model import Ridge  # noqa: E402
from sklearn.metrics import mean_absolute_error, r2_score, root_mean_squared_error  # noqa: E402
from sklearn.model_selection import KFold, cross_validate, train_test_split  # noqa: E402
from sklearn.pipeline import Pipeline  # noqa: E402
from sklearn.preprocessing import OneHotEncoder, StandardScaler  # noqa: E402

ARTIFACTS = ROOT / "artifacts"
RANDOM_STATE = 42
N_FOLDS = 5


# --------------------------------------------------------------------------- #
# Pipeline building blocks
# --------------------------------------------------------------------------- #
def make_preprocessor(scale_numeric: bool) -> ColumnTransformer:
    """One-hot the categoricals; pass through (or scale) the numeric/boolean cols.

    ``infrequent_if_exist`` buckets rare/unseen categories, which keeps the model
    robust to whatever the NLP block hands us at inference time.
    """
    cat = OneHotEncoder(handle_unknown="infrequent_if_exist", min_frequency=25, sparse_output=False)
    num = StandardScaler() if scale_numeric else "passthrough"
    return ColumnTransformer(
        [
            ("cat", cat, CATEGORICAL_COLS),
            ("num", num, BOOL_COLS + NUMERIC_COLS),
        ],
        remainder="drop",
    )


def log_target(reg) -> TransformedTargetRegressor:
    """Wrap a regressor so it trains on log1p(salary) and predicts in dollars."""
    return TransformedTargetRegressor(regressor=reg, func=np.log1p, inverse_func=np.expm1)


def build_models() -> dict[str, Pipeline]:
    """Return the candidate models (each a full preprocessing+estimator pipeline)."""
    return {
        "baseline_mean": Pipeline(
            [("prep", make_preprocessor(scale_numeric=False)),
             ("model", DummyRegressor(strategy="mean"))]
        ),
        "ridge": Pipeline(
            [("prep", make_preprocessor(scale_numeric=True)),
             ("model", log_target(Ridge(alpha=1.0, random_state=RANDOM_STATE)))]
        ),
        "random_forest": Pipeline(
            [("prep", make_preprocessor(scale_numeric=False)),
             ("model", log_target(RandomForestRegressor(
                 n_estimators=300, min_samples_leaf=2,
                 n_jobs=-1, random_state=RANDOM_STATE)))]
        ),
        "hist_gbr": Pipeline(
            [("prep", make_preprocessor(scale_numeric=False)),
             ("model", log_target(HistGradientBoostingRegressor(
                 max_iter=400, learning_rate=0.06, max_leaf_nodes=31,
                 l2_regularization=1.0, random_state=RANDOM_STATE)))]
        ),
    }


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _round(x, n: int = 2):
    return None if x is None or (isinstance(x, float) and np.isnan(x)) else round(float(x), n)


def error_analysis(model: Pipeline, X_test: pd.DataFrame, y_test: pd.Series) -> dict:
    """Residual diagnostics on the held-out set."""
    pred = model.predict(X_test)
    resid = pred - y_test.to_numpy()
    abs_err = np.abs(resid)

    by_title = (
        pd.DataFrame({"title": X_test["job_title_short"].to_numpy(), "abs_err": abs_err})
        .groupby("title")["abs_err"].agg(["count", "mean"]).sort_values("mean", ascending=False)
    )
    quint = pd.qcut(y_test, 5, duplicates="drop")
    by_quint = (
        pd.DataFrame({"q": quint.astype(str).to_numpy(), "abs_err": abs_err})
        .groupby("q")["abs_err"].mean()
    )
    return {
        "mean_abs_error": _round(abs_err.mean()),
        "median_abs_error": _round(np.median(abs_err)),
        "mean_signed_error_bias": _round(resid.mean()),
        "abs_error_by_title": {k: _round(v) for k, v in by_title["mean"].items()},
        "count_by_title": {k: int(v) for k, v in by_title["count"].items()},
        "abs_error_by_salary_quintile": {k: _round(v) for k, v in by_quint.items()},
    }


# --------------------------------------------------------------------------- #
# Training entry points
# --------------------------------------------------------------------------- #
def train_salary_model() -> None:
    print("Loading salary data ...")
    df = load_salary_df()
    X = build_feature_frame(df)
    X[BOOL_COLS] = X[BOOL_COLS].astype(int)
    y = df[TARGET_COL].astype(float)
    print(f"  rows={len(X):,} | features={len(FEATURE_COLS)} "
          f"({len(CATEGORICAL_COLS)} cat + {len(BOOL_COLS)} bool + {len(NUMERIC_COLS)} num)")

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=RANDOM_STATE
    )
    print(f"  train={len(X_train):,} | test={len(X_test):,}")

    cv = KFold(n_splits=N_FOLDS, shuffle=True, random_state=RANDOM_STATE)
    scoring = {
        "rmse": "neg_root_mean_squared_error",
        "mae": "neg_mean_absolute_error",
        "r2": "r2",
    }

    print(f"\nComparing {len(build_models())} models with {N_FOLDS}-fold CV "
          f"(metrics in USD/year) ...\n")
    cv_results: dict[str, dict] = {}
    rows = []
    for name, pipe in build_models().items():
        res = cross_validate(pipe, X_train, y_train, cv=cv, scoring=scoring, n_jobs=-1)
        rmse = -res["test_rmse"]
        mae = -res["test_mae"]
        r2 = res["test_r2"]
        cv_results[name] = {
            "rmse_mean": _round(rmse.mean()), "rmse_std": _round(rmse.std()),
            "mae_mean": _round(mae.mean()), "mae_std": _round(mae.std()),
            "r2_mean": _round(r2.mean(), 4), "r2_std": _round(r2.std(), 4),
        }
        rows.append((name, rmse.mean(), mae.mean(), r2.mean()))

    table = pd.DataFrame(rows, columns=["model", "CV_RMSE", "CV_MAE", "CV_R2"]).set_index("model")
    print(table.round(2).to_string())

    # Best non-baseline model by CV-RMSE.
    ranked = sorted(
        [(n, v["rmse_mean"]) for n, v in cv_results.items() if n != "baseline_mean"],
        key=lambda t: t[1],
    )
    best_name = ranked[0][0]
    print(f"\nBest model by CV-RMSE: {best_name}")

    # Refit best on full train, evaluate on held-out test.
    best = build_models()[best_name]
    best.fit(X_train, y_train)
    pred = best.predict(X_test)
    holdout = {
        "model": best_name,
        "rmse": _round(root_mean_squared_error(y_test, pred)),
        "mae": _round(mean_absolute_error(y_test, pred)),
        "r2": _round(r2_score(y_test, pred), 4),
        "n_test": int(len(y_test)),
    }

    base = build_models()["baseline_mean"]
    base.fit(X_train, y_train)
    bpred = base.predict(X_test)
    baseline_holdout = {
        "rmse": _round(root_mean_squared_error(y_test, bpred)),
        "mae": _round(mean_absolute_error(y_test, bpred)),
        "r2": _round(r2_score(y_test, bpred), 4),
    }
    uplift = _round(100 * (1 - holdout["rmse"] / baseline_holdout["rmse"]), 1)

    print(f"\nHeld-out test ({best_name}): "
          f"RMSE={holdout['rmse']:,.0f}  MAE={holdout['mae']:,.0f}  R2={holdout['r2']}")
    print(f"Held-out test (baseline) : RMSE={baseline_holdout['rmse']:,.0f}  "
          f"MAE={baseline_holdout['mae']:,.0f}  R2={baseline_holdout['r2']}")
    print(f"RMSE reduction vs. baseline: {uplift}%")

    ea = error_analysis(best, X_test, y_test)
    print(f"\nError analysis: MAE={ea['mean_abs_error']:,.0f}  "
          f"bias(mean signed)={ea['mean_signed_error_bias']:,.0f}")
    print("  worst titles (MAE):",
          dict(list(ea["abs_error_by_title"].items())[:3]))

    # --- Persist artifacts (refit on ALL data for the shipped model) --------- #
    ARTIFACTS.mkdir(exist_ok=True)
    final = build_models()[best_name]
    final.fit(X, y)
    # compress=3: RandomForest tree arrays are highly redundant, so this shrinks
    # the artifact ~5x (≈190 MB → <40 MB) with no change to predictions — keeping
    # it under GitHub's 100 MB hard limit and friendly for the HF Space.
    joblib.dump(final, ARTIFACTS / "salary_model.joblib", compress=3)

    metadata = {
        "task": "salary_regression",
        "target": TARGET_COL,
        "target_unit": "USD per year",
        "target_transform": "log1p / expm1",
        "model": best_name,
        "feature_cols": FEATURE_COLS,
        "categorical_cols": CATEGORICAL_COLS,
        "bool_cols": BOOL_COLS,
        "numeric_cols": NUMERIC_COLS,
        "skill_vocab": SKILL_VOCAB,
        "defaults": DEFAULTS,
        "title_categories": sorted(df["job_title_short"].dropna().unique().tolist()),
        "n_rows_total": int(len(X)),
        "trained_at": date.today().isoformat(),
        "sklearn_version": sklearn.__version__,
    }
    (ARTIFACTS / "salary_metadata.json").write_text(
        json.dumps(metadata, indent=2, ensure_ascii=False), encoding="utf-8"
    )

    metrics = {
        "cv": {"folds": N_FOLDS, "metric_unit": "USD/year", "models": cv_results},
        "holdout": holdout,
        "baseline_holdout": baseline_holdout,
        "rmse_reduction_vs_baseline_pct": uplift,
        "error_analysis": ea,
    }
    (ARTIFACTS / "salary_metrics.json").write_text(
        json.dumps(metrics, indent=2, ensure_ascii=False), encoding="utf-8"
    )

    print(f"\nSaved -> {ARTIFACTS / 'salary_model.joblib'}")
    print(f"Saved -> {ARTIFACTS / 'salary_metadata.json'}")
    print(f"Saved -> {ARTIFACTS / 'salary_metrics.json'}")


def main() -> None:
    print("#" * 72 + "\n# ML block: salary regression\n" + "#" * 72)
    train_salary_model()

    # NLP block: the supervised résumé↔JD fit classifier (writes fit_* artifacts).
    # Imported (not run as __main__) so its custom transformers pickle correctly.
    print("\n" + "#" * 72 + "\n# NLP block: résumé↔JD fit classifier\n" + "#" * 72)
    from src.fit_model import train_fit_classifier

    train_fit_classifier()
    print("\nALL ARTIFACTS WRITTEN. DONE.")


if __name__ == "__main__":
    main()