analysis/datasets.py

"""
Датасеты и вспомогательные структуры данных для KADID-10k. Взяты из репозитория PatchSAE.
"""

import gzip
import re
import json
import os
from functools import lru_cache
from pathlib import Path
from typing import List, Optional, Sequence

import numpy as np
import pandas as pd
import torch
from PIL import Image, ImageColor
from scipy.special import softmax
from torch.utils.data import Dataset
from torchvision import transforms

from log_config import get_logger

logger = get_logger(__name__)

distortion_types_mapping = {
    1:  "gaussian_blur",
    2:  "lens_blur",
    3:  "motion_blur",
    4:  "color_diffusion",
    5:  "color_shift",
    6:  "color_quantization",
    7:  "color_saturation_1",
    8:  "color_saturation_2",
    9:  "jpeg2000",
    10: "jpeg",
    11: "white_noise",
    12: "white_noise_color_component",
    13: "impulse_noise",
    14: "multiplicative_noise",
    15: "denoise",
    16: "brighten",
    17: "darken",
    18: "mean_shift",
    19: "jitter",
    20: "non_eccentricity_patch",
    21: "pixelate",
    22: "quantization",
    23: "color_block",
    24: "high_sharpen",
    25: "contrast_change",
}

available_distortions = {
    "gaussian_blur":              "blur",
    "lens_blur":                  "blur",
    "motion_blur":                "blur",
    "color_diffusion":            "color_distortion",
    "color_shift":                "color_distortion",
    "color_quantization":         "color_distortion",
    "color_saturation_1":         "color_distortion",
    "color_saturation_2":         "color_distortion",
    "jpeg2000":                   "jpeg",
    "jpeg":                       "jpeg",
    "white_noise":                "noise",
    "white_noise_color_component": "noise",
    "impulse_noise":              "noise",
    "multiplicative_noise":       "noise",
    "denoise":                    "noise",
    "brighten":                   "brightness_change",
    "darken":                     "brightness_change",
    "mean_shift":                 "brightness_change",
    "jitter":                     "spatial_distortion",
    "non_eccentricity_patch":     "spatial_distortion",
    "pixelate":                   "spatial_distortion",
    "quantization":               "spatial_distortion",
    "color_block":                "spatial_distortion",
    "high_sharpen":               "sharpness_contrast",
    "contrast_change":            "sharpness_contrast",
}

distortion_groups = {
    "blur":              ["gaussian_blur", "lens_blur", "motion_blur"],
    "color_distortion":   ["color_diffusion", "color_shift", "color_quantization",
                           "color_saturation_1", "color_saturation_2"],
    "jpeg":               ["jpeg2000", "jpeg"],
    "noise":              ["white_noise", "white_noise_color_component", "impulse_noise",
                           "multiplicative_noise", "denoise"],
    "brightness_change":  ["brighten", "darken", "mean_shift"],
    "spatial_distortion": ["jitter", "non_eccentricity_patch", "pixelate",
                           "quantization", "color_block"],
    "sharpness_contrast":  ["high_sharpen", "contrast_change"],
}

QGROUND_DISTORTION_TYPES = {
    'jitter': np.array(ImageColor.getrgb('#4b54e1')),
    'noise': np.array(ImageColor.getrgb('#93fff0')),
    'overexposure': np.array(ImageColor.getrgb('#cde55d')),
    'blur': np.array(ImageColor.getrgb('#e45c5c')),
    'low light': np.array(ImageColor.getrgb('#35e344')),
}

distortion_types_mapping_qground = {
    0: 'background',
    1: 'jitter',
    2: 'noise',
    3: 'overexposure',
    4: 'blur',
    5: 'low light',
}

available_distortions_qground = {
    'background': 'background',
    'jitter': 'jitter',
    'noise': 'noise',
    'overexposure': 'overexposure',
    'blur': 'blur',
    'low light': 'low light',
}

SRGROUND_CLASS_ORDER = (
    'no_distortion',
    'blur',
    'jitter',
    'lowlight',
    'noise',
    'overexposure',
    'sr_artifact',
)

SRGROUND_DISTORTION_TYPES = {
    'blur': np.array(ImageColor.getrgb('#e45c5c')),
    'jitter': np.array(ImageColor.getrgb('#4b54e1')),
    'lowlight': np.array(ImageColor.getrgb('#35e344')),
    'noise': np.array(ImageColor.getrgb('#93fff0')),
    'overexposure': np.array(ImageColor.getrgb('#cde55d')),
    'sr_artifact': np.array(ImageColor.getrgb('#c000a0')),
}

distortion_types_mapping_srground = {
    0: 'background',
    1: 'blur',
    2: 'jitter',
    3: 'lowlight',
    4: 'noise',
    5: 'overexposure',
    6: 'sr_artifact',
}

available_distortions_srground = {
    'background': 'background',
    'blur': 'blur',
    'jitter': 'jitter',
    'lowlight': 'lowlight',
    'noise': 'noise',
    'overexposure': 'overexposure',
    'sr_artifact': 'sr_artifact',
}

SRGROUND_SR_ARTIFACT_THRESHOLD = 0.3


def _load_npy_gz(path: Path) -> np.ndarray:
    with gzip.open(path, 'rb') as handle:
        return np.load(handle)


def _real_distortion_labels(
    real_maps: np.ndarray,
    prominences: Optional[object] = None,
) -> np.ndarray:
    real_maps = np.asarray(real_maps, dtype=np.float64)
    real_prom = np.array(prominences)[:-1, None, None]
    real_prob = softmax(real_maps, axis=0) * real_prom
    return np.argmax(real_prob, axis=0).astype(np.uint8)


