"""Real-IAD data loading with multi-view (sample) grouping.

Official annotation JSONs define train/test. Filenames:
    OK/S0003/usb_0003_OK_C1_20231019140654.jpg
    NG/BX/S0001/usb_0001_NG_BX_C1_20231020105442.jpg
A physical sample = the image's parent directory; every sample has 5 views (C1..C5).
"""
from __future__ import annotations

import json
import os
from dataclasses import dataclass, field

from PIL import Image

Image.MAX_IMAGE_PIXELS = None


@dataclass
class Sample:
    """One physical object instance with its 5 camera views."""
    sample_id: str          # e.g. "usb/OK/S0003" or "usb/NG/BX/S0001"
    category: str
    label: int              # 0 = OK, 1 = anomaly
    anomaly_class: str      # "OK", "BX", ...
    image_paths: list = field(default_factory=list)   # absolute paths, sorted by camera
    mask_paths: list = field(default_factory=list)

    @property
    def n_views(self):
        return len(self.image_paths)


def _load_split_json(category: str, json_root: str) -> dict:
    path = os.path.join(json_root, f"{category}.json")
    with open(path) as f:
        return json.load(f)


def load_samples(category: str, image_root: str, json_root: str, split: str) -> list:
    """Return samples for one category/split, grouped by parent directory."""
    d = _load_split_json(category, json_root)
    base = os.path.join(image_root, category)
    groups = {}
    for rec in d[split]:
        rel = rec["image_path"]                     # e.g. "OK/S0003/xxx_C1_t.jpg"
        parent = os.path.dirname(rel)               # "OK/S0003"
        key = f"{category}/{parent}"
        g = groups.setdefault(key, Sample(
            sample_id=key, category=category,
            label=0 if rec["anomaly_class"] == "OK" else 1,
            anomaly_class=rec["anomaly_class"],
        ))
        g.image_paths.append(os.path.join(base, rel))
        mp = rec.get("mask_path")
        g.mask_paths.append(os.path.join(base, mp) if mp else None)
    out = list(groups.values())
    for s in out:
        order = sorted(range(s.n_views), key=lambda i: s.image_paths[i])
        s.image_paths = [s.image_paths[i] for i in order]
        s.mask_paths = [s.mask_paths[i] for i in order]
    return out


def subsample_samples(samples: list, n_ok: int, n_ng: int, seed: int = 0) -> list:
    """Deterministic stratified pick of samples (keeps all 5 views)."""
    import random
    rng = random.Random(seed)
    ok = [s for s in samples if s.label == 0]
    ng = [s for s in samples if s.label == 1]
    rng.shuffle(ok)
    rng.shuffle(ng)
    return ok[:n_ok] + ng[:n_ng]


def split_calibration(samples: list, calib_frac: float = 0.25, seed: int = 0):
    """Stratified sample-level split into (calibration, eval)."""
    import random
    rng = random.Random(seed)
    ok = [s for s in samples if s.label == 0]
    ng = [s for s in samples if s.label == 1]
    rng.shuffle(ok)
    rng.shuffle(ng)
    n_ok, n_ng = int(len(ok) * calib_frac), int(len(ng) * calib_frac)
    calib = ok[:n_ok] + ng[:n_ng]
    ev = ok[n_ok:] + ng[n_ng:]
    return calib, ev