"""
Статистики по SAE-признакам и их визуализация (bar charts).
"""

from typing import List, Optional, Tuple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from analysis.features.feature_indexing import FeatureMatrix


def compute_feature_stats(features: FeatureMatrix) -> pd.DataFrame:
    """
    Вычисляет статистики для каждого признака SAE.

    Параметры
    ----------
    features : CSR activations; ``column_feature_ids[j]`` = global SAE id

    Возвращает
    ----------
    DataFrame с колонками: feature_id (global SAE id), mean, frequency, max, mean_acts
        mean      — средняя активация по всем патчам
        frequency — доля патчей с ненулевой активацией
        max       — максимальная активация среди всех патчей
        mean_acts — средняя активация по ненулевым патчам
    """
    codes = features.codes
    mat = codes if codes.dtype == np.float32 else codes.astype(np.float32)

    n_rows = mat.shape[0]
    means = np.asarray(mat.mean(axis=0)).ravel()

    freqs = np.asarray(mat.getnnz(axis=0), dtype=np.float32).ravel() / float(n_rows)

    maxes = np.asarray(mat.tocsc().max(axis=0).toarray()).ravel()
    mean_acts = np.where(freqs > 0, means / freqs, 0.0)

    col_ids = [int(fid) for fid in features.column_feature_ids]
    return pd.DataFrame({
        'feature_id': col_ids,
        'mean':       means,
        'frequency':  freqs,
        'max':        maxes,
        'mean_acts':  mean_acts,
    })


def get_top_features(
    stats: pd.DataFrame,
    top_k: int = 20,
    criterion: str = 'mean',
    min_mean_acts: Optional[float] = None,
) -> List[int]:
    """
    Возвращает индексы top-K признаков по выбранному критерию.

    Параметры
    ----------
    stats         : DataFrame из compute_feature_stats
    top_k         : число топ-признаков
    criterion     : 'mean' | 'frequency' | 'max' | 'mean_acts'
    min_mean_acts : предварительный фильтр по mean_acts

    Возвращает
    ----------
    List[int] — feature_id в порядке убывания критерия
    """
    assert criterion in ('mean', 'frequency', 'max', 'mean_acts'), \
        f"criterion must be one of 'mean', 'frequency', 'max', 'mean_acts', got {criterion!r}"
    filtered = stats
    if min_mean_acts is not None:
        filtered = stats[stats['mean_acts'] >= min_mean_acts]
    return filtered.nlargest(top_k, criterion)['feature_id'].tolist()


def plot_top_features(
    stats: pd.DataFrame,
    top_k: int = 20,
    criterion: str = 'mean',
    figsize: Tuple[int, int] = (14, 4),
    min_mean_acts: Optional[float] = None,
) -> None:
    """
    Bar-chart топ-K признаков по выбранному критерию.
    Над каждым баром подписывается frequency (доля ненулевых патчей).
    """
    top_ids = get_top_features(stats, top_k=top_k, criterion=criterion,
                               min_mean_acts=min_mean_acts)
    top_stats = stats.set_index('feature_id').loc[top_ids]

    fig, ax = plt.subplots(figsize=figsize)
    bars = ax.bar(range(len(top_ids)), top_stats[criterion].values, color='steelblue')

    for bar, freq_val in zip(bars, top_stats['frequency'].values):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            bar.get_height(),
            f'{freq_val:.3f}',
            ha='center', va='bottom', fontsize=6, color='black', rotation=90,
        )

    ax.set_xticks(range(len(top_ids)))
    ax.set_xticklabels([str(i) for i in top_ids], rotation=90, fontsize=8)
    ax.set_xlabel('Feature ID')
    ax.set_ylabel(criterion)
    ax.set_title(f'Top-{top_k} features by {criterion}')
    plt.tight_layout()
    plt.show()