def _sr_artifact_labels(
    sr_maps: np.ndarray,
    prominences: Optional[object] = None,
    threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
) -> np.ndarray:
    sr_maps = np.asarray(sr_maps, dtype=np.float64)
    if sr_maps.ndim == 3:
        if sr_maps.shape[0] == 1:
            sr_maps = sr_maps[0]
        else:
            raise ValueError(f'Expected SR artifact maps with shape (1, H, W) or (H, W), got {sr_maps.shape}')

    sr_prom = prominences[-1]
    return np.where(sr_maps * sr_prom >= threshold, 6, 0).astype(np.uint8)


def merge_srground_masks(
    annot_rd: np.ndarray | None,
    annot_sr: np.ndarray | None,
) -> np.ndarray:
    """Merge real-distortion and SR-artifact label maps (0=background, 1..6=classes).

    If ``annot_sr`` is None, returns ``annot_rd`` unchanged (no SR merge).
    """
    if annot_sr is None:
        if annot_rd is None:
            raise ValueError('merge_srground_masks requires annot_rd when annot_sr is None')
        return annot_rd.astype(np.uint8, copy=False)
    if annot_rd is None:
        return annot_sr.astype(np.uint8, copy=False)
    return np.where(annot_sr == 6, 6, annot_rd).astype(np.uint8)


def label2rgb_srground(mask_label: np.ndarray) -> np.ndarray:
    mask_rgb = np.zeros(mask_label.shape + (3,), dtype=np.uint8)
    for label_id, dist_name in distortion_types_mapping_srground.items():
        if label_id == 0:
            continue
        mask_rgb[mask_label == label_id] = SRGROUND_DISTORTION_TYPES[dist_name]
    return mask_rgb


SRGROUND_LEGEND_LABELS = {
    'blur': 'blur',
    'jitter': 'jitter',
    'lowlight': 'low light',
    'noise': 'noise',
    'overexposure': 'overexposure',
    'sr_artifact': 'SR artifact',
}


def _parse_prominences(value: object) -> np.ndarray | None:
    """Convert a prominences field from JSON/CSV into a 1D float array.

    Does not read any files — callers load ``prominences`` from ``srground_train.json``
    (or another source) and pass the cell value here.
    """
    if value is None:
        return None
    if isinstance(value, float) and np.isnan(value):
        return None
    if isinstance(value, str):
        text = value.strip()
        if not text:
            return None
        try:
            value = json.loads(text)
        except json.JSONDecodeError:
            return None
    try:
        return np.asarray(value, dtype=np.float64)
    except (TypeError, ValueError):
        return None


def _center_crop_label_map(mask_label: np.ndarray, crop_size: int, reference_size: tuple[int, int]) -> np.ndarray:
    """Center-crop a label map to match heatmap preprocessing (reference_size is W×H)."""
    width, height = reference_size
    mask_img = Image.fromarray(mask_label.astype(np.uint8), mode='L')
    if mask_label.shape[:2] != (height, width):
        mask_img = mask_img.resize((width, height), resample=Image.NEAREST)
    return np.asarray(transforms.CenterCrop(int(crop_size))(mask_img), dtype=np.uint8)


@lru_cache(maxsize=8)
def _srground_train_dataframe(datasets_root: str) -> pd.DataFrame:
    from analysis.config import dataset_images_root as _dataset_images_root

    root = Path(_dataset_images_root(datasets_root, 'SRGround'))
    return pd.read_json(root / 'srground_train.json')


def _resolve_datasets_root(datasets_root: str | None) -> str:
    if datasets_root is not None:
        return str(datasets_root)
    from analysis.config import load_sae_vis_config

    return str(load_sae_vis_config().DATASETS_ROOT)


def srground_image_key(path: str | Path, *, datasets_root: str | None = None) -> str:
    """Normalize a path to ``image_path`` keys used in ``srground_train.json`` (relative to SRGround root)."""
    from analysis.config import dataset_images_root

    root = _resolve_datasets_root(datasets_root)
    path_obj = Path(path)
    sr_root = Path(dataset_images_root(root, 'SRGround'))
    if path_obj.is_absolute():
        return _to_relative_dataset_path(path_obj, sr_root)
    return path_obj.as_posix()


@lru_cache(maxsize=8)
def srground_prominences_index(datasets_root: str | None = None) -> dict[str, np.ndarray]:
    """Cached ``image_path`` → prominences array from ``srground_train.json``."""
    root = _resolve_datasets_root(datasets_root)
    df = _srground_train_dataframe(root)
    out: dict[str, np.ndarray] = {}
    for image_path, raw_prom in zip(df['image_path'].astype(str), df['prominences']):
        prom = _parse_prominences(raw_prom)
        if prom is not None:
            out[str(image_path)] = prom
    return out


def _image_rel_from_meta_row(meta_row: pd.Series | None) -> str | None:
    if meta_row is None:
        return None
    for column in ('image_path', 'distorted_img_path'):
        if column not in meta_row:
            continue
        value = meta_row.get(column)
        if value is None or (isinstance(value, float) and np.isnan(value)):
            continue
        text = str(value).strip()
        if text:
            return text
    return None


def srground_rows_for_image_paths(
    image_paths: Sequence[str],
    *,
    datasets_root: str | None = None,
) -> pd.DataFrame:
    """Subset of ``srground_train.json`` for the given ``image_path`` keys."""
    if not image_paths:
        return pd.DataFrame()

    root = _resolve_datasets_root(datasets_root)
    keys = {srground_image_key(path, datasets_root=root) for path in image_paths if path}
    if not keys:
        return pd.DataFrame()

    df = _srground_train_dataframe(root)
    return df[df['image_path'].astype(str).isin(keys)].copy()


