scripts/train_price_model.py

"""Train and evaluate structured used-car price prediction models."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.append(str(PROJECT_ROOT))

import joblib
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer, TransformedTargetRegressor
from sklearn.ensemble import HistGradientBoostingRegressor, RandomForestRegressor
from sklearn.impute import SimpleImputer
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error, mean_absolute_percentage_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler

from app.price_features import (
    BINARY_FEATURES,
    CATEGORICAL_FEATURES,
    MODEL_FEATURES,
    NUMERIC_FEATURES,
    PRICE_MAX_EUR,
    PRICE_MIN_EUR,
    apply_category_grouping,
    clean_and_engineer_features,
    fit_category_grouping,
    get_model_frame,
    read_price_dataset,
)


DEFAULT_INPUT = PROJECT_ROOT / "data/raw/autoscout24_dataset_20251108.csv"
MODEL_PATH = PROJECT_ROOT / "models/price_model.joblib"
METRICS_PATH = PROJECT_ROOT / "reports/price_model_metrics.json"
REPORT_PATH = PROJECT_ROOT / "reports/price_model_report.md"
PREDICTIONS_PATH = PROJECT_ROOT / "reports/price_model_test_predictions.csv"
FEATURE_IMPORTANCE_PATH = PROJECT_ROOT / "reports/price_model_feature_importance.csv"
ERROR_ANALYSIS_PATH = PROJECT_ROOT / "reports/price_model_error_analysis.md"

RANDOM_STATE = 42


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--sample-rows", type=int, default=None)
    parser.add_argument(
        "--fast",
        action="store_true",
        help="Use faster model settings for local iteration.",
    )
    return parser.parse_args()


def split_data(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Split data into train, validation, and test sets: 70/15/15."""
    train_val, test = train_test_split(
        df,
        test_size=0.15,
        random_state=RANDOM_STATE,
        stratify=df["make_grouped"],
    )
    validation_fraction_of_train_val = 0.15 / 0.85
    train, validation = train_test_split(
        train_val,
        test_size=validation_fraction_of_train_val,
        random_state=RANDOM_STATE,
        stratify=train_val["make_grouped"],
    )
    return train, validation, test


def build_ridge_model() -> Pipeline:
    """Build a Ridge regression baseline with one-hot encoded categories."""
    numeric_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="median")),
            ("scaler", StandardScaler()),
        ]
    )
    binary_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="constant", fill_value=0, keep_empty_features=True)),
        ]
    )
    categorical_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="most_frequent")),
            (
                "onehot",
                OneHotEncoder(handle_unknown="ignore", min_frequency=25),
            ),
        ]
    )
    preprocessor = ColumnTransformer(
        transformers=[
            ("numeric", numeric_transformer, NUMERIC_FEATURES),
            ("binary", binary_transformer, BINARY_FEATURES),
            ("categorical", categorical_transformer, CATEGORICAL_FEATURES),
        ]
    )
    ridge = TransformedTargetRegressor(
        regressor=Ridge(alpha=2.0),
        func=np.log1p,
        inverse_func=np.expm1,
    )
    return Pipeline(steps=[("preprocessor", preprocessor), ("model", ridge)])


def build_tree_preprocessor() -> ColumnTransformer:
    """Build compact preprocessing for tree models."""
    numeric_transformer = Pipeline(
        steps=[("imputer", SimpleImputer(strategy="median"))]
    )
    binary_transformer = Pipeline(
        steps=[("imputer", SimpleImputer(strategy="constant", fill_value=0, keep_empty_features=True))]
    )
    categorical_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="most_frequent")),
            (
                "ordinal",
                OrdinalEncoder(
                    handle_unknown="use_encoded_value",
                    unknown_value=-1,
                    encoded_missing_value=-1,
                ),
            ),
        ]
    )
    return ColumnTransformer(
        transformers=[
            ("numeric", numeric_transformer, NUMERIC_FEATURES),
            ("binary", binary_transformer, BINARY_FEATURES),
            ("categorical", categorical_transformer, CATEGORICAL_FEATURES),
        ],
        verbose_feature_names_out=False,
    )


