SummerProject2026/scripts/train_baselines.py

"""Train leakage-safe baselines from the private inspection episode mart.

This script deliberately accepts only an explicit, prior-only feature schema.
It never queries countydata, never writes row-level predictions, and writes all
model artifacts beneath ``artifacts/private`` or ``data/private``.

The locked 2025 partition is not evaluated unless ``--evaluate-locked`` is
provided explicitly.
"""

from __future__ import annotations

import argparse
import csv
import gzip
import hashlib
import json
import math
import os
import statistics
import sys
import warnings
from collections import Counter
from dataclasses import dataclass
from datetime import date, datetime, timezone
from pathlib import Path
from typing import Dict, Iterator, List, Mapping, Optional, Sequence, Tuple


PROJECT_ROOT = Path(__file__).resolve().parents[1]
MODEL_VERSION = "baseline_v1"

ID_COLUMNS = (
    "vehicle_token",
    "vehicle_bucket",
    "is_vin_audit",
    "episode_number",
)
TARGET_COLUMNS = (
    "episode_start",
    "first_outcome",
    "target_nonpass",
    "eligible_returning_target",
    "temporal_partition",
)
NUMERIC_FEATURES = (
    "vehicle_age",
    "prior_episode_count",
    "prior_total_attempt_count",
    "prior_attempt_count",
    "days_since_prior_episode",
    "days_since_prior_adverse",
    "prior_nonpass_rate",
)
CATEGORICAL_FEATURES = (
    "public_county",
    "target_season",
    "prior_first_outcome",
    "prior_final_outcome",
    "last_observed_make",
    "last_observed_model",
)
AUDIT_COLUMNS = ("source_era", "target_outcome_label_source")
REQUIRED_COLUMNS = (
    ID_COLUMNS + TARGET_COLUMNS + NUMERIC_FEATURES + CATEGORICAL_FEATURES + AUDIT_COLUMNS
)

PASS_OUTCOMES = {"pass", "p"}
NONPASS_OUTCOMES = {"fail", "f", "reject", "abort"}
MISSING_TEXT = {"", "null", "none", "na", "n/a", "<null>"}
ALLOWED_TARGET_LABEL_SOURCES = {"overall_result", "utah_obd_proxy"}

# This trainer is intentionally binary-only. The Utah OBD proxy establishes a
# pass/non-pass label but is not approved to synthesize pass/fail/reject/abort
# classes. Any future four-class trainer must fail closed on utah_obd_proxy.
TARGET_CONTRACT = {
    "task": "binary_pass_vs_nonpass",
    "allowed_label_sources": sorted(ALLOWED_TARGET_LABEL_SOURCES),
    "utah_obd_proxy_binary_approved": True,
    "utah_obd_proxy_four_class_approved": False,
}
PLATT_CALIBRATION_CONFIG = {
    "method": "sigmoid_platt_scaling",
    "solver": "liblinear",
    "penalty": "l2",
    "c": 1000.0,
}

# These fields are forbidden even when they are not selected as model features.
# Their presence indicates that the input is not the approved episode mart.
FORBIDDEN_EXACT_COLUMNS = {
    "vin",
    "plate",
    "zip",
    "zipcode",
    "station",
    "station_id",
    "internal_event_id",
    "overall_result",
    "obd_result",
    "result_reason",
    "dtc_count",
    "mil",
    "readiness",
    "raw_json",
    "eventual_outcome",
    "eventually_passed",
    "target_attempt_count",
    "target_final_outcome",
}
FORBIDDEN_COLUMN_FRAGMENTS = (
    "license_plate",
    "certificate",
    "calibration_id",
    "visual_check",
    "readiness_monitor",
    "target_dtc",
    "target_pid",
    "target_obd",
    "target_mil",
)
FORBIDDEN_UNKNOWN_COLUMN_TOKENS = {
    "address",
    "certificate",
    "dtc",
    "email",
    "ip",
    "mil",
    "obd",
    "owner",
    "pid",
    "plate",
    "readiness",
    "session",
    "station",
    "user",
    "vin",
    "zip",
}

SPLIT_BOUNDS = {
    "train": (date(2016, 1, 1), date(2023, 1, 1)),
    "tune": (date(2023, 1, 1), date(2024, 1, 1)),
    "calibrate": (date(2024, 1, 1), date(2025, 1, 1)),
    "locked_test": (date(2025, 1, 1), date(2026, 1, 1)),
    "shadow": (date(2026, 1, 1), date(2027, 1, 1)),
}
MART_PARTITION_NAMES = {
    "train": {"train"},
    "tune": {"tune", "validation"},
    "calibrate": {"calibrate", "calibration"},
    "locked_test": {"test", "locked_test", "locked-test"},
    "shadow": {"shadow", "drift", "monitor"},
}


class BaselineError(Exception):
    """Base class for expected, user-facing baseline errors."""


class SchemaError(BaselineError):
    """Raised when the private mart schema is unsafe or incomplete."""


class DataValidationError(BaselineError):
    """Raised when a mart row violates the modeling contract."""


class DependencyError(BaselineError):
    """Raised when optional modeling dependencies are unavailable."""


@dataclass(frozen=True)
class EpisodeRow:
    """One eligible returning-vehicle target episode."""

    vehicle_token: str
    vehicle_bucket: int
    is_vin_audit: bool
    episode_number: int
    episode_start: datetime
    partition: str
    target: int
    source_era: Optional[str]
    target_outcome_label_source: str
    prior_first_outcome: Optional[str]
    numeric: Mapping[str, Optional[float]]
    categorical: Mapping[str, Optional[str]]

    @property
    def audit_vehicle(self) -> bool:
        return self.is_vin_audit


@dataclass
class MartData:
    """Validated rows and non-sensitive cohort counts."""

    rows_by_partition: Dict[str, List[EpisodeRow]]
    input_rows: int
    ineligible_rows: int
    eligible_rows: int