def srground_prominences_by_image_paths(
    image_paths: Sequence[str],
    *,
    dataset_root: str | Path | None = None,
    datasets_root: str | None = None,
) -> dict[str, np.ndarray]:
    """Look up prominences for paths (absolute or SRGround-relative) via cached index."""
    if not image_paths:
        return {}

    if datasets_root is None and dataset_root is not None:
        datasets_root = parent_datasets_root(dataset_root)
    root = _resolve_datasets_root(datasets_root)
    index = srground_prominences_index(root)
    keys = {srground_image_key(path, datasets_root=root) for path in image_paths if path}
    return {key: index[key] for key in keys if key in index}


def _meta_row_path(meta_row: pd.Series, column: str) -> str | None:
    if column not in meta_row:
        return None
    value = meta_row.get(column)
    if value is None or (isinstance(value, float) and np.isnan(value)):
        return None
    text = str(value).strip()
    return text or None


def parent_datasets_root(dataset_root: str | Path) -> str:
    """Parent directory that contains dataset folders (e.g. ``Kadid`` for ``Kadid/SRGround``)."""
    return str(Path(dataset_root).resolve().parent)


def _resolve_dataset_file_path(rel_path: str, *, dataset_root: str | Path) -> Path:
    path = Path(str(rel_path).strip())
    if path.is_absolute():
        return path
    return Path(dataset_root) / path


def infer_spatial_mask_dataset(meta_row: pd.Series | None) -> str | None:
    """Return ``QGround`` or ``SRGround`` when meta row carries spatial mask fields."""
    if meta_row is None:
        return None
    if _meta_row_path(meta_row, 'real_distortions_ann_path') or _meta_row_path(meta_row, 'sr_artifacts_ann_path'):
        return 'SRGround'
    if _meta_row_path(meta_row, 'mask_path'):
        return 'QGround'
    return None


def qground_mask_rgb_for_meta_row(
    meta_row: pd.Series | None,
    *,
    dataset_root: str | Path,
) -> np.ndarray | None:
    """Load QGround RGB segmentation mask PNG referenced from cache metadata."""
    if meta_row is None:
        return None

    mask_rel = _meta_row_path(meta_row, 'mask_path')
    if not mask_rel:
        return None

    mask_path = _resolve_dataset_file_path(mask_rel, dataset_root=dataset_root)
    if not mask_path.is_file():
        return None
    return np.asarray(Image.open(mask_path).convert('RGB'), dtype=np.uint8)


def _resize_label_map_to_image(
    mask_label: np.ndarray | None,
    reference_size: tuple[int, int],
) -> np.ndarray | None:
    if mask_label is None:
        return None
    width, height = reference_size
    if mask_label.shape[:2] == (height, width):
        return mask_label.astype(np.uint8, copy=False)
    mask_image = Image.fromarray(mask_label.astype(np.uint8), mode='L')
    mask_image = mask_image.resize((width, height), resample=Image.NEAREST)
    return np.asarray(mask_image, dtype=np.uint8)


def srground_label_map_for_meta_row(
    meta_row: pd.Series | None,
    *,
    dataset_root: str | Path,
    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
    prominences: np.ndarray | None = None,
) -> np.ndarray | None:
    """Build a full-frame SRGround label map from paths stored in cache metadata."""
    if meta_row is None:
        return None

    image_rel = _image_rel_from_meta_row(meta_row)
    real_rel = _meta_row_path(meta_row, 'real_distortions_ann_path')
    sr_rel = _meta_row_path(meta_row, 'sr_artifacts_ann_path')
    if not real_rel and not sr_rel:
        return None

    if prominences is None and image_rel:
        prominences = srground_prominences_by_image_paths(
            [image_rel],
            dataset_root=dataset_root,
        ).get(image_rel)

    annot_rd = None
    annot_sr = None
    if real_rel:
        real_path = _resolve_dataset_file_path(real_rel, dataset_root=dataset_root)
        if real_path.is_file():
            real_maps = _load_npy_gz(real_path)
            prom = prominences
            if prom is None:
                prom = np.ones(int(real_maps.shape[0]) + 1, dtype=np.float64)
            annot_rd = _real_distortion_labels(real_maps, prom)

    if sr_rel:
        sr_path = _resolve_dataset_file_path(sr_rel, dataset_root=dataset_root)
        if sr_path.is_file():
            sr_maps = _load_npy_gz(sr_path)
            prom = prominences
            if prom is None:
                prom = np.ones(6, dtype=np.float64)
            annot_sr = _sr_artifact_labels(
                sr_maps,
                prom,
                threshold=sr_artifact_threshold,
            )

    if annot_rd is None and annot_sr is None:
        return None

    reference_size = (annot_rd if annot_rd is not None else annot_sr).shape[1], (
        annot_rd if annot_rd is not None else annot_sr
    ).shape[0]
    if image_rel:
        image_path = _resolve_dataset_file_path(image_rel, dataset_root=dataset_root)
        if image_path.is_file():
            with Image.open(image_path) as image:
                reference_size = image.size

    annot_rd = _resize_label_map_to_image(annot_rd, reference_size)
    annot_sr = _resize_label_map_to_image(annot_sr, reference_size)
    return merge_srground_masks(annot_rd, annot_sr)


