"""Core TrainingJob class — orchestrates the full training pipeline."""

from __future__ import annotations

import gc
import os
import tempfile
from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import polars as pl

from haute._logging import get_logger
from haute.modelling._algorithms import (
    ALGORITHM_REGISTRY,
    IterationCallback,
    _malloc_trim,
    resolve_loss_function,
)
from haute.modelling._metrics import compute_metrics
from haute.modelling._split import (
    PARTITION_HOLDOUT,
    PARTITION_TRAIN,
    PARTITION_VALIDATION,
    SplitConfig,
    split_mask,
)

logger = get_logger(component="training_job")

_MODEL_EXT_MAP: dict[str, str] = {"catboost": ".cbm", "glm": ".rsglm"}


@dataclass
class TrainResult:
    """Result of a training job."""

    metrics: dict[str, float]
    feature_importance: list[dict[str, Any]]
    model_path: str
    train_rows: int
    test_rows: int  # validation rows (kept as test_rows for backward compat)
    features: list[str]
    cat_features: list[str]
    holdout_rows: int = 0
    holdout_metrics: dict[str, float] = field(default_factory=dict)
    diagnostics_set: str = "validation"  # "train" | "validation" | "holdout"
    best_iteration: int | None = None
    loss_history: list[dict[str, float]] = field(default_factory=list)
    double_lift: list[dict[str, Any]] = field(default_factory=list)
    shap_summary: list[dict[str, Any]] = field(default_factory=list)
    feature_importance_loss: list[dict[str, Any]] = field(default_factory=list)
    cv_results: dict[str, Any] | None = None
    ave_per_feature: list[dict[str, Any]] = field(default_factory=list)
    residuals_histogram: list[dict[str, Any]] = field(default_factory=list)
    residuals_stats: dict[str, float] = field(default_factory=dict)
    actual_vs_predicted: list[dict[str, float]] = field(default_factory=list)
    lorenz_curve: list[dict[str, float]] = field(default_factory=list)
    lorenz_curve_perfect: list[dict[str, float]] = field(default_factory=list)
    pdp_data: list[dict[str, Any]] = field(default_factory=list)
    # GLM-specific (empty for CatBoost)
    glm_coefficients: list[dict[str, Any]] = field(default_factory=list)
    glm_relativities: list[dict[str, Any]] = field(default_factory=list)
    glm_fit_statistics: dict[str, float] = field(default_factory=dict)
    glm_regularization_path: dict[str, Any] | None = None


@dataclass
class _PreparedData:
    """Intermediate result from the data-preparation phase."""

    data_path: str
    owns_tmp: bool
    features: list[str]
    cat_features: list[str]
    total_rows: int


@dataclass
class _SplitResult:
    """Intermediate result from the train/validation/holdout split phase."""

    split_path: str
    owns_tmp: bool
    n_train: int
    n_validation: int
    n_holdout: int


@dataclass
class _TrainModelResult:
    """Intermediate result from the model-fitting phase."""

    model: Any
    algo: Any
    fit_result: Any
    fit_params: dict[str, Any]


@dataclass
class _MetricsResult:
    """Intermediate result from the metrics/evaluation phase."""

    metrics: dict[str, float]
    holdout_metrics: dict[str, float]
    diagnostics_set: str  # "train" | "validation" | "holdout"
    importance: list[dict[str, Any]]
    double_lift: list[dict[str, Any]]
    shap_summary: list[dict[str, float]]
    feature_importance_loss: list[dict[str, Any]]
    cv_results: dict[str, Any] | None
    ave_per_feature: list[dict[str, Any]]
    residuals_histogram: list[dict[str, Any]]
    residuals_stats: dict[str, float]
    actual_vs_predicted: list[dict[str, float]]
    lorenz_curve: list[dict[str, float]]
    lorenz_curve_perfect: list[dict[str, float]]
    pdp_data: list[dict[str, Any]]
    # GLM-specific
    glm_coefficients: list[dict[str, Any]] = field(default_factory=list)
    glm_relativities: list[dict[str, Any]] = field(default_factory=list)
    glm_fit_statistics: dict[str, float] = field(default_factory=dict)
    glm_regularization_path: dict[str, Any] | None = None


