analysis/cache_utils.py

"""
Утилиты для кэширования и загрузки активаций SAE.
"""

from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union

import numpy as np
import pandas as pd
import scipy.sparse as sp
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm.auto import tqdm

from analysis.models import _iqa_activations
from log_config import get_logger

logger = get_logger(__name__)


def _df_mb(df: pd.DataFrame) -> float:
    """Объём памяти DataFrame в МБ."""
    return df.memory_usage(deep=True).sum() / 1024 ** 2


def _sparse_mb(mat: sp.csr_matrix) -> float:
    """Объём памяти CSR-матрицы (данные + индексы) в МБ."""
    return (mat.data.nbytes + mat.indices.nbytes + mat.indptr.nbytes) / 1024 ** 2


def _cache_paths(base_path: str) -> Tuple[str, str, str]:
    """
    Возвращает пути к трём файлам кэша из базового пути.

    Пример:
        'cache/kadid_acts.feather'
          → ('cache/kadid_acts_meta.feather', 'cache/kadid_acts_codes.npz',
              'cache/kadid_acts_steps.npz')
    """
    p = Path(base_path)
    stem = p.stem.removesuffix('.feather')
    return (
        str(p.parent / f'{stem}_meta.feather'),
        str(p.parent / f'{stem}_codes.npz'),
        str(p.parent / f'{stem}_steps.npz'),
    )


def _pristine_cache_paths(base_path: str) -> Tuple[str, str, str]:
    """Return paths for pristine cache files derived from base cache path."""
    meta_path, codes_path, steps_path = _cache_paths(base_path)
    return (
        meta_path.replace('.feather', '_pristine.feather'),
        codes_path.replace('.npz', '_pristine.npz'),
        steps_path.replace('.npz', '_pristine.npz'),
    )


def load_parquet_cache(cache_path: Optional[str], *, label: str = 'cache') -> Optional[pd.DataFrame]:
    """Load cached parquet table if present."""
    if cache_path is None:
        return None

    cache = Path(cache_path)
    if not cache.exists():
        return None

    logger.debug('[cache] Loading %s from %s', label, cache)
    return pd.read_parquet(cache)


def save_parquet_cache(df: pd.DataFrame, cache_path: Optional[str], *, label: str = 'cache') -> None:
    """Persist a dataframe to parquet cache if path is provided."""
    if cache_path is None:
        return

    cache = Path(cache_path)
    cache.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(cache)
    logger.debug('[cache] Saved %s to %s', label, cache)


_PATCH_LABEL_META_KEYS = frozenset({'dist_type', 'dist_group'})


def _patch_labels_to_dist_meta(
    patch_labels: np.ndarray,
    *,
    label_to_dist_type: Dict[int, str],
    label_to_dist_group: Dict[str, str],
) -> Tuple[List[str], List[str]]:
    flat_labels = patch_labels.reshape(-1)
    dist_types = [label_to_dist_type.get(int(label_id), 'background') for label_id in flat_labels]
    dist_groups = [label_to_dist_group.get(dist_type, dist_type) for dist_type in dist_types]
    return dist_types, dist_groups


