"""Real-IAD D3 (CVPR 2025) adapter: RGB / PS / XYZ / PCD as a 4-agent swarm.

Each D3 sample directory holds one file per modality:
    RGBL05 (jpg)  - color image
    PS    (jpg)   - photometric-stereo surface map (pseudo-3D normals)
    XYZ   (tiff)  - surface height/coordinate map
    PCD   (pcd)   - registered point cloud
Splits per category: train=200, test=150, validation=150 (all-labeled).
The fusion layer treats the four modality scores exactly like camera views:
agents that "lie" (miss the defect or cry wolf) are inverted into information.
"""
from __future__ import annotations

import json
import os
from dataclasses import dataclass, field

MODALITIES = ("rgb", "ps", "xyz", "pcd")


@dataclass
class D3Sample:
    sample_id: str
    category: str
    label: int
    anomaly_class: str
    paths: dict = field(default_factory=dict)   # modality -> abs path

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


def load_d3_samples(category: str, image_root: str, json_root: str, split: str,
                    modalities=MODALITIES) -> list:
    """modalities: subset of MODALITIES; pcd requires open3d at scoring time."""
    path = os.path.join(json_root, f"{category}.json")
    with open(path) as f:
        d = json.load(f)
    base = os.path.join(image_root, category)
    keymap = {"rgb": "image_path", "ps": "ps_path", "xyz": "xyz_path", "pcd": "pcd_path"}
    samples = {}
    for rec in d[split]:
        parent = os.path.dirname(rec["image_path"])
        key = f"{category}/{parent}"
        s = samples.setdefault(key, D3Sample(sample_id=key, category=category,
                                             label=0 if rec["anomaly_class"] == "OK" else 1,
                                             anomaly_class=rec["anomaly_class"]))
        for m in modalities:
            rel = rec.get(keymap[m])
            if rel:
                s.paths[m] = os.path.join(base, rel)
    return list(samples.values())