def annotation_mask_rgb_for_meta_row(
    meta_row: pd.Series | None,
    *,
    dataset_root: str | Path,
    dataset: str | None = None,
    crop_size: int = 512,
    full_frame: bool = True,
    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
    prominences: np.ndarray | None = None,
) -> np.ndarray | None:
    """RGB annotation mask for QGround / SRGround visualization (QGround-style full frame by default)."""
    
    if dataset_root is None:
        from analysis.config import load_sae_vis_config
        cfg = load_sae_vis_config()
        datasets_root = str(cfg.DATASETS_ROOT)
        dataset_root = Path(datasets_root) / dataset

    dataset_name = dataset or infer_spatial_mask_dataset(meta_row)
    if dataset_name == 'QGround':
        return qground_mask_rgb_for_meta_row(meta_row, dataset_root=dataset_root)
    if dataset_name == 'SRGround':
        label_map = srground_label_map_for_meta_row(
            meta_row,
            dataset_root=dataset_root,
            sr_artifact_threshold=sr_artifact_threshold,
            prominences=prominences,
        )
        if label_map is None:
            return None
        if not full_frame:
            image_rel = _image_rel_from_meta_row(meta_row)
            reference_size = (label_map.shape[1], label_map.shape[0])
            if image_rel:
                image_path = _resolve_dataset_file_path(image_rel, dataset_root=dataset_root)
                if image_path.is_file():
                    with Image.open(image_path) as image:
                        reference_size = image.size
            label_map = _center_crop_label_map(label_map, crop_size, reference_size)
        return label2rgb_srground(label_map)
    return None


def get_srground_rgb_mask(
    meta_row: pd.Series | None,
    *,
    dataset_root: str | Path,
    crop_size: int = 224,
    full_frame: bool = True,
    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
) -> np.ndarray | None:
    """RGB SRGround mask for one cache meta row (prominences from ``srground_train.json``)."""
    return annotation_mask_rgb_for_meta_row(
        meta_row,
        dataset_root=dataset_root,
        dataset='SRGround',
        crop_size=crop_size,
        full_frame=full_frame,
        sr_artifact_threshold=sr_artifact_threshold,
    )


def srground_mask_rgb_for_meta_row(
    meta_row: pd.Series | None,
    *,
    dataset_root: str | Path,
    crop_size: int = 224,
    sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
) -> np.ndarray | None:
    """Center-cropped SRGround mask preview (alias of :func:`get_srground_rgb_mask`)."""
    return get_srground_rgb_mask(
        meta_row,
        dataset_root=dataset_root,
        crop_size=crop_size,
        full_frame=False,
        sr_artifact_threshold=sr_artifact_threshold,
    )


def _rgb2label_qground(mask_rgb: np.ndarray) -> np.ndarray:
    mask_label = np.zeros(mask_rgb.shape[:2], dtype=np.uint8)
    for label, rgb_code in enumerate(QGROUND_DISTORTION_TYPES.values(), start=1):
        matches = np.isclose(mask_rgb, rgb_code, rtol=0.2, atol=20).all(axis=-1)
        mask_label[matches] = label
    return mask_label


def _label2rgb_qground(mask_label: np.ndarray) -> np.ndarray:
    mask_rgb = np.zeros(mask_label.shape + (3,), dtype=np.uint8)
    for label, rgb_code in enumerate(QGROUND_DISTORTION_TYPES.values(), start=1):
        mask_rgb[mask_label == label] = rgb_code
    return mask_rgb


def _infer_kadid_original_path(distorted_path: Path) -> Path | None:
    match = re.search(r'(I\d+)_\d+_\d+\.png$', distorted_path.name)
    if match:
        return distorted_path.with_name(f'{match.group(1)}.png')
    return None


def _to_relative_dataset_path(path: Path, root: Path) -> str:
    try:
        return path.relative_to(root).as_posix()
    except ValueError:
        return path.as_posix()


class Kadid10kDataset(Dataset):
    """
    KADID-10k dataset. При семплинге применяется RandomCrop.
    """

    def __init__(
        self,
        root: str,
        crop_size: int = 224,
        min_distortion_level: int = 1,
        transform=None,
    ):
        self.root = Path(root)
        self.crop_size = crop_size
        self.mos_range = (1, 5)
        self.min_distortion_level = int(min_distortion_level)

        if not (1 <= self.min_distortion_level <= 5):
            raise ValueError(
                f"min_distortion_level must be in [1, 5], got {self.min_distortion_level}"
            )

        if transform is None:
            self.transform = transforms.Compose([
                transforms.RandomCrop(self.crop_size),
                transforms.ToTensor(),
            ])
        else:
            self.transform = transform

        scores_csv = pd.read_csv(self.root / "dmos.csv")
        scores_csv = scores_csv[["dist_img", "dmos"]]

        self.images = np.array([
            self.root / "images" / el
            for el in scores_csv["dist_img"].values.tolist()
        ])
        self.mos = np.array(scores_csv["dmos"].values.tolist())

        self.distortion_types = []
        self.distortion_groups = []
        self.distortion_levels = []

        for image in self.images:
            match = re.search(r'I\d+_(\d+)_(\d+)\.png$', str(image))
            dist_type = distortion_types_mapping[int(match.group(1))]
            self.distortion_types.append(dist_type)
            self.distortion_groups.append(available_distortions[dist_type])
            self.distortion_levels.append(int(match.group(2)))

        self.distortion_types = np.array(self.distortion_types)
        self.distortion_groups = np.array(self.distortion_groups)
        self.distortion_levels = np.array(self.distortion_levels)

        if self.min_distortion_level > 1:
            keep_mask = self.distortion_levels >= self.min_distortion_level
            self.images = self.images[keep_mask]
            self.mos = self.mos[keep_mask]
            self.distortion_types = self.distortion_types[keep_mask]
            self.distortion_groups = self.distortion_groups[keep_mask]
            self.distortion_levels = self.distortion_levels[keep_mask]

    def __len__(self) -> int:
        return len(self.images)

    def __getitem__(self, index: int) -> dict:
        img = Image.open(self.images[index]).convert("RGB")
        img = self.transform(img)
        original_path = _infer_kadid_original_path(Path(self.images[index]))
        return {
            "img":        img,
            "mos":        float(self.mos[index]),
            "dist_type":  self.distortion_types[index],
            "dist_group": self.distortion_groups[index],
            "dist_level": int(self.distortion_levels[index]),
            "distorted_img_path": _to_relative_dataset_path(Path(self.images[index]), self.root),
            "original_img_path": _to_relative_dataset_path(original_path, self.root) if original_path is not None else None,
        }