def _process_dataloader(
    dataloader,
    iqa,
    sae,
    layer_name,
    scaling_factor,
    device,
    patches_per_image,
    patch_grid_shape,
    meta_keys,
    max_batches,
    max_memory_gb,
    add_patch_mask_stats,
    show_progress_bars: bool = True,
    *,
    label_to_dist_type: Optional[Dict[int, str]] = None,
    label_to_dist_group: Optional[Dict[str, str]] = None,
):
    all_sparse_codes = []
    all_meta = []
    all_sparse_steps = []

    n_patches_known = patches_per_image
    patch_grid_known = patch_grid_shape
    image_offset = 0

    for batch_i, batch in enumerate(tqdm(dataloader, desc='Caching activations', disable=not show_progress_bars)):
        if max_batches is not None and batch_i >= max_batches:
            break

        imgs = batch['images'].to(device)
        B = imgs.shape[0]

        with torch.no_grad():
            iqa(imgs)
            acts = _iqa_activations[layer_name].to(device)
            acts = acts * scaling_factor
            enc_out = sae.get_acts(acts)

            if isinstance(enc_out, tuple):
                codes, activation_steps = enc_out
            else:
                codes, activation_steps = enc_out, None

        codes_np = codes.cpu().float().numpy()

        if activation_steps is None:
            steps_np = np.zeros_like(codes_np, dtype=np.int32)
        else:
            steps_np = activation_steps.cpu().numpy().astype(np.int32)

        if n_patches_known is None:
            n_patches_known = codes_np.shape[0] // B
            logger.info('Detected %s patches per image', n_patches_known)

        P = n_patches_known

        use_patch_label_meta = (
            label_to_dist_type is not None
            and label_to_dist_group is not None
        )
        meta = {}
        for k in meta_keys:
            if k in batch:
                if use_patch_label_meta and k in _PATCH_LABEL_META_KEYS:
                    continue
                vals = batch[k]
                meta[k] = [v for v in vals for _ in range(P)]

        meta['patch_idx'] = list(range(P)) * B
        meta['image_idx'] = [image_offset + i for i in range(B) for _ in range(P)]

        if add_patch_mask_stats and 'masks' in batch:
            masks = batch['masks'].to(device=device, dtype=torch.float32)
            if patch_grid_known is not None:
                grid_h, grid_w = patch_grid_known
                if grid_h * grid_w == P:
                    mask_labels = masks.to(dtype=torch.int64)
                    max_label = int(mask_labels.max().item())
                    max_cov = None

                    if max_label <= 0:
                        patch_labels = torch.zeros((B, grid_h, grid_w), device=device, dtype=torch.int64)
                    else:
                        class_coverages = []
                        for label_id in range(1, max_label + 1):
                            label_cov = F.adaptive_avg_pool2d(
                                (mask_labels == label_id).to(dtype=torch.float32),
                                (grid_h, grid_w),
                            )
                            class_coverages.append(label_cov)

                        coverages = torch.cat(class_coverages, dim=1)  # (B, classes, H, W)
                        max_cov, max_idx = coverages.max(dim=1)
                        patch_labels = torch.where(
                            max_cov > 0,
                            max_idx.to(dtype=torch.int64) + 1,
                            torch.zeros_like(max_idx, dtype=torch.int64),
                        )

                    patch_labels_np = patch_labels.reshape(B, P).cpu().numpy().astype(np.int16)
                    patch_is_dist = (patch_labels_np > 0).astype(np.int8)
                    patch_coverage_np = max_cov.reshape(B, P).cpu().numpy() if max_label > 0 else patch_is_dist.astype(np.float32)

                    meta['patch_mask_label'] = patch_labels_np.reshape(-1).tolist()
                    meta['patch_mask_coverage'] = patch_coverage_np.reshape(-1).tolist()
                    meta['patch_is_distorted'] = patch_is_dist.reshape(-1).tolist()

                    if use_patch_label_meta:
                        dist_types, dist_groups = _patch_labels_to_dist_meta(
                            patch_labels_np,
                            label_to_dist_type=label_to_dist_type,
                            label_to_dist_group=label_to_dist_group,
                        )
                        if 'dist_type' in meta_keys:
                            meta['dist_type'] = dist_types
                        if 'dist_group' in meta_keys:
                            meta['dist_group'] = dist_groups

        all_meta.append(pd.DataFrame(meta))
        all_sparse_codes.append(sp.csr_matrix(codes_np))
        all_sparse_steps.append(sp.csr_matrix(steps_np))

        image_offset += B

    meta_df = pd.concat(all_meta, ignore_index=True)
    codes_csr = sp.vstack(all_sparse_codes, format='csr')
    steps_csr = sp.vstack(all_sparse_steps, format='csr')

    return meta_df, codes_csr, steps_csr