@dataclass
class FeatureEncoder:
    """Train-only numeric imputation/scaling and categorical rare grouping."""

    min_category_count: int
    numeric_medians: Dict[str, float]
    numeric_means: Dict[str, float]
    numeric_scales: Dict[str, float]
    kept_categories: Dict[str, set]
    vectorizer: object

    @classmethod
    def fit_transform(
        cls,
        rows: Sequence[EpisodeRow],
        min_category_count: int,
        np_module: object,
        dict_vectorizer_class: object,
    ) -> Tuple["FeatureEncoder", object]:
        if not rows:
            raise DataValidationError("Cannot fit preprocessing on zero rows")
        if min_category_count < 1:
            raise DataValidationError("min_category_count must be at least 1")

        numeric_medians: Dict[str, float] = {}
        numeric_means: Dict[str, float] = {}
        numeric_scales: Dict[str, float] = {}
        for name in NUMERIC_FEATURES:
            observed = [row.numeric[name] for row in rows if row.numeric[name] is not None]
            if not observed:
                raise DataValidationError(
                    "Training data has no observed values for numeric feature " + name
                )
            median = float(statistics.median(observed))
            imputed = [median if row.numeric[name] is None else row.numeric[name] for row in rows]
            mean = float(sum(imputed) / len(imputed))
            variance = float(sum((value - mean) ** 2 for value in imputed) / len(imputed))
            numeric_medians[name] = median
            numeric_means[name] = mean
            numeric_scales[name] = math.sqrt(variance) if variance > 0.0 else 1.0

        kept_categories: Dict[str, set] = {}
        for name in CATEGORICAL_FEATURES:
            counts = Counter(_category_value(row.categorical[name]) for row in rows)
            kept_categories[name] = {
                value for value, count in counts.items() if count >= min_category_count
            }

        vectorizer = dict_vectorizer_class(dtype=np_module.float64, sparse=True, sort=True)
        encoder = cls(
            min_category_count=min_category_count,
            numeric_medians=numeric_medians,
            numeric_means=numeric_means,
            numeric_scales=numeric_scales,
            kept_categories=kept_categories,
            vectorizer=vectorizer,
        )
        encoded = vectorizer.fit_transform([encoder.feature_dict(row) for row in rows])
        return encoder, encoded

    def feature_dict(self, row: EpisodeRow) -> Dict[str, object]:
        values: Dict[str, object] = {}
        for name in NUMERIC_FEATURES:
            raw = row.numeric[name]
            missing = raw is None
            imputed = self.numeric_medians[name] if missing else raw
            values["num__" + name] = (
                float(imputed) - self.numeric_means[name]
            ) / self.numeric_scales[name]
            values["missing__" + name] = 1.0 if missing else 0.0

        for name in CATEGORICAL_FEATURES:
            category = _category_value(row.categorical[name])
            if category not in self.kept_categories[name]:
                category = "<rare>"
            values["cat__" + name] = category
        return values

    def transform(self, rows: Sequence[EpisodeRow]) -> object:
        return self.vectorizer.transform([self.feature_dict(row) for row in rows])

    def artifact_state(self) -> Dict[str, object]:
        return {
            "min_category_count": self.min_category_count,
            "numeric_medians": self.numeric_medians,
            "numeric_means": self.numeric_means,
            "numeric_scales": self.numeric_scales,
            "kept_categories": {
                name: sorted(values) for name, values in self.kept_categories.items()
            },
            "vectorizer": self.vectorizer,
        }


@dataclass
class TrainedBaselines:
    training_prevalence: float
    encoder: FeatureEncoder
    logistic_model: object
    platt_model: object
    selected_c: float
    selected_n_iter: int
    platt_n_iter: int
    max_iter: int
    tuning_results: List[Dict[str, object]]


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--mart",
        type=Path,
        required=True,
        help="Private episode mart (.csv, .csv.gz, or .parquet)",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=PROJECT_ROOT / "artifacts/private/baselines" / MODEL_VERSION,
    )
    parser.add_argument("--seed", type=int, default=20260715)
    parser.add_argument("--min-category-count", type=int, default=100)
    parser.add_argument("--calibration-bins", type=int, default=10)
    parser.add_argument("--c-grid", default="0.03,0.1,0.3,1.0,3.0")
    parser.add_argument("--max-iter", type=int, default=5000)
    parser.add_argument(
        "--evaluate-locked",
        action="store_true",
        help="Explicitly unlock final metric calculation on the 2025 partition",
    )
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args(argv)


def parse_c_grid(value: str) -> List[float]:
    try:
        result = [float(part.strip()) for part in value.split(",") if part.strip()]
    except ValueError as exc:
        raise DataValidationError("--c-grid must contain positive numbers") from exc
    if not result or any(not math.isfinite(item) or item <= 0.0 for item in result):
        raise DataValidationError("--c-grid must contain positive finite numbers")
    return result


def _category_value(value: Optional[str]) -> str:
    return "<missing>" if value is None else value


def _normalize_column(name: str) -> str:
    return name.strip().lower()


def validate_schema(columns: Sequence[str]) -> None:
    normalized = [_normalize_column(name) for name in columns]
    if any(not name for name in normalized):
        raise SchemaError("The mart contains a blank column name")
    if len(normalized) != len(set(normalized)):
        raise SchemaError("The mart contains duplicate column names")

    missing = sorted(set(REQUIRED_COLUMNS) - set(normalized))
    if missing:
        raise SchemaError("The mart is missing required columns: " + ", ".join(missing))

    forbidden = []
    for name in normalized:
        if name in FORBIDDEN_EXACT_COLUMNS or any(
            fragment in name for fragment in FORBIDDEN_COLUMN_FRAGMENTS
        ):
            forbidden.append(name)
            continue
        if name not in REQUIRED_COLUMNS and (
            set(name.split("_")) & FORBIDDEN_UNKNOWN_COLUMN_TOKENS
        ):
            forbidden.append(name)
    if forbidden:
        raise SchemaError(
            "The input contains forbidden identifier/leakage columns: "
            + ", ".join(sorted(forbidden))
        )


def _parse_bool(value: object, column: str, row_number: int) -> bool:
    normalized = str(value).strip().lower()
    if normalized in {"true", "t", "1", "yes", "y"}:
        return True
    if normalized in {"false", "f", "0", "no", "n"}:
        return False
    raise DataValidationError(
        "Row {} has invalid boolean in {}: {!r}".format(row_number, column, value)
    )


