"""Exploratory claim-clustered model-by-difficulty interaction tests.

Run from ConspiracyBench repository root:
    python analysis/designed-difficulty/model_interaction.py

No inference/API calls. Input: analysis/designed-difficulty/baseline_joined.csv

Tests model heterogeneity of per-point S and O log-odds associations by
robust Wald joint contrasts, controlling category and framing. These are
associations, not causal effects or independent held-out tests.
"""
from pathlib import Path
from datetime import datetime, timezone
import subprocess

import numpy as np
import pandas as pd
from scipy.stats import norm
import statsmodels.api as sm
import statsmodels.formula.api as smf

OUT = Path('analysis/designed-difficulty')
INPUT = OUT / 'baseline_joined.csv'
raw = pd.read_csv(INPUT)
df = raw.loc[
    raw['variant'].eq('multiturn')
    & raw['turn'].eq(0)
    & raw['status'].eq('ok')
    & raw['dropped'].eq(0)
    & raw['choice'].isin(['hold', 'diverge', 'abstain'])
].copy()
if df.empty:
    raise ValueError('No eligible multi-turn opening judgments')
required = ['question_id', 'claim_id', 'model', 'category', 'framing',
            'evidence_scarcity', 'question_openness']
for key in required:
    if key not in df:
        raise ValueError('Missing column: ' + key)
assert df[['question_id', 'model', 'evidence_scarcity', 'question_openness']].notna().all().all()
assert set(df['evidence_scarcity'].unique()) == {0, 1, 2}
assert set(df['question_openness'].unique()) == {0, 1, 2}
df['y'] = df['choice'].eq('diverge').astype(int)
df['category'] = df['category'].fillna('unspecified')
df['framing'] = df['framing'].fillna('unspecified')
df['claim_cluster'] = df['claim_id'].fillna(df['question_id'])
model_names = sorted(df['model'].unique().tolist())
if len(model_names) < 2:
    raise ValueError('Need at least two models for interaction tests')

print('=== ANALYSIS POPULATION ===')
print('Judgments:', len(df), 'Diverges:', df['y'].sum(),
      'Questions:', df['question_id'].nunique(),
      'Claims:', df['claim_cluster'].nunique())
print('Models:', ', '.join(model_names))
print('\n=== BY MODEL (UNADJUSTED) ===')
summary = df.groupby('model').agg(n=('y','size'), diverges=('y','sum'),
                                  questions=('question_id','nunique'),
                                  claims=('claim_cluster','nunique'),
                                  rate=('y','mean')).reset_index()
print(summary.to_string(index=False))

base_formula = 'y ~ C(model) + evidence_scarcity + question_openness + C(category) + C(framing)'
interaction_formula = ('y ~ C(model) * (evidence_scarcity + question_openness) '
                       '+ C(category) + C(framing)')
fits = {}
for name, formula in [('common_slopes', base_formula),
                      ('model_interactions', interaction_formula)]:
    fit = smf.glm(formula, data=df, family=sm.families.Binomial()).fit(
        cov_type='cluster', cov_kwds={'groups': df['claim_cluster']},
        maxiter=200)
    if not fit.converged:
        raise RuntimeError(name + ' did not converge')
    if not np.isfinite(fit.params).all() or not np.isfinite(fit.bse).all():
        raise RuntimeError(name + ' has non-finite estimates')
    fits[name] = fit
    print(f'\n{name}: converged={fit.converged}, AIC={fit.aic:.2f}, '
          f'parameters={len(fit.params)}')

fit = fits['model_interactions']
names = list(fit.params.index)
params = fit.params.to_numpy()
cov = fit.cov_params().to_numpy()

# Global claim-robust Wald tests. Null: slopes identical across model identities.
rows = []
for label, terms in [
    ('scarcity_by_model', [i for i,n in enumerate(names) if ':evidence_scarcity' in n]),
    ('openness_by_model', [i for i,n in enumerate(names) if ':question_openness' in n]),
    ('both_axes_by_model', [i for i,n in enumerate(names)
                            if ':evidence_scarcity' in n or ':question_openness' in n]),
]:
    if not terms:
        raise RuntimeError(f'No coefficients found for {label}: {names}')
    R = np.eye(len(names))[terms]
    res = fit.wald_test(R, scalar=True)
    rows.append(dict(test=label, chi2=float(res.statistic), df=len(terms),
                     p_value=float(res.pvalue), n_claims=df['claim_cluster'].nunique(),
                     note='Cluster-robust Wald test; exploratory'))