def collect_and_cache(
    dataloader: DataLoader,
    iqa: torch.nn.Module,
    sae,
    layer_name: str,
    output_path: str,
    scaling_factor: float = 1.0,
    patches_per_image: Optional[int] = None,
    patch_grid_shape: Optional[Tuple[int, int]] = None,
    meta_keys: Sequence[str] = (
        'dist_type',
        'dist_group',
        'dist_level',
        'mos',
        'distorted_img_path',
        'original_img_path',
        'sample_id',
    ),
    device: str = 'cuda',
    max_batches: Optional[int] = None,
    max_memory_gb: Optional[float] = None,
    add_patch_mask_stats: bool = True,
    pristine_dataloader: Optional[DataLoader] = None,
    show_progress_bars: bool = True,
    label_to_dist_type: Optional[Dict[int, str]] = None,
    label_to_dist_group: Optional[Dict[str, str]] = None,
) -> Tuple[pd.DataFrame, sp.csr_matrix]:

    meta_df, codes_csr, steps_csr = _process_dataloader(
        dataloader=dataloader,
        iqa=iqa,
        sae=sae,
        layer_name=layer_name,
        scaling_factor=scaling_factor,
        device=device,
        patches_per_image=patches_per_image,
        patch_grid_shape=patch_grid_shape,
        meta_keys=meta_keys,
        max_batches=max_batches,
        max_memory_gb=max_memory_gb,
        add_patch_mask_stats=add_patch_mask_stats,
        show_progress_bars=show_progress_bars,
        label_to_dist_type=label_to_dist_type,
        label_to_dist_group=label_to_dist_group,
    )
    sparse_mb = _sparse_mb(codes_csr)
    steps_mb = _sparse_mb(steps_csr)

    logger.info('Activations: shape=%s, %.1f МБ (sparse)', codes_csr.shape, sparse_mb)
    logger.info('Activation steps: shape=%s, %.1f МБ (sparse)', steps_csr.shape, steps_mb)

    meta_path, codes_path, steps_path = _cache_paths(output_path)

    meta_df.to_feather(meta_path)
    sp.save_npz(codes_path, codes_csr)
    sp.save_npz(steps_path, steps_csr)

    logger.info('Saved metadata (%s rows) -> %s', len(meta_df), meta_path)
    logger.info('Saved activations -> %s', codes_path)
    logger.info('Saved activation steps -> %s', steps_path)
    logger.info('  Metadata:    %.1f МБ', _df_mb(meta_df))
    logger.info('  Activations: %.1f МБ', sparse_mb)
    logger.info('  Steps:       %.1f МБ', steps_mb)

    if pristine_dataloader is not None:
        logger.info('Processing pristine dataset...')

        pristine_meta, pristine_codes, pristine_steps = _process_dataloader(
            dataloader=pristine_dataloader,
            iqa=iqa,
            sae=sae,
            layer_name=layer_name,
            scaling_factor=scaling_factor,
            device=device,
            patches_per_image=patches_per_image,
            patch_grid_shape=patch_grid_shape,
            meta_keys=meta_keys,
            max_batches=max_batches,
            max_memory_gb=max_memory_gb,
            add_patch_mask_stats=False,
            show_progress_bars=show_progress_bars,
        )
        pristine_sparse_mb = _sparse_mb(pristine_codes)
        pristine_steps_mb = _sparse_mb(pristine_steps)

        pristine_meta_path = meta_path.replace(".feather", "_pristine.feather")
        pristine_codes_path = codes_path.replace(".npz", "_pristine.npz")
        pristine_steps_path = steps_path.replace(".npz", "_pristine.npz")

        pristine_meta.to_feather(pristine_meta_path)
        sp.save_npz(pristine_codes_path, pristine_codes)
        sp.save_npz(pristine_steps_path, pristine_steps)

        logger.info(
            'Pristine activations: shape=%s, %.1f МБ',
            pristine_codes.shape,
            pristine_sparse_mb,
        )
        logger.info(
            'Pristine steps: shape=%s, %.1f МБ',
            pristine_steps.shape,
            pristine_steps_mb,
        )

        logger.info('Saved pristine metadata (%s rows) -> %s', len(pristine_meta), pristine_meta_path)
        logger.info('Saved pristine activations -> %s', pristine_codes_path)
        logger.info('Saved pristine activation steps -> %s', pristine_steps_path)

    return meta_df, codes_csr