class LocalKadidPresavedDataset(Dataset):
    """Presaved KADID dataset with local distortions and binary masks.

    Expects a dataset root directory produced by the local-distortion presave script.
    The root must contain index.csv.
    Required columns:
        distorted_img_path, mask_path
    Optional metadata columns are returned as-is in each sample.
    """

    def __init__(
        self,
        root: str,
        crop_size: Optional[int] = 224,
    ):
        self.root = Path(root)
        self.index_path = self.root / "index.csv"
        self.index_dir = self.root
        if not self.index_path.exists():
            raise FileNotFoundError(f"index.csv not found in local_kadid root: {self.root}")
        self.df = pd.read_csv(self.index_path)

        required_cols = {"distorted_img_path", "mask_path"}
        missing = required_cols - set(self.df.columns)
        if missing:
            raise ValueError(f"Index is missing columns: {sorted(missing)}")

        self.image_to_tensor = transforms.ToTensor()
        self.mask_to_tensor = transforms.ToTensor()
        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None

    def __len__(self) -> int:
        return len(self.df)

    def _resolve_path(self, value: str) -> Path:
        p = Path(value)
        return p if p.is_absolute() else (self.index_dir / p)

    def __getitem__(self, index: int) -> dict:
        row = self.df.iloc[index]

        img_path = self._resolve_path(str(row["distorted_img_path"]))
        mask_path = self._resolve_path(str(row["mask_path"]))

        img = Image.open(img_path).convert("RGB")
        mask = Image.open(mask_path).convert("L")

        img = self.image_to_tensor(img)
        mask = self.mask_to_tensor(mask)
        mask = (mask > 0.5).to(img.dtype)

        if self.crop is not None:
            img = self.crop(img)
            mask = self.crop(mask)

        sample = {
            "img": img,
            "mask": mask,
        }

        for key, value in row.items():
            if key == "mask_path":
                continue
            if isinstance(value, np.generic):
                value = value.item()
            if key in ("distorted_img_path", "original_img_path"):
                value = _to_relative_dataset_path(Path(value), self.root)
            sample[key] = value

        return sample


def kadid_collate_fn(batch: List[dict]) -> dict:
    """
    Collate для Kadid10kDataset.

    Возвращает:
        images: Tensor (B, C, H, W)
        + все остальные ключи как списки длины B
    """
    images = torch.stack([item["img"] for item in batch], dim=0)
    collated: dict = {"images": images}
    for key in batch[0]:
        if key == "img":
            continue
        collated[key] = [item[key] for item in batch]
    return collated


def local_kadid_collate_fn(batch: List[dict]) -> dict:
    """Collate for LocalKadidPresavedDataset.

    Returns:
        images: Tensor (B, C, H, W)
        masks:  Tensor (B, 1, H, W)
        + remaining keys as lists with length B
    """
    images = torch.stack([item["img"] for item in batch], dim=0)
    masks = torch.stack([item["mask"] for item in batch], dim=0)

    collated: dict = {
        "images": images,
        "masks": masks,
    }
    for key in batch[0]:
        if key in ("img", "mask"):
            continue
        collated[key] = [item[key] for item in batch]
    return collated


