"""Train the BioPatch clinical risk model.

Generates 5 000 literature-calibrated synthetic snapshots
(:mod:`ai.synthetic`), trains a :class:`RandomForestClassifier`, evaluates
on a held-out 20% split, and persists the model bundle to
``backend/ai/artifacts/rf_v1.joblib`` together with the training-set
summary statistics used at inference time for the XAI breakdown.

Run from the ``backend/`` directory::

    python -m scripts.seed_clinical_data
"""

from __future__ import annotations

import csv
import logging
from pathlib import Path

import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split

from ai.features import FEATURE_ORDER
from ai.synthetic import SEVERITY_CLASSES, generate_dataset

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("seed_clinical_data")


ROOT = Path(__file__).resolve().parent.parent
ARTIFACTS_DIR = ROOT / "ai" / "artifacts"
TRAINING_CSV = ARTIFACTS_DIR / "training_data.csv"
MODEL_PATH   = ARTIFACTS_DIR / "rf_v1.joblib"


def _dump_csv(samples: list[dict], path: Path) -> None:
    cols = FEATURE_ORDER + ["severity_class", "severity_label"]
    with path.open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=cols)
        w.writeheader()
        for s in samples:
            w.writerow({k: s[k] for k in cols})


def train(n_samples: int = 5000, seed: int = 42) -> dict:
    ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)

    logger.info("Generating %d calibrated synthetic samples (seed=%d)", n_samples, seed)
    samples = generate_dataset(n_samples, seed=seed)
    _dump_csv(samples, TRAINING_CSV)
    logger.info("Training data → %s", TRAINING_CSV)

    X = np.array([[s[f] for f in FEATURE_ORDER] for s in samples], dtype=np.float32)
    y = np.array([s["severity_class"] for s in samples], dtype=np.int32)

    X_tr, X_te, y_tr, y_te = train_test_split(
        X, y, test_size=0.20, random_state=seed, stratify=y,
    )

    logger.info("Training RandomForestClassifier (n_estimators=200, max_depth=14)")
    clf = RandomForestClassifier(
        n_estimators=200,
        max_depth=14,
        min_samples_leaf=4,
        class_weight="balanced",
        random_state=seed,
        n_jobs=-1,
    )
    clf.fit(X_tr, y_tr)
    acc = float(clf.score(X_te, y_te))
    logger.info("Held-out accuracy: %.3f", acc)

    pred = clf.predict(X_te)
    target_names = [SEVERITY_CLASSES[i] for i in sorted(set(y))]
    report = classification_report(
        y_te, pred, target_names=target_names, digits=3, zero_division=0
    )
    logger.info("Classification report:\n%s", report)

    baseline_mean = {f: float(X_tr[:, i].mean()) for i, f in enumerate(FEATURE_ORDER)}
    baseline_std  = {f: float(max(X_tr[:, i].std(),  0.05)) for i, f in enumerate(FEATURE_ORDER)}

    importance = dict(zip(FEATURE_ORDER, clf.feature_importances_.tolist()))
    logger.info("Feature importance: %s",
                ", ".join(f"{k}={v:.3f}" for k, v in importance.items()))

    bundle = {
        "clf": clf,
        "version": "rf-v1",
        "feature_order": FEATURE_ORDER,
        "baseline_mean": baseline_mean,
        "baseline_std":  baseline_std,
        "test_accuracy": acc,
        "n_samples": n_samples,
        "seed": seed,
        "classes": SEVERITY_CLASSES,
    }
    joblib.dump(bundle, MODEL_PATH)
    logger.info("Saved model bundle → %s (%.0f KB)", MODEL_PATH, MODEL_PATH.stat().st_size / 1024)
    return bundle


if __name__ == "__main__":
    train()