def build_random_forest_model(fast: bool = False) -> Pipeline:
    estimator = RandomForestRegressor(
        n_estimators=60 if fast else 120,
        max_depth=20 if fast else 26,
        min_samples_leaf=3,
        n_jobs=-1,
        random_state=RANDOM_STATE,
    )
    model = TransformedTargetRegressor(
        regressor=estimator,
        func=np.log1p,
        inverse_func=np.expm1,
    )
    return Pipeline(steps=[("preprocessor", build_tree_preprocessor()), ("model", model)])


def build_hist_gradient_boosting_model(fast: bool = False) -> Pipeline:
    estimator = HistGradientBoostingRegressor(
        max_iter=150 if fast else 260,
        learning_rate=0.08,
        max_leaf_nodes=31,
        l2_regularization=0.05,
        random_state=RANDOM_STATE,
    )
    model = TransformedTargetRegressor(
        regressor=estimator,
        func=np.log1p,
        inverse_func=np.expm1,
    )
    return Pipeline(steps=[("preprocessor", build_tree_preprocessor()), ("model", model)])


def evaluate_model(model: Pipeline, x: pd.DataFrame, y: pd.Series) -> dict[str, float]:
    predictions = np.clip(model.predict(x), PRICE_MIN_EUR, PRICE_MAX_EUR)
    rmse = mean_squared_error(y, predictions) ** 0.5
    return {
        "mae": float(mean_absolute_error(y, predictions)),
        "rmse": float(rmse),
        "r2": float(r2_score(y, predictions)),
        "mape": float(mean_absolute_percentage_error(y, predictions)),
    }


def format_metrics(metrics: dict[str, float]) -> str:
    return (
        f"MAE EUR {metrics['mae']:,.0f}, "
        f"RMSE EUR {metrics['rmse']:,.0f}, "
        f"R2 {metrics['r2']:.3f}, "
        f"MAPE {metrics['mape']:.3f}"
    )


def get_feature_importance(model: Pipeline) -> pd.DataFrame:
    """Extract feature importance from tree-based final model if available."""
    preprocessor = model.named_steps["preprocessor"]
    transformed_names = list(preprocessor.get_feature_names_out())
    regressor = model.named_steps["model"].regressor_

    if hasattr(regressor, "feature_importances_"):
        importance = regressor.feature_importances_
    else:
        return pd.DataFrame(columns=["feature", "importance"])

    return (
        pd.DataFrame({"feature": transformed_names, "importance": importance})
        .sort_values("importance", ascending=False)
        .reset_index(drop=True)
    )


def make_report(
    metrics: dict[str, dict[str, dict[str, float]]],
    best_model_name: str,
    feature_importance: pd.DataFrame,
    train_rows: int,
    validation_rows: int,
    test_rows: int,
) -> str:
    lines = [
        "# Price Model Training Report",
        "",
        "## Data Split",
        "",
        f"- Train rows: {train_rows:,}",
        f"- Validation rows: {validation_rows:,}",
        f"- Test rows: {test_rows:,}",
        "",
        "## Model Comparison",
        "",
        "| Model | Split | MAE (EUR) | RMSE (EUR) | R2 | MAPE |",
        "|---|---|---:|---:|---:|---:|",
    ]
    for model_name, split_metrics in metrics.items():
        for split_name, values in split_metrics.items():
            lines.append(
                f"| {model_name} | {split_name} | "
                f"{values['mae']:,.0f} | {values['rmse']:,.0f} | "
                f"{values['r2']:.3f} | {values['mape']:.3f} |"
            )

    lines.extend(
        [
            "",
            f"Best model selected by validation MAE: **{best_model_name}**.",
            "",
            "## Top Feature Importances",
            "",
        ]
    )
    if feature_importance.empty:
        lines.append("Feature importance is not available for the selected model.")
    else:
        lines.extend(["| Feature | Importance |", "|---|---:|"])
        for _, row in feature_importance.head(20).iterrows():
            lines.append(f"| {row['feature']} | {row['importance']:.4f} |")

    lines.extend(
        [
            "",
            "## Interpretation",
            "",
            "The comparison uses a linear baseline and non-linear tree-based models. "
            "The final model is selected by validation MAE because the application "
            "needs accurate absolute price recommendations in EUR. Test-set metrics "
            "are reported only after model selection.",
            "",
            "## Known Difficult Cases",
            "",
            "- Rare makes and models with few comparable listings.",
            "- Very expensive or highly customized vehicles.",
            "- Very new vehicles with extremely low mileage.",
            "- Listings with missing or inconsistent vehicle history fields.",
            "- Cases where visible damage is not reflected in the structured price dataset.",
        ]
    )
    return "\n".join(lines)


