"""CT abnormality classification: a folder of .nii.gz volumes -> predictions JSON.

    python predict_classification.py --weights ct_classification \
        --input scans/ --output results.json
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

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

from fusion import FusionPolicy
from model import build_from_config, preprocess_ct


def case_name(path: Path) -> str:
    return path.name[:-len(".nii.gz")] if path.name.endswith(".nii.gz") else path.stem


def load_models(weights_dir: Path, device):
    spec = json.loads((weights_dir / "streams.json").read_text())
    models, cfgs = [], []
    for stream in spec["streams"]:
        sd = load_file(str(weights_dir / stream["weights_file"]))
        models.append(build_from_config(stream["config"], sd, len(spec["labels"])).to(device))
        cfgs.append(stream["config"])
    return spec["labels"], models, cfgs


@torch.inference_mode()
def predict_logits(path: Path, models, cfgs, device) -> np.ndarray:
    """(K, M) raw logits of the M models for one volume."""
    img = nib.load(str(path))
    volume = np.asanyarray(img.dataobj).astype(np.float32)
    affine = np.asarray(img.affine, dtype=np.float32)
    grids, out = {}, []
    for model, cfg in zip(models, cfgs):
        key = (cfg["img_size"], cfg["xy_spacing"], cfg["z_spacing"], cfg["max_slices"])
        if key not in grids:                       # the two 1.5 mm models share one grid
            grids[key] = preprocess_ct(volume, affine, img_size=cfg["img_size"], xy_spacing=cfg["xy_spacing"],
                                       z_spacing=cfg["z_spacing"], max_slices=cfg["max_slices"])
        x = grids[key].unsqueeze(0).to(device)
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            out.append(model(x).float().cpu().numpy()[0])
    return np.stack(out, axis=-1)


def challenge_json(names, labels, probs):
    return [{"outputs": [{"value": {
        "name": "Generated probabilities", "type": "Abnormality Classification",
        "version": {"major": 1, "minor": 0},
        "predictions": [{"input_image_name": n,
                         "probabilities": {l: float(p) for l, p in zip(labels, row)}}
                        for n, row in zip(names, probs)]}}]}]


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)
    ap.add_argument("--save-logits", type=Path, help="optional .npz with raw logits per case")
    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")
    assert policy.labels == labels

    paths = sorted(args.input.glob("*.nii.gz"))
    names = [case_name(p) for p in paths]
    logits = []
    for i, p in enumerate(paths, 1):
        logits.append(predict_logits(p, models, cfgs, device))
        print(f"[{i}/{len(paths)}] {names[i - 1]}", flush=True)
    logits = np.stack(logits)
    _, emitted = policy(logits)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(challenge_json(names, labels, emitted)))
    if args.save_logits:
        np.savez(args.save_logits, names=np.array(names), logits=logits, emitted=emitted)


if __name__ == "__main__":
    main()