class TrainingJob:
    """Orchestrates model training: split, fit, evaluate, optionally log to MLflow.

    Parameters
    ----------
    name : str
        Name for the model / training run.
    data : str | pl.DataFrame
        Path to parquet file, or a DataFrame directly.
    target : str
        Target column name.
    weight : str | None
        Optional weight/exposure column name.
    exclude : list[str] | None
        Columns to exclude from features. Everything not in
        {target, weight, *exclude} is automatically a feature.
    algorithm : str
        Algorithm key from ALGORITHM_REGISTRY (default: "catboost").
    task : str
        "regression" or "classification".
    params : dict | None
        Algorithm-specific hyperparameters.
    split : dict | SplitConfig | None
        Split configuration (strategy, test_size, seed, etc.).
    metrics : list[str] | None
        Metrics to compute (default: ["gini", "rmse"]).
    mlflow_experiment : str | None
        MLflow experiment path. If set and mlflow is importable, logs the run.
    model_name : str | None
        Optional MLflow registered model name.
    output_dir : str
        Directory to save the model file.
    """

    def __init__(
        self,
        *,
        name: str,
        data: str | pl.DataFrame,
        target: str,
        weight: str | None = None,
        exclude: list[str] | None = None,
        algorithm: str = "catboost",
        task: str = "regression",
        params: dict[str, Any] | None = None,
        split: dict[str, Any] | SplitConfig | None = None,
        metrics: list[str] | None = None,
        mlflow_experiment: str | None = None,
        model_name: str | None = None,
        output_dir: str = "outputs",
        loss_function: str | None = None,
        variance_power: float | None = None,
        offset: str | None = None,
        monotone_constraints: dict[str, int] | None = None,
        feature_weights: dict[str, float] | None = None,
        cv_folds: int | None = None,
    ) -> None:
        self.name = name
        self._data: str | pl.DataFrame | None = data
        self.target = target
        self.weight = weight
        self.exclude = exclude or []
        self.algorithm = algorithm
        self.task = task
        self.params = params or {}
        self.metrics = metrics or (["gini", "rmse"] if task == "regression" else ["auc", "logloss"])
        self.mlflow_experiment = mlflow_experiment
        self.model_name = model_name
        self.output_dir = output_dir
        self.loss_function = loss_function
        self.variance_power = variance_power
        self.offset = offset
        self.monotone_constraints = monotone_constraints
        self.feature_weights = feature_weights
        self.cv_folds = cv_folds

        # Parse split config
        if isinstance(split, SplitConfig):
            self.split_config = split
        elif isinstance(split, dict):
            self.split_config = SplitConfig(**split)
        else:
            self.split_config = SplitConfig()

    def run(
        self,
        progress: Callable[[str, float], None] | None = None,
        on_iteration: IterationCallback | None = None,
    ) -> TrainResult:
        """Execute the full training pipeline.

        Parameters
        ----------
        progress : callable | None
            Optional callback ``(message, fraction)`` for progress reporting.
        on_iteration : callable | None
            Optional callback ``(iteration, total, metrics_dict)`` called
            after each training iteration.

        Returns
        -------
        TrainResult
            Metrics, feature importances, model path, and split sizes.
        """

        def _report(msg: str, frac: float) -> None:
            if progress:
                progress(msg, frac)

        from haute.modelling._algorithms import _mem_checkpoint

        _mem_checkpoint("training run START")

        _report("Loading data", 0.0)
        prepared = self._prepare_data(_report)

        # GLM: narrow features to only the terms the user selected.
        # CatBoost uses all features; GLM should only carry the columns
        # referenced by its terms dict so we don't build a massive
        # design matrix or load unnecessary columns from parquet.
        if self.algorithm == "glm":
            glm_terms = self.params.get("terms", {})
            if glm_terms:
                term_names = set(glm_terms.keys())
                missing = term_names - set(prepared.features)
                if missing:
                    raise ValueError(
                        f"GLM terms reference columns not found in training data: "
                        f"{sorted(missing)}. Available columns: {prepared.features[:20]}"
                        + ("..." if len(prepared.features) > 20 else "")
                    )
                prepared = _PreparedData(
                    data_path=prepared.data_path,
                    owns_tmp=prepared.owns_tmp,
                    features=[f for f in prepared.features if f in term_names],
                    cat_features=[f for f in prepared.cat_features if f in term_names],
                    total_rows=prepared.total_rows,
                )
                _report(
                    f"GLM: using {len(prepared.features)} term features "
                    f"({len(prepared.cat_features)} categorical)",
                    0.12,
                )
                if not prepared.features:
                    raise ValueError(
                        "GLM: no valid features remaining after matching terms to data columns. "
                        "Check that your factor names match column names in the training data."
                    )

        _report("Splitting data", 0.15)
        split_result = self._split_data(prepared, _report)

        _report("Training model", 0.2)
        train_result = self._train_model(
            split_result,
            prepared.features,
            prepared.cat_features,
            on_iteration,
            _report,
        )

        _report("Evaluating model", 0.7)
        metrics_result = self._compute_metrics(
            split_result,
            prepared.features,
            prepared.cat_features,
            train_result,
            _report,
        )

        _report("Saving model", 0.9)
        model_path = self._save_artifacts(train_result)

        result = TrainResult(
            metrics=metrics_result.metrics,
            feature_importance=metrics_result.importance,
            model_path=str(model_path),
            train_rows=split_result.n_train,
            test_rows=split_result.n_validation,
            holdout_rows=split_result.n_holdout,
            holdout_metrics=metrics_result.holdout_metrics,
            diagnostics_set=metrics_result.diagnostics_set,
            features=prepared.features,
            cat_features=prepared.cat_features,
            best_iteration=train_result.fit_result.best_iteration,
            loss_history=train_result.fit_result.loss_history,
            double_lift=metrics_result.double_lift,
            shap_summary=metrics_result.shap_summary,
            feature_importance_loss=metrics_result.feature_importance_loss,
            cv_results=metrics_result.cv_results,
            ave_per_feature=metrics_result.ave_per_feature,
            residuals_histogram=metrics_result.residuals_histogram,
            residuals_stats=metrics_result.residuals_stats,
            actual_vs_predicted=metrics_result.actual_vs_predicted,
            lorenz_curve=metrics_result.lorenz_curve,
            lorenz_curve_perfect=metrics_result.lorenz_curve_perfect,
            pdp_data=metrics_result.pdp_data,
            glm_coefficients=metrics_result.glm_coefficients,
            glm_relativities=metrics_result.glm_relativities,
            glm_fit_statistics=metrics_result.glm_fit_statistics,
            glm_regularization_path=metrics_result.glm_regularization_path,
        )

        if self.mlflow_experiment:
            self._log_to_mlflow(result)

        _report("Done", 1.0)
        return result

    # ------------------------------------------------------------------
    # Pipeline sub-methods
    # ------------------------------------------------------------------

    def _prepare_data(
        self,
        _report: Callable[[str, float], None],
    ) -> _PreparedData:
        """Load data, validate columns, clean null targets, and derive features."""
        from haute.modelling._algorithms import _mem_checkpoint

        owns_tmp = False
        data_path: str | None = None

        if isinstance(self._data, str) and self._data.endswith(".parquet"):
            # Route already sunk the LazyFrame to disk -- no collect needed
            data_path = self._data
            self._data = None
            _mem_checkpoint(f"using on-disk parquet: {data_path}")
        else:
            # DataFrame or LazyFrame: collect, write to temp parquet, free
            df = self._load_data()
            self._data = None
            _mem_checkpoint(f"data loaded ({len(df):,} rows)")

            if self.target in df.columns:
                null_count = df[self.target].null_count()
                if null_count is not None and null_count > 0:
                    _mem_checkpoint(f"target has {null_count:,} null rows (will be cleaned)")

            with tempfile.NamedTemporaryFile(
                suffix=".parquet",
                prefix="haute_split_",
                delete=False,
            ) as f:
                data_path = f.name
            owns_tmp = True
            try:
                df.write_parquet(data_path)
            except BaseException:
                os.unlink(data_path)
                raise
            del df
            gc.collect()
            _malloc_trim()
            _mem_checkpoint("wrote temp parquet, freed df")

        # Validate schema from parquet metadata (cheap, no data loaded)
        _report("Validating columns", 0.05)
        from haute._polars_utils import read_parquet_metadata

        pq_meta = read_parquet_metadata(Path(data_path))
        if pq_meta["row_count"] == 0:
            raise ValueError("DataFrame is empty — cannot train on zero rows")
        schema_lf = pl.scan_parquet(data_path)
        schema_df = schema_lf.head(0).collect()
        self._validate_columns(schema_df)

        # Drop null targets -- re-write parquet without nulls if needed
        null_count = (
            pl.scan_parquet(data_path).select(pl.col(self.target).is_null().sum()).collect().item()
        )
        if null_count is not None and null_count > 0:
            clean_lf = pl.scan_parquet(data_path).filter(pl.col(self.target).is_not_null())
            with tempfile.NamedTemporaryFile(
                suffix=".parquet",
                prefix="haute_clean_",
                delete=False,
            ) as f:
                clean_path = f.name
            from haute._polars_utils import safe_sink

            safe_sink(clean_lf, clean_path)
            # Swap: delete old temp, use new one
            if owns_tmp:
                os.unlink(data_path)
            data_path = clean_path
            owns_tmp = True
            _mem_checkpoint(f"dropped {null_count:,} null-target rows")
            # Re-read row count from cleaned file so split mask matches
            pq_meta = read_parquet_metadata(Path(data_path))

        # Derive features from schema
        features, cat_features = self._derive_features(schema_df)
        del schema_df, schema_lf
        _report(f"Using {len(features)} features ({len(cat_features)} categorical)", 0.1)

        return _PreparedData(
            data_path=data_path,
            owns_tmp=owns_tmp,
            features=features,
            cat_features=cat_features,
            total_rows=pq_meta["row_count"],
        )

    def _split_data(
        self,
        prepared: _PreparedData,
        _report: Callable[[str, float], None],
    ) -> _SplitResult:
        """Compute train/validation/holdout split mask and write split parquet."""
        from haute.modelling._algorithms import _mem_checkpoint

        data_path = prepared.data_path
        owns_tmp = prepared.owns_tmp
        total_rows = prepared.total_rows

        # Compute mask -- for temporal/group we need a small scan
        mask_df = None
        if self.split_config.strategy in ("temporal", "group"):
            col = self.split_config.date_column or self.split_config.group_column
            mask_df = pl.scan_parquet(data_path).select(col).collect()
        mask = split_mask(total_rows, self.split_config, df=mask_df)
        del mask_df
        n_train = int((mask == PARTITION_TRAIN).sum())
        n_validation = int((mask == PARTITION_VALIDATION).sum())
        n_holdout = int((mask == PARTITION_HOLDOUT).sum())
        _mem_checkpoint(
            f"split mask (train={n_train:,} val={n_validation:,} holdout={n_holdout:,})"
        )

        # Write split parquet: original data + _partition column.
        # Lazy scan avoids doubling peak memory; collect + write is a single pass.
        with tempfile.NamedTemporaryFile(
            suffix=".parquet",
            prefix="haute_split_",
            delete=False,
        ) as f:
            split_path = f.name
        pl.scan_parquet(data_path).with_columns(mask).collect().write_parquet(split_path)
        del mask
        gc.collect()
        _malloc_trim()

        # Free the original temp parquet if we own it
        if owns_tmp and data_path != split_path:
            os.unlink(data_path)
        _mem_checkpoint("wrote split parquet")

        return _SplitResult(
            split_path=split_path,
            owns_tmp=True,
            n_train=n_train,
            n_validation=n_validation,
            n_holdout=n_holdout,
        )

    def _train_model(
        self,
        split_result: _SplitResult,
        features: list[str],
        cat_features: list[str],
        on_iteration: IterationCallback | None,
        _report: Callable[[str, float], None],
    ) -> _TrainModelResult:
        """Build train/eval pools (or DataFrames for GLM), fit the model."""
        from haute.modelling._algorithms import _mem_checkpoint

        data_path = split_result.split_path
        has_validation = split_result.n_validation > 0

        # Look up algorithm
        algo_cls = ALGORITHM_REGISTRY.get(self.algorithm)
        if algo_cls is None:
            raise ValueError(
                f"Unknown algorithm: {self.algorithm}. Available: {list(ALGORITHM_REGISTRY.keys())}"
            )
        algo = algo_cls()

        # Resolve loss function and inject into params
        fit_params = {**self.params}

        # GLM: pack all GLM-specific config into fit_params for the algorithm
        is_glm = self.algorithm == "glm"
        if not is_glm:
            resolved_loss = resolve_loss_function(
                self.loss_function,
                self.task,
                self.variance_power,
            )
            if resolved_loss:
                fit_params["loss_function"] = resolved_loss

        # Read train partition
        _report("Loading training data", 0.2)
        train_df = (
            self._scan_with_columns(data_path, features)
            .filter(pl.col("_partition") == PARTITION_TRAIN)
            .drop("_partition")
            .collect()
        )
        _mem_checkpoint(f"read train partition ({len(train_df):,} rows)")

        eval_df = None
        if has_validation:
            _report("Loading validation data", 0.25)
            eval_df = (
                self._scan_with_columns(data_path, features)
                .filter(pl.col("_partition") == PARTITION_VALIDATION)
                .drop("_partition")
                .collect()
            )
            _mem_checkpoint(f"read validation partition ({len(eval_df):,} rows)")

        if is_glm:
            # GLM: pass DataFrames directly (no Pool conversion needed)
            _report("Fitting GLM", 0.3)
            fit_result = algo.fit(
                train_df,
                features,
                cat_features,
                self.target,
                self.weight,
                fit_params,
                self.task,
                on_iteration=on_iteration,
                eval_df=eval_df,
                offset=self.offset,
                monotone_constraints=self.monotone_constraints,
                feature_weights=self.feature_weights,
            )
            _mem_checkpoint("glm algo.fit() returned")
            del train_df, eval_df
            gc.collect()
            _malloc_trim()
        else:
            # CatBoost: build memory-efficient pools, then fit
            from haute.modelling._algorithms import _build_pool

            train_y = train_df[self.target].cast(pl.Float64).to_numpy()
            train_w = train_df[self.weight].cast(pl.Float64).to_numpy() if self.weight else None
            train_baseline = (
                train_df[self.offset].cast(pl.Float64).to_numpy() if self.offset else None
            )
            train_features_df = train_df.select(features)
            del train_df
            gc.collect()
            _malloc_trim()
            _mem_checkpoint("extracted labels, freed train_df")

            train_pool = _build_pool(
                train_features_df,
                features,
                cat_features,
                y=train_y,
                w=train_w,
                baseline=train_baseline,
            )
            del train_features_df, train_y, train_w, train_baseline
            gc.collect()
            _malloc_trim()
            _mem_checkpoint("train pool built")

            eval_pool = None
            if eval_df is not None:
                _report("Building eval pool", 0.25)
                val_y = eval_df[self.target].cast(pl.Float64).to_numpy()
                val_w = eval_df[self.weight].cast(pl.Float64).to_numpy() if self.weight else None
                val_baseline = (
                    eval_df[self.offset].cast(pl.Float64).to_numpy() if self.offset else None
                )
                val_features_df = eval_df.select(features)
                del eval_df
                gc.collect()
                _malloc_trim()

                eval_pool = _build_pool(
                    val_features_df,
                    features,
                    cat_features,
                    y=val_y,
                    w=val_w,
                    baseline=val_baseline,
                )
                del val_features_df, val_y, val_w, val_baseline
                gc.collect()
                _malloc_trim()
                _mem_checkpoint("eval pool built")

            _report("Training model", 0.3)
            fit_result = algo.fit(
                None,
                features,
                cat_features,
                self.target,
                self.weight,
                fit_params,
                self.task,
                on_iteration=on_iteration,
                offset=self.offset,
                monotone_constraints=self.monotone_constraints,
                feature_weights=self.feature_weights,
                pool=train_pool,
                eval_pool=eval_pool,
            )
            _mem_checkpoint("algo.fit() returned")
            del train_pool, eval_pool
            gc.collect()
            _malloc_trim()
            _mem_checkpoint("del pools")

        return _TrainModelResult(
            model=fit_result.model,
            algo=algo,
            fit_result=fit_result,
            fit_params=fit_params,
        )

    def _glm_select_columns(self, features: list[str]) -> list[str] | None:
        """Column subset needed for GLM parquet reads, or ``None`` for CatBoost.

        GLM only needs the terms + target + weight + offset columns.
        Returning ``None`` means "read all columns" (CatBoost path).
        """
        if self.algorithm != "glm":
            return None
        needed = set(features)
        needed.add(self.target)
        if self.weight:
            needed.add(self.weight)
        if self.offset:
            needed.add(self.offset)
        return sorted(needed)

    def _scan_with_columns(self, data_path: str, features: list[str]) -> pl.LazyFrame:
        """Scan parquet with optional GLM column projection (includes _partition)."""
        scan = pl.scan_parquet(data_path)
        glm_columns = self._glm_select_columns(features)
        if glm_columns is not None:
            scan = scan.select([*glm_columns, "_partition"])
        return scan

    def _read_partition(
        self,
        data_path: str,
        partition: int,
        columns: list[str] | None = None,
    ) -> pl.DataFrame:
        """Read a single partition from the split parquet.

        If *columns* is given, only those columns (plus ``_partition`` for
        filtering) are loaded — Polars pushes the projection into the
        parquet reader so unused columns never touch RAM.
        """
        scan = pl.scan_parquet(data_path)
        if columns is not None:
            # Always need _partition for the filter; drop it after
            select_cols = columns if "_partition" in columns else [*columns, "_partition"]
            scan = scan.select(select_cols)
        return scan.filter(pl.col("_partition") == partition).drop("_partition").collect()

    def _compute_metrics(
        self,
        split_result: _SplitResult,
        features: list[str],
        cat_features: list[str],
        train_result: _TrainModelResult,
        _report: Callable[[str, float], None],
    ) -> _MetricsResult:
        """Evaluate model: metrics on validation + holdout, diagnostics on best available set.

        Memory-optimised: each partition is read at most once.  The diagnostics
        partition (holdout > validation > train) is read once and used for both
        its per-partition metrics *and* all diagnostic plots.  A separate read
        is only done for validation when holdout also exists.
        """
        from haute.modelling._algorithms import _build_pool, _mem_checkpoint
        from haute.modelling._metrics import (
            compute_actual_vs_predicted,
            compute_ave_per_feature,
            compute_double_lift,
            compute_lorenz_curve,
            compute_pdp,
            compute_residuals_histogram,
        )

        data_path = split_result.split_path
        algo = train_result.algo
        model = train_result.model
        need_cv = bool(self.cv_folds and self.cv_folds > 1)

        has_validation = split_result.n_validation > 0
        has_holdout = split_result.n_holdout > 0
        glm_columns = self._glm_select_columns(features)

        # Feature importance (doesn't need eval data)
        importance = algo.feature_importance(model)
        sorted_features = [fi["feature"] for fi in importance if fi["feature"] in features]
        sorted_features += [f for f in features if f not in sorted_features]

        # ── Determine which set to use for diagnostics (holdout > validation > train) ──
        if has_holdout:
            diag_partition = PARTITION_HOLDOUT
            diagnostics_set = "holdout"
        elif has_validation:
            diag_partition = PARTITION_VALIDATION
            diagnostics_set = "validation"
        else:
            diag_partition = PARTITION_TRAIN
            diagnostics_set = "train"

        # ── Read the diagnostics partition ONCE — metrics + all diagnostics ──
        _report("Computing diagnostics", 0.8)
        diag_df = self._read_partition(data_path, diag_partition, columns=glm_columns)
        _mem_checkpoint(f"read {diagnostics_set} partition for diagnostics ({len(diag_df):,} rows)")
        y_true = diag_df[self.target].to_numpy()
        y_pred = algo.predict(model, diag_df, features)
        w = diag_df[self.weight].to_numpy() if self.weight else None

        # Primary metrics from the diagnostics set
        vp = self.variance_power
        metrics = compute_metrics(
            y_true,
            y_pred,
            w,
            self.metrics,
            variance_power=vp,
        )

        # When holdout is present, diagnostics were computed on holdout.
        # Also compute validation metrics separately so both are available.
        holdout_metrics: dict[str, float] = {}
        if diagnostics_set == "holdout":
            holdout_metrics = metrics
            # Compute validation metrics separately if a validation set exists
            if has_validation:
                val_df = self._read_partition(
                    data_path,
                    PARTITION_VALIDATION,
                    columns=glm_columns,
                )
                val_y_true = val_df[self.target].to_numpy()
                val_y_pred = algo.predict(model, val_df, features)
                val_w = val_df[self.weight].to_numpy() if self.weight else None
                metrics = compute_metrics(
                    val_y_true,
                    val_y_pred,
                    val_w,
                    self.metrics,
                    variance_power=vp,
                )
                del val_df

        # Double-lift
        double_lift = compute_double_lift(y_true, y_pred, w)

        # AvE per feature
        ave_per_feature = compute_ave_per_feature(
            diag_df,
            sorted_features,
            cat_features,
            y_true,
            y_pred,
            w,
        )

        # SHAP + LossFunctionChange importance
        _report("Computing SHAP values", 0.85)
        shap_summary: list[dict[str, float]] = []
        feature_importance_loss: list[dict[str, Any]] = []
        if hasattr(algo, "shap_summary"):
            try:
                shap_summary = algo.shap_summary(model, diag_df, features, cat_features)
            except Exception as exc:
                logger.warning("shap_summary_failed", error=str(exc))
        if hasattr(algo, "feature_importance_typed"):
            try:
                _diag_pool = _build_pool(
                    diag_df,
                    features,
                    cat_features,
                    target=self.target,
                )
                feature_importance_loss = algo.feature_importance_typed(
                    model,
                    _diag_pool,
                    "LossFunctionChange",
                )
                del _diag_pool
            except Exception as exc:
                logger.warning("feature_importance_loss_failed", error=str(exc))

        # Residuals, scatter, Lorenz
        _report("Computing diagnostics", 0.86)
        residuals_histogram, residuals_stats = compute_residuals_histogram(y_true, y_pred, w)
        actual_vs_predicted = compute_actual_vs_predicted(y_true, y_pred, w)
        lorenz_model, lorenz_perfect = compute_lorenz_curve(y_true, y_pred, w)

        # PDP
        _report("Computing partial dependence", 0.87)
        pdp_data = compute_pdp(model, algo, diag_df, sorted_features, cat_features)

        # ── GLM-specific diagnostics ──
        glm_coefficients: list[dict[str, Any]] = []
        glm_relativities: list[dict[str, Any]] = []
        glm_fit_statistics: dict[str, float] = {}
        glm_regularization_path: dict[str, Any] | None = None

        if hasattr(algo, "coefficients_table"):
            try:
                glm_coefficients = algo.coefficients_table(model)
            except Exception as exc:
                logger.warning("glm_coefficients_failed", error=str(exc))
        if hasattr(algo, "relativities"):
            try:
                glm_relativities = algo.relativities(model)
            except Exception as exc:
                logger.warning("glm_relativities_failed", error=str(exc))
        if hasattr(algo, "fit_statistics"):
            try:
                glm_fit_statistics = algo.fit_statistics(model)
            except Exception as exc:
                logger.warning("glm_fit_statistics_failed", error=str(exc))
        if hasattr(model, "regularization_path") and model.regularization_path:
            try:
                rp = model.regularization_path
                glm_regularization_path = {
                    "selected_alpha": float(getattr(rp, "selected_alpha", 0)),
                    "n_nonzero": int(model.n_nonzero()) if hasattr(model, "n_nonzero") else 0,
                }
            except Exception as exc:
                logger.warning("glm_regularization_path_failed", error=str(exc))

        del diag_df
        gc.collect()

        # ── Cross-validation ──
        cv_results: dict[str, Any] | None = None
        if need_cv and hasattr(algo, "cross_validate"):
            _report("Running cross-validation", 0.88)
            try:
                _cv_df = (
                    self._scan_with_columns(data_path, features)
                    .filter(pl.col("_partition") == PARTITION_TRAIN)
                    .drop("_partition")
                    .collect()
                )
                cv_results = algo.cross_validate(
                    _cv_df,
                    features,
                    cat_features,
                    self.target,
                    self.weight,
                    train_result.fit_params,
                    self.task,
                    n_folds=self.cv_folds,
                )
                del _cv_df
                gc.collect()
            except Exception as exc:
                logger.warning("cv_failed", error=str(exc))

        # Clean up split parquet
        if split_result.owns_tmp and os.path.exists(data_path):
            os.unlink(data_path)

        return _MetricsResult(
            metrics=metrics,
            holdout_metrics=holdout_metrics,
            diagnostics_set=diagnostics_set,
            importance=importance,
            double_lift=double_lift,
            shap_summary=shap_summary,
            feature_importance_loss=feature_importance_loss,
            cv_results=cv_results,
            ave_per_feature=ave_per_feature,
            residuals_histogram=residuals_histogram,
            residuals_stats=residuals_stats,
            actual_vs_predicted=actual_vs_predicted,
            lorenz_curve=lorenz_model,
            lorenz_curve_perfect=lorenz_perfect,
            pdp_data=pdp_data,
            glm_coefficients=glm_coefficients,
            glm_relativities=glm_relativities,
            glm_fit_statistics=glm_fit_statistics,
            glm_regularization_path=glm_regularization_path,
        )

    def _save_artifacts(self, train_result: _TrainModelResult) -> Path:
        """Save the trained model to disk and return the model path."""
        ext = _MODEL_EXT_MAP.get(self.algorithm, ".model")
        model_path = Path(self.output_dir) / f"{self.name}{ext}"
        train_result.algo.save(train_result.model, model_path)
        return model_path

    # ------------------------------------------------------------------
    # Utility methods (unchanged)
    # ------------------------------------------------------------------

    def _load_data(self) -> pl.DataFrame:
        """Load data from path or use directly if already a DataFrame."""
        if isinstance(self._data, pl.DataFrame):
            return self._data
        if isinstance(self._data, pl.LazyFrame):
            return self._data.collect()
        if self._data is None:
            raise RuntimeError("Training data has already been consumed")
        from haute._io import read_source

        path = Path(self._data)
        if not path.exists():
            raise FileNotFoundError(f"Data file not found: {path}")
        return read_source(str(path)).collect()

    def _validate_columns(self, df: pl.DataFrame) -> None:
        """Validate that required columns exist in the DataFrame."""
        if self.target not in df.columns:
            raise ValueError(f"Target column '{self.target}' not found. Available: {df.columns}")
        if self.weight and self.weight not in df.columns:
            raise ValueError(f"Weight column '{self.weight}' not found. Available: {df.columns}")
        if self.offset and self.offset not in df.columns:
            raise ValueError(f"Offset column '{self.offset}' not found. Available: {df.columns}")
        # Excluded columns may already have been projected out during
        # pipeline execution — only flag genuinely unknown columns.
        available = set(df.columns)
        non_feature_cols = {self.target}
        if self.weight:
            non_feature_cols.add(self.weight)
        if self.offset:
            non_feature_cols.add(self.offset)
        dropped = [
            col for col in self.exclude if col not in available and col not in non_feature_cols
        ]
        if dropped:
            logger.debug("exclude_columns_already_dropped", count=len(dropped))

    def _derive_features(self, df: pl.DataFrame) -> tuple[list[str], list[str]]:
        """Derive feature list: all columns minus {target, weight, *exclude}.

        Also detects categorical features from Polars dtype.
        """
        non_features = {self.target}
        if self.weight:
            non_features.add(self.weight)
        if self.offset:
            non_features.add(self.offset)
        non_features.update(self.exclude)

        features = [c for c in df.columns if c not in non_features]
        if not features:
            raise ValueError(
                "No feature columns remaining after excluding "
                f"{non_features}. Check your target/weight/exclude settings."
            )

        # Detect categorical features from Polars dtype
        cat_features = []
        for col in features:
            dtype = df[col].dtype
            if dtype in (pl.Utf8, pl.Categorical, pl.String):
                cat_features.append(col)

        return features, cat_features

    def _log_to_mlflow(self, result: TrainResult) -> None:
        """Log training run to MLflow (conditional import).

        Delegates to the standalone ``log_experiment()`` function so the
        same logic is reused by the "Log to MLflow" button in the UI.
        """
        try:
            from haute.modelling._mlflow_log import log_experiment
        except ImportError:
            return
        from haute.modelling._result_types import (
            ModelCardMetadata,
            ModelDiagnostics,
        )

        if not self.mlflow_experiment:
            return

        diagnostics = ModelDiagnostics(
            feature_importance=result.feature_importance,
            shap_summary=result.shap_summary,
            feature_importance_loss=result.feature_importance_loss,
            double_lift=result.double_lift,
            loss_history=result.loss_history,
            cv_results=result.cv_results,
            ave_per_feature=result.ave_per_feature,
            residuals_histogram=result.residuals_histogram,
            residuals_stats=result.residuals_stats,
            actual_vs_predicted=result.actual_vs_predicted,
            lorenz_curve=result.lorenz_curve,
            glm_coefficients=result.glm_coefficients,
            glm_relativities=result.glm_relativities,
            glm_fit_statistics=result.glm_fit_statistics,
            glm_regularization_path=result.glm_regularization_path,
            lorenz_curve_perfect=result.lorenz_curve_perfect,
            pdp_data=result.pdp_data,
            holdout_metrics=result.holdout_metrics,
            diagnostics_set=result.diagnostics_set,
        )
        metadata = ModelCardMetadata(
            algorithm=self.algorithm,
            task=self.task,
            train_rows=result.train_rows,
            test_rows=result.test_rows,
            holdout_rows=result.holdout_rows,
            features=result.features,
            split_config=asdict(self.split_config) if self.split_config else {},
            best_iteration=result.best_iteration,
        )

        log_experiment(
            experiment_name=self.mlflow_experiment,
            run_name=self.name,
            metrics=result.metrics,
            params={
                "algorithm": self.algorithm,
                "task": self.task,
                "target": self.target,
                "weight": self.weight or "",
                "split_strategy": self.split_config.strategy,
                "validation_size": self.split_config.validation_size,
                "holdout_size": self.split_config.holdout_size,
                **{f"param_{k}": v for k, v in self.params.items()},
            },
            diagnostics=diagnostics,
            metadata=metadata,
            model_path=result.model_path or None,
            model_name=self.model_name,
        )
