"""MR report generation: a folder of study zips -> predictions.json with one report per study.

    python predict_report.py --weights mr_report_generation \
        --input studies/ --output reports.json
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from safetensors.torch import load_file

from implications import closure_matrix, select_with_implications
from model import build_model
from predict_classification import discover
from preprocess import process_files, process_zip
from report import ReportWriter
from tensors import to_model_input


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", type=Path, required=True)
    ap.add_argument("--input", type=Path, required=True)
    ap.add_argument("--output", type=Path, required=True)
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    cfg = json.loads((args.weights / "config.json").read_text())
    rc = json.loads((args.weights / "report_config.json").read_text())
    net = build_model(cfg["arch"], in_ch=cfg["in_channels"], num_classes=len(cfg["labels"]))
    net.load_state_dict(load_file(str(args.weights / "model.safetensors")))
    net = net.to(device).eval()

    order = rc["order"]                                  # the 32 scored diagnoses
    idx = [cfg["labels"].index(c) for c in order]
    thr = np.array([rc["thresholds"][c] for c in order])
    closure = closure_matrix(order)
    writer = ReportWriter(rc["template_lines"], rc["bullets"], order,
                          header_findings=rc.get("header_findings", True),
                          header_impression=rc.get("header_impression", True),
                          normal_line=rc.get("normal_line", ""),
                          short_template=rc.get("short_template_lines"),
                          short_from=int(rc.get("short_from", 0)))
    fallback = writer.write(set(rc.get("fallback_positives", [])))   # for a study that cannot be decoded

    reports = []
    for i, (name, study) in enumerate(discover(args.input), 1):
        geo = cfg["geometry"]
        per, n_used = process_zip(study, (geo,)) if isinstance(study, Path) else process_files(study, (geo,))
        if n_used == 0:
            text = fallback
        else:
            vol, presence = per[geo]
            x = to_model_input(torch.from_numpy(vol)[None].to(device),
                               torch.from_numpy(presence.astype(np.float32))[None].to(device))
            with torch.no_grad():
                probs = torch.sigmoid(net(x).float())[0].cpu().numpy().astype(np.float64)[idx]
            chosen = select_with_implications(probs[None], thr, int(rc.get("max_bullets", 0)), closure)[0]
            text = writer.write({c for c, k in zip(order, chosen) if k})
        reports.append({"input_image_name": f"{name}.nii.gz", "report": text})
        print(f"[{i}] {name}", flush=True)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps({"generated_reports": reports}))


if __name__ == "__main__":
    main()