"""Feature selection strategies for automated SAE analysis reports."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple

import numpy as np
import pandas as pd

from analysis.features.feature_indexing import FeatureMatrix
from analysis.metrics.correlations import compute_distortion_correlations
from analysis.metrics.mutual_information import compute_distortion_mutual_information
from analysis.metrics.paired_deltas import build_paired_delta_tables
from analysis.metrics.roc_auc import compute_distortion_roc_auc

from log_config import get_logger

logger = get_logger(__name__)


Category = Literal['type', 'group']
DeltaMode = Literal['abs', 'signed', 'relative']
Level = Literal['patch', 'image']


def _metric_table_key(metric: str, category: Category) -> str:
    return f'{metric}_{category}_df'


def _paired_delta_table_key(delta_mode: DeltaMode, category: Category) -> str:
    return f'paired_{delta_mode}_dist_{category}_df'


def _group_col(category: Category) -> str:
    return 'dist_group' if category == 'group' else 'dist_type'


@dataclass(frozen=True)
class SelectorResult:
    """Output bundle produced by selector execution."""

    selected_features: pd.DataFrame
    metric_tables: Dict[str, pd.DataFrame]
    selection_stats: Dict[str, object]
    selector_name: str
    selector_description: str
    relation_label: str


@dataclass(frozen=True)
class SelectorInputs:
    """Typed selector inputs with required core data and optional pristine data."""

    meta_df: pd.DataFrame
    features: FeatureMatrix
    dataset: str
    selector_top_k: int = 20
    selector_params: Mapping[str, object] = field(default_factory=dict)
    cache_paths: Any | None = None
    pristine_meta_df: pd.DataFrame | None = None
    pristine_features: FeatureMatrix | None = None
    feature_weight_norms: Mapping[int, float] | None = None


def _match_metric_df_columns(corr_df: pd.DataFrame, feature_ids: Sequence[int]) -> List[object]:
    id_set = {int(x) for x in feature_ids}
    matched_columns: List[object] = []
    for col in corr_df.columns:
        try:
            col_id = int(col)
        except (TypeError, ValueError):
            continue
        if col_id in id_set:
            matched_columns.append(col)
    return matched_columns


def _build_feature_importance_table(
    importance_series: pd.Series,
    top_k: int,
    metric_name: str,
) -> pd.DataFrame:
    top_series = importance_series.nlargest(top_k)
    result = pd.DataFrame({
        'feature_id': [int(fid) for fid in top_series.index.tolist()],
        'importance_score': [float(x) for x in top_series.values.tolist()],
    })
    result['importance_rank'] = range(1, len(result) + 1)
    result['importance_metric'] = metric_name
    return result[['feature_id', 'importance_score', 'importance_rank', 'importance_metric']]


def _collect_table_stats(
    stats: Dict[str, object],
    table_key: str,
    table_df: pd.DataFrame,
    feature_ids: Sequence[int],
) -> None:
    matched_cols = _match_metric_df_columns(table_df, feature_ids)
    if not matched_cols:
        return

    table_values = table_df.loc[:, matched_cols]
    stats[f'{table_key}_matched_cols'] = len(matched_cols)
    stats[f'{table_key}_mean_abs'] = float(table_values.abs().values.mean())
    stats[f'{table_key}_max_abs'] = float(table_values.abs().values.max())


class BaseFeatureSelector(ABC):
    """Common interface for feature selectors."""

    name: str
    description: str
    relation_label: str = 'feature relevance'

    @abstractmethod
    def compute_metric_tables(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> Dict[str, pd.DataFrame]:
        """Compute metric tables required by the selector."""

    @abstractmethod
    def compute_importance(
        self,
        metric_tables: Dict[str, pd.DataFrame],
        inputs: SelectorInputs | None = None,
        feature_ids: Sequence[int] | None = None,
    ) -> Tuple[pd.Series, str]:
        """Return per-feature importance scores and metric name."""

    def execute(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> SelectorResult:
        """Run full selector pipeline and return selected features + computed measures."""
        metric_tables = self.compute_metric_tables(inputs, feature_ids=feature_ids)
        importance_series, metric_name = self.compute_importance(
            metric_tables=metric_tables,
            inputs=inputs,
            feature_ids=feature_ids,
        )
        top_k = int(inputs.selector_top_k)
        selected_features = _build_feature_importance_table(
            importance_series=importance_series,
            top_k=top_k,
            metric_name=metric_name,
        )

        selected_features = selected_features.copy()
        selected_features['feature_id'] = selected_features['feature_id'].astype(int)
        selected_features['importance_score'] = selected_features['importance_score'].astype(float)
        selected_features['importance_rank'] = selected_features['importance_rank'].astype(int)

        selected_features = selected_features.drop_duplicates(
            subset=['feature_id'], keep='first').reset_index(drop=True)
        if selected_features.empty:
            raise ValueError('Selector returned empty feature list')

        selected_features['importance_rank'] = range(1, len(selected_features) + 1)
        stats = self.selection_stats(selected_features, metric_tables)
        return SelectorResult(
            selected_features=selected_features,
            metric_tables=metric_tables,
            selection_stats=stats,
            selector_name=self.name,
            selector_description=self.description,
            relation_label=self.relation_label,
        )

    def select_features(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> pd.DataFrame:
        """Template method: normalize and return top-k with importance stats."""
        metric_tables = self.compute_metric_tables(inputs, feature_ids=feature_ids)
        importance_series, metric_name = self.compute_importance(
            metric_tables=metric_tables,
            inputs=inputs,
            feature_ids=feature_ids,
        )

        top_k = int(inputs.selector_top_k)
        selected_features = _build_feature_importance_table(
            importance_series=importance_series,
            top_k=top_k,
            metric_name=metric_name,
        )

        selected_features = selected_features.copy()
        selected_features['feature_id'] = selected_features['feature_id'].astype(int)
        selected_features['importance_score'] = selected_features['importance_score'].astype(float)
        selected_features['importance_rank'] = selected_features['importance_rank'].astype(int)

        selected_features = selected_features.drop_duplicates(
            subset=['feature_id'], keep='first').reset_index(drop=True)
        if selected_features.empty:
            raise ValueError('Selector returned empty feature list')

        selected_features['importance_rank'] = range(1, len(selected_features) + 1)
        return selected_features

    def selection_stats(
        self,
        selected_features: pd.DataFrame,
        metric_tables: Mapping[str, pd.DataFrame],
    ) -> Dict[str, object]:
        feature_ids = selected_features['feature_id'].tolist()
        stats: Dict[str, object] = {
            'selector_name': self.name,
            'selector_description': self.description,
            'n_selected': len(feature_ids),
        }
        if not selected_features.empty:
            stats['mean_importance_score'] = float(selected_features['importance_score'].mean())
            stats['max_importance_score'] = float(selected_features['importance_score'].max())
            stats['min_importance_score'] = float(selected_features['importance_score'].min())
        if feature_ids:
            numeric_ids = [int(x) for x in feature_ids]
            stats['min_feature_id'] = min(numeric_ids)
            stats['max_feature_id'] = max(numeric_ids)

        for corr_key in ('corr_type_df', 'corr_group_df'):
            corr_df = metric_tables.get(corr_key)
            if isinstance(corr_df, pd.DataFrame) and len(feature_ids) > 0:
                matched_cols = _match_metric_df_columns(corr_df, feature_ids)
                if matched_cols:
                    abs_corr = corr_df.loc[:, matched_cols].abs()
                    stats[f'{corr_key}_matched_cols'] = len(matched_cols)
                    stats[f'{corr_key}_mean_abs_corr'] = float(abs_corr.values.mean())
                    stats[f'{corr_key}_max_abs_corr'] = float(abs_corr.values.max())

        for mi_key in ('mi_type_df', 'mi_group_df'):
            mi_df = metric_tables.get(mi_key)
            if isinstance(mi_df, pd.DataFrame) and len(feature_ids) > 0:
                matched_cols = _match_metric_df_columns(mi_df, feature_ids)
                if matched_cols:
                    mi_values = mi_df.loc[:, matched_cols]
                    stats[f'{mi_key}_matched_cols'] = len(matched_cols)
                    stats[f'{mi_key}_mean'] = float(mi_values.values.mean())
                    stats[f'{mi_key}_max'] = float(mi_values.values.max())

        for auc_key in ('auc_type_df', 'auc_group_df'):
            auc_df = metric_tables.get(auc_key)
            if isinstance(auc_df, pd.DataFrame) and len(feature_ids) > 0:
                matched_cols = _match_metric_df_columns(auc_df, feature_ids)
                if matched_cols:
                    auc_values = auc_df.loc[:, matched_cols]
                    stats[f'{auc_key}_matched_cols'] = len(matched_cols)
                    stats[f'{auc_key}_mean'] = float(auc_values.values.mean())
                    stats[f'{auc_key}_max'] = float(auc_values.values.max())
                    stats[f'{auc_key}_min'] = float(auc_values.values.min())

        for key, value in metric_tables.items():
            if not (isinstance(key, str) and key.startswith('paired_') and key.endswith('_df')):
                continue
            if isinstance(value, pd.DataFrame) and len(feature_ids) > 0:
                _collect_table_stats(stats, key, value, feature_ids)

        return stats


@dataclass(frozen=True)
class TopKAbsCorrSelector(BaseFeatureSelector):
    """Top-k by max absolute correlation across category rows."""

    name: str
    description: str
    category: Category = 'type'
    relation_label: str = 'high correlation'

    @property
    def corr_key(self) -> str:
        return _metric_table_key('corr', self.category)

    def compute_metric_tables(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> Dict[str, pd.DataFrame]:
        selector_params = dict(inputs.selector_params)
        level = str(selector_params.get('level', 'patch'))
        binarize = bool(selector_params.get('binarize', False))
        group_col = _group_col(self.category)

        cache_path = selector_params.get('cache_path')
        if cache_path is None:
            cache_paths = inputs.cache_paths
            if cache_paths is not None:
                attr_map = {
                    ('corr_type_df', 'patch'): 'corr_type_patch_cache_path',
                    ('corr_group_df', 'patch'): 'corr_group_patch_cache_path',
                    ('corr_type_df', 'image'): 'corr_type_cache_path',
                    ('corr_group_df', 'image'): 'corr_group_cache_path',
                }
                cache_attr = attr_map.get((self.corr_key, level))
                if cache_attr is not None and hasattr(cache_paths, cache_attr):
                    cache_path = getattr(cache_paths, cache_attr)

        corr_df = compute_distortion_correlations(
            meta=inputs.meta_df,
            features=inputs.features,
            group_col=group_col,
            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
            cache_path=cache_path,
        )
        return {self.corr_key: corr_df}

    def compute_importance(
        self,
        metric_tables: Dict[str, pd.DataFrame],
        inputs: SelectorInputs | None = None,
        feature_ids: Sequence[int] | None = None,
    ) -> Tuple[pd.Series, str]:
        corr_df = metric_tables[self.corr_key]
        if feature_ids is not None:
            corr_df = corr_df[[int(fid) for fid in feature_ids]]
        importance_series = corr_df.abs().max(axis=0)
        return importance_series, f'max_abs_corr:{self.corr_key}'


@dataclass(frozen=True)
class TopKMutualInfoSelector(BaseFeatureSelector):
    """Top-k by max mutual information across category rows."""

    name: str
    description: str
    category: Category = 'type'
    relation_label: str = 'High mutual information'

    @property
    def mi_key(self) -> str:
        return _metric_table_key('mi', self.category)

    def compute_metric_tables(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> Dict[str, pd.DataFrame]:
        selector_params = dict(inputs.selector_params)
        level = str(selector_params.get('level', 'patch'))
        binarize = bool(selector_params.get('binarize', False))
        n_bins = int(selector_params.get('n_bins', 16))
        feature_chunk_size = int(selector_params.get('feature_chunk_size', 512))
        group_col = _group_col(self.category)

        cache_path = selector_params.get('cache_path')
        if cache_path is None:
            cache_paths = inputs.cache_paths
            if cache_paths is not None:
                attr_map = {
                    ('mi_type_df', 'patch'): 'mi_type_patch_cache_path',
                    ('mi_group_df', 'patch'): 'mi_group_patch_cache_path',
                    ('mi_type_df', 'image'): 'mi_type_cache_path',
                    ('mi_group_df', 'image'): 'mi_group_cache_path',
                }
                cache_attr = attr_map.get((self.mi_key, level))
                if cache_attr is not None and hasattr(cache_paths, cache_attr):
                    cache_path = getattr(cache_paths, cache_attr)

        mi_df = compute_distortion_mutual_information(
            meta=inputs.meta_df,
            features=inputs.features,
            group_col=group_col,
            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
            binarize=binarize,
            level=level,
            n_bins=n_bins,
            feature_chunk_size=feature_chunk_size,
            cache_path=cache_path,
        )
        return {self.mi_key: mi_df}

    def compute_importance(
        self,
        metric_tables: Dict[str, pd.DataFrame],
        inputs: SelectorInputs | None = None,
        feature_ids: Sequence[int] | None = None,
    ) -> Tuple[pd.Series, str]:
        mi_df = metric_tables[self.mi_key]
        if feature_ids is not None:
            mi_df = mi_df[[int(fid) for fid in feature_ids]]
        importance_series = mi_df.max(axis=0)
        return importance_series, f'max_mi:{self.mi_key}'


@dataclass(frozen=True)
class TopKRocAucSelector(BaseFeatureSelector):
    """Top-k by one-vs-rest ROC-AUC across distortion categories."""

    name: str
    description: str
    category: Category = 'type'
    level: Level = 'patch'
    relation_label: str = 'High ROC-AUC distinguishability'

    @property
    def auc_key(self) -> str:
        return _metric_table_key('auc', self.category)

    def compute_metric_tables(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> Dict[str, pd.DataFrame]:
        selector_params = dict(inputs.selector_params)
        level = str(self.level)
        feature_chunk_size = int(selector_params.get('feature_chunk_size', 512))
        group_col = _group_col(self.category)

        cache_path = selector_params.get('cache_path')
        if cache_path is None:
            cache_paths = inputs.cache_paths
            if cache_paths is not None:
                attr_map = {
                    ('auc_type_df', 'patch'): 'auc_type_patch_cache_path',
                    ('auc_group_df', 'patch'): 'auc_group_patch_cache_path',
                    ('auc_type_df', 'image'): 'auc_type_cache_path',
                    ('auc_group_df', 'image'): 'auc_group_cache_path',
                }
                cache_attr = attr_map.get((self.auc_key, level))
                if cache_attr is not None and hasattr(cache_paths, cache_attr):
                    cache_path = getattr(cache_paths, cache_attr)

        auc_df = compute_distortion_roc_auc(
            meta=inputs.meta_df,
            features=inputs.features,
            group_col=group_col,
            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
            level=level,
            feature_chunk_size=feature_chunk_size,
            cache_path=cache_path,
        )
        return {self.auc_key: auc_df}

    def compute_importance(
        self,
        metric_tables: Dict[str, pd.DataFrame],
        inputs: SelectorInputs | None = None,
        feature_ids: Sequence[int] | None = None,
    ) -> Tuple[pd.Series, str]:
        auc_df = metric_tables[self.auc_key]
        if feature_ids is not None:
            auc_df = auc_df[[int(fid) for fid in feature_ids]]

        # Symmetric separability score: high both for direct and inverse predictors.
        direct_auc = auc_df.max(axis=0)
        inverse_auc = 1.0 - auc_df.min(axis=0)
        importance_series = pd.concat([direct_auc, inverse_auc], axis=1).max(axis=1)
        return importance_series, f'symmetric_auc:{self.auc_key}'


@dataclass(frozen=True)
class TopKPairedDeltaSelector(BaseFeatureSelector):
    """Top-k by original-vs-distorted paired activation deltas."""

    name: str
    description: str
    delta_mode: DeltaMode = 'relative'
    category: Category = 'type'
    relation_label: str = 'высокая парная дельта активаций'

    @property
    def delta_key(self) -> str:
        return _paired_delta_table_key(self.delta_mode, self.category)

    def compute_metric_tables(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> Dict[str, pd.DataFrame]:
        if inputs.pristine_meta_df is None or inputs.pristine_features is None:
            raise ValueError(
                f"Selector '{self.name}' requires 'pristine_meta_df' and 'pristine_features' for paired deltas."
            )

        selector_params = dict(inputs.selector_params)
        level = str(selector_params.get('level', 'patch'))
        ranking_mode = str(selector_params.get('ranking_mode', 'default')).strip().lower()
        if ranking_mode not in {'default', 'weighted_average'}:
            raise ValueError(
                f"Selector '{self.name}': unsupported ranking_mode={ranking_mode!r}. "
                "Allowed: {'default', 'weighted_average'}."
            )
        delta_mode = self.delta_mode
        group_col = _group_col(self.category)

        cache_prefix = selector_params.get('cache_prefix')
        if cache_prefix is None:
            cache_paths = inputs.cache_paths
            dataset_cache_dir = getattr(cache_paths, 'dataset_cache_dir', None) if cache_paths is not None else None
            if dataset_cache_dir is not None:
                cache_prefix = str(Path(dataset_cache_dir) / f'paired_delta_{level}')

        tables = build_paired_delta_tables(
            distorted_meta=inputs.meta_df,
            distorted=inputs.features,
            original_meta=inputs.pristine_meta_df,
            original=inputs.pristine_features,
            group_cols=(group_col,),
            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
            delta_mode=delta_mode,
            level=level,
            cache_prefix=cache_prefix,
        )
        if self.delta_key not in tables:
            available = ', '.join(sorted(tables.keys()))
            raise ValueError(
                f"Selector '{self.name}' expected table '{self.delta_key}', available: {available}"
            )

        metric_table = tables[self.delta_key]
        if ranking_mode == 'weighted_average':
            if not inputs.feature_weight_norms:
                raise ValueError(
                    f"Selector '{self.name}' (ranking_mode='weighted_average') requires per-feature decoder weight norms."
                )

            norm_by_feature_id: Dict[int, float] = {}
            for raw_feature_id, raw_norm in inputs.feature_weight_norms.items():
                try:
                    feature_id = int(raw_feature_id)
                    norm_by_feature_id[feature_id] = float(raw_norm)
                except (TypeError, ValueError):
                    continue

            if not norm_by_feature_id:
                raise ValueError(
                    f"Selector '{self.name}' (ranking_mode='weighted_average') received empty/invalid decoder weight norms."
                )

            weight_series = pd.Series(index=metric_table.columns, dtype='float64')
            missing_columns: List[object] = []
            for column in metric_table.columns:
                try:
                    weight_series.loc[column] = norm_by_feature_id[int(column)]
                except (TypeError, ValueError, KeyError):
                    missing_columns.append(column)

            if missing_columns:
                missing_preview = ', '.join(str(x) for x in missing_columns[:8])
                raise ValueError(
                    f"Selector '{self.name}' has no decoder norms for {len(missing_columns)} features "
                    f"(first: {missing_preview})."
                )

            weighted_abs = metric_table.abs().mul(weight_series, axis=1)
            col_min = weighted_abs.min(axis=0)
            col_max = weighted_abs.max(axis=0)
            denom = col_max - col_min
            normalized = weighted_abs.subtract(col_min, axis=1)
            nonzero_denom = denom > 0
            if nonzero_denom.any():
                normalized.loc[:, nonzero_denom] = normalized.loc[:, nonzero_denom].div(denom[nonzero_denom], axis=1)
            if (~nonzero_denom).any():
                normalized.loc[:, ~nonzero_denom] = 0.0
            metric_table = normalized.astype('float32', copy=False)

        return {self.delta_key: metric_table}

    def compute_importance(
        self,
        metric_tables: Dict[str, pd.DataFrame],
        inputs: SelectorInputs | None = None,
        feature_ids: Sequence[int] | None = None,
    ) -> Tuple[pd.Series, str]:
        delta_df = metric_tables[self.delta_key]
        if feature_ids is not None:
            delta_df = delta_df[[int(fid) for fid in feature_ids]]

        ranking_mode = 'default'
        if inputs is not None:
            ranking_mode = str(inputs.selector_params.get('ranking_mode', 'default')).strip().lower()

        if ranking_mode == 'weighted_average':
            importance_series = delta_df.mean(axis=0)
            return importance_series, f'mean_weighted_delta_norm01:{self.delta_key}'

        delta_abs_mean = delta_df.abs().mean(axis=0)
        return delta_abs_mean, f'mean_abs_delta:{self.delta_key}'


@dataclass(frozen=True)
class TopKIoUSelector(BaseFeatureSelector):
    """Top-k by median IoU with distortion masks."""

    name: str
    description: str
    category: Category = 'type'
    relation_label: str = 'высокая пространственная локализация'

    @property
    def iou_key(self) -> str:
        return _metric_table_key('iou', self.category)

    def compute_metric_tables(
        self,
        inputs: SelectorInputs,
        feature_ids: Sequence[int] | None = None,
    ) -> Dict[str, pd.DataFrame]:
        from analysis.metrics.iou_utils import compute_iou_per_distortion_type_and_feature

        selector_params = dict(inputs.selector_params)
        iou_key = self.iou_key
        group_col = _group_col(self.category)

        cache_path = selector_params.get('cache_path')
        if cache_path is None:
            cache_paths = inputs.cache_paths
            if cache_paths is not None:
                attr_map = {
                    ('iou_type_df',): 'iou_type_cache_path',
                    ('iou_group_df',): 'iou_group_cache_path',
                }
                cache_attr = attr_map.get((iou_key,))
                if cache_attr is not None and hasattr(cache_paths, cache_attr):
                    cache_path = getattr(cache_paths, cache_attr)

        # Infer spatial shape from patch counts
        # Assume square grid: n_patches = height * width
        n_samples_total = inputs.features.codes.shape[0]
        n_unique_images = inputs.meta_df['image_idx'].nunique()
        patches_per_image = n_samples_total // n_unique_images if n_unique_images > 0 else 49
        spatial_dim = int(np.sqrt(patches_per_image))
        
        # Allow override via selector_params
        if 'spatial_shape' in selector_params:
            spatial_shape = tuple(selector_params['spatial_shape'])
        else:
            spatial_shape = (spatial_dim, spatial_dim)

        feature_batch_size = int(selector_params.get('feature_batch_size', 1024))

        iou_df = compute_iou_per_distortion_type_and_feature(
            features=inputs.features,
            meta_df=inputs.meta_df,
            spatial_shape=spatial_shape,
            group_col=group_col,
            global_feature_ids=[int(fid) for fid in feature_ids] if feature_ids is not None else None,
            feature_batch_size=feature_batch_size,
            cache_path=cache_path,
            dataset=inputs.dataset,
        )

        return {iou_key: iou_df}

    def compute_importance(
        self,
        metric_tables: Dict[str, pd.DataFrame],
        inputs: SelectorInputs | None = None,
        feature_ids: Sequence[int] | None = None,
    ) -> Tuple[pd.Series, str]:
        iou_key = self.iou_key
        iou_df = metric_tables[iou_key]
        if iou_df.empty:
            # Return empty series if no IoU data
            if feature_ids is not None:
                feature_ids_int = [int(fid) for fid in feature_ids]
                importance_series = pd.Series(0.0, index=feature_ids_int)
            else:
                importance_series = pd.Series(dtype=float)
            return importance_series, f'max_median_iou:{iou_key}'
        
        if feature_ids is not None:
            iou_df = iou_df[[int(fid) for fid in feature_ids]]

        # Take max IoU across all distortion types (each row is a distortion type)
        importance_series = iou_df.max(axis=0)
        return importance_series, f'max_median_iou:{iou_key}'


def _parse_category(params: Mapping[str, object], *, default: Category = 'type') -> Category:
    value = str(params.get('category', default)).strip().lower()
    if value not in ('type', 'group'):
        raise ValueError(f"Invalid category {value!r}; expected 'type' or 'group'")
    return value  # type: ignore[return-value]


def _parse_level(params: Mapping[str, object], *, default: Level = 'patch') -> Level:
    value = str(params.get('level', default)).strip().lower()
    if value not in ('patch', 'image'):
        raise ValueError(f"Invalid level {value!r}; expected 'patch' or 'image'")
    return value  # type: ignore[return-value]


def _parse_delta_mode(params: Mapping[str, object], *, default: DeltaMode = 'relative') -> DeltaMode:
    value = str(params.get('delta_mode', default)).strip().lower()
    if value not in ('abs', 'signed', 'relative'):
        raise ValueError(f"Invalid delta_mode {value!r}; expected 'abs', 'signed', or 'relative'")
    return value  # type: ignore[return-value]


def canonical_selector_run_key(name: str, params: Mapping[str, object]) -> str:
    """Stable run/cache key derived from base selector name + params."""
    params = dict(params or {})
    base = str(name).strip()
    if base == 'topk_abs_corr':
        return base if _parse_category(params) == 'type' else f'{base}_group'
    if base == 'topk_mi':
        return base if _parse_category(params) == 'type' else f'{base}_group'
    if base == 'topk_auc':
        return f"{base}_{_parse_category(params)}_{_parse_level(params)}"
    if base == 'paired_delta':
        return f"{base}_{_parse_delta_mode(params)}_{_parse_category(params)}"
    if base == 'topk_iou':
        return f'{base}_type' if _parse_category(params) == 'type' else f'{base}_group'
    return base


def instantiate_selector(name: str, params: Mapping[str, object] | None = None) -> BaseFeatureSelector:
    """Build a selector instance from registry name and config params."""
    params = dict(params or {})
    base = str(name).strip()
    run_name = canonical_selector_run_key(base, params)
    if base == 'topk_abs_corr':
        category = _parse_category(params)
        scope = 'dist_type' if category == 'type' else 'dist_group'
        return TopKAbsCorrSelector(
            name=run_name,
            description=f'Top-k features by max absolute {scope} correlation.',
            category=category,
        )
    if base == 'topk_mi':
        category = _parse_category(params)
        scope = 'dist_type' if category == 'type' else 'dist_group'
        return TopKMutualInfoSelector(
            name=run_name,
            description=f'Top-k features by max {scope} mutual information.',
            category=category,
        )
    if base == 'topk_auc':
        category = _parse_category(params)
        level = _parse_level(params)
        scope = 'dist_type' if category == 'type' else 'dist_group'
        return TopKRocAucSelector(
            name=run_name,
            description=f'Top-k features by ROC-AUC at {scope} level ({level}).',
            category=category,
            level=level,
        )
    if base == 'paired_delta':
        category = _parse_category(params)
        delta_mode = _parse_delta_mode(params)
        scope = 'dist_type' if category == 'type' else 'dist_group'
        return TopKPairedDeltaSelector(
            name=run_name,
            description=f'Top-k features by {delta_mode} paired delta at {scope} level.',
            delta_mode=delta_mode,
            category=category,
        )
    if base == 'topk_iou':
        category = _parse_category(params)
        scope = 'dist_type' if category == 'type' else 'dist_group'
        return TopKIoUSelector(
            name=run_name,
            description=f'Top-k features by median IoU with distortion regions at {scope} level.',
            category=category,
        )
    available = ', '.join(sorted(SELECTOR_REGISTRY))
    raise ValueError(f"Unknown selector '{base}'. Available: {available}")


SELECTOR_REGISTRY: Tuple[str, ...] = (
    'topk_abs_corr',
    'topk_mi',
    'topk_auc',
    'paired_delta',
    'topk_iou',
)


def get_selector_registry() -> Dict[str, BaseFeatureSelector]:
    """Return default instances for each base selector (default params)."""
    return {name: instantiate_selector(name, {}) for name in SELECTOR_REGISTRY}


def load_selector(selector_name: str, params: Mapping[str, object] | None = None) -> BaseFeatureSelector:
    """Load built-in selector by name and optional params."""
    if not str(selector_name).strip():
        raise ValueError('selector_name is required')
    return instantiate_selector(selector_name, params)


def run_selector(
    selector: BaseFeatureSelector,
    inputs: SelectorInputs,
    feature_ids: Sequence[int] | None = None,
) -> pd.DataFrame:
    """Compatibility wrapper around selector.select_features."""
    selected_features = selector.select_features(inputs, feature_ids=feature_ids)
    required_cols = {'feature_id', 'importance_score', 'importance_rank', 'importance_metric'}
    missing_cols = required_cols.difference(selected_features.columns)
    if missing_cols:
        missing = ', '.join(sorted(missing_cols))
        raise ValueError(f'Selector result is missing required columns: {missing}')
    return selected_features


def run_selector_with_metrics(
    selector: BaseFeatureSelector,
    inputs: SelectorInputs,
    feature_ids: Sequence[int] | None = None,
) -> SelectorResult:
    """Run selector and return selected features plus computed metric tables."""
    result = selector.execute(inputs, feature_ids=feature_ids)
    required_cols = {'feature_id', 'importance_score', 'importance_rank', 'importance_metric'}
    missing_cols = required_cols.difference(result.selected_features.columns)
    if missing_cols:
        missing = ', '.join(sorted(missing_cols))
        raise ValueError(f'Selector result is missing required columns: {missing}')
    return result


def selected_feature_ids(selected_features: pd.DataFrame) -> List[int]:
    """Convenience helper to get ordered feature ids from selector output."""
    return [int(x) for x in selected_features['feature_id'].tolist()]


def _normalize_selector_configs(selector_configs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    if not selector_configs:
        raise ValueError("selector_configs must contain at least one selector config")

    normalized: List[Dict[str, Any]] = []
    for selector_cfg in selector_configs:
        if not isinstance(selector_cfg, dict):
            raise ValueError("Each selector config must be a dict: {'name': str, 'params': dict, 'top_k': int}")

        selector_name = str(selector_cfg.get('name', '')).strip()
        if not selector_name:
            raise ValueError("Each selector config must include non-empty 'name'")

        if selector_name not in SELECTOR_REGISTRY:
            available = ', '.join(sorted(SELECTOR_REGISTRY))
            raise ValueError(f"Unknown selector '{selector_name}'. Available: {available}")

        if 'params' not in selector_cfg:
            raise ValueError(f"Selector '{selector_name}': 'params' is required and must be a dict")
        selector_params = selector_cfg['params']
        if selector_params is None:
            selector_params = {}
        if not isinstance(selector_params, dict):
            raise ValueError(f"Selector '{selector_name}': 'params' must be a dict")

        instantiate_selector(selector_name, selector_params)

        if 'top_k' not in selector_cfg:
            raise ValueError(f"Selector '{selector_name}': 'top_k' is required and must be positive")
        selector_top_k = int(selector_cfg['top_k'])
        if selector_top_k <= 0:
            raise ValueError(f"Selector '{selector_name}': 'top_k' must be positive")

        selector_feature_ids = selector_cfg.get('feature_ids')
        if selector_feature_ids is not None:
            if not isinstance(selector_feature_ids, (list, tuple)):
                raise ValueError(f"Selector '{selector_name}': 'feature_ids' must be list/tuple or null")
            selector_feature_ids = [int(fid) for fid in selector_feature_ids]

        normalized.append({
            'name': selector_name,
            'params': selector_params,
            'top_k': selector_top_k,
            'feature_ids': selector_feature_ids,
        })

    return normalized


def run_selectors_from_configs(
    *,
    selector_configs: List[Dict[str, Any]],
    meta_df: Any,
    features: FeatureMatrix,
    pristine_meta_df: Any,
    pristine_features: FeatureMatrix | None,
    dataset: str,
    sae_checkpoint_path: str,
    cache_dir: str,
    cache_paths: Any = None,
    selector_registry: Optional[Dict[str, BaseFeatureSelector]] = None,
    feature_weight_norms: Optional[Mapping[int, float]] = None,
) -> Dict[str, SelectorResult]:
    """Execute configured selectors and return full per-run selector outputs."""
    from analysis.models import extract_decoder_weight_norms

    normalized_configs = _normalize_selector_configs(selector_configs)

    needs_pristine_inputs = any(
        selector_cfg['name'] == 'paired_delta'
        for selector_cfg in normalized_configs
    )
    if not needs_pristine_inputs:
        pristine_meta_df = None
        pristine_features = None

    selector_entries = []
    for selector_cfg in normalized_configs:
        base_name = selector_cfg['name']
        selector_params = dict(selector_cfg.get('params', {}))
        if selector_registry is not None and base_name in selector_registry:
            selector = selector_registry[base_name]
        else:
            selector = instantiate_selector(base_name, selector_params)
        run_key = canonical_selector_run_key(base_name, selector_params)
        selector_entries.append((run_key, selector, selector_cfg))

    needs_weighted_delta = False
    for selector_cfg in normalized_configs:
        selector_params = selector_cfg.get('params', {})
        if selector_params and str(selector_params.get('ranking_mode', 'default')) == 'weighted_average':
            needs_weighted_delta = True
            break

    resolved_weight_norms: Optional[Mapping[int, float]] = feature_weight_norms
    if needs_weighted_delta and not resolved_weight_norms:
        logger.info('[run] Loading SAE decoder norms for weighted_average ranking mode...')
        resolved_weight_norms = extract_decoder_weight_norms(
            checkpoint_path=sae_checkpoint_path,
            cache_dir=cache_dir,
        )
        logger.info('[run] Loaded decoder norms for %s SAE features.', len(resolved_weight_norms))

    selector_run_key_counts: Dict[str, int] = {}
    for run_key, _, _ in selector_entries:
        selector_run_key_counts[run_key] = selector_run_key_counts.get(run_key, 0) + 1

    selector_results: Dict[str, SelectorResult] = {}
    for selector_idx, (run_key, selector, selector_cfg) in enumerate(selector_entries, start=1):
        selector_run_key = run_key
        if selector_run_key_counts[run_key] > 1:
            selector_run_key = f'{run_key}__{selector_idx}'

        selector_inputs = SelectorInputs(
            meta_df=meta_df,
            features=features,
            dataset=dataset,
            pristine_meta_df=pristine_meta_df,
            pristine_features=pristine_features,
            feature_weight_norms=resolved_weight_norms,
            cache_paths=cache_paths,
            selector_top_k=int(selector_cfg['top_k']),
            selector_params=dict(selector_cfg.get('params', {})),
        )
        selector_feature_ids = selector_cfg.get('feature_ids')
        selector_results[selector_run_key] = selector.execute(
            selector_inputs,
            feature_ids=selector_feature_ids,
        )

    return selector_results