experiments/analyze_study.py

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd
from scipy.stats import spearmanr

from experiments.common import ARTIFACT_DIR
from featurelens.stats import paired_bootstrap_difference_ci, paired_sign_flip_pvalue

POLICIES = ('final_token', 'max_feature_activation')


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Aggregate held-out association, stability, and causal evidence by concept.'
    )
    parser.add_argument('--artifact-dir', type=Path, default=ARTIFACT_DIR)
    parser.add_argument('--seed', type=int, default=42)
    return parser.parse_args()


def selected_features(catalog: pd.DataFrame) -> pd.DataFrame:
    scored = catalog.copy()
    scored['activation_contrast'] = (
        scored['activation_rate_pos'].astype(float) - scored['activation_rate_neg'].astype(float)
    )
    ordered = scored.sort_values(
        ['concept', 'train_auroc', 'activation_contrast'],
        ascending=[True, False, False],
    )
    return ordered.groupby('concept', as_index=False).first()


def causal_path(artifact_dir: Path, policy: str) -> Path:
    explicit = artifact_dir / f'causal_results_{policy if policy == "final_token" else "max_active"}.csv'
    if explicit.exists():
        return explicit
    if policy == 'final_token':
        legacy = artifact_dir / 'causal_results.csv'
        if legacy.exists():
            return legacy
    return explicit


def _task_level_specificity(
    frame: pd.DataFrame,
    *,
    effect_column: str,
    seed: int,
    active_only: bool = False,
) -> dict[str, float | list[float]]:
    if frame.empty:
        return _empty_specificity()

    sae_rows = frame[frame['condition'] == 'sae_feature'].copy()
    if active_only:
        active_column = (
            'feature_active_at_intervention'
            if 'feature_active_at_intervention' in sae_rows.columns
            else 'feature_activation'
        )
        if active_column == 'feature_activation':
            active_tasks = sae_rows.loc[sae_rows[active_column].astype(float) > 0.0, 'task_id'].unique()
        else:
            active_tasks = sae_rows.loc[sae_rows[active_column].astype(float) > 0.0, 'task_id'].unique()
        frame = frame[frame['task_id'].isin(active_tasks)].copy()
        sae_rows = frame[frame['condition'] == 'sae_feature'].copy()

    sae = (
        sae_rows.assign(
            _abs_effect=lambda data: np.abs(pd.to_numeric(data[effect_column], errors='coerce'))
        )
        .groupby('task_id', as_index=False)
        .agg(sae_abs_effect=('_abs_effect', 'mean'), sae_signed_effect=(effect_column, 'mean'))
    )
    random_rows = frame[frame['condition'] == 'random_norm_matched'].assign(
        _abs_effect=lambda data: np.abs(pd.to_numeric(data[effect_column], errors='coerce'))
    )
    if 'intervention' in random_rows.columns:
        random = (
            random_rows.groupby(['task_id', 'intervention'], as_index=False)['_abs_effect']
            .mean()
            .groupby('task_id', as_index=False)['_abs_effect']
            .mean()
            .rename(columns={'_abs_effect': 'random_abs_effect'})
        )
    else:
        random = (
            random_rows.groupby('task_id', as_index=False)['_abs_effect']
            .mean()
            .rename(columns={'_abs_effect': 'random_abs_effect'})
        )
    paired = sae.merge(random, on='task_id', how='inner')
    if paired.empty:
        return _empty_specificity()

    sae_abs = paired['sae_abs_effect'].to_numpy(dtype=float)
    sae_signed = paired['sae_signed_effect'].to_numpy(dtype=float)
    random_abs = paired['random_abs_effect'].to_numpy(dtype=float)
    ci_low, ci_high = paired_bootstrap_difference_ci(sae_abs, random_abs, seed=seed)
    return {
        'sae_abs_mean': float(np.mean(sae_abs)),
        'sae_signed_mean': float(np.mean(sae_signed)),
        'random_abs_mean': float(np.mean(random_abs)),
        'specificity_ratio': float(np.mean(sae_abs) / max(float(np.mean(random_abs)), 1e-12)),
        'paired_advantage': float(np.mean(sae_abs - random_abs)),
        'paired_advantage_ci_95': [float(ci_low), float(ci_high)],
        'paired_sign_flip_pvalue': float(paired_sign_flip_pvalue(sae_abs, random_abs, seed=seed + 1)),
        'n_tasks': int(len(paired)),
    }