def build_activation_cache(
    *,
    dataset: str,
    cache_path: str,
    checkpoint_path: str,
    dataset_root: str,
    layer_num: int,
    iqa_metric: str,
    swin_num: int,
    device: str,
    batch_size: int,
    num_workers: int,
    crop_size: int,
    scaling_factor: float = 1.0,
    min_distortion_level: Optional[int] = None,
    max_batches: Optional[int] = None,
    max_memory_gb: Optional[float] = None,
    add_patch_mask_stats: bool = True,
    include_pristine: bool = True,
    show_progress_bars: bool = True,
    srground_include_sr_artifact: bool = False,
) -> Dict[str, Any]:
    """Build activation cache end-to-end for KADID/local-KADID datasets."""
    from .datasets import (
        Kadid10kDataset,
        KadidPristineDataset,
        LocalKadidPresavedDataset,
        LocalKadidPristineDataset,
        QGroundDataset,
        SRGroundSmallDataset,
        available_distortions_qground,
        available_distortions_srground,
        distortion_types_mapping_qground,
        distortion_types_mapping_srground,
        kadid_collate_fn,
        kadid_pristine_collate_fn,
        local_kadid_collate_fn,
        local_kadid_pristine_collate_fn,
        qground_collate_fn,
        srground_collate_fn,
    )
    from .models import _iqa_activation_grids, load_iqa_model, load_sae, read_sae_config

    if min_distortion_level is not None and not (1 <= min_distortion_level <= 5):
        raise ValueError('min_distortion_level must be in [1, 5]')

    label_to_dist_type = None
    label_to_dist_group = None

    if dataset == 'local_kadid':
        data = LocalKadidPresavedDataset(root=dataset_root, crop_size=crop_size)
        collate_fn = local_kadid_collate_fn
        meta_keys = [
            'dist_type',
            'dist_group',
            'dist_level',
            'mos',
            'local_dist_type',
            'local_dist_level',
            'mask_shape',
            'mask_coverage',
            'sample_id',
            'distorted_img_path',
            'original_img_path',
        ]
        pristine_data = LocalKadidPristineDataset(root=dataset_root, crop_size=crop_size) if include_pristine else None
        pristine_collate = local_kadid_pristine_collate_fn
    elif dataset in {'kadid10k', 'kadid'}:
        data = Kadid10kDataset(
            root=dataset_root,
            crop_size=crop_size,
            min_distortion_level=min_distortion_level or 1,
        )
        collate_fn = kadid_collate_fn
        meta_keys = [
            'dist_type',
            'dist_group',
            'dist_level',
            'mos',
            'distorted_img_path',
            'original_img_path',
        ]
        pristine_data = KadidPristineDataset(root=dataset_root, crop_size=crop_size) if include_pristine else None
        pristine_collate = kadid_pristine_collate_fn
    elif dataset == 'QGround':
        data = QGroundDataset(
            root=dataset_root,
            split='test',
            crop_size=crop_size,
        )
        collate_fn = qground_collate_fn
        meta_keys = [
            'dist_type',
            'dist_group',
            'dist_level',
            'mos',
            'mask_coverage',
            'qground_ann_id',
            'sample_id',
            'distorted_img_path',
            'original_img_path',
            'image_path',
            'mask_path',
            'split',
        ]
        label_to_dist_type = distortion_types_mapping_qground
        label_to_dist_group = available_distortions_qground
        pristine_data = None
        pristine_collate = None
    elif dataset == 'SRGround':
        data = SRGroundSmallDataset(
            root=dataset_root,
            json_path='srground_train.json',
            crop_size=crop_size,
            include_sr_artifact=srground_include_sr_artifact,
        )
        collate_fn = srground_collate_fn
        meta_keys = [
            'dist_type',
            'dist_group',
            'dist_level',
            'mos',
            'mask_coverage',
            'sample_id',
            'distorted_img_path',
            'image_path',
            'real_distortions_ann_path',
            'sr_artifacts_ann_path',
        ]
        label_to_dist_type = distortion_types_mapping_srground
        label_to_dist_group = available_distortions_srground
        pristine_data = None
        pristine_collate = None
    else:
        raise ValueError(f'Unsupported dataset: {dataset}')

    loader = DataLoader(
        data,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collate_fn,
    )

    pristine_loader = None
    if pristine_data is not None:
        pristine_loader = DataLoader(
            pristine_data,
            batch_size=batch_size,
            shuffle=False,
            num_workers=num_workers,
            collate_fn=pristine_collate,
        )

    iqa_model, layer_name = load_iqa_model(
        layer_num=layer_num,
        device=device,
        iqa_metric=iqa_metric,
        swin_num=swin_num,
    )
    dtype = torch.float16 if iqa_metric == 'qalign' else torch.float32
    sae_cfg = read_sae_config(checkpoint_path)
    sae_model = load_sae(checkpoint_path, device=device, dtype=dtype, sae_config=sae_cfg)

    with torch.no_grad():
        dummy = torch.rand(1, 3, crop_size, crop_size, device=device).clamp(0, 1)
        iqa_model(dummy)

    if layer_name not in _iqa_activations or layer_name not in _iqa_activation_grids:
        raise RuntimeError(f'Cannot infer activation grid for layer {layer_name}')

    patch_grid_shape = _iqa_activation_grids[layer_name]
    patches_per_image = patch_grid_shape[0] * patch_grid_shape[1]

    Path(cache_path).parent.mkdir(parents=True, exist_ok=True)
    collect_and_cache(
        dataloader=loader,
        iqa=iqa_model,
        sae=sae_model,
        layer_name=layer_name,
        output_path=cache_path,
        scaling_factor=scaling_factor,
        patches_per_image=patches_per_image,
        patch_grid_shape=patch_grid_shape,
        meta_keys=meta_keys,
        device=device,
        max_batches=max_batches,
        max_memory_gb=max_memory_gb,
        add_patch_mask_stats=add_patch_mask_stats,
        pristine_dataloader=pristine_loader,
        show_progress_bars=show_progress_bars,
        label_to_dist_type=label_to_dist_type,
        label_to_dist_group=label_to_dist_group,
    )

    return {
        'layer_name': layer_name,
        'patch_grid_shape': patch_grid_shape,
        'patches_per_image': patches_per_image,
        'sae_config': sae_cfg,
    }