def _parse_int(value: object, column: str, row_number: int) -> int:
    try:
        parsed = int(str(value).strip())
    except (TypeError, ValueError) as exc:
        raise DataValidationError(
            "Row {} has invalid integer in {}: {!r}".format(row_number, column, value)
        ) from exc
    return parsed


def _parse_optional_float(value: object, column: str, row_number: int) -> Optional[float]:
    if value is None or str(value).strip().lower() in MISSING_TEXT:
        return None
    try:
        parsed = float(str(value).strip())
    except (TypeError, ValueError) as exc:
        raise DataValidationError(
            "Row {} has invalid numeric value in {}: {!r}".format(
                row_number, column, value
            )
        ) from exc
    if not math.isfinite(parsed):
        raise DataValidationError(
            "Row {} has a non-finite value in {}".format(row_number, column)
        )
    return parsed


def _parse_optional_category(value: object) -> Optional[str]:
    if value is None:
        return None
    normalized = str(value).strip().lower()
    return None if normalized in MISSING_TEXT else normalized


def _parse_timestamp(value: object, row_number: int) -> datetime:
    text = str(value).strip()
    if not text:
        raise DataValidationError("Row {} has a blank episode_start".format(row_number))
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        return datetime.fromisoformat(text)
    except ValueError as exc:
        raise DataValidationError(
            "Row {} has invalid episode_start: {!r}".format(row_number, value)
        ) from exc


def _binary_outcome(value: object, column: str, row_number: int) -> int:
    normalized = str(value).strip().lower()
    if normalized in PASS_OUTCOMES:
        return 0
    if normalized in NONPASS_OUTCOMES:
        return 1
    raise DataValidationError(
        "Row {} has unrecognized {}: {!r}".format(row_number, column, value)
    )


def _target_nonpass(value: object, row_number: int) -> int:
    parsed = _parse_int(value, "target_nonpass", row_number)
    if parsed not in (0, 1):
        raise DataValidationError(
            "Row {} target_nonpass must be 0 or 1".format(row_number)
        )
    return parsed


def partition_for_timestamp(timestamp: datetime) -> Optional[str]:
    target_date = timestamp.date()
    for name, (start, end) in SPLIT_BOUNDS.items():
        if start <= target_date < end:
            return name
    return None


def _validate_declared_partition(
    declared: object, derived: str, row_number: int
) -> None:
    value = str(declared).strip().lower()
    if value not in MART_PARTITION_NAMES[derived]:
        raise DataValidationError(
            "Row {} declares temporal_partition={!r}, but episode_start belongs to {}".format(
                row_number, declared, derived
            )
        )


def _normalize_record(record: Mapping[str, object]) -> Dict[str, object]:
    return {_normalize_column(str(key)): value for key, value in record.items()}


def episode_from_record(record: Mapping[str, object], row_number: int) -> Optional[EpisodeRow]:
    row = _normalize_record(record)
    eligible = _parse_bool(
        row["eligible_returning_target"], "eligible_returning_target", row_number
    )
    if not eligible:
        return None

    token = str(row["vehicle_token"]).strip().lower()
    if len(token) != 64:
        raise DataValidationError(
            "Row {} vehicle_token is not a 64-character HMAC token".format(row_number)
        )
    try:
        int(token, 16)
    except ValueError as exc:
        raise DataValidationError(
            "Row {} vehicle_token is not hexadecimal".format(row_number)
        ) from exc

    bucket = _parse_int(row["vehicle_bucket"], "vehicle_bucket", row_number)
    if not 0 <= bucket <= 99:
        raise DataValidationError(
            "Row {} vehicle_bucket must be between 0 and 99".format(row_number)
        )
    is_vin_audit = _parse_bool(row["is_vin_audit"], "is_vin_audit", row_number)
    if is_vin_audit != (bucket < 10):
        raise DataValidationError(
            "Row {} is_vin_audit disagrees with the mart's bucket<10 contract".format(
                row_number
            )
        )

    episode_number = _parse_int(row["episode_number"], "episode_number", row_number)
    if episode_number <= 1:
        raise DataValidationError(
            "Row {} is marked eligible but is not a returning episode".format(row_number)
        )

    episode_start = _parse_timestamp(row["episode_start"], row_number)
    partition = partition_for_timestamp(episode_start)
    if partition is None or partition == "shadow":
        # Historical context and post-cutoff rows are not supervised here. Shadow
        # monitoring is intentionally outside this locked baseline trainer.
        return None
    _validate_declared_partition(row["temporal_partition"], partition, row_number)

    target = _binary_outcome(row["first_outcome"], "first_outcome", row_number)
    provided_target = _target_nonpass(row["target_nonpass"], row_number)
    if target != provided_target:
        raise DataValidationError(
            "Row {} target_nonpass disagrees with first_outcome".format(row_number)
        )

    numeric = {
        name: _parse_optional_float(row[name], name, row_number)
        for name in NUMERIC_FEATURES
    }
    prior_episode_count = numeric["prior_episode_count"]
    if prior_episode_count is None or prior_episode_count < 1.0:
        raise DataValidationError(
            "Row {} is marked eligible but prior_episode_count is below 1".format(
                row_number
            )
        )
    rate = numeric["prior_nonpass_rate"]
    if rate is not None and not 0.0 <= rate <= 1.0:
        raise DataValidationError(
            "Row {} prior_nonpass_rate is outside [0, 1]".format(row_number)
        )
    for name in (
        "prior_episode_count",
        "prior_total_attempt_count",
        "prior_attempt_count",
        "days_since_prior_episode",
        "days_since_prior_adverse",
    ):
        value = numeric[name]
        if value is not None and value < 0.0:
            raise DataValidationError(
                "Row {} {} is negative".format(row_number, name)
            )

    categorical = {
        name: _parse_optional_category(row[name]) for name in CATEGORICAL_FEATURES
    }
    source_era = _parse_optional_category(row["source_era"])
    if source_era is None:
        raise DataValidationError(
            "Row {} is missing source_era required for drift auditing".format(row_number)
        )
    label_source = _parse_optional_category(row["target_outcome_label_source"])
    if label_source not in ALLOWED_TARGET_LABEL_SOURCES:
        raise DataValidationError(
            "Row {} has unapproved target_outcome_label_source: {!r}".format(
                row_number, row["target_outcome_label_source"]
            )
        )
    if label_source == "utah_obd_proxy" and source_era != "utah":
        raise DataValidationError(
            "Row {} uses utah_obd_proxy outside source_era='utah'".format(row_number)
        )
    return EpisodeRow(
        vehicle_token=token,
        vehicle_bucket=bucket,
        is_vin_audit=is_vin_audit,
        episode_number=episode_number,
        episode_start=episode_start,
        partition=partition,
        target=target,
        source_era=source_era,
        target_outcome_label_source=label_source,
        prior_first_outcome=categorical["prior_first_outcome"],
        numeric=numeric,
        categorical=categorical,
    )