def make_error_analysis(predictions: pd.DataFrame) -> str:
    """Create markdown error analysis from full test-set predictions."""
    analyzed = predictions.copy()
    analyzed["percentage_error"] = analyzed["absolute_error"] / analyzed["price"]
    analyzed["price_band"] = pd.cut(
        analyzed["price"],
        bins=[0, 10_000, 20_000, 40_000, 80_000, 150_000, 300_000],
        labels=[
            "0-10k",
            "10k-20k",
            "20k-40k",
            "40k-80k",
            "80k-150k",
            "150k-300k",
        ],
        include_lowest=True,
    )

    by_band = (
        analyzed.groupby("price_band", observed=False)
        .agg(
            count=("absolute_error", "size"),
            mae=("absolute_error", "mean"),
            median_ae=("absolute_error", "median"),
            mape=("percentage_error", "mean"),
        )
        .reset_index()
    )
    by_make = (
        analyzed.groupby("make")
        .agg(
            count=("absolute_error", "size"),
            mae=("absolute_error", "mean"),
            median_ae=("absolute_error", "median"),
            mape=("percentage_error", "mean"),
        )
        .query("count >= 100")
        .sort_values("mae", ascending=False)
        .head(15)
        .reset_index()
    )
    worst_cases = analyzed.sort_values("absolute_error", ascending=False).head(20)

    lines = [
        "# Price Model Error Analysis",
        "",
        "## Error by Price Band",
        "",
        "| Price band | Count | MAE (EUR) | Median AE (EUR) | MAPE |",
        "|---|---:|---:|---:|---:|",
    ]
    for _, row in by_band.iterrows():
        lines.append(
            f"| {row['price_band']} | {row['count']:,.0f} | "
            f"{row['mae']:,.0f} | {row['median_ae']:,.0f} | {row['mape']:.3f} |"
        )

    lines.extend(
        [
            "",
            "## Highest-MAE Makes with at Least 100 Test Listings",
            "",
            "| Make | Count | MAE (EUR) | Median AE (EUR) | MAPE |",
            "|---|---:|---:|---:|---:|",
        ]
    )
    for _, row in by_make.iterrows():
        lines.append(
            f"| {row['make']} | {row['count']:,.0f} | "
            f"{row['mae']:,.0f} | {row['median_ae']:,.0f} | {row['mape']:.3f} |"
        )

    lines.extend(
        [
            "",
            "## Worst Absolute Errors",
            "",
            "| Make | Model | Year | Mileage km | Actual EUR | Predicted EUR | Absolute Error EUR |",
            "|---|---|---:|---:|---:|---:|---:|",
        ]
    )
    for _, row in worst_cases.iterrows():
        lines.append(
            f"| {row['make']} | {row['model']} | {row['production_year']:.0f} | "
            f"{row['mileage_km']:,.0f} | {row['price']:,.0f} | "
            f"{row['predicted_price']:,.0f} | {row['absolute_error']:,.0f} |"
        )

    lines.extend(
        [
            "",
            "## Interpretation",
            "",
            "Absolute errors increase for high-price vehicles because expensive cars have a wider value range and often depend on details that are not fully captured in the structured fields. Percentage errors are more informative for comparing cheap and expensive vehicles. Rare or highly customized vehicles remain difficult because the model has fewer comparable listings.",
        ]
    )
    return "\n".join(lines)