class QGroundDataset(Dataset):
    """QGround dataset stored locally as JSON index files plus image/mask folders.

    Expected layout:
        root/
          qground_train.json
          qground_test.json
          images/
          masks/
    """

    def __init__(
        self,
        root: str,
        split: str = 'test',
        json_path: Optional[str] = None,
        crop_size: Optional[int] = 224,
        annotation_index: int = 0,
        transform=None,
    ):
        self.root = Path(root)
        self.images_root = self.root / 'images'
        self.masks_root = self.root / 'masks'
        self.split = str(split).strip().lower()
        self.annotation_index = int(annotation_index)

        if json_path is None:
            candidates = [
                self.root / f'qground_{self.split}.json',
                self.root / f'QGround_{self.split}.json',
            ]
            self.json_path = next((path for path in candidates if path.exists()), candidates[0])
        else:
            self.json_path = Path(json_path)
            if not self.json_path.is_absolute():
                self.json_path = self.root / self.json_path

        if not self.json_path.exists():
            raise FileNotFoundError(f'QGround split file not found: {self.json_path}')

        with self.json_path.open('r', encoding='utf-8') as handle:
            raw_samples = json.load(handle)

        if not isinstance(raw_samples, list):
            raise ValueError(f'QGround JSON must contain a list of samples: {self.json_path}')

        self.samples = []
        for sample in raw_samples:
            if not isinstance(sample, dict):
                continue

            ann_list = sample.get('ann_list') or []
            if isinstance(ann_list, dict):
                ann_list = [ann_list]
            if not ann_list:
                continue

            ann = ann_list[min(self.annotation_index, len(ann_list) - 1)]
            image_rel = sample.get('image')
            mask_rel = ann.get('segmentation_mask')
            if not image_rel or not mask_rel:
                continue

            self.samples.append({
                'sample_id': sample.get('id'),
                'image_rel': image_rel,
                'mask_rel': mask_rel,
                'ann_id': ann.get('id'),
                'quality_description': ann.get('quality_description'),
            })

        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None
        self.image_to_tensor = transforms.ToTensor() if transform is None else transform
        self.image_paths = [self._resolve_path(self.images_root, sample['image_rel']) for sample in self.samples]

    def __len__(self) -> int:
        return len(self.samples)

    def _resolve_path(self, base_dir: Path, rel_path: str) -> Path:
        rel = Path(str(rel_path))
        candidates = [
            base_dir / rel,
            self.root / rel,
            base_dir / rel.name,
            self.root / rel.name,
        ]
        for candidate in candidates:
            if candidate.exists():
                return candidate
        return candidates[0]

    def __getitem__(self, index: int) -> dict:
        sample = self.samples[index]

        img_path = self._resolve_path(self.images_root, sample['image_rel'])
        mask_path = self._resolve_path(self.masks_root, sample['mask_rel'])

        image = Image.open(img_path).convert('RGB')
        mask_image = Image.open(mask_path).convert('RGB')

        if self.crop is not None:
            image = self.crop(image)
            mask_image = self.crop(mask_image)

        image = self.image_to_tensor(image)
        if isinstance(image, Image.Image):
            image = transforms.ToTensor()(image)

        mask_rgb = np.asarray(mask_image, dtype=np.uint8)
        mask_label = _rgb2label_qground(mask_rgb)
        mask = torch.from_numpy(mask_label.astype(np.float32)).unsqueeze(0)

        return {
            'img': image,
            'mask': mask,
            'mos': float('nan'),
            'dist_level': 5,
            'mask_coverage': float((mask_label > 0).mean()),
            'qground_ann_id': sample['ann_id'],
            'sample_id': str(sample['sample_id'] or Path(sample['image_rel']).stem),
            'distorted_img_path': _to_relative_dataset_path(img_path, self.root),
            'original_img_path': '',  # no reference images in QGround
            'image_path': _to_relative_dataset_path(img_path, self.root),
            'mask_path': _to_relative_dataset_path(mask_path, self.root),
            'split': self.split,
        }


def qground_collate_fn(batch: List[dict]) -> dict:
    """Collate for QGroundDataset.

    Returns:
        images: Tensor (B, C, H, W)
        masks: Tensor (B, 1, H, W)
        + remaining keys as lists with length B
    """
    images = torch.stack([item['img'] for item in batch], dim=0)
    masks = torch.stack([item['mask'] for item in batch], dim=0)

    collated: dict = {
        'images': images,
        'masks': masks,
    }
    for key in batch[0]:
        if key in ('img', 'mask'):
            continue
        collated[key] = [item[key] for item in batch]
    return collated