def _csv_records(path: Path) -> Tuple[List[str], Iterator[Mapping[str, object]]]:
    opener = gzip.open if path.name.lower().endswith(".gz") else open
    with opener(path, "rt", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        if reader.fieldnames is None:
            raise SchemaError("The mart has no header row")
        columns = list(reader.fieldnames)

    def iterator() -> Iterator[Mapping[str, object]]:
        with opener(path, "rt", encoding="utf-8", newline="") as handle:
            row_reader = csv.DictReader(handle)
            if row_reader.fieldnames is None:  # pragma: no cover - file changed mid-run
                raise SchemaError("The mart header disappeared while reading")
            if list(row_reader.fieldnames) != columns:
                raise SchemaError("The mart header changed while reading")
            for record in row_reader:
                if None in record:
                    raise SchemaError("A CSV mart row has more fields than its header")
                yield record

    return columns, iterator()


def _parquet_records(path: Path) -> Tuple[List[str], Iterator[Mapping[str, object]]]:
    try:
        import duckdb
    except ImportError as exc:
        raise DependencyError(
            "Reading Parquet requires duckdb; use the CSV/CSV.GZ mart or install duckdb"
        ) from exc
    connection = duckdb.connect(database=":memory:")
    try:
        cursor = connection.execute("SELECT * FROM read_parquet(?)", [str(path)])
        columns = [description[0] for description in cursor.description]
    except Exception as exc:
        connection.close()
        raise SchemaError("DuckDB could not read the Parquet mart") from exc
    connection.close()

    def iterator() -> Iterator[Mapping[str, object]]:
        row_connection = duckdb.connect(database=":memory:")
        try:
            row_cursor = row_connection.execute(
                "SELECT * FROM read_parquet(?)", [str(path)]
            )
            while True:
                rows = row_cursor.fetchmany(65_536)
                if not rows:
                    return
                for values in rows:
                    yield dict(zip(columns, values))
        except Exception as exc:
            raise DataValidationError("DuckDB failed while scanning the Parquet mart") from exc
        finally:
            row_connection.close()

    return columns, iterator()


def mart_records(path: Path) -> Tuple[List[str], Iterator[Mapping[str, object]]]:
    name = path.name.lower()
    if name.endswith(".parquet"):
        return _parquet_records(path)
    if name.endswith(".csv") or name.endswith(".csv.gz"):
        return _csv_records(path)
    raise SchemaError("The mart must be .csv, .csv.gz, or .parquet")


def load_mart(path: Path) -> MartData:
    if not path.is_file():
        raise SchemaError("Mart file does not exist: " + str(path))
    columns, records = mart_records(path)
    validate_schema(columns)

    rows_by_partition = {name: [] for name in ("train", "tune", "calibrate", "locked_test")}
    input_rows = 0
    ineligible_rows = 0
    for row_number, record in enumerate(records, start=2):
        input_rows += 1
        episode = episode_from_record(record, row_number)
        if episode is None:
            ineligible_rows += 1
            continue
        rows_by_partition[episode.partition].append(episode)

    eligible_rows = sum(len(rows) for rows in rows_by_partition.values())
    if eligible_rows == 0:
        raise DataValidationError("The mart contains no eligible supervised episodes")
    return MartData(
        rows_by_partition=rows_by_partition,
        input_rows=input_rows,
        ineligible_rows=ineligible_rows,
        eligible_rows=eligible_rows,
    )


def _require_two_classes(rows: Sequence[EpisodeRow], purpose: str) -> None:
    classes = {row.target for row in rows}
    if classes != {0, 1}:
        raise DataValidationError(
            "{} requires both pass and non-pass rows; observed classes={}".format(
                purpose, sorted(classes)
            )
        )


def _non_audit(rows: Sequence[EpisodeRow]) -> List[EpisodeRow]:
    return [row for row in rows if not row.audit_vehicle]


def _targets(rows: Sequence[EpisodeRow]) -> List[int]:
    return [row.target for row in rows]


def require_ml_dependencies() -> Tuple[object, object, object, object, object]:
    try:
        import joblib
        import numpy as np
        import sklearn
        from sklearn.exceptions import ConvergenceWarning
        from sklearn.feature_extraction import DictVectorizer
        from sklearn.linear_model import LogisticRegression
    except ImportError as exc:
        raise DependencyError(
            "Training requires numpy, scikit-learn, and joblib in the active environment"
        ) from exc
    return np, sklearn, DictVectorizer, LogisticRegression, ConvergenceWarning


def fit_with_convergence_capture(
    model: object,
    features: object,
    targets: object,
    convergence_warning_class: object,
) -> List[str]:
    """Fit while capturing only convergence warnings and replaying all others."""
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always", convergence_warning_class)
        model.fit(features, targets)

    convergence_messages: List[str] = []
    for warning in caught:
        if issubclass(warning.category, convergence_warning_class):
            convergence_messages.append(str(warning.message))
            continue
        # record=True redirects every warning that passed the caller's active
        # filters. Replay unrelated warnings rather than silently swallowing
        # them as a side effect of convergence auditing.
        warnings.warn_explicit(
            message=warning.message,
            category=warning.category,
            filename=warning.filename,
            lineno=warning.lineno,
        )
    return convergence_messages


def model_n_iter(model: object, np_module: object) -> Optional[int]:
    values = getattr(model, "n_iter_", None)
    if values is None:
        return None
    array = np_module.asarray(values)
    if array.size == 0 or not np_module.isfinite(array).all():
        return None
    return int(array.max())


def model_parameters_finite(model: object, np_module: object) -> bool:
    coefficient = getattr(model, "coef_", None)
    intercept = getattr(model, "intercept_", None)
    if coefficient is None or intercept is None:
        return False
    return bool(
        np_module.isfinite(np_module.asarray(coefficient)).all()
        and np_module.isfinite(np_module.asarray(intercept)).all()
    )


def probability_array_finite(probabilities: object, np_module: object) -> bool:
    values = np_module.asarray(probabilities)
    return bool(
        values.size > 0
        and np_module.isfinite(values).all()
        and (values >= 0.0).all()
        and (values <= 1.0).all()
    )


def train_baselines(
    mart: MartData,
    c_grid: Sequence[float],
    min_category_count: int,
    seed: int,
    max_iter: int,
) -> TrainedBaselines:
    (
        np,
        _sklearn,
        dict_vectorizer_class,
        logistic_class,
        convergence_warning_class,
    ) = require_ml_dependencies()
    train_rows = _non_audit(mart.rows_by_partition["train"])
    tune_rows = _non_audit(mart.rows_by_partition["tune"])
    calibrate_rows = _non_audit(mart.rows_by_partition["calibrate"])
    _require_two_classes(train_rows, "Training")
    _require_two_classes(tune_rows, "Tuning")
    _require_two_classes(calibrate_rows, "Calibration")

    train_targets = np.asarray(_targets(train_rows), dtype=np.int8)
    tune_targets = _targets(tune_rows)
    calibrate_targets = np.asarray(_targets(calibrate_rows), dtype=np.int8)
    training_prevalence = float(train_targets.mean())

    encoder, x_train = FeatureEncoder.fit_transform(
        train_rows,
        min_category_count=min_category_count,
        np_module=np,
        dict_vectorizer_class=dict_vectorizer_class,
    )
    x_tune = encoder.transform(tune_rows)

    candidates: List[Tuple[float, float, float, object, int]] = []
    tuning_results: List[Dict[str, object]] = []
    for c_value in c_grid:
        model = logistic_class(
            C=float(c_value),
            penalty="l2",
            solver="saga",
            class_weight=None,
            max_iter=max_iter,
            random_state=seed,
        )
        convergence_messages = fit_with_convergence_capture(
            model,
            x_train,
            train_targets,
            convergence_warning_class,
        )
        n_iter = model_n_iter(model, np)
        coefficients_finite = model_parameters_finite(model, np)
        probability_values = None
        probabilities_finite = False
        if coefficients_finite:
            probability_values = model.predict_proba(x_tune)[:, 1]
            probabilities_finite = probability_array_finite(probability_values, np)

        converged = (
            not convergence_messages
            and n_iter is not None
            and n_iter < max_iter
        )
        eligible_for_selection = (
            converged and coefficients_finite and probabilities_finite
        )
        result: Dict[str, object] = {
            "c": float(c_value),
            "converged": converged,
            "n_iter": n_iter,
            "convergence_warning": (
                " | ".join(convergence_messages) if convergence_messages else None
            ),
            "finite_coefficients": coefficients_finite,
            "finite_probabilities": probabilities_finite,
            "eligible_for_selection": eligible_for_selection,
            "brier": None,
            "average_precision": None,
        }
        if eligible_for_selection and probability_values is not None:
            probabilities = probability_values.tolist()
            metrics = binary_metrics(tune_targets, probabilities)
            brier = float(metrics["brier"])
            average_precision = float(metrics["average_precision"])
            result["brier"] = brier
            result["average_precision"] = average_precision
            candidates.append(
                (brier, -average_precision, float(c_value), model, int(n_iter))
            )
        tuning_results.append(result)

    if not candidates:
        raise DataValidationError(
            "No logistic candidate converged with finite coefficients and "
            "probabilities; increase --max-iter or inspect feature scaling"
        )

    _best_brier, _best_negative_ap, selected_c, selected_model, selected_n_iter = min(
        candidates, key=lambda item: (item[0], item[1], item[2])
    )
    if not model_parameters_finite(selected_model, np):
        raise DataValidationError("Selected logistic model has non-finite coefficients")
    selected_tune_probabilities = selected_model.predict_proba(x_tune)[:, 1]
    if not probability_array_finite(selected_tune_probabilities, np):
        raise DataValidationError("Selected logistic model has non-finite probabilities")

    # Platt scaling is fitted only on the dedicated 2024 calibration partition.
    x_calibrate = encoder.transform(calibrate_rows)
    decision_scores = selected_model.decision_function(x_calibrate).reshape(-1, 1)
    if not bool(np.isfinite(decision_scores).all()):
        raise DataValidationError("Selected logistic model has non-finite decision scores")
    # C=1000 with liblinear is a deliberately finite approximation to an
    # unregularized sigmoid calibrator. It is stable for this Python 3.9 /
    # scikit-learn 1.6 / NumPy 2.0 stack, where lbfgs emitted numeric warnings.
    platt_model = logistic_class(
        C=PLATT_CALIBRATION_CONFIG["c"],
        penalty=PLATT_CALIBRATION_CONFIG["penalty"],
        solver=PLATT_CALIBRATION_CONFIG["solver"],
        class_weight=None,
        max_iter=max_iter,
        random_state=seed,
    )
    platt_warnings = fit_with_convergence_capture(
        platt_model,
        decision_scores,
        calibrate_targets,
        convergence_warning_class,
    )
    platt_n_iter = model_n_iter(platt_model, np)
    if platt_warnings or platt_n_iter is None or platt_n_iter >= max_iter:
        if platt_warnings:
            detail = " | ".join(platt_warnings)
        elif platt_n_iter is None:
            detail = "missing n_iter_"
        else:
            detail = "n_iter reached max_iter without a ConvergenceWarning"
        raise DataValidationError("Platt calibration did not converge: " + detail)
    if not model_parameters_finite(platt_model, np):
        raise DataValidationError("Platt calibration has non-finite coefficients")
    calibrated_probabilities = platt_model.predict_proba(decision_scores)[:, 1]
    if not probability_array_finite(calibrated_probabilities, np):
        raise DataValidationError("Platt calibration has non-finite probabilities")

    return TrainedBaselines(
        training_prevalence=training_prevalence,
        encoder=encoder,
        logistic_model=selected_model,
        platt_model=platt_model,
        selected_c=selected_c,
        selected_n_iter=selected_n_iter,
        platt_n_iter=platt_n_iter,
        max_iter=max_iter,
        tuning_results=tuning_results,
    )


def literal_previous_probabilities(
    rows: Sequence[EpisodeRow], training_prevalence: float
) -> Tuple[List[float], int]:
    probabilities: List[float] = []
    fallback_count = 0
    for row in rows:
        outcome = row.prior_first_outcome
        if outcome in PASS_OUTCOMES:
            probabilities.append(0.0)
        elif outcome in NONPASS_OUTCOMES:
            probabilities.append(1.0)
        else:
            probabilities.append(training_prevalence)
            fallback_count += 1
    return probabilities, fallback_count


def logistic_probabilities(
    trained: TrainedBaselines, rows: Sequence[EpisodeRow], calibrated: bool
) -> List[float]:
    matrix = trained.encoder.transform(rows)
    if calibrated:
        decision = trained.logistic_model.decision_function(matrix).reshape(-1, 1)
        return trained.platt_model.predict_proba(decision)[:, 1].tolist()
    return trained.logistic_model.predict_proba(matrix)[:, 1].tolist()


def _validate_probabilities(targets: Sequence[int], probabilities: Sequence[float]) -> None:
    if len(targets) != len(probabilities) or not targets:
        raise DataValidationError("Targets and probabilities must have equal nonzero length")
    if any(target not in (0, 1) for target in targets):
        raise DataValidationError("Metrics require binary targets")
    for probability in probabilities:
        if not math.isfinite(probability) or not 0.0 <= probability <= 1.0:
            raise DataValidationError("Predicted probabilities must be finite and in [0, 1]")


def average_precision(targets: Sequence[int], probabilities: Sequence[float]) -> Optional[float]:
    positives = sum(targets)
    if positives == 0:
        return None
    ordered = sorted(zip(probabilities, targets), key=lambda item: item[0], reverse=True)
    true_positives = 0
    false_positives = 0
    previous_recall = 0.0
    result = 0.0
    index = 0
    while index < len(ordered):
        score = ordered[index][0]
        group_positive = 0
        group_total = 0
        while index < len(ordered) and ordered[index][0] == score:
            group_positive += ordered[index][1]
            group_total += 1
            index += 1
        true_positives += group_positive
        false_positives += group_total - group_positive
        recall = true_positives / positives
        precision = true_positives / (true_positives + false_positives)
        result += (recall - previous_recall) * precision
        previous_recall = recall
    return result


def roc_auc(targets: Sequence[int], probabilities: Sequence[float]) -> Optional[float]:
    positives = sum(targets)
    negatives = len(targets) - positives
    if positives == 0 or negatives == 0:
        return None
    ordered = sorted(zip(probabilities, targets), key=lambda item: item[0], reverse=True)
    true_positives = 0
    false_positives = 0
    previous_tpr = 0.0
    previous_fpr = 0.0
    area = 0.0
    index = 0
    while index < len(ordered):
        score = ordered[index][0]
        group_positive = 0
        group_total = 0
        while index < len(ordered) and ordered[index][0] == score:
            group_positive += ordered[index][1]
            group_total += 1
            index += 1
        true_positives += group_positive
        false_positives += group_total - group_positive
        tpr = true_positives / positives
        fpr = false_positives / negatives
        area += (fpr - previous_fpr) * (tpr + previous_tpr) / 2.0
        previous_tpr = tpr
        previous_fpr = fpr
    return area


def top_capacity_metrics(
    targets: Sequence[int], probabilities: Sequence[float], capacity: float
) -> Dict[str, Optional[float]]:
    if not 0.0 < capacity <= 1.0:
        raise DataValidationError("Capacity must be in (0, 1]")
    selected_count = capacity * len(targets)
    ordered = sorted(zip(probabilities, targets), key=lambda item: item[0], reverse=True)
    remaining = selected_count
    selected_positive = 0.0
    index = 0
    while index < len(ordered) and remaining > 0.0:
        score = ordered[index][0]
        group_targets: List[int] = []
        while index < len(ordered) and ordered[index][0] == score:
            group_targets.append(ordered[index][1])
            index += 1
        fraction = min(1.0, remaining / len(group_targets))
        selected_positive += fraction * sum(group_targets)
        remaining -= fraction * len(group_targets)
    positives = sum(targets)
    return {
        "precision": selected_positive / selected_count,
        "capture": None if positives == 0 else selected_positive / positives,
    }


def binary_metrics(targets: Sequence[int], probabilities: Sequence[float]) -> Dict[str, object]:
    _validate_probabilities(targets, probabilities)
    epsilon = 1e-15
    brier = sum(
        (probability - target) ** 2
        for target, probability in zip(targets, probabilities)
    ) / len(targets)
    log_loss = -sum(
        target * math.log(min(1.0 - epsilon, max(epsilon, probability)))
        + (1 - target)
        * math.log(min(1.0 - epsilon, max(epsilon, 1.0 - probability)))
        for target, probability in zip(targets, probabilities)
    ) / len(targets)
    top_5 = top_capacity_metrics(targets, probabilities, 0.05)
    top_10 = top_capacity_metrics(targets, probabilities, 0.10)
    return {
        "average_precision": average_precision(targets, probabilities),
        "roc_auc": roc_auc(targets, probabilities),
        "brier": brier,
        "log_loss": log_loss,
        "top_5_precision": top_5["precision"],
        "top_5_capture": top_5["capture"],
        "top_10_precision": top_10["precision"],
        "top_10_capture": top_10["capture"],
    }


def calibration_bins(
    targets: Sequence[int], probabilities: Sequence[float], bin_count: int
) -> List[Dict[str, object]]:
    _validate_probabilities(targets, probabilities)
    if bin_count < 2:
        raise DataValidationError("Calibration requires at least two bins")
    bin_count = min(bin_count, len(targets))
    ordered = sorted(zip(probabilities, targets), key=lambda item: item[0])
    base_size, extra = divmod(len(ordered), bin_count)
    result: List[Dict[str, object]] = []
    start = 0
    for bin_index in range(bin_count):
        size = base_size + (1 if bin_index < extra else 0)
        values = ordered[start : start + size]
        start += size
        result.append(
            {
                "bin": bin_index + 1,
                "count": size,
                "mean_probability": sum(item[0] for item in values) / size,
                "observed_nonpass_rate": sum(item[1] for item in values) / size,
                "min_probability": values[0][0],
                "max_probability": values[-1][0],
            }
        )
    return result


def evaluation_cohorts(
    rows: Sequence[EpisodeRow], include_audit_breakout: bool
) -> List[Tuple[str, List[EpisodeRow]]]:
    if not include_audit_breakout:
        return [("non_audit", _non_audit(rows))]
    return [
        ("all", list(rows)),
        ("non_audit", _non_audit(rows)),
        ("never_fit_audit", [row for row in rows if row.audit_vehicle]),
    ]


def evaluate_models(
    mart: MartData,
    trained: TrainedBaselines,
    evaluate_locked: bool,
    bin_count: int,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]]]:
    metric_rows: List[Dict[str, object]] = []
    calibration_rows: List[Dict[str, object]] = []
    partitions = ["train", "tune", "calibrate"]
    if evaluate_locked:
        partitions.append("locked_test")

    for partition in partitions:
        source_rows = mart.rows_by_partition[partition]
        for cohort_name, rows in evaluation_cohorts(
            source_rows, include_audit_breakout=(partition == "locked_test")
        ):
            if not rows:
                continue
            targets = _targets(rows)
            model_probabilities: List[Tuple[str, List[float], int]] = []
            model_probabilities.append(
                (
                    "training_prevalence",
                    [trained.training_prevalence] * len(rows),
                    0,
                )
            )
            previous, fallback_count = literal_previous_probabilities(
                rows, trained.training_prevalence
            )
            model_probabilities.append(
                ("previous_episode_literal", previous, fallback_count)
            )
            model_probabilities.append(
                ("logistic_raw", logistic_probabilities(trained, rows, calibrated=False), 0)
            )
            if partition in {"calibrate", "locked_test"}:
                model_probabilities.append(
                    (
                        "logistic_platt",
                        logistic_probabilities(trained, rows, calibrated=True),
                        0,
                    )
                )

            for model_name, probabilities, fallback in model_probabilities:
                metrics = binary_metrics(targets, probabilities)
                metric_rows.append(
                    {
                        "model": model_name,
                        "partition": partition,
                        "cohort": cohort_name,
                        "episodes": len(rows),
                        "vehicles": len({row.vehicle_token for row in rows}),
                        "nonpass": sum(targets),
                        "previous_outcome_fallbacks": fallback,
                        **metrics,
                    }
                )
                for bin_values in calibration_bins(targets, probabilities, bin_count):
                    calibration_rows.append(
                        {
                            "model": model_name,
                            "partition": partition,
                            "cohort": cohort_name,
                            **bin_values,
                        }
                    )
    return metric_rows, calibration_rows


