"""MR abnormality classification: a folder of study zips -> predictions JSON.

    python predict_classification.py --weights mr_classification \
        --input studies/ --output results.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 model import build_model
from preprocess import process_files, process_zip
from tensors import to_model_input


def discover(root: Path):
    """Study zips, or pre-extracted <study>/**/atlas*/*.nii.gz folders."""
    zips = sorted(root.glob("*.zip"))
    if zips:
        return [(z.stem, z) for z in zips]
    out = []
    for d in sorted(p for p in root.iterdir() if p.is_dir()):
        vols = sorted(p for p in d.rglob("*.nii.gz")
                      if "/atlas" in str(p) and "/atlas_seg/" not in str(p) and "mask" not in str(p))
        if vols:
            out.append((d.name, vols))
    return out


def load_ensemble(weights_dir: Path, device):
    cfg = json.loads((weights_dir / "config.json").read_text())
    models = []
    for m in cfg["members"]:
        net = build_model(m["arch"], in_ch=m["in_channels"], num_classes=len(cfg["labels"]))
        net.load_state_dict(load_file(str(weights_dir / m["weights_file"])))
        models.append((net.to(device).eval(), m["geometry"]))
    return cfg, models


@torch.no_grad()
def ensemble_probabilities(study, models, device) -> np.ndarray | None:
    """Mean over members of sigmoid(mean of logits for the volume and its left-right mirror)."""
    geos = sorted({g for _, g in models})
    per, n_used = process_zip(study, geos) if isinstance(study, Path) else process_files(study, geos)
    if n_used == 0:
        return None
    inputs = {}
    for g, (vol, presence) in per.items():
        x = to_model_input(torch.from_numpy(vol)[None].to(device),
                           torch.from_numpy(presence.astype(np.float32))[None].to(device))
        inputs[g] = (x, torch.flip(x, dims=[2]))
    acc = 0.0
    for net, g in models:
        x, xf = inputs[g]
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            logits = (net(x).float() + net(xf).float()) / 2.0
        acc = acc + torch.sigmoid(logits)[0].cpu().numpy().astype(np.float64)
    return acc / len(models)


def move_threshold_to_half(p: np.ndarray, t: np.ndarray) -> np.ndarray:
    p = np.clip(p, 1e-9, 1 - 1e-9)
    t = np.clip(t, 1e-6, 1 - 1e-6)
    num = p * (1.0 - t)
    return np.clip(num / np.maximum(num + (1.0 - p) * t, 1e-15), 1e-9, 1 - 1e-9)


def apply_positive_rates(probs: np.ndarray, rates: np.ndarray) -> np.ndarray:
    n = probs.shape[0]
    thr = np.empty(probs.shape[1])
    for c in range(probs.shape[1]):
        k = max(1, min(n - 1, int(round(float(rates[c]) * n))))
        s = np.sort(probs[:, c])[::-1]
        thr[c] = (s[k - 1] + s[k]) / 2.0
    return move_threshold_to_half(probs, np.clip(thr, 1e-6, 1 - 1e-6))


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, models = load_ensemble(args.weights, device)
    labels = cfg["labels"]
    rates = np.array([cfg["positive_rates"][l] for l in labels])

    names, rows, failed = [], [], []
    for i, (name, study) in enumerate(discover(args.input), 1):
        p = ensemble_probabilities(study, models, device)
        names.append(f"{name}.nii.gz")
        rows.append(p)
        if p is None:
            failed.append(len(rows) - 1)
        print(f"[{i}] {name}{' (no usable series)' if p is None else ''}", flush=True)

    real = [i for i in range(len(rows)) if rows[i] is not None]
    if len(real) < 2:
        raise RuntimeError("need at least 2 decoded studies")
    probs = np.full((len(rows), len(labels)), 0.5)
    probs[real] = apply_positive_rates(np.stack([rows[i] for i in real]), rates)

    preds = [{"input_image_name": n, "probabilities": {l: float(v) for l, v in zip(labels, row)}}
             for n, row in zip(names, probs)]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps([{"outputs": [{"value": {
        "name": "Generated probabilities", "type": "Abnormality Classification",
        "version": {"major": 1, "minor": 0}, "predictions": preds}}]}]))


if __name__ == "__main__":
    main()