"""MR-RATE atlas study -> 4-channel uint8 volume + 4 presence flags.

Every series of a study is registered to the MNI152 1 mm grid (193x229x193), so
the series are voxel-aligned. Each MR type becomes one channel: the mean of up
to four of its series, each scaled to [0, 1] first. The same code builds the
training cache and runs at inference.
"""
from __future__ import annotations

import gzip
import io
import re
import zipfile

import nibabel as nib
import numpy as np
from scipy.ndimage import zoom

MODALITIES = ("t1w", "t2w", "flair", "swi")     # channel order is part of the model
SRC_SHAPE = (193, 229, 193)
MAX_PER_MODALITY = 4

# d15: head box cropped from the atlas, linear zoom to 1.5 mm
# d2:  near-full field of view, exact 2x2x2 block average to 2 mm
GEOMETRIES = {
    "d15": {"offset": (12, 16, 6), "crop": (168, 192, 156), "out": (112, 128, 104), "mode": "zoom"},
    "d2": {"offset": (0, 2, 0), "crop": (192, 224, 192), "out": (96, 112, 96), "mode": "block"},
}

_MOD_RE = re.compile(r"_atlas_([a-z0-9]+)-", re.IGNORECASE)
_MOD_ALT = {m: re.compile(rf"(?:^|[^a-z]){m}(?:[^a-z]|$)", re.IGNORECASE) for m in MODALITIES}


def modality_of(name: str) -> str | None:
    stem = name.rsplit("/", 1)[-1]
    m = _MOD_RE.search(stem)
    tok = m.group(1).lower() if m else None
    if tok in MODALITIES:
        return tok
    return next((mod for mod, rx in _MOD_ALT.items() if rx.search(stem)), None)


def normalize(vol: np.ndarray) -> np.ndarray:
    """Min-max to [0, 1] between the 1st and 99.5th percentile of non-zero voxels.

    Registration leaves exact zeros outside the head, so non-zero voxels stand in for a head mask.
    """
    fg = vol[vol > 0]
    if fg.size < 1000:
        return np.zeros_like(vol)
    lo, hi = np.percentile(fg, (1.0, 99.5))
    if not np.isfinite(hi) or hi <= lo:
        lo, hi = float(fg.min()), float(fg.max())
        if hi <= lo:
            return np.zeros_like(vol)
    return np.clip((vol - lo) / (hi - lo), 0.0, 1.0)


def _crop(vol: np.ndarray, geometry: str) -> np.ndarray:
    g = GEOMETRIES[geometry]
    offset, shape = g["offset"], g["crop"]
    if vol.shape == SRC_SHAPE:
        return vol[tuple(slice(o, o + c) for o, c in zip(offset, shape))]
    if vol.shape == shape:
        return vol
    out = np.zeros(shape, dtype=vol.dtype)            # unexpected grid: centre crop or pad
    src, dst = [], []
    for s, t in zip(vol.shape, shape):
        o = abs(s - t) // 2
        src.append(slice(o, o + t) if s >= t else slice(0, s))
        dst.append(slice(0, t) if s >= t else slice(o, o + s))
    out[tuple(dst)] = vol[tuple(src)]
    return out


def resample(vol: np.ndarray, geometry: str) -> np.ndarray:
    g = GEOMETRIES[geometry]
    v = _crop(vol, geometry)
    d, h, w = g["out"]
    if g["mode"] == "block":
        return v.reshape(d, 2, h, 2, w, 2).mean(axis=(1, 3, 5))
    out = zoom(v, tuple(o / c for o, c in zip(g["out"], v.shape)), order=1, prefilter=False)
    if out.shape != g["out"]:                          # zoom can be off by one
        pad = np.zeros(g["out"], dtype=out.dtype)
        sl = tuple(slice(0, min(a, b)) for a, b in zip(out.shape, g["out"]))
        pad[sl] = out[sl]
        out = pad
    return out


def _read_normalized(read, name):
    try:
        vol = read(name)
    except Exception:
        return None
    if vol is None:
        return None
    if vol.ndim == 4:
        vol = vol[..., 0]
    return normalize(vol) if vol.ndim == 3 else None


def process_study(names, read, geometries=("d15",)):
    """Series names + reader -> ({geometry: (volume uint8 [4, *out], presence uint8 [4])}, n_series_used)."""
    cache = {}

    def load(n):
        if n not in cache:
            cache[n] = _read_normalized(read, n)
        return cache[n]

    per_mod, unmatched, n_used = {}, [], 0
    for name in sorted(names):
        mod = modality_of(name)
        if mod is None:
            unmatched.append(name)
            continue
        bucket = per_mod.setdefault(mod, [])
        if len(bucket) < MAX_PER_MODALITY and load(name) is not None:
            bucket.append(name)
            n_used += 1
    if n_used == 0 and unmatched:                      # unrecognised names: use them as T1w rather than nothing
        for name in unmatched[:MAX_PER_MODALITY]:
            if load(name) is not None:
                per_mod.setdefault(MODALITIES[0], []).append(name)
                n_used += 1

    out = {}
    for geo in geometries:
        vol = np.zeros((len(MODALITIES), *GEOMETRIES[geo]["out"]), dtype=np.uint8)
        presence = np.zeros(len(MODALITIES), dtype=np.uint8)
        for i, mod in enumerate(MODALITIES):
            series = [resample(cache[n], geo) for n in per_mod.get(mod, [])]
            if series:
                acc = series[0] if len(series) == 1 else np.mean(series, axis=0)
                vol[i] = np.clip(np.rint(acc * 255.0), 0, 255).astype(np.uint8)
                presence[i] = 1
        out[geo] = (vol, presence)
    return out, n_used


def _load_nii_bytes(raw: bytes) -> np.ndarray:
    fh = nib.FileHolder(fileobj=gzip.GzipFile(fileobj=io.BytesIO(raw)))
    img = nib.Nifti1Image.from_file_map({"header": fh, "image": fh})
    return np.asanyarray(img.dataobj, dtype=np.float32)


def process_zip(path, geometries=("d15",)):
    """A study zip (training layout <id>/atlas_img/, challenge layout <id>/atlas/)."""
    with zipfile.ZipFile(path) as zf:
        names = [n for n in zf.namelist() if n.endswith((".nii.gz", ".nii")) and "/atlas" in n
                 and "/atlas_seg/" not in n and "mask" not in n]
        return process_study(names, lambda n: _load_nii_bytes(zf.read(n)), geometries)


def process_files(paths, geometries=("d15",)):
    return process_study([str(p) for p in paths],
                         lambda p: np.asanyarray(nib.load(p).dataobj, dtype=np.float32), geometries)