class SRGroundSmallDataset(Dataset):
    """
    Dataset wrapper for SRGround JSON indexes such as `srground_train.json`.

    Expects entries with fields like `image_path`, `real_distortions_ann_path`,
    `sr_artifacts_ann_path`, `has_markup`, `prominences`.
    """

    def __init__(
        self,
        root: str,
        json_path: Optional[str] = None,
        require_markup: bool = True,
        require_sr: bool = True,
        allowed_methods: Optional[List[str]] = ['DiT4SR_x2'],
        crop_size: Optional[int] = None,
        sr_artifact_threshold: float = SRGROUND_SR_ARTIFACT_THRESHOLD,
        include_sr_artifact: bool = False,
        transform=None,
    ):
        self.root = Path(root)
        self.sr_artifact_threshold = float(sr_artifact_threshold)
        self.include_sr_artifact = bool(include_sr_artifact)
        self.json_path = self.root / 'srground_train.json'
        df = pd.read_json(self.json_path)

        if require_markup:
            df = df[df['has_markup'].fillna(False).astype(bool)]
        
        df = df[~df['image_path'].str.contains('blur')]

        df = df[df['image_path'].notna() & (df['image_path'].astype(str) != '')]

        if require_sr:
            df = df[df['image_path'].astype(str).str.contains('@SR@', na=False)]

        if allowed_methods is not None:
            method_pattern = '|'.join(re.escape(method) for method in allowed_methods)
            df = df[df['image_path'].astype(str).str.contains(method_pattern, na=False)]

        df = df.assign(sample_id=df['image_path'].astype(str).map(lambda path: Path(path).stem))
        df = df.reset_index(drop=True)


        self.df = df.copy()
        self.image_to_tensor = transforms.ToTensor() if transform is None else transform
        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None
        self.image_paths = [self.root / Path(path) for path in self.df['image_path'].tolist()]

    def __len__(self) -> int:
        return len(self.df)

    def _resolve_path(self, rel_path: Optional[str]) -> Path:
        return self.root / Path(str(rel_path))

    def _load_mask_labels(self, sample: dict) -> tuple[Optional[np.ndarray], Optional[np.ndarray]]:
        prominences = sample.get('prominences')
        annot_rd = None
        annot_sr = None

        real_path = sample.get('real_distortions_ann_path')
        real_ann_path = self._resolve_path(real_path)
        if real_ann_path.exists():
            annot_rd = _real_distortion_labels(_load_npy_gz(real_ann_path), prominences)

        if self.include_sr_artifact:
            sr_path = sample.get('sr_artifacts_ann_path')
            sr_ann_path = self._resolve_path(sr_path)
            if sr_ann_path.exists():
                annot_sr = _sr_artifact_labels(
                    _load_npy_gz(sr_ann_path),
                    prominences,
                    threshold=self.sr_artifact_threshold,
                )

        if annot_rd is None and annot_sr is None:
            return None, None
        return annot_rd, annot_sr

    def _resize_mask(self, mask_label: Optional[np.ndarray], image: Image.Image) -> Optional[np.ndarray]:
        if mask_label is None:
            return None
        if mask_label.shape[:2] == image.size[::-1]:
            return mask_label

        mask_image = Image.fromarray(mask_label.astype(np.uint8), mode='L')
        mask_image = mask_image.resize(image.size, resample=Image.NEAREST)
        return np.asarray(mask_image, dtype=np.uint8)

    def __getitem__(self, index: int) -> dict:
        sample = self.df.iloc[index].to_dict()
        img_path = self.image_paths[index]

        image = Image.open(img_path).convert('RGB')
        annot_rd, annot_sr = self._load_mask_labels(sample)
        annot_rd = self._resize_mask(annot_rd, image)
        annot_sr = self._resize_mask(annot_sr, image)

        if self.crop is not None:
            image = self.crop(image)
            if annot_rd is not None:
                mask_image = Image.fromarray(annot_rd.astype(np.uint8), mode='L')
                annot_rd = np.asarray(self.crop(mask_image), dtype=np.uint8)
            if annot_sr is not None:
                mask_image = Image.fromarray(annot_sr.astype(np.uint8), mode='L')
                annot_sr = np.asarray(self.crop(mask_image), dtype=np.uint8)

        img_tensor = self.image_to_tensor(image)
        mask_hw = (image.height, image.width)

        if annot_rd is None and annot_sr is None:
            mask_label = np.zeros(mask_hw, dtype=np.uint8)
        else:
            mask_label = merge_srground_masks(annot_rd, annot_sr)
        mask = torch.from_numpy(mask_label.astype(np.float32)).unsqueeze(0)
        mask_rd = (
            torch.from_numpy(annot_rd.astype(np.float32)).unsqueeze(0)
            if annot_rd is not None
            else torch.zeros((1, *mask_hw), dtype=torch.float32)
        )
        mask_sr = (
            torch.from_numpy(annot_sr.astype(np.float32)).unsqueeze(0)
            if annot_sr is not None
            else torch.zeros((1, *mask_hw), dtype=torch.float32)
        )
        mask_coverage = float((mask > 0).float().mean().item())

        real_ann_path = self._resolve_path(sample.get('real_distortions_ann_path'))
        sr_ann_path = None
        if self.include_sr_artifact:
            candidate = self._resolve_path(sample.get('sr_artifacts_ann_path'))
            if candidate.exists():
                sr_ann_path = candidate

        return {
            'img': img_tensor,
            'mask': mask,
            'mask_rd': mask_rd,
            'mask_sr': mask_sr,
            'mos': float('nan'),
            'dist_level': 5,
            'mask_coverage': mask_coverage,
            'prominences': sample.get('prominences'),
            'has_markup': sample.get('has_markup', False),
            'sr_artifacts_ann_path': (
                _to_relative_dataset_path(sr_ann_path, self.root) if sr_ann_path is not None else None
            ),
            'real_distortions_ann_path': _to_relative_dataset_path(real_ann_path, self.root),
            'sample_id': sample.get('sample_id'),
            'distorted_img_path': _to_relative_dataset_path(img_path, self.root),
            'image_path': _to_relative_dataset_path(img_path, self.root),
            'mask_path': _to_relative_dataset_path(real_ann_path, self.root),
        }


def srground_collate_fn(batch: List[dict]) -> dict:
    """Collate for SRGroundSmallDataset.

    Returns:
        images: Tensor (B, C, H, W)
        masks: Tensor (B, 1, H, W)
        + remaining keys as lists length B
    """
    images = torch.stack([item['img'] for item in batch], dim=0)
    masks = torch.stack([item['mask'] for item in batch], dim=0)

    collated: dict = {
        'images': images,
        'masks': masks,
    }
    for key in batch[0]:
        if key in ('img', 'mask'):
            continue
        collated[key] = [item[key] for item in batch]
    return collated


class KadidPristineDataset(Dataset):
    """
    KADID-10k pristine (reference) images dataset. 
    Возвращает только оригинальные изображения без искажений с RandomCrop.
    """

    def __init__(
        self,
        root: str,
        crop_size: int = 224,
        transform=None,
    ):
        self.root = Path(root)
        self.crop_size = crop_size

        if transform is None:
            self.transform = transforms.Compose([
                transforms.RandomCrop(self.crop_size),
                transforms.ToTensor(),
            ])
        else:
            self.transform = transform

        # Find all original images (I{number}.png format, no suffixes)
        images_dir = self.root / "images"
        pristine_pattern = re.compile(r'^I\d+\.png$')

        self.images = sorted([
            images_dir / fname
            for fname in os.listdir(images_dir)
            if pristine_pattern.match(fname)
        ])

        if len(self.images) == 0:
            raise ValueError(f"No pristine images found in {images_dir}")

        logger.info('Found %s pristine images', len(self.images))

    def __len__(self) -> int:
        return len(self.images)

    def __getitem__(self, index: int) -> dict:
        img_path = self.images[index]
        img = Image.open(img_path).convert("RGB")
        img = self.transform(img)
        img_rel_path = _to_relative_dataset_path(Path(img_path), self.root)

        return {
            "img": img,
            "mos": 5.0,  # maximum quality for pristine images
            "dist_type": "pristine",
            "dist_group": "pristine",
            "dist_level": 0,  # distortion level = 0
            "distorted_img_path": img_rel_path,
            "original_img_path": img_rel_path,  # self-reference
            "sample_id": img_path.stem,  # e.g. "I01"
        }