def manifest_audit_counts(
    mart: MartData, evaluate_locked: bool
) -> Tuple[
    Dict[str, Dict[str, object]],
    Dict[str, Dict[str, Dict[str, object]]],
    Dict[str, Dict[str, int]],
]:
    """Return non-sensitive counts while withholding locked labels by default."""
    partition_counts: Dict[str, Dict[str, object]] = {}
    source_era_counts: Dict[str, Dict[str, Dict[str, object]]] = {}
    label_source_counts: Dict[str, Dict[str, int]] = {}
    for partition, rows in mart.rows_by_partition.items():
        locked_and_closed = partition == "locked_test" and not evaluate_locked
        partition_values: Dict[str, object] = {
            "episodes": len(rows),
            "audit_episodes": sum(row.audit_vehicle for row in rows),
        }
        if locked_and_closed:
            partition_values["labels_withheld"] = True
        else:
            partition_values["nonpass"] = sum(row.target for row in rows)
        partition_counts[partition] = partition_values

        partition_eras: Dict[str, Dict[str, object]] = {}
        for row in rows:
            era = row.source_era
            if era is None:  # guarded during mart validation
                continue
            values = partition_eras.setdefault(era, {"episodes": 0})
            values["episodes"] = int(values["episodes"]) + 1
            if not locked_and_closed:
                values["nonpass"] = int(values.get("nonpass", 0)) + row.target
        source_era_counts[partition] = partition_eras
        label_source_counts[partition] = dict(
            sorted(Counter(row.target_outcome_label_source for row in rows).items())
        )
    return partition_counts, source_era_counts, label_source_counts


