#!/usr/bin/env python3
"""Calibration-first modeling scaffold for INP-era lab-to-field prediction.

The script expects a leakage-reviewed panel CSV and produces validation metrics.
It uses blocked origin holdout when origin is available. Install scikit-learn/pandas in your execution environment.
"""
from __future__ import annotations
import argparse
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pandas as pd

from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.metrics import brier_score_loss, roc_auc_score, average_precision_score, mean_absolute_error
from sklearn.model_selection import GroupShuffleSplit
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.calibration import CalibratedClassifierCV, calibration_curve

NUMERIC_FEATURES = [
    'lighthouse_score', 'lab_lcp_ms', 'lab_tbt_ms', 'lab_cls', 'lab_fcp_ms',
    'lab_speed_index_ms', 'lab_total_bytes', 'p75_ttfb', 'p75_rtt'
]
CATEGORICAL_FEATURES = ['form_factor', 'lcp_resource_type']
TARGET = 'cwv_good_primary'


def expected_calibration_error(y_true: np.ndarray, y_prob: np.ndarray, bins: int = 15) -> float:
    edges = np.linspace(0, 1, bins + 1)
    ece = 0.0
    for lo, hi in zip(edges[:-1], edges[1:]):
        mask = (y_prob >= lo) & (y_prob < hi if hi < 1 else y_prob <= hi)
        if mask.sum() == 0:
            continue
        ece += (mask.sum() / len(y_prob)) * abs(y_true[mask].mean() - y_prob[mask].mean())
    return float(ece)


def build_pipeline() -> Pipeline:
    pre = ColumnTransformer([
        ('num', Pipeline([('impute', SimpleImputer(strategy='median')), ('scale', StandardScaler())]), NUMERIC_FEATURES),
        ('cat', Pipeline([('impute', SimpleImputer(strategy='most_frequent')), ('onehot', OneHotEncoder(handle_unknown='ignore'))]), CATEGORICAL_FEATURES),
    ], remainder='drop')
    base = HistGradientBoostingClassifier(max_iter=250, learning_rate=0.05, l2_regularization=0.05)
    return Pipeline([('pre', pre), ('model', base)])


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--panel', required=True, help='Leakage-reviewed CSV panel')
    ap.add_argument('--out', required=True, help='Output metrics JSON path')
    args = ap.parse_args()

    df = pd.read_csv(args.panel)
    missing = [c for c in [TARGET] + NUMERIC_FEATURES + CATEGORICAL_FEATURES + ['origin'] if c not in df.columns]
    if missing:
        raise SystemExit(f'Missing required columns: {missing}')

    df = df.dropna(subset=[TARGET, 'origin'])
    X = df[NUMERIC_FEATURES + CATEGORICAL_FEATURES]
    y = df[TARGET].astype(int).to_numpy()
    groups = df['origin'].astype(str).to_numpy()

    splitter = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=20260622)
    train_idx, test_idx = next(splitter.split(X, y, groups))

    pipe = build_pipeline()
    # CalibratedClassifierCV with cv=3 is used only inside the training block; origin blocking is preserved for final test.
    clf = CalibratedClassifierCV(pipe, method='isotonic', cv=3)
    clf.fit(X.iloc[train_idx], y[train_idx])
    prob = clf.predict_proba(X.iloc[test_idx])[:, 1]

    metrics = {
        'n_train': int(len(train_idx)),
        'n_test': int(len(test_idx)),
        'brier': brier_score_loss(y[test_idx], prob),
        'auroc': roc_auc_score(y[test_idx], prob),
        'auprc': average_precision_score(y[test_idx], prob),
        'ece_15_bins': expected_calibration_error(y[test_idx], prob, bins=15),
        'base_rate_test': float(y[test_idx].mean()),
        'access_note': 'No CrUX target metric may be used as a feature for itself; verify NUMERIC_FEATURES before each run.'
    }
    Path(args.out).write_text(pd.Series(metrics).to_json(indent=2), encoding='utf-8')


if __name__ == '__main__':
    main()