class LocalKadidPristineDataset(Dataset):
    """
    Pristine KADID dataset.

    Ожидает корневую директорию с index.csv.
    Обязательные колонки: original_img_path
    Опциональные: mask_path (если есть маска области без искажений)
    """

    def __init__(
        self,
        root: str,
        crop_size: Optional[int] = 224,
    ):
        self.root = Path(root)
        self.index_path = self.root / "index.csv"
        self.index_dir = self.root

        if not self.index_path.exists():
            raise FileNotFoundError(f"index.csv not found in pristine root: {self.root}")

        df = pd.read_csv(self.index_path)

        required_cols = {"original_img_path"}

        missing = required_cols - set(df.columns)
        if missing:
            raise ValueError(f"Index is missing columns: {sorted(missing)}")

        self.images = sorted(df['original_img_path'].unique())
        self.image_to_tensor = transforms.ToTensor()
        self.crop = transforms.CenterCrop(crop_size) if crop_size is not None else None

    def __len__(self) -> int:
        return len(self.images)

    def _resolve_path(self, value: str) -> Path:
        p = Path(value)
        return p if p.is_absolute() else (self.index_dir / p)

    def __getitem__(self, index: int) -> dict:
        img_path = self._resolve_path(str(self.images[index]))

        img = Image.open(img_path).convert("RGB")
        img = self.image_to_tensor(img)
        img_rel_path = _to_relative_dataset_path(Path(img_path), self.root)

        sample = {
            "img": img,
            "dist_type": "pristine",
            "dist_group": "pristine",
            "dist_level": 0,
            "distorted_img_path": img_rel_path,
            "original_img_path": img_rel_path,
            "sample_id": img_path.stem,
        }

        if self.crop is not None:
            sample["img"] = self.crop(sample["img"])

        return sample


def kadid_pristine_collate_fn(batch: List[dict]) -> dict:
    """
    Collate для KadidPristineDataset.

    Возвращает:
        images: Tensor (B, C, H, W)
        + все остальные ключи как списки длины B
    """
    images = torch.stack([item["img"] for item in batch], dim=0)
    collated: dict = {"images": images}

    for key in batch[0]:
        if key == "img":
            continue
        collated[key] = [item[key] for item in batch]

    return collated


def local_kadid_pristine_collate_fn(batch: List[dict]) -> dict:
    """
    Collate для LocalKadidPristinePresavedDataset.

    Возвращает:
        images: Tensor (B, C, H, W)
        masks: Tensor (B, 1, H, W) если has_masks=True
        + остальные ключи как списки длины B
    """
    images = torch.stack([item["img"] for item in batch], dim=0)

    collated: dict = {"images": images}

    # Include masks in the batch when present
    if "mask" in batch[0]:
        masks = torch.stack([item["mask"] for item in batch], dim=0)
        collated["masks"] = masks

    for key in batch[0]:
        if key in ("img", "mask"):
            continue
        collated[key] = [item[key] for item in batch]

    return collated


def resolve_dataset_image_path(
    dataset: str,
    path_from_meta: str,
    datasets_root: str | None = None,
) -> Path:
    """Resolve a path stored in activation-cache metadata to an absolute file path."""
    from analysis.config import dataset_images_root, load_sae_vis_config

    root = Path(
        datasets_root
        if datasets_root is not None
        else load_sae_vis_config().DATASETS_ROOT
    )
    dataset_root = Path(dataset_images_root(str(root), dataset))
    path = Path(path_from_meta)
    if path.is_absolute():
        parts = Path(path).parts
        i = parts.index(dataset)
        suffix = Path(*parts[i:])
        path = dataset_root / suffix
    else:
        path = dataset_root / path
    return path


def dataset_image_paths(
    dataset: str,
    dataset_root: str,
    crop_size: int,
    min_distortion_level: int,
) -> List[str]:
    if dataset == 'kadid10k':
        ds = Kadid10kDataset(
            dataset_root,
            crop_size=crop_size,
            min_distortion_level=min_distortion_level,
        )
        return [_to_relative_dataset_path(Path(p), ds.root) for p in ds.images]

    if dataset == 'local_kadid':
        ds = LocalKadidPresavedDataset(dataset_root, crop_size=crop_size)
        return [_to_relative_dataset_path(Path(str(p)), ds.root) for p in ds.df['distorted_img_path'].tolist()]

    if dataset == 'QGround':
        ds = QGroundDataset(dataset_root, split='test', crop_size=crop_size)
        return [_to_relative_dataset_path(Path(p), ds.root) for p in ds.image_paths]

    if dataset == 'SRGround':
        ds = SRGroundSmallDataset(dataset_root, json_path='srground_train.json', allowed_methods=['DiT4SR_x2'])
        return [_to_relative_dataset_path(Path(str(p)), ds.root) for p in ds.df['image_path'].tolist()]

    raise ValueError(f'Unsupported dataset: {dataset}')