"""Full-volume inference for every trained configuration.

Reads each of the 888 CT volumes exactly once and runs every configuration
whose fold model covers that scan, so the HDD cost does not grow with the size
of the experiment matrix.

Usage
-----
    python scripts/04_predict.py                  # everything trained but not predicted
    python scripts/04_predict.py --groups main
    python scripts/04_predict.py --force
"""

from __future__ import annotations

import argparse
import sys
import time
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from luna_rev import config as cfg
from luna_rev import splits
from luna_rev.predict import candidates_path, predict_folds
from luna_rev.train import is_complete, weights_path


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--groups", nargs="*", default=None)
    ap.add_argument("--exp", nargs="*", default=None)
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--top-k", type=int, default=150,
                    help="candidates kept per scan (FROC only needs 8 FP/scan)")
    args = ap.parse_args()

    experiments = cfg.experiments_for_group(*args.groups) if args.groups else list(cfg.ALL_EXPERIMENTS)
    if args.exp:
        experiments = [e for e in experiments if e.name in set(args.exp)]

    ready, skipped = [], []
    for e in experiments:
        missing = [k for k in cfg.folds_for(e) if not is_complete(e, k)]
        (skipped if missing else ready).append((e, missing))
    for e, missing in skipped:
        print(f"  skip {e.name}: {len(missing)} fold(s) not trained yet -> {missing}")

    todo = [e for e, _ in ready if args.force or not candidates_path(e.name).exists()]
    if not todo:
        print("Nothing to predict.")
        return 0

    # Group configurations by the fold set they need, so each pass over the
    # scans covers exactly the configurations that use those folds.
    by_folds: dict[tuple[int, ...], list] = defaultdict(list)
    for e in todo:
        by_folds[tuple(cfg.folds_for(e))].append(e)

    all_folds = {f.index: f for f in splits.get_folds("official")}
    t0 = time.time()
    for fold_key, exps in sorted(by_folds.items(), key=lambda kv: -len(kv[0])):
        print(f"\n=== {len(exps)} configuration(s) over folds {list(fold_key)} ===")
        for e in exps:
            print(f"    {e.name}")
        predict_folds(exps, [all_folds[k] for k in fold_key],
                      top_k_per_scan=args.top_k, force=args.force)

    print(f"\nInference finished in {(time.time() - t0) / 3600:.2f} h")
    for e in todo:
        p = candidates_path(e.name)
        print(f"  {e.name:28s} {p.stat().st_size / 1e6:6.1f} MB")
    return 0


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