def _is_private_project_path(path: Path) -> bool:
    resolved = path.resolve()
    private_roots = (
        (PROJECT_ROOT / "data/private").resolve(),
        (PROJECT_ROOT / "artifacts/private").resolve(),
    )
    for root in private_roots:
        try:
            resolved.relative_to(root)
            return True
        except ValueError:
            continue
    return False


def require_private_path(path: Path, label: str) -> Path:
    resolved = path.resolve()
    if not _is_private_project_path(resolved):
        raise DataValidationError(
            "{} must remain under data/private or artifacts/private".format(label)
        )
    return resolved


def _atomic_text(path: Path, content: str) -> None:
    partial = path.with_name(path.name + ".partial")
    partial.write_text(content, encoding="utf-8")
    os.replace(partial, path)


def _atomic_csv(path: Path, rows: Sequence[Mapping[str, object]]) -> None:
    if not rows:
        raise DataValidationError("Refusing to write an empty CSV artifact")
    partial = path.with_name(path.name + ".partial")
    columns = list(rows[0].keys())
    with partial.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=columns, extrasaction="raise")
        writer.writeheader()
        writer.writerows(rows)
    os.replace(partial, path)


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            block = handle.read(1024 * 1024)
            if not block:
                break
            digest.update(block)
    return digest.hexdigest()


