train.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Training and experiment entry point for Track A.

This script owns model fitting and evaluation:
    1. Load the labelled Phase 1 training scenarios.
    2. Cross-validate the LightGBM template classifier and option selector.
    3. Save fold metrics, overall metrics, OOF predictions, train/val split
       manifests, full train/val split JSON files, and training logs under a
       timestamped experiment directory.
    4. Train final models on all labelled data.
    5. Save the model bundle in the experiment folder and copy it to --out.

Shared feature extraction, prediction-time selection, and bundle IO live in
src/model_core.py. Runtime inference orchestration lives in main.py.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
import time
from collections import Counter
from contextlib import redirect_stderr, redirect_stdout
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np
import pandas as pd
from lightgbm import LGBMClassifier
from sklearn.feature_extraction import DictVectorizer
from sklearn.model_selection import KFold, StratifiedKFold
from sklearn.pipeline import Pipeline

from src.model_core import (
    MODEL_BUNDLE_VERSION,
    answer_to_template,
    build_prediction_context,
    extract_scenario_features,
    iou_score,
    get_options,
    load_json,
    option_feature_dict,
    parse_answer,
    predict_labels,
    predict_labels_batch,
    save_model_bundle,
    template_to_actions,
)


class Tee:
    def __init__(self, *streams):
        self.streams = streams

    def write(self, data: str) -> None:
        for stream in self.streams:
            stream.write(data)
            stream.flush()

    def flush(self) -> None:
        for stream in self.streams:
            stream.flush()


def json_safe(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): json_safe(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_safe(v) for v in value]
    if isinstance(value, np.generic):
        return value.item()
    return value


def get_last_version(results_dir: Path) -> int:
    versions = []
    if not results_dir.exists():
        return 0
    for child in results_dir.iterdir():
        if not child.is_dir():
            continue
        prefix = child.name.split("_", 1)[0]
        if prefix.isdigit():
            versions.append(int(prefix))
    return sorted(versions)[-1] if versions else 0


def make_experiment_dir(root: Path, name: str = "") -> Path:
    root.mkdir(parents=True, exist_ok=True)
    clean_name = name.strip()[:10]
    last_version = get_last_version(root)
    base = (
        f"{last_version + 1:02d}_{clean_name}"
        if clean_name
        else f"{last_version + 1:02d}"
    )
    exp_dir = root / base
    suffix = 2
    while exp_dir.exists():
        exp_dir = root / f"{base}_{suffix}"
        suffix += 1
    exp_dir.mkdir(parents=True)
    return exp_dir


def scenario_id(s: Dict[str, Any], index: int) -> str:
    return str(s.get("scenario_id") or s.get("ID") or f"train_{index}")


def write_json(path: Path, obj: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(json_safe(obj), ensure_ascii=False, indent=2), encoding="utf-8"
    )


def fit_template_model(
    train: List[Dict[str, Any]],
    seed: int = 42,
    n_jobs: int = 1,
    boost_rounds: int = 2000,
) -> Pipeline:
    started = time.time()
    print(f"Building template features for {len(train)} scenarios...")
    X = [extract_scenario_features(s) for s in train]
    y = [answer_to_template(s) for s in train]
    print(
        f"Template features built in {time.time() - started:.1f}s; fitting LightGBM for {boost_rounds} rounds..."
    )
    clf = LGBMClassifier(
        objective="multiclass",
        n_estimators=boost_rounds,
        learning_rate=0.05,
        num_leaves=20,
        path_smooth=10,
        feature_fraction=0.8,
        bagging_fraction=0.8,
        bagging_freq=5,
        min_child_samples=20,
        class_weight="balanced",
        random_state=seed,
        n_jobs=n_jobs,
        verbosity=-1,
        force_col_wise=True,
    )
    model = Pipeline([("vec", DictVectorizer(sparse=False)), ("clf", clf)])
    model.fit(X, y)
    print(f"Template model fitted in {time.time() - started:.1f}s total.")
    return model