def load_cache(
    path: str,
    return_activation_steps: bool = False,
    min_distortion_level: Optional[int] = None,
    max_distortion_level: Optional[int] = None,
) -> Union[
    Tuple[pd.DataFrame, sp.csr_matrix],
    Tuple[pd.DataFrame, sp.csr_matrix, sp.csr_matrix],
]:
    """Загружает кэш активаций SAE из раздельных файлов.

    Если ``return_activation_steps=True``, дополнительно возвращает CSR-матрицу
    порядка активаций, где значение ``0`` означает отсутствие активации,
    а ``k>0`` соответствует шагу ``k`` в pursuit.
    """
    meta_path, codes_path, steps_path = _cache_paths(path)
    meta = pd.read_feather(meta_path)
    codes = sp.load_npz(codes_path)

    logger.debug(
        'Loaded from %s: %s rows × %s cols',
        meta_path,
        meta.shape[0],
        meta.shape[1],
    )
    logger.debug('Loaded from %s: shape=%s, dtype=%s', codes_path, codes.shape, codes.dtype)
    logger.debug('  Metadata:    %.1f МБ', _df_mb(meta))
    logger.debug('  Activations: %.1f МБ (sparse)', _sparse_mb(codes))

    keep_idx: Optional[np.ndarray] = None
    if min_distortion_level is not None or max_distortion_level is not None:
        if 'dist_level' not in meta.columns:
            raise ValueError('Cannot filter by distortion level: metadata has no "dist_level" column')

        min_level = 1 if min_distortion_level is None else int(min_distortion_level)
        max_level = 5 if max_distortion_level is None else int(max_distortion_level)
        if min_level > max_level:
            raise ValueError(
                f'Invalid distortion-level range: min_distortion_level={min_level} > max_distortion_level={max_level}'
            )
        if 'Ground' not in path:
            keep_mask = (meta['dist_level'] >= min_level) & (meta['dist_level'] <= max_level)
            keep_idx = np.flatnonzero(keep_mask.to_numpy())
        else: 
            keep_mask = (meta['dist_level'] >= -1000) 
            keep_idx = np.flatnonzero(keep_mask.to_numpy())  # Temporary workaround -- fix later
 
    if return_activation_steps:
        if Path(steps_path).exists():
            steps = sp.load_npz(steps_path)
            if steps.shape != codes.shape:
                raise ValueError(
                    f'Steps cache shape mismatch: expected {codes.shape}, got {steps.shape}'
                )
            logger.info('Loaded from %s: shape=%s, dtype=%s', steps_path, steps.shape, steps.dtype)
        else:
            logger.warning('No steps cache found. Using all-zero activation steps.')
            steps = sp.csr_matrix(codes.shape, dtype=np.int32)

        if keep_idx is not None:
            meta = meta.iloc[keep_idx].reset_index(drop=True)
            codes = codes[keep_idx]
            steps = steps[keep_idx]
            logger.info(
                'Applied dist_level filter [%s, %s] -> %s rows kept',
                min_level,
                max_level,
                meta.shape[0],
            )

        logger.info('  Steps:       %.1f МБ (sparse)', _sparse_mb(steps))
        return meta, codes, steps

    if keep_idx is not None:
        meta = meta.iloc[keep_idx].reset_index(drop=True)
        codes = codes[keep_idx]
        logger.info(
            'Applied dist_level filter [%s, %s] -> %s rows kept',
            min_level,
            max_level,
            meta.shape[0],
        )

    return meta, codes


