"""Run in the existing Gate 8/9 Colab: %run gate9_audit.py
Requires the original orientation-level dataframe `d`; never edits it.
Writes only to a newly created timestamped audit directory.
"""
from pathlib import Path
from datetime import datetime, timezone
import json
import numpy as np
import pandas as pd


def audit(d, root=Path('/content/drive/MyDrive/proteus')):
    x = d.copy(deep=True)
    required = ['WT_name', 'm1', 'm2', 'ddG_1', 'ddG_2', 'dmg', 'prt',
                'physical_event', 'isotonic_heldout_event']
    assert set(required) <= set(x), set(required) - set(x)
    for c in ['physical_event', 'isotonic_heldout_event']:
        assert x[c].notna().all() and x[c].isin([True, False, 0, 1]).all(), c
    x['robust'] = x.physical_event.astype(bool) & x.isotonic_heldout_event.astype(bool)
    # Reproduce Gate 9A, including its original default relative tolerance.
    f = np.isclose(x.dmg, x.ddG_1, atol=1e-6) & np.isclose(x.prt, x.ddG_2, atol=1e-6)
    r = np.isclose(x.dmg, x.ddG_2, atol=1e-6) & np.isclose(x.prt, x.ddG_1, atol=1e-6)
    valid = f ^ r
    x['damaged_mut'] = np.where(f, x.m1, x.m2)
    x['partner_mut'] = np.where(f, x.m2, x.m1)
    x = x.loc[valid & (x.dmg < -1)].copy()
    key = ['WT_name', 'partner_mut', 'damaged_mut']
    # Conflicting duplicates must be reviewed, never silently maximized.
    for c in ['robust', 'physical_event', 'isotonic_heldout_event', 'prt']:
        assert x.groupby(key)[c].nunique(dropna=False).max() == 1, ('conflicting duplicates', c)
    e = x.drop_duplicates(key).copy()
    for prefix in ['partner', 'damaged']:
        p = e[prefix + '_mut'].str.extract(r'^([A-Z])(\d+)([A-Z])$')
        assert p.notna().all().all()
        e[prefix + '_site'] = p[1].astype(int)
    out = Path(root) / ('gate9_audit_' + datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S%fZ'))
    out.mkdir(parents=True, exist_ok=False)
    e.to_csv(out / 'eligible_orientation_evidence.csv', index=False)
    summary = e.groupby(['WT_name', 'partner_mut']).agg(
        tested=('robust', 'size'), rescued=('robust', 'sum'),
        damaged_sites=('damaged_site', 'nunique'),
        partner_min=('prt', 'min'), partner_max=('prt', 'max'),
        partner_median=('prt', 'median')).reset_index()
    summary.to_csv(out / 'candidate_summary.csv', index=False)
    triad = e.loc[(e.WT_name == '2K28.pdb') & e.partner_mut.isin(['E16F', 'E16W', 'E16Y'])]
    matrix = triad.pivot(index='damaged_mut', columns='partner_mut', values='robust')
    matrix.to_csv(out / 'E16_background_matrix.csv')
    site = e.loc[(e.WT_name == '2K28.pdb') & (e.partner_site == 16)]
    site.groupby('partner_mut').robust.agg(['size', 'sum', 'mean']).to_csv(out / 'E16_all_substitutions.csv')

    def tail(k, probabilities):
        dist = np.array([1.])
        for p in probabilities:
            dist = np.convolve(dist, [1-p, p])
        return float(dist[k:].sum())

    # Reproduce the existing conditional test exactly.
    for label, cols in [('bg', ['WT_name', 'damaged_mut']),
                        ('site', ['WT_name', 'damaged_mut', 'partner_site'])]:
        e['p_' + label] = e.groupby(cols).robust.transform('mean')
    rows = []
    blocks = {k: g for k, g in e.groupby(['WT_name', 'partner_site', 'damaged_site'])}
    for (protein, partner), g in e.groupby(['WT_name', 'partner_mut']):
        if len(g) < 5:
            continue
        observed = int(g.robust.sum())
        # Sensitivity null: exchange whole substitution columns within each
        # damaged-site/partner-site block, preserving correlation across the
        # damaged amino-acid variants. Use only fully observed comparators.
        # This is a different, explicitly stated null, not a universally valid
        # replacement. Separate blocks are still assumed independent.
        dist = np.array([1.])
        alternatives = []
        for (ps, ds), sub in g.groupby(['partner_site', 'damaged_site']):
            b = blocks[(protein, ps, ds)].pivot(index='damaged_mut', columns='partner_mut', values='robust')
            b = b.reindex(sub.damaged_mut).dropna(axis=1)
            assert partner in b
            scores = b.astype(int).sum(axis=0).to_numpy()
            alternatives.append(len(scores))
            pmf = np.bincount(scores, minlength=len(sub)+1) / len(scores)
            dist = np.convolve(dist, pmf)
        rows.append(dict(WT_name=protein, partner_mut=partner, tested=len(g), rescued=observed,
                         damaged_sites=g.damaged_site.nunique(),
                         expected_site=g.p_site.sum(), p_bg=tail(observed, g.p_bg),
                         p_site=tail(observed, g.p_site),
                         p_column_sensitivity=float(dist[observed:].sum()),
                         full_comparators_per_block=json.dumps(alternatives)))
    result = pd.DataFrame(rows)
    for c in ['p_bg', 'p_site', 'p_column_sensitivity']:
        result[c + '_bonf'] = np.minimum(1, result[c] * len(result))
    result.to_csv(out / 'conditional_pvalue_audit.csv', index=False)
    print('Eligible observations:', len(e), 'rescue rate:', e.robust.mean())
    print('Tested hypotheses:', len(result))
    print(matrix.to_string())
    print(summary.sort_values('rescued', ascending=False).head(20).to_string(index=False))
    print(result.sort_values('p_site').head(26).to_string(index=False))
    print('Evidence columns available for recovery decomposition:',
          [c for c in d if any(s in c.lower() for s in ['gain', 'recover', 'epist', 'epsilon', 'correct'])])
    print('Saved new audit files:', out)
    return out


if __name__ == '__main__':
    if 'd' not in globals():
        raise RuntimeError('Run with %run -i gate9_audit.py in the existing Colab containing d.')
    audit(d)