def _empty_specificity() -> dict[str, float | list[float]]:
    return {
        'sae_abs_mean': float('nan'),
        'sae_signed_mean': float('nan'),
        'random_abs_mean': float('nan'),
        'specificity_ratio': float('nan'),
        'paired_advantage': float('nan'),
        'paired_advantage_ci_95': [float('nan'), float('nan')],
        'paired_sign_flip_pvalue': float('nan'),
        'n_tasks': 0,
    }


def _paired_specificity(
    frame: pd.DataFrame,
    *,
    effect_column: str,
    task_column: str = 'task_id',
    seed: int,
) -> dict[str, float | list[float]]:
    """Backward-compatible public helper using task-level inference."""
    if task_column != 'task_id':
        frame = frame.rename(columns={task_column: 'task_id'})
    return _task_level_specificity(frame, effect_column=effect_column, seed=seed)


def _safe_spearman(x: pd.Series, y: pd.Series) -> dict[str, float | int]:
    a = pd.to_numeric(x, errors='coerce').to_numpy(dtype=float)
    b = pd.to_numeric(y, errors='coerce').to_numpy(dtype=float)
    mask = np.isfinite(a) & np.isfinite(b)
    if int(mask.sum()) < 3:
        return {'rho': float('nan'), 'pvalue': float('nan'), 'n': int(mask.sum())}
    if np.unique(a[mask]).size < 2 or np.unique(b[mask]).size < 2:
        return {'rho': float('nan'), 'pvalue': float('nan'), 'n': int(mask.sum())}
    result = spearmanr(a[mask], b[mask])
    return {'rho': float(result.statistic), 'pvalue': float(result.pvalue), 'n': int(mask.sum())}


def _coverage(sae_rows: pd.DataFrame, column: str, fallback: str | None = None) -> float:
    if column in sae_rows.columns:
        return float(pd.to_numeric(sae_rows[column], errors='coerce').fillna(0).astype(float).gt(0).mean())
    if fallback and fallback in sae_rows.columns:
        return float(pd.to_numeric(sae_rows[fallback], errors='coerce').fillna(0).astype(float).gt(0).mean())
    return float('nan')


def _policy_summary(frame: pd.DataFrame, policy: str, seed: int) -> dict:
    target = _task_level_specificity(frame, effect_column='target_mean_logprob_delta', seed=seed)
    target_active = _task_level_specificity(
        frame, effect_column='target_mean_logprob_delta', seed=seed + 7, active_only=True
    )
    js = _task_level_specificity(frame, effect_column='js_divergence', seed=seed + 17)
    js_active = _task_level_specificity(
        frame, effect_column='js_divergence', seed=seed + 29, active_only=True
    )
    sae_rows = frame[frame['condition'] == 'sae_feature']
    return {
        'position_policy': policy,
        'tasks': int(target['n_tasks']),
        'feature_active_at_intervention_rate': _coverage(
            sae_rows, 'feature_active_at_intervention', fallback='feature_activation'
        ),
        'feature_active_at_final_token_rate': _coverage(
            sae_rows, 'feature_active_at_final_token', fallback='feature_activation'
        ),
        'feature_active_anywhere_rate': _coverage(
            sae_rows, 'feature_active_anywhere', fallback='feature_activation'
        ),
        'target_sae_abs_mean': target['sae_abs_mean'],
        'target_random_abs_mean': target['random_abs_mean'],
        'target_specificity_ratio': target['specificity_ratio'],
        'target_paired_advantage': target['paired_advantage'],
        'target_paired_ci_low': float(target['paired_advantage_ci_95'][0]),
        'target_paired_ci_high': float(target['paired_advantage_ci_95'][1]),
        'target_sign_flip_pvalue': target['paired_sign_flip_pvalue'],
        'active_target_sae_abs_mean': target_active['sae_abs_mean'],
        'active_target_random_abs_mean': target_active['random_abs_mean'],
        'active_target_specificity_ratio': target_active['specificity_ratio'],
        'active_target_paired_advantage': target_active['paired_advantage'],
        'active_target_paired_ci_low': float(target_active['paired_advantage_ci_95'][0]),
        'active_target_paired_ci_high': float(target_active['paired_advantage_ci_95'][1]),
        'active_target_sign_flip_pvalue': target_active['paired_sign_flip_pvalue'],
        'active_tasks': int(target_active['n_tasks']),
        'js_sae_mean': js['sae_abs_mean'],
        'js_random_mean': js['random_abs_mean'],
        'js_specificity_ratio': js['specificity_ratio'],
        'active_js_specificity_ratio': js_active['specificity_ratio'],
    }