def load_pristine_cache(
    path: str,
    return_activation_steps: bool = False,
) -> Union[
    Tuple[pd.DataFrame, sp.csr_matrix],
    Tuple[pd.DataFrame, sp.csr_matrix, sp.csr_matrix],
]:
    """Load pristine (original-image) activation cache saved by collect_and_cache."""
    meta_path, codes_path, steps_path = _pristine_cache_paths(path)
    meta = pd.read_feather(meta_path)
    codes = sp.load_npz(codes_path)

    logger.info(
        'Loaded pristine from %s: %s rows × %s cols',
        meta_path,
        meta.shape[0],
        meta.shape[1],
    )
    logger.info('Loaded pristine from %s: shape=%s, dtype=%s', codes_path, codes.shape, codes.dtype)
    logger.info('  Metadata:    %.1f МБ', _df_mb(meta))
    logger.info('  Activations: %.1f МБ (sparse)', _sparse_mb(codes))

    if return_activation_steps:
        if Path(steps_path).exists():
            steps = sp.load_npz(steps_path)
            if steps.shape != codes.shape:
                raise ValueError(
                    f'Pristine steps cache shape mismatch: expected {codes.shape}, got {steps.shape}'
                )
            logger.info(
                'Loaded pristine from %s: shape=%s, dtype=%s',
                steps_path,
                steps.shape,
                steps.dtype,
            )
        else:
            logger.warning('No pristine steps cache found. Using all-zero activation steps.')
            steps = sp.csr_matrix(codes.shape, dtype=np.int32)

        logger.info('  Steps:       %.1f МБ (sparse)', _sparse_mb(steps))
        return meta, codes, steps

    return meta, codes


def ensure_cache_ready(
    cache_path: str,
    *,
    force_recache: bool = False,
    build_cache_if_missing: bool = True,
    load_cache_kwargs: Optional[Dict[str, Any]] = None,
    build_cache_fn: Optional[Callable[[], None]] = None,
) -> None:
    """Проверяет доступность кэша и при необходимости собирает его.

    Поведение:
    - пытается загрузить кэш через ``load_cache``;
    - если кэш отсутствует или выставлен ``force_recache=True``, запускает сборку;
    - если сборка отключена, пробрасывает ``FileNotFoundError``.
    """
    needs_rebuild = force_recache
    if not needs_rebuild:
        try:
            load_cache(cache_path, **(load_cache_kwargs or {}))
            return
        except FileNotFoundError:
            needs_rebuild = True

    if not needs_rebuild:
        return

    if not build_cache_if_missing:
        raise FileNotFoundError(
            f'Activation cache not found at {cache_path}, and build is disabled. '
            'Use --build-cache-if-missing or provide existing cache files.'
        )

    if build_cache_fn is None:
        raise ValueError(
            'build_cache_fn must be provided when cache rebuild is required '
            '(missing cache or force_recache=True).'
        )

    logger.debug('[cache] Building activation cache...')
    build_cache_fn()


