"""Recompute a run's val/test probabilities in dataset order, on one GPU.

train.py evaluates under DDP, which returns rows in rank-interleaved order with
a padded tail. Combining runs needs rows aligned by study id, so this script
re-scores val and test on one GPU with a sequential loader and saves the uids.
"""

from __future__ import annotations

import argparse
import os
import sys

import numpy as np
import torch
from torch.utils.data import DataLoader

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from dataset import AtlasDataset, ShardIndex, to_model_input  # noqa: E402
from model import build_model  # noqa: E402
from train import build_splits, load_tables  # noqa: E402


@torch.no_grad()
def infer(model, loader, device, tta: bool = True):
    P = []
    for xu, _y, pres in loader:
        x = to_model_input(xu.to(device, non_blocking=True),
                           pres.to(device, non_blocking=True))
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            lg = model(x).float()
            if tta:
                lg = (lg + model(torch.flip(x, dims=[2])).float()) / 2.0
        P.append(torch.sigmoid(lg).cpu())
    return torch.cat(P).numpy().astype(np.float64)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True)
    ap.add_argument("--cache", required=True)
    ap.add_argument("--labels", required=True)
    ap.add_argument("--splits", required=True)
    ap.add_argument("--batch-size", type=int, default=48)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--no-tta", action="store_true")
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ck = torch.load(os.path.join(args.run, "best.pt"), map_location="cpu", weights_only=False)
    model = build_model(ck["arch"], num_classes=len(ck["labels"]),
                        in_ch=int(ck["model"]["stem.0.weight"].shape[1]))
    model.load_state_dict(ck["model"])
    model = model.to(device).eval()

    index = ShardIndex(args.cache)
    labels, split, cols, _scored = load_tables(args.labels, args.splits)
    if list(cols) != list(ck["labels"]):
        raise SystemExit("label columns differ from the checkpoint's")
    sp = build_splits(index, labels, split, len(cols))

    for name in ("val", "test"):
        ids, y = sp[name]
        ds = AtlasDataset(index, ids, y, ck.get("geometry", "d15"), train=False)
        dl = DataLoader(ds, batch_size=args.batch_size, shuffle=False,
                        num_workers=args.workers, pin_memory=True)
        probs = infer(model, dl, device, tta=not args.no_tta)
        uids = np.array([index.uids[i] for i in ids])
        assert len(probs) == len(uids) == len(y)
        out = os.path.join(args.run, f"{name}_probs_fixed.npz")
        np.savez_compressed(out, probs=probs, y=y, uids=uids)
        print(f"[{os.path.basename(args.run)}] {name}: {len(uids)} studies -> {out}",
              flush=True)
    return 0


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