#!/usr/bin/env python3
"""Adversarial regression tests for T^GPT v0.8.6 MNAR/P0 instrumentation repair."""
from __future__ import annotations
import importlib.util, random
from pathlib import Path
HERE=Path(__file__).resolve().parent; DATA=HERE.parent/'data'
def load(name,fn):
    spec=importlib.util.spec_from_file_location(name,HERE/fn); mod=importlib.util.module_from_spec(spec); assert spec.loader is not None; spec.loader.exec_module(mod); return mod
def close(a,b,tol=1e-9): return abs(a-b)<=tol
def expect_raises(fn,contains=None):
    try: fn()
    except Exception as e:
        if contains and contains not in str(e): raise AssertionError((contains,str(e)))
        return
    raise AssertionError('expected exception')
def make_24_distractors(aid,tokens,fp_near=0,fp_mid=0,fp_far=0):
    rows=[]
    for s,nfp in [('NEAR',fp_near),('MID',fp_mid),('FAR',fp_far)]:
        for i in range(8): rows.append({'aggregate_id':aid,'candidate_cluster_id':f'D_{s}_{i}','candidate_type':'DISTRACTOR','distractor_stratum':s,'represented':1 if i<nfp else 0,'aggregate_output_tokens':tokens,'condition_blind_id':aid})
    return rows
