from __future__ import annotations

import glob
from dataclasses import dataclass
from pathlib import Path

import numpy as np


@dataclass(frozen=True)
class MeasuredAIFPool:
    aif: np.ndarray
    time: np.ndarray
    sources: tuple[str, ...]

    @property
    def n_items(self) -> int:
        return int(self.aif.shape[0])

    @property
    def n_frames(self) -> int:
        return int(self.aif.shape[1])

    def sample(self, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray, str]:
        idx = int(rng.integers(0, self.n_items))
        return self.aif[idx], self.time[idx], self.sources[idx]


def resolve_measured_files(
    entries: list[str] | tuple[str, ...], *, base_dir: Path | None
) -> list[Path]:
    paths: list[Path] = []
    for entry in entries:
        raw = str(entry)
        path = Path(raw)
        if not path.is_absolute() and base_dir is not None:
            path = base_dir / path

        if any(ch in str(path) for ch in "*?[]"):
            matched = [Path(p) for p in sorted(glob.glob(str(path)))]
            paths.extend(matched)
            continue

        if path.is_dir():
            paths.extend(sorted(path.glob("*.npz")))
            continue

        paths.append(path)

    return paths


def _load_single_aif_file(path: Path) -> tuple[np.ndarray, np.ndarray]:
    if not path.exists():
        raise FileNotFoundError(f"Measured AIF file not found: {path}")

    with np.load(path) as data:
        if "AIF" not in data or "TIME" not in data:
            raise KeyError(
                f"{path} must contain 'AIF' and 'TIME' arrays (npz keys)."
            )
        aif = np.asarray(data["AIF"], dtype=np.float32).reshape(-1)
        time = np.asarray(data["TIME"], dtype=np.float32).reshape(-1)

    if aif.size < 2:
        raise ValueError(f"{path} has too few frames. Expected at least 2.")
    if aif.shape != time.shape:
        raise ValueError(
            f"{path} has mismatched shapes: AIF={aif.shape}, TIME={time.shape}."
        )
    if not np.isfinite(aif).all() or not np.isfinite(time).all():
        raise ValueError(f"{path} contains non-finite values.")
    if np.any(time[1:] <= time[:-1]):
        raise ValueError(f"{path} TIME must be strictly increasing.")

    return aif, time


def load_measured_aif_pool(
    entries: list[str] | tuple[str, ...], *, base_dir: Path | None = None
) -> MeasuredAIFPool:
    files = resolve_measured_files(entries, base_dir=base_dir)
    if not files:
        raise ValueError(
            "No measured AIF files resolved from config 'aif.measured_files'."
        )

    aif_list: list[np.ndarray] = []
    time_list: list[np.ndarray] = []
    sources: list[str] = []

    expected_frames: int | None = None
    for path in files:
        aif, time = _load_single_aif_file(path)
        n_frames = int(aif.shape[0])
        if expected_frames is None:
            expected_frames = n_frames
        elif n_frames != expected_frames:
            raise ValueError(
                "All measured AIF/TIME vectors must have the same frame count. "
                f"Expected {expected_frames}, got {n_frames} for {path}."
            )

        aif_list.append(aif)
        time_list.append(time)
        sources.append(str(path))

    aif_arr = np.stack(aif_list, axis=0).astype(np.float32, copy=False)
    time_arr = np.stack(time_list, axis=0).astype(np.float32, copy=False)
    return MeasuredAIFPool(aif=aif_arr, time=time_arr, sources=tuple(sources))

