"""CT report generation: a folder of .nii.gz volumes -> challenge report JSON.

    python predict_report.py --weights ct_classification \
        --writer ct_report_generation/report_writer.json \
        --input scans/ --output reports.json
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import numpy as np
import torch

from fusion import FusionPolicy
from predict_classification import case_name, load_models, predict_logits
from report_writer import ReportWriter

UUID = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", re.I)


def report_name(path: Path) -> str:
    """The report track names a case by the trailing UUID in its file name, if there is one."""
    stem = case_name(path)
    return next((part for part in reversed(stem.split("_")) if UUID.fullmatch(part)), stem)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", type=Path, required=True)
    ap.add_argument("--writer", 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")
    labels, models, cfgs = load_models(args.weights, device)
    policy = FusionPolicy(args.weights / "fusion_policy.json")
    writer = ReportWriter(args.writer)
    assert policy.labels == labels == writer.labels

    paths = sorted(args.input.rglob("*.nii.gz"))
    names = [report_name(p) for p in paths]
    assert len(set(names)) == len(names), "duplicate case names"
    logits = np.stack([predict_logits(p, models, cfgs, device) for p in paths])
    _, emitted = policy(logits)
    reports = [{"input_image_name": n, "report": writer.render(row)} for n, row in zip(names, emitted)]
    out = [{"outputs": [{"value": {"name": "Generated reports", "type": "Report generation",
                                   "version": "1.0", "generated_reports": reports}}]}]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(out, ensure_ascii=False))


if __name__ == "__main__":
    main()