def write_artifacts(
    output_dir: Path,
    mart_path: Path,
    mart: MartData,
    trained: TrainedBaselines,
    metric_rows: Sequence[Mapping[str, object]],
    calibration_rows: Sequence[Mapping[str, object]],
    evaluate_locked: bool,
    seed: int,
    overwrite: bool,
    sklearn_version: str,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    expected = (
        output_dir / "model.joblib",
        output_dir / "metrics.json",
        output_dir / "calibration_bins.csv",
        output_dir / "manifest.json",
    )
    existing = [path for path in expected if path.exists()]
    if existing and not overwrite:
        raise DataValidationError(
            "Refusing to overwrite existing artifacts: "
            + ", ".join(path.name for path in existing)
        )

    try:
        import joblib
    except ImportError as exc:
        raise DependencyError("Writing the fitted model requires joblib") from exc

    model_path = output_dir / "model.joblib"
    partial_model = model_path.with_name(model_path.name + ".partial")
    joblib.dump(
        {
            "model_version": MODEL_VERSION,
            "target_contract": TARGET_CONTRACT,
            "platt_calibration_config": PLATT_CALIBRATION_CONFIG,
            "numeric_features": NUMERIC_FEATURES,
            "categorical_features": CATEGORICAL_FEATURES,
            "training_prevalence": trained.training_prevalence,
            "selected_c": trained.selected_c,
            "convergence": {
                "max_iter": trained.max_iter,
                "selected_n_iter": trained.selected_n_iter,
                "platt_n_iter": trained.platt_n_iter,
            },
            "encoder": trained.encoder.artifact_state(),
            "logistic_model": trained.logistic_model,
            "platt_model": trained.platt_model,
        },
        partial_model,
    )
    os.replace(partial_model, model_path)

    _atomic_text(
        output_dir / "metrics.json",
        json.dumps(list(metric_rows), indent=2, sort_keys=True, allow_nan=False) + "\n",
    )
    _atomic_csv(output_dir / "calibration_bins.csv", calibration_rows)

    partition_counts, source_era_counts, label_source_counts = manifest_audit_counts(
        mart, evaluate_locked=evaluate_locked
    )
    manifest = {
        "model_version": MODEL_VERSION,
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "classification": "private_model_artifact_no_row_predictions",
        "input_file": mart_path.name,
        "input_sha256": _sha256(mart_path),
        "metrics_sha256": _sha256(output_dir / "metrics.json"),
        "input_rows": mart.input_rows,
        "eligible_rows": mart.eligible_rows,
        "ineligible_or_out_of_scope_rows": mart.ineligible_rows,
        "partition_counts": partition_counts,
        "source_era_audit_counts": source_era_counts,
        "target_label_source_counts": label_source_counts,
        "target_contract": TARGET_CONTRACT,
        "platt_calibration_config": PLATT_CALIBRATION_CONFIG,
        "split_bounds": {
            name: {"start_inclusive": start.isoformat(), "end_exclusive": end.isoformat()}
            for name, (start, end) in SPLIT_BOUNDS.items()
            if name != "shadow"
        },
        "never_fit_audit_rule": "mart is_vin_audit; validated as vehicle_bucket < 10",
        "locked_test_evaluated": evaluate_locked,
        "numeric_features": NUMERIC_FEATURES,
        "categorical_features": CATEGORICAL_FEATURES,
        "selected_c": trained.selected_c,
        "convergence": {
            "max_iter": trained.max_iter,
            "selected_n_iter": trained.selected_n_iter,
            "platt_n_iter": trained.platt_n_iter,
            "artifact_acceptance_requires_convergence": True,
            "artifact_acceptance_requires_finite_coefficients_and_probabilities": True,
        },
        "tuning_results": trained.tuning_results,
        "seed": seed,
        "python_version": sys.version,
        "scikit_learn_version": sklearn_version,
        "artifacts": [path.name for path in expected],
    }
    _atomic_text(
        output_dir / "manifest.json",
        json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n",
    )


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    os.umask(0o077)
    try:
        mart_path = require_private_path(args.mart, "--mart")
        output_dir = require_private_path(args.output_dir, "--output-dir")
        c_grid = parse_c_grid(args.c_grid)
        if args.max_iter < 1:
            raise DataValidationError("--max-iter must be positive")
        if args.calibration_bins < 2:
            raise DataValidationError("--calibration-bins must be at least 2")
        mart = load_mart(mart_path)
        trained = train_baselines(
            mart,
            c_grid=c_grid,
            min_category_count=args.min_category_count,
            seed=args.seed,
            max_iter=args.max_iter,
        )
        metric_rows, calibration_rows = evaluate_models(
            mart,
            trained,
            evaluate_locked=args.evaluate_locked,
            bin_count=args.calibration_bins,
        )
        (
            _np,
            sklearn,
            _dict_vectorizer,
            _logistic,
            _convergence_warning,
        ) = require_ml_dependencies()
        write_artifacts(
            output_dir=output_dir,
            mart_path=mart_path,
            mart=mart,
            trained=trained,
            metric_rows=metric_rows,
            calibration_rows=calibration_rows,
            evaluate_locked=args.evaluate_locked,
            seed=args.seed,
            overwrite=args.overwrite,
            sklearn_version=sklearn.__version__,
        )
    except BaselineError as exc:
        print("Baseline training failed: {}".format(exc), file=sys.stderr)
        return 2

    print("Wrote private baseline artifacts to {}".format(output_dir))
    if args.evaluate_locked:
        print("The explicitly unlocked 2025 test metrics were evaluated.")
    else:
        print("The 2025 locked test was not evaluated.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())