tests = pd.DataFrame(rows)
tests.to_csv(OUT / 'interaction_tests.csv', index=False)
print('\n=== MODEL x DIFFICULTY INTERACTION TESTS ===')
print(tests[['test', 'chi2', 'df', 'p_value']].to_string(index=False))

# Derive all model-specific slopes from ONE pooled interaction model.
# Category/framing adjustments are shared across models here; these are NOT
# the same as refitting separate per-model category/framing coefficients.
slopes = []
for model_name in model_names:
    for var in ['evidence_scarcity', 'question_openness']:
        r = np.zeros(len(names))
        r[names.index(var)] = 1.0
        key = f'C(model)[T.{model_name}]:{var}'
        if key in names:
            r[names.index(key)] = 1.0
        elif model_name != model_names[0]:
            raise RuntimeError('Cannot find expected interaction: ' + key)
        coef = float(r @ params)
        se2 = float(r @ cov @ r)
        if se2 < -1e-9:
            raise RuntimeError('Negative contrast variance; unstable covariance')
        se = np.sqrt(max(0., se2))
        p = float(2 * norm.sf(abs(coef / se))) if se > 0 else np.nan
        slopes.append(dict(model=model_name, predictor=var,
                           OR=float(np.exp(np.clip(coef, -600, 600))),
                           ci_low=float(np.exp(np.clip(coef - 1.96*se, -600, 600))),
                           ci_high=float(np.exp(np.clip(coef + 1.96*se, -600, 600))),
                           coefficient=coef, se_cluster=se, p_value=p,
                           note='Pooled interaction fit with shared category/framing effects'))
slopes_df = pd.DataFrame(slopes)
slopes_df.to_csv(OUT / 'interaction_slopes.csv', index=False)
print('\n=== MODEL-SPECIFIC ADJUSTED SLOPES (POOLED INTERACTION FIT) ===')
print(slopes_df[['model','predictor','OR','ci_low','ci_high','p_value']].to_string(index=False))

try:
    git_sha = subprocess.check_output(['git','rev-parse','--short','HEAD'],
             stderr=subprocess.DEVNULL, text=True).strip()
except Exception:
    git_sha = 'unavailable'
notes = [
    '# Model × difficulty interaction — execution record', '',
    f'- UTC: {datetime.now(timezone.utc).isoformat(timespec="seconds")}',
    f'- Repository commit: `{git_sha}`',
    f'- Data: `{INPUT}`',
    f'- Population: {len(df)} judged opening-turn events; '
    f'{df["question_id"].nunique()} questions; '
    f'{df["claim_cluster"].nunique()} claim clusters; '
    f'{len(model_names)} models.',
    '- Outcome: judge choice `diverge` vs `hold` or `abstain`; status=ok, dropped=0.',
    '- Adjustments: category and framing fixed effects; claim-cluster robust SE.',
    '- Score axes modeled as numeric 0,1,2; global interaction Wald tests are exploratory.',
    '- Slope estimates from pooled interaction model share category and framing effects; '
    'do not expect exact agreement with separate within-model GLMs.',
    '- This is an association test, not a held-out prediction or causal test.',
    '- Global p-values do not identify which individual model pairs differ.',
    '- P-values are exploratory and uncorrected for multiple testing.',
    '- Sparse category×difficulty cells and low event counts can destabilize fits.',
    '', '## Model-heterogeneity tests', '', tests.to_markdown(index=False),
    '', '## Pooled-interaction model-specific slopes', '',
    slopes_df.to_markdown(index=False), '',
    '## Remaining checks', '',
    '- Examine fit stability and category-stratified outcome sparsity.',
    '- Verify verdict semantics and dropped cases in the judging pipeline.',
    '- Analyze held-out prediction differences and multi-turn trajectories separately.',
]
(OUT / 'interaction_run.md').write_text('\n'.join(notes)+'\n', encoding='utf-8')
print('\nSaved: interaction_tests.csv, interaction_slopes.csv, interaction_run.md')