def main():
    rec=load('rec','compute_recovery_metrics_v0.8.6.py'); agg=load('agg','evaluate_aggregation_v0.8.6.py'); plan=load('plan','plan_task_precision_v0.8.6.py'); cov=load('cov','coverage_baseline_v0.8.6.py')
    rng=random.Random(123); cats=[f'C{i}' for i in range(1,11)]; weights=[.25,.18,.14,.11,.09,.07,.06,.04,.035,.025]; rel={c:1 for c in cats}
    for j in range(30):
        A=rng.choices(cats,weights,k=128); B=['UNMAPPED' if rng.random()<.20 else rng.choices(cats,weights,k=1)[0] for _ in range(128)]
        o=rec.one_task({'task_id':f'MCAR{j}','family_id':'F1_POLYSEMY_UNDERSPECIFICATION','relevance':rel,'reference_A':A,'contracted_B':B},100,50,9000+j); assert o['measurement_status']=='MEASUREMENT_INSUFFICIENT'; assert not o['shift']['PRIMARY_CONDITIONAL_ENDPOINT_AVAILABLE']
    A=['C1']*40+['C2']*24+['C3']*24+['C4']*40; B=['C1']*40+['C2']*19+['UNMAPPED']*5+['C3']*24+['C4']*40
    mnar=rec.one_task({'task_id':'MNAR','family_id':'F1_POLYSEMY_UNDERSPECIFICATION','relevance':{'C1':1,'C2':1,'C3':1,'C4':1},'reference_A':A,'contracted_B':B,'unmapped_affinity':{'B':['C2']*5}},300,50,77)
    assert mnar['measurement_status']=='OK'; assert mnar['partial_identification']['TOTAL_DEFICIT_IDENTIFICATION_INTERVAL'][0]<=1e-12; assert not mnar['shift']['CONFIRMATORY_TOTAL_SEMANTIC_CLAIM_AVAILABLE']; assert mnar['unmapped_affinity_diagnostic']['arms']['B']['counts']['C2']==5
    ex=[]
    for j in range(40):
        A=rng.choices(cats,weights,k=128); B=rng.choices(cats,weights,k=128); o=rec.one_task({'task_id':f'NULL{j}','family_id':'F1_POLYSEMY_UNDERSPECIFICATION','relevance':rel,'reference_A':A,'contracted_B':B},120,50,12000+j); ex.append(o['shift']['SEMANTIC_DEFICIT_EXCESS_CONDITIONAL'])
    assert abs(sum(ex)/len(ex))<.03
    expect_raises(lambda: rec.one_task({'task_id':'UNEQ','family_id':'F1_POLYSEMY_UNDERSPECIFICATION','relevance':{'C1':1,'C2':1},'reference_A':['C1']*32+['C2']*32,'reference_A_prime':['C1']*32+['C2']*32,'contracted_B':['C1']*16+['C2']*16},100,50,7),'N_A == N_A_prime == N_B')
    expect_raises(lambda: rec.one_task({'task_id':'SMALL','family_id':'F1_POLYSEMY_UNDERSPECIFICATION','relevance':{'C1':1},'reference_A':['C1']*31,'contracted_B':['C1']*31},100,50,7),'MIN_N_PER_ARM')
    pairs=[]
    for i in range(64):
        base='C1' if i<40 else 'C2'; pairs.append({'pair_id':str(i),'base':base,'recovery':base,'control':'UNMAPPED' if i<12 else base})
    ro=rec.one_task({'task_id':'RMAP','family_id':'F2_CAUSAL_ALTERNATIVES','relevance':{'C1':1,'C2':1},'reference_A':['C1']*40+['C2']*24,'paired_probes':pairs},100,100,9); assert not ro['recovery']['PRIMARY_ENDPOINT_AVAILABLE']; assert 'COMPLETE_CASE_RECOVERY_LIFT_SENSITIVITY' not in ro['recovery']
    true=[{'id':'C1','weight':.7},{'id':'C2','weight':.3}]; anns=[]
    for cid,rep in [('C1',1),('C2',1)]: anns.append({'aggregate_id':'A1','candidate_cluster_id':cid,'candidate_type':'TRUE_POOL','represented':rep,'aggregate_output_tokens':90,'condition_blind_id':'A1'})
    anns+=make_24_distractors('A1',90,fp_near=3); a1=agg.analyze({'max_output_tokens':100,'true_pool_clusters':true,'annotations':anns})['aggregates'][0]; assert a1['measurement_status']=='MEASUREMENT_INSUFFICIENT' and a1['AGG_COVERAGE_ADJUSTED_PRIMARY_NEAR_FPR'] is None
    ok=[]
    for cid,rep in [('C1',1),('C2',0)]: ok.append({'aggregate_id':'A2','candidate_cluster_id':cid,'candidate_type':'TRUE_POOL','represented':rep,'aggregate_output_tokens':90,'condition_blind_id':'A2'})
    ok+=make_24_distractors('A2',90,fp_near=1); a2=agg.analyze({'max_output_tokens':100,'true_pool_clusters':true,'annotations':ok})['aggregates'][0]; assert a2['measurement_status']=='OK'
    collision=list(ok); collision.append({'aggregate_id':'A2','candidate_cluster_id':'C1','candidate_type':'DISTRACTOR','distractor_stratum':'FAR','represented':0,'aggregate_output_tokens':90,'condition_blind_id':'A2'}); expect_raises(lambda: agg.analyze({'max_output_tokens':100,'true_pool_clusters':true,'annotations':collision}),'collides')
    b=cov.solve({'token_budget':240,'clusters':[{'id':'C1','weight':.4},{'id':'C2','weight':.35},{'id':'C3','weight':.25}],'candidates':[{'id':'S1','token_cost':100,'covers':{'C1':1}},{'id':'S2','token_cost':120,'covers':{'C2':1}},{'id':'S3','token_cost':180,'covers':{'C1':.5,'C3':1}}]}); assert b['selected_ids']==['S1','S2'] and close(b['objective_normalized'],.75)
    families=sorted(plan.EXPECTED_FAMILIES); rows=[]; p0=[]
    for fi,f in enumerate(families):
        for j in range(4): rows.append({'task_id':f'{fi}-{j}','family_id':f,'effect':.04*fi+(-.03,-.01,.02,.04)[j],'n_per_task':64,'measurement_status':'OK'}); p0.append({'task_id':f'P0-{fi}-{j}','family_id':f,'measurement_status':'OK' if j<3 else 'MEASUREMENT_INSUFFICIENT'})
    valid={'experiment_id':'E1_STAGE','planning_started_utc':'2026-09-12T06:00:00Z','planning_tasks':rows,'p0_registry':p0,'confirmation_n_per_task':64,'bootstrap_reps':2500,'seed':77}; po=plan.plan(valid,DATA/'minimum-effects-v0.8.6.json'); assert close(po['meoi_artifact']['minimum_relevant_effect'],.05)
    print('T_GPT_V0_8_6_MNAR_P0_SELFTEST_OK')
if __name__=='__main__': main()