def main() -> None:
    args = parse_args()
    artifact_dir = args.artifact_dir
    catalog = pd.read_csv(artifact_dir / 'feature_catalog.csv')
    paraphrase = pd.read_csv(artifact_dir / 'stability.csv')
    stability_path = artifact_dir / 'selection_stability.csv'
    selection_stability = pd.read_csv(stability_path) if stability_path.exists() else pd.DataFrame()

    causal_by_policy: dict[str, pd.DataFrame] = {}
    for policy in POLICIES:
        path = causal_path(artifact_dir, policy)
        if path.exists():
            frame = pd.read_csv(path)
            if 'position_policy' not in frame.columns:
                frame['position_policy'] = policy
            causal_by_policy[policy] = frame

    if 'final_token' not in causal_by_policy or 'max_feature_activation' not in causal_by_policy:
        missing = [policy for policy in POLICIES if policy not in causal_by_policy]
        raise SystemExit(f'Missing causal position policies: {missing}')

    selected = selected_features(catalog)
    position_rows: list[dict] = []
    for policy_idx, policy in enumerate(POLICIES):
        overall = _policy_summary(causal_by_policy[policy], policy, args.seed + 1000 * policy_idx)
        overall['concept'] = '__all__'
        position_rows.append(overall)
        for concept_idx, concept in enumerate(sorted(selected['concept'].astype(str).unique())):
            subset = causal_by_policy[policy][causal_by_policy[policy]['concept'] == concept].copy()
            row = _policy_summary(
                subset,
                policy,
                args.seed + 1000 * policy_idx + 100 * (concept_idx + 1),
            )
            row['concept'] = concept
            position_rows.append(row)

    position_summary = pd.DataFrame(position_rows)
    position_summary.to_csv(artifact_dir / 'causal_position_summary.csv', index=False)

    rows: list[dict] = []
    for concept_idx, selected_row in selected.sort_values('concept').reset_index(drop=True).iterrows():
        concept = str(selected_row['concept'])
        layer = int(selected_row['layer'])
        feature_id = int(selected_row['feature_id'])
        para_rows = paraphrase[
            (paraphrase['concept'] == concept) & (paraphrase['layer'].astype(int) == layer)
        ]
        selection_row = pd.DataFrame()
        if not selection_stability.empty:
            selection_row = selection_stability[
                (selection_stability['concept'] == concept)
                & (selection_stability['layer'].astype(int) == layer)
                & (selection_stability['feature_id'].astype(int) == feature_id)
            ]

        base = {
            'concept': concept,
            'layer': layer,
            'feature_id': feature_id,
            'train_auroc': float(selected_row['train_auroc']),
            'heldout_auroc': float(selected_row['auroc']),
            'heldout_f1': float(selected_row['f1']),
            'activation_rate_pos_train': float(selected_row['activation_rate_pos']),
            'activation_rate_neg_train': float(selected_row['activation_rate_neg']),
            'candidate_resample_support': (
                float(selection_row.iloc[0]['resample_support']) if not selection_row.empty else 0.0
            ),
            'candidate_median_resample_rank': (
                float(selection_row.iloc[0]['median_resample_rank'])
                if not selection_row.empty
                else float('nan')
            ),
            'mean_paraphrase_topk_jaccard': (
                float(para_rows['topk_jaccard'].mean()) if not para_rows.empty else float('nan')
            ),
            'mean_paraphrase_sparse_cosine': (
                float(para_rows['sparse_cosine'].mean()) if not para_rows.empty else float('nan')
            ),
        }
        for policy in POLICIES:
            policy_frame = causal_by_policy[policy]
            subset = policy_frame[policy_frame['concept'] == concept].copy()
            summary = _policy_summary(
                subset,
                policy,
                args.seed + 10_000 + 1000 * POLICIES.index(policy) + 100 * concept_idx,
            )
            prefix = 'final' if policy == 'final_token' else 'max_active'
            for key, value in summary.items():
                if key in {'position_policy'}:
                    continue
                base[f'{prefix}_{key}'] = value
        base['target_specificity_gain_max_vs_final'] = (
            float(base['max_active_target_specificity_ratio'])
            - float(base['final_target_specificity_ratio'])
        )
        base['js_specificity_gain_max_vs_final'] = (
            float(base['max_active_js_specificity_ratio'])
            - float(base['final_js_specificity_ratio'])
        )
        rows.append(base)

    study = pd.DataFrame(rows)
    study_path = artifact_dir / 'study_feature_summary.csv'
    study.to_csv(study_path, index=False)

    correlations = {
        'heldout_auroc_vs_max_active_target_specificity': _safe_spearman(
            study['heldout_auroc'], study['max_active_target_specificity_ratio']
        ),
        'heldout_auroc_vs_max_active_js_specificity': _safe_spearman(
            study['heldout_auroc'], study['max_active_js_specificity_ratio']
        ),
        'heldout_f1_vs_max_active_target_specificity': _safe_spearman(
            study['heldout_f1'], study['max_active_target_specificity_ratio']
        ),
        'candidate_resample_support_vs_max_active_target_specificity': _safe_spearman(
            study['candidate_resample_support'], study['max_active_target_specificity_ratio']
        ),
    }

    overall = position_summary[position_summary['concept'] == '__all__'].set_index('position_policy')
    final_row = overall.loc['final_token']
    max_row = overall.loc['max_feature_activation']
    most_predictive = study.sort_values('heldout_auroc', ascending=False).iloc[0]
    most_target_specific = study.sort_values('max_active_target_specificity_ratio', ascending=False).iloc[0]
    most_js_specific = study.sort_values('max_active_js_specificity_ratio', ascending=False).iloc[0]

    summary = {
        'n_concepts': int(len(study)),
        'selected_feature_pooling': 'prompt-wide max SAE activation across non-padding prompt tokens',
        'dense_probe_pooling': 'final prompt token residual',
        'primary_causal_position_policy': 'max_feature_activation',
        'causal_statistical_unit': 'causal task; ablation and amplification are averaged within task before paired inference',
        'median_selected_feature_resample_support': float(study['candidate_resample_support'].median()),
        'final_token_feature_coverage': float(final_row['feature_active_at_intervention_rate']),
        'max_active_feature_coverage': float(max_row['feature_active_at_intervention_rate']),
        'final_token_target_specificity_ratio': float(final_row['target_specificity_ratio']),
        'max_active_target_specificity_ratio': float(max_row['target_specificity_ratio']),
        'final_token_target_paired_advantage': float(final_row['target_paired_advantage']),
        'max_active_target_paired_advantage': float(max_row['target_paired_advantage']),
        'final_token_target_paired_ci_95': [
            float(final_row['target_paired_ci_low']), float(final_row['target_paired_ci_high'])
        ],
        'max_active_target_paired_ci_95': [
            float(max_row['target_paired_ci_low']), float(max_row['target_paired_ci_high'])
        ],
        'final_token_target_sign_flip_pvalue': float(final_row['target_sign_flip_pvalue']),
        'max_active_target_sign_flip_pvalue': float(max_row['target_sign_flip_pvalue']),
        'most_predictive_concept': {
            'concept': str(most_predictive['concept']),
            'heldout_auroc': float(most_predictive['heldout_auroc']),
        },
        'highest_max_active_target_specificity': {
            'concept': str(most_target_specific['concept']),
            'ratio': float(most_target_specific['max_active_target_specificity_ratio']),
        },
        'highest_max_active_js_specificity': {
            'concept': str(most_js_specific['concept']),
            'ratio': float(most_js_specific['max_active_js_specificity_ratio']),
        },
        'correlations': correlations,
        'guardrail': (
            'Max-active causal positions are selected from SAE activation only, never from behavioral outcome. '
            'Cross-concept Spearman correlations are descriptive because the study contains seven concepts.'
        ),
    }
    (artifact_dir / 'study_summary.json').write_text(json.dumps(summary, indent=2), encoding='utf-8')
    print(f'Wrote {artifact_dir / "causal_position_summary.csv"}')
    print(f'Wrote {study_path}')
    print(f'Wrote {artifact_dir / "study_summary.json"}')


if __name__ == '__main__':
    main()