#!/usr/bin/env python3
"""
Retrain and save the chosen theme classifiers.

classify.py runs the full grid and keeps only scores and predictions. This script takes, for each
of the 15 modality combinations and each of the 5 folds, the configuration that validation picked
(results/classification/phaseb/runs.csv), trains it again with exactly the same fold split,
validation split, standardization, seed and settings, and saves it.

Output: <out>/<combination>/fold_<k>/
  model.pt      state_dict (PyTorch); rebuild with classify.build_model(run, modalities, in_dims, model_cfg)
  config.json   modalities, method and settings, input sizes, model config, validation / test
                macro-F1 (and the test macro-F1 of the original GPU run for comparison)
  norm.npz      per-modality mean and std of this fold's training songs (inputs must be
                standardized with them: (x - mean) / std)
<out>/index.csv one row per saved model

Combinations whose chosen configuration is late fusion get only config.json (it averages the
saved single-modality models of the same fold).

Usage (from the repository root; GPU recommended):
  python classification/save_models.py --out meet_models/classification
On Kaggle (no arguments): SET_ENV="SAVE_JOBS=1" EMBED_FILES="classification/classify.py
classification/config.yaml results/classification/phaseb/runs.csv" push_to_kaggle.sh ...
"""

import argparse
import json
import os
import sys
from itertools import combinations

import numpy as np
import pandas as pd
import torch
from joblib import Parallel, delayed

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, "classification")          # Kaggle: classify.py is written there (EMBED_FILES)
import classify as C  # noqa: E402


def as_bool(v):
    return str(v).strip().lower() in ("true", "1", "1.0")


def fold_job(k, chosen, out, config_path):
    if not torch.cuda.is_available():
        torch.set_num_threads(2)
    cfg = C.load_config(config_path, [])
    meta, audio, lyrics, folds, groups, y, soft = C.load_data(cfg)
    test = np.flatnonzero(folds == k)
    tr, va = C.inner_split(np.flatnonzero(folds != k), y, groups, cfg)
    split = {"train": tr, "val": va, "test": test}
    T = C.fold_tensors(k, audio, lyrics, split)
    sources = {**audio, "lyrics": lyrics[k] if lyrics.ndim == 3 else lyrics}
    mods_all = cfg["modalities"]
    rows = []
    for n in range(1, len(mods_all) + 1):
        for mods in combinations(mods_all, n):
            mods = list(mods)
            name = "+".join(mods)
            row = chosen[(chosen.combo == name) & (chosen.fold == k)].iloc[0]
            d = os.path.join(out, name, f"fold_{k}")
            os.makedirs(d, exist_ok=True)
            info = {"combination": name, "modalities": mods, "fold": k, "method": row["method"],
                    "themes": C.THEMES, "lyrics_encoder": f"theme-tuned BhashaEmbed, fold {k}",
                    "original_val_f1": float(row["val_f1"]), "original_test_f1": float(row["test_f1"])}
            if row["method"] == "late":
                info["note"] = "late fusion: average the class probabilities of the single-modality models " \
                               f"{', '.join(f'{m}/fold_{k}' for m in mods)}"
                json.dump(info, open(os.path.join(d, "config.json"), "w"), indent=2)
                rows.append({**info, "test_f1": np.nan})
                continue
            run = {"method": row["method"], "lr": float(row["lr"]), "soft": as_bool(row["soft"]),
                   "moddrop": float(row["moddrop"]) if pd.notna(row["moddrop"]) else 0.0,
                   "aux": as_bool(row["aux"])}
            r = C.train_eval(run, mods, T, y, soft, split, cfg)
            tf1, tacc = C.scores(y[test], r["pt"])
            torch.save(r["model"].state_dict(), os.path.join(d, "model.pt"))
            np.savez(os.path.join(d, "norm.npz"),
                     **{f"{m}_mean": sources[m][tr].mean(0) for m in mods},
                     **{f"{m}_std": np.where(sources[m][tr].std(0) < 1e-8, 1.0, sources[m][tr].std(0)) for m in mods})
            info.update({"run": run, "in_dims": [int(T[m]["train"].shape[1]) for m in mods],
                         "model_config": cfg["model"], "training": cfg["training"], "seed": cfg["seed"],
                         "val_f1": float(r["val_f1"]), "test_f1": float(tf1), "test_accuracy": float(tacc),
                         "epochs": int(r["epochs"]), "n_train": len(tr), "n_val": len(va), "n_test": len(test)})
            json.dump(info, open(os.path.join(d, "config.json"), "w"), indent=2)
            rows.append({k2: v for k2, v in info.items() if not isinstance(v, (dict, list))})
            print(f"fold {k} {name:30s} {run['method']:5s} test {tf1:.4f} (original {row['test_f1']:.4f})", flush=True)
    return rows


def main():
    ap = argparse.ArgumentParser(description="Retrain and save the chosen classifiers")
    ap.add_argument("--out", default=os.environ.get("SAVE_OUT", "meet_models/classification"))
    ap.add_argument("--config", default="classification/config.yaml")
    ap.add_argument("--runs", default="results/classification/phaseb/runs.csv")
    ap.add_argument("--jobs", type=int, default=int(os.environ.get("SAVE_JOBS", 5)),
                    help="folds trained in parallel (1 on a single GPU)")
    args = ap.parse_args()
    chosen = C.pick_by_val(pd.read_csv(args.runs))
    out = Parallel(n_jobs=args.jobs)(delayed(fold_job)(k, chosen, args.out, args.config) for k in range(5))
    df = pd.DataFrame([r for rows in out for r in rows]).sort_values(["combination", "fold"])
    df.to_csv(os.path.join(args.out, "index.csv"), index=False)
    s = df.groupby("combination")[["test_f1", "original_test_f1"]].mean().sort_values("test_f1", ascending=False)
    print("\nmean test macro-F1 per combination (retrained on this machine vs original run):")
    print(s.round(4).to_string())


if __name__ == "__main__":
    main()