def zero_codes_outside_activation_steps(
    codes_csr: sp.csr_matrix,
    activation_steps_csr: sp.csr_matrix,
    activation_steps_to_keep: List[int],
) -> sp.csr_matrix:
    """Обнуляет активации, шаг появления которых не входит в allow-list.

    Параметры
    ----------
    codes_csr : CSR-матрица активаций SAE.
    activation_steps_csr : CSR-матрица шагов активаций (0 = не активирован).
    activation_steps_to_keep : список шагов, которые нужно сохранить.

    Возвращает
    ----------
    CSR-матрицу той же формы, где вне указанных шагов значения занулены.
    Если список шагов пуст, возвращается исходная матрица без изменений.
    """
    if not activation_steps_to_keep:
        return codes_csr

    if codes_csr.shape != activation_steps_csr.shape:
        raise ValueError(
            f'Codes/steps shape mismatch: {codes_csr.shape} vs {activation_steps_csr.shape}'
        )

    keep_steps = sorted({int(step) for step in activation_steps_to_keep})
    if any(step <= 0 for step in keep_steps):
        raise ValueError('activation_steps_to_keep must contain only positive integers')

    codes_coo = codes_csr.tocoo(copy=False)
    steps_coo = activation_steps_csr.tocoo(copy=False)

    # Steps matrix stores indices for nonzero entries of codes, so coordinates must match.
    if (
        codes_coo.nnz != steps_coo.nnz
        or not np.array_equal(codes_coo.row, steps_coo.row)
        or not np.array_equal(codes_coo.col, steps_coo.col)
    ):
        raise ValueError('Codes and steps matrices must have the same sparsity pattern. Something weird is going on.')
    else:
        steps_for_codes = steps_coo.data

    keep_mask = np.isin(np.asarray(steps_for_codes), np.asarray(keep_steps, dtype=np.int32))

    filtered = sp.coo_matrix(
        (codes_coo.data[keep_mask], (codes_coo.row[keep_mask], codes_coo.col[keep_mask])),
        shape=codes_csr.shape,
        dtype=codes_csr.dtype,
    )
    return filtered.tocsr()


def ensure_activation_cache(
    dataset: str,
    acts_cache_path: str,
    dataset_root: str,
    min_distortion_level: int,
    params: dict,
    include_pristine_cache: Optional[bool] = None,
) -> None:
    """Build distorted+pristine activation cache if missing."""
    cache_filter_min = int(min_distortion_level) if dataset == 'kadid10k' else None
    if include_pristine_cache is None:
        needs_pristine_cache = dataset in {'kadid10k', 'local_kadid'}
    else:
        needs_pristine_cache = bool(include_pristine_cache)

    try:
        load_cache(
            acts_cache_path,
            return_activation_steps=True,
            min_distortion_level=cache_filter_min,
            max_distortion_level=params.get('KADID_MAX_DISTORTION_LEVEL') if dataset == 'kadid10k' else None,
        )
        if needs_pristine_cache:
            load_pristine_cache(acts_cache_path, return_activation_steps=True)
        return
    except FileNotFoundError:
        pass

    logger.info('[run] Activation cache not found for %s. Building cache...', acts_cache_path)
    build_activation_cache(
        dataset=dataset,
        cache_path=acts_cache_path,
        checkpoint_path=params.get('SAE_CHECKPOINT_PATH'),
        dataset_root=dataset_root,
        layer_num=params.get('LAYER_NUM'),
        iqa_metric=params.get('IQA_METRIC'),
        swin_num=params.get('SWIN_NUM'),
        device=params.get('DEVICE'),
        batch_size=params.get('BATCH_SIZE'),
        num_workers=params.get('NUM_WORKERS'),
        crop_size=params.get('CROP_SIZE'),
        scaling_factor=params.get('SCALING_FACTOR'),
        min_distortion_level=min_distortion_level,
        max_batches=None,
        max_memory_gb=30.0,
        add_patch_mask_stats=True,
        include_pristine=needs_pristine_cache,
        srground_include_sr_artifact=bool(params.get('SRGROUND_INCLUDE_SR_ARTIFACT', False)),
    )
    logger.info('[run] Activation cache build completed.')