def build_selector_rows(
    train: List[Dict[str, Any]],
) -> Tuple[List[Dict[str, float]], List[int], List[Dict[str, Any]]]:
    X, y, meta = [], [], []
    for s in train:
        context = build_prediction_context(s)
        opts = context["options"]
        ans_ids = set(parse_answer(s.get("answer", "")))
        template = answer_to_template(s)
        for cid, label in opts.items():
            action = context["option_actions"].get(cid, "other")
            feats = option_feature_dict(s, cid, label, template, action, context)
            X.append(feats)
            y.append(1 if cid in ans_ids else 0)
            meta.append(
                {
                    "scenario_id": s.get("scenario_id"),
                    "cid": cid,
                    "action": action,
                    "template": template,
                    "in_template": action in set(template_to_actions(template)),
                }
            )
    return X, y, meta


def fit_selector_model(
    train: List[Dict[str, Any]],
    seed: int = 42,
    n_jobs: int = 1,
    boost_rounds: int = 2000,
) -> Pipeline:
    started = time.time()
    print(f"Building selector rows for {len(train)} scenarios...")
    X, y, _ = build_selector_rows(train)
    print(
        f"Selector rows built: {len(X)} rows in {time.time() - started:.1f}s; fitting LightGBM for {boost_rounds} rounds..."
    )
    clf = LGBMClassifier(
        objective="binary",
        n_estimators=boost_rounds,
        learning_rate=0.05,
        num_leaves=20,
        path_smooth=10,
        feature_fraction=0.8,
        bagging_fraction=0.8,
        bagging_freq=5,
        min_child_samples=20,
        class_weight="balanced",
        random_state=seed,
        n_jobs=n_jobs,
        verbosity=-1,
        force_col_wise=True,
    )
    model = Pipeline([("vec", DictVectorizer(sparse=False)), ("clf", clf)])
    model.fit(X, y)
    print(f"Selector model fitted in {time.time() - started:.1f}s total.")
    return model


def save_split_artifacts(
    exp_dir: Path,
    fold: int,
    train: List[Dict[str, Any]],
    tr_idx: np.ndarray,
    va_idx: np.ndarray,
) -> Tuple[Path, Path]:
    split_dir = exp_dir / "splits" / f"fold_{fold}"
    split_dir.mkdir(parents=True, exist_ok=True)

    def rows(indices: np.ndarray, split: str) -> List[Dict[str, Any]]:
        out = []
        for idx in indices:
            s = train[int(idx)]
            out.append(
                {
                    "fold": fold,
                    "split": split,
                    "row_index": int(idx),
                    "scenario_id": scenario_id(s, int(idx)),
                    "answer": s.get("answer", ""),
                    "template": answer_to_template(s),
                }
            )
        return out

    train_rows = rows(tr_idx, "train")
    val_rows = rows(va_idx, "val")
    pd.DataFrame(train_rows).to_csv(split_dir / "train_manifest.csv", index=False)
    pd.DataFrame(val_rows).to_csv(split_dir / "val_manifest.csv", index=False)

    train_json = split_dir / "train.json"
    val_json = split_dir / "val.json"
    write_json(train_json, [train[int(i)] for i in tr_idx])
    write_json(val_json, [train[int(i)] for i in va_idx])
    return train_json, val_json


def cross_validate_and_save(
    train: List[Dict[str, Any]],
    exp_dir: Path,
    folds: int,
    seed: int,
    n_jobs: int,
    template_boost_rounds: int,
    selector_boost_rounds: int,
) -> Dict[str, Any]:
    for child in ("metrics", "predictions", "splits"):
        (exp_dir / child).mkdir(parents=True, exist_ok=True)

    templates = [answer_to_template(s) for s in train]
    counts = Counter(templates)
    split_y = [t if counts[t] >= folds else "__rare__" for t in templates]
    split_counts = Counter(split_y)
    usable_folds = min(folds, len(train))
    min_split_count = min(split_counts.values()) if split_counts else 0
    if min_split_count >= 2:
        usable_folds = min(usable_folds, min_split_count)
        splitter = StratifiedKFold(
            n_splits=usable_folds, shuffle=True, random_state=seed
        )
        splits = splitter.split(np.zeros(len(train)), split_y)
        split_strategy = "stratified"
    else:
        usable_folds = min(usable_folds, max(2, len(train)))
        splitter = KFold(n_splits=usable_folds, shuffle=True, random_state=seed)
        splits = splitter.split(np.zeros(len(train)))
        split_strategy = "kfold"

    metrics: List[Dict[str, Any]] = []
    oof_rows: List[Dict[str, Any]] = []
    split_manifest: List[Dict[str, Any]] = []

    print(f"Template classes: {len(counts)}")
    print("Top templates:")
    for k, v in counts.most_common(20):
        print(f"  {k}: {v}")

    if usable_folds != folds:
        print(
            f"Adjusted CV folds from {folds} to {usable_folds} because of rare classes."
        )
    print(f"CV split strategy: {split_strategy}")

    for fold, (tr_idx, va_idx) in enumerate(splits, start=1):
        fold_start = time.time()
        print(f"\nTraining fold {fold}/{folds}...")
        train_json, val_json = save_split_artifacts(
            exp_dir, fold, train, tr_idx, va_idx
        )
        split_manifest.append(
            {
                "fold": fold,
                "train_rows": int(len(tr_idx)),
                "val_rows": int(len(va_idx)),
                "train_json": str(train_json),
                "val_json": str(val_json),
            }
        )

        tr = [train[int(i)] for i in tr_idx]
        va = [train[int(i)] for i in va_idx]
        template_model = fit_template_model(
            tr,
            seed + fold,
            n_jobs=n_jobs,
            boost_rounds=template_boost_rounds,
        )
        selector_model = fit_selector_model(
            tr,
            seed + fold,
            n_jobs=n_jobs,
            boost_rounds=selector_boost_rounds,
        )

        fold_ious: List[float] = []
        fold_tpl: List[float] = []
        pred_start = time.time()
        print(f"Scoring {len(va)} validation scenarios...")
        predictions = predict_labels_batch(template_model, selector_model, va)
        for local_i, s, (pred, dbg) in zip(va_idx, va, predictions):
            truth = parse_answer(s.get("answer", ""))
            iou = iou_score(pred, truth)
            template_ok = 1.0 if dbg["template"] == answer_to_template(s) else 0.0
            fold_ious.append(iou)
            fold_tpl.append(template_ok)
            oof_rows.append(
                {
                    "fold": fold,
                    "row_index": int(local_i),
                    "scenario_id": scenario_id(s, int(local_i)),
                    "truth": "|".join(truth),
                    "prediction": "|".join(pred),
                    "iou": float(iou),
                    "true_template": answer_to_template(s),
                    "pred_template": dbg["template"],
                    "template_prob": float(dbg["template_prob"]),
                    "template_correct": bool(template_ok),
                }
            )
        print(f"Validation scoring completed in {time.time() - pred_start:.1f}s.")

        row = {
            "fold": fold,
            "train_rows": int(len(tr_idx)),
            "val_rows": int(len(va_idx)),
            "iou_mean": float(np.mean(fold_ious)),
            "iou_std": float(np.std(fold_ious)),
            "template_acc": float(np.mean(fold_tpl)),
            "duration_seconds": round(time.time() - fold_start, 3),
        }
        metrics.append(row)
        print(
            f"Fold {fold}: IoU={row['iou_mean']:.4f} "
            f"template_acc={row['template_acc']:.4f} "
            f"duration={row['duration_seconds']:.1f}s"
        )

        pd.DataFrame(metrics).to_csv(
            exp_dir / "metrics" / "fold_metrics.csv", index=False
        )
        pd.DataFrame(oof_rows).to_csv(
            exp_dir / "predictions" / "oof_predictions.csv", index=False
        )
        write_json(exp_dir / "splits" / "split_manifest.json", split_manifest)

    summary = {
        "requested_folds": folds,
        "folds": usable_folds,
        "split_strategy": split_strategy,
        "seed": seed,
        "n_jobs": n_jobs,
        "template_boost_rounds": template_boost_rounds,
        "selector_boost_rounds": selector_boost_rounds,
        "iou_mean": float(np.mean([m["iou_mean"] for m in metrics])),
        "iou_std": float(np.std([m["iou_mean"] for m in metrics])),
        "template_acc_mean": float(np.mean([m["template_acc"] for m in metrics])),
        "template_acc_std": float(np.std([m["template_acc"] for m in metrics])),
        "fold_metrics": metrics,
    }
    write_json(exp_dir / "metrics" / "overall_metrics.json", summary)
    pd.DataFrame([summary]).drop(columns=["fold_metrics"]).to_csv(
        exp_dir / "metrics" / "overall_metrics.csv", index=False
    )
    return summary


def run_training(args: argparse.Namespace) -> None:
    started = time.time()
    exp_dir = make_experiment_dir(args.experiments_root, args.experiment_name)
    log_path = exp_dir / "training.log"

    config = {
        "train_path": args.train_path,
        "out": args.out,
        "experiments_root": str(args.experiments_root),
        "experiment_dir": str(exp_dir),
        "experiment_name": args.experiment_name,
        "folds": args.folds,
        "seed": args.seed,
        "n_jobs": args.n_jobs,
        "template_boost_rounds": args.template_boost_rounds,
        "selector_boost_rounds": args.selector_boost_rounds,
        "selector_training_scope": "all_options",
        "selector_prediction_mode": "candidate_ranker",
        "model_bundle_version": MODEL_BUNDLE_VERSION,
    }
    write_json(exp_dir / "config.json", config)
    (exp_dir / "experiment.txt").write_text(
        f"{exp_dir.name}\n{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n",
        encoding="utf-8",
    )

    with log_path.open("w", encoding="utf-8") as log_f, redirect_stdout(
        Tee(sys.stdout, log_f)
    ), redirect_stderr(Tee(sys.stderr, log_f)):
        print(f"Experiment directory: {exp_dir}")
        print(f"Training log: {log_path}")
        print(f"Config: {json.dumps(config, indent=2)}")

        train = load_json(args.train_path)
        print(f"Loaded train={len(train)} from {args.train_path}")

        cv_summary = cross_validate_and_save(
            train,
            exp_dir,
            args.folds,
            args.seed,
            args.n_jobs,
            args.template_boost_rounds,
            args.selector_boost_rounds,
        )

        print("\nTraining full template model...")
        template_model = fit_template_model(
            train,
            args.seed,
            n_jobs=args.n_jobs,
            boost_rounds=args.template_boost_rounds,
        )
        print("Training full selector model...")
        selector_model = fit_selector_model(
            train,
            args.seed,
            n_jobs=args.n_jobs,
            boost_rounds=args.selector_boost_rounds,
        )

        templates = Counter(answer_to_template(s) for s in train)
        metadata = {
            "version": MODEL_BUNDLE_VERSION,
            "train_path": args.train_path,
            "train_rows": len(train),
            "seed": args.seed,
            "n_jobs": args.n_jobs,
            "template_boost_rounds": args.template_boost_rounds,
            "selector_boost_rounds": args.selector_boost_rounds,
            "selector_training_scope": "all_options",
            "selector_prediction_mode": "candidate_ranker",
            "template_classes": len(templates),
            "top_templates": templates.most_common(20),
            "cv": cv_summary,
            "experiment_dir": str(exp_dir),
            "trained_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "duration_seconds": round(time.time() - started, 3),
        }

        bundle_path = exp_dir / "models" / "model_v4_bundle.pkl"
        save_model_bundle(str(bundle_path), template_model, selector_model, metadata)
        write_json(exp_dir / "metadata.json", metadata)
        Path(args.out).parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(bundle_path, args.out)

        print(f"\nSaved experiment model bundle: {bundle_path}")
        print(f"Copied model bundle to: {args.out}")
        print(f"Saved metadata: {exp_dir / 'metadata.json'}")
        print(f"Saved metrics: {exp_dir / 'metrics'}")
        print(f"Saved splits: {exp_dir / 'splits'}")
        print(f"Total duration: {metadata['duration_seconds']:.1f}s")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Cross-validate, train, and save the Track A v4 ML model bundle."
    )
    parser.add_argument("--train_path", default="data/Phase_1/train.json")
    parser.add_argument("--out", default="models/model_v4_bundle.pkl")
    parser.add_argument(
        "--experiments_root", type=Path, default=Path("results/experiments")
    )
    parser.add_argument("--experiment_name", default="lgbm_v4")
    parser.add_argument("--folds", type=int, default=5)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--n_jobs",
        type=int,
        default=-1,
        help="Parallel jobs for LightGBM.",
    )
    parser.add_argument("--template_boost_rounds", type=int, default=2000)
    parser.add_argument("--selector_boost_rounds", type=int, default=2000)
    args = parser.parse_args()
    run_training(args)


if __name__ == "__main__":
    main()