def train_models(input_path: Path, sample_rows: int | None = None, fast: bool = False) -> None:
    print("Loading price dataset...")
    raw = read_price_dataset(input_path, sample_rows=sample_rows)
    print(f"Raw shape: {raw.shape}")

    print("Cleaning and engineering features...")
    cleaned = clean_and_engineer_features(raw)
    print(f"Cleaned shape: {cleaned.shape}")

    preliminary_grouping = fit_category_grouping(cleaned)
    grouped = apply_category_grouping(cleaned, preliminary_grouping)
    train_df, validation_df, test_df = split_data(grouped)

    grouping_metadata = fit_category_grouping(train_df)
    train_df = apply_category_grouping(train_df, grouping_metadata)
    validation_df = apply_category_grouping(validation_df, grouping_metadata)
    test_df = apply_category_grouping(test_df, grouping_metadata)

    x_train = get_model_frame(train_df)
    y_train = train_df["price"]
    x_validation = get_model_frame(validation_df)
    y_validation = validation_df["price"]
    x_test = get_model_frame(test_df)
    y_test = test_df["price"]

    candidates: dict[str, Pipeline] = {
        "ridge_regression": build_ridge_model(),
        "random_forest": build_random_forest_model(fast=fast),
        "hist_gradient_boosting": build_hist_gradient_boosting_model(fast=fast),
    }

    metrics: dict[str, dict[str, dict[str, float]]] = {}
    fitted_models: dict[str, Pipeline] = {}

    for model_name, model in candidates.items():
        print(f"Training {model_name}...")
        model.fit(x_train, y_train)
        fitted_models[model_name] = model
        metrics[model_name] = {
            "validation": evaluate_model(model, x_validation, y_validation),
        }
        print(f"{model_name} validation: {format_metrics(metrics[model_name]['validation'])}")

    best_model_name = min(
        metrics,
        key=lambda name: metrics[name]["validation"]["mae"],
    )
    best_model = fitted_models[best_model_name]

    print(f"Best validation model: {best_model_name}")
    for model_name, model in fitted_models.items():
        metrics[model_name]["test"] = evaluate_model(model, x_test, y_test)
        print(f"{model_name} test: {format_metrics(metrics[model_name]['test'])}")

    feature_importance = get_feature_importance(best_model)

    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)

    artifact: dict[str, Any] = {
        "model": best_model,
        "model_name": best_model_name,
        "grouping_metadata": grouping_metadata,
        "metrics": metrics,
        "feature_columns": MODEL_FEATURES,
    }
    joblib.dump(artifact, MODEL_PATH, compress=3)

    METRICS_PATH.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    report = make_report(
        metrics=metrics,
        best_model_name=best_model_name,
        feature_importance=feature_importance,
        train_rows=len(train_df),
        validation_rows=len(validation_df),
        test_rows=len(test_df),
    )
    REPORT_PATH.write_text(report, encoding="utf-8")
    feature_importance.to_csv(FEATURE_IMPORTANCE_PATH, index=False)

    best_predictions = np.clip(best_model.predict(x_test), PRICE_MIN_EUR, PRICE_MAX_EUR)
    prediction_results = test_df[
        ["make", "model", "production_year", "mileage_km", "fuel_category", "price"]
    ].copy()
    prediction_results["predicted_price"] = best_predictions
    prediction_results["absolute_error"] = (
        prediction_results["price"] - prediction_results["predicted_price"]
    ).abs()
    prediction_results["selected_model"] = best_model_name
    prediction_results.to_csv(PREDICTIONS_PATH, index=False)
    ERROR_ANALYSIS_PATH.write_text(
        make_error_analysis(prediction_results),
        encoding="utf-8",
    )

    print(f"Saved model artifact to {MODEL_PATH}")
    print(f"Saved report to {REPORT_PATH}")


def main() -> None:
    args = parse_args()
    train_models(args.input, sample_rows=args.sample_rows, fast=args.fast)


if __name__ == "__main__":
    main()