"""One locally frozen matrix; fixture feedback is disjoint and recorded."""
import argparse, copy, csv, datetime, hashlib, itertools, json, math, pathlib, statistics, sys, tempfile, time
import circuit as c
import task
ROOT=pathlib.Path(__file__).resolve().parent
SOURCES=['circuit.py','task.py','experiment.py','run_job.py','config.json','protocol.md','development-history.md']

def now():return datetime.datetime.now(datetime.timezone.utc).isoformat()
def load(name):return json.loads((ROOT/name).read_text())
def save(path,obj):path.write_text(json.dumps(obj,indent=2,sort_keys=True)+'\n')
def digest(path):return hashlib.sha256(path.read_bytes()).hexdigest()
def settings(cfg):return list(itertools.product(cfg['sizes']['m'],cfg['sizes']['r'],cfg['sizes']['d']))
def key(s):return 'm%d-r%d-d%d'%tuple(s)

def simulate(art,stream,trace=False):
    rt=c.Runtime(art);rt.trace_enabled=trace; rows=[]
    targets,ages=task.oracle(stream)
    cpu=time.process_time();wall=time.perf_counter()
    for i,e in enumerate(stream):
        end=(stream[i+1]['t']-1) if i+1<len(stream) else e['t']+47
        out=rt.transact(e,end)
        if e['kind']=='Q':rows.append(dict(event=e,age=ages[e['id']],target=targets[e['id']],**task.score(e,targets[e['id']],out)))
    return rt,rows,dict(execution_cpu_seconds=time.process_time()-cpu,execution_wall_seconds=time.perf_counter()-wall,
        sent_messages=rt.s['sent'],communication_payload_bytes=rt.s['logical_comm_bytes'],
        peak_buffer_messages=rt.s['peak_buffer'],state_serialized_bytes=rt.max_serialized_state,state_python_object_bytes=rt.max_python_state,
        logical_execution_ticks=rt.s['clock']+1,active_state_boundaries=len(rt.trajectory))

def replay(art,stream,original):
    # An actual round trip through serialized bytes, including zero state/buffers.
    compiled=json.loads(c.canonical(art))
    rt,rows,stats=simulate(compiled,stream)
    if original.trajectory!=rt.trajectory or original.s!=rt.s:raise AssertionError('exact compiled replay defect')
    return c.sha(original.trajectory)

def retime_for_fixture(art):
    """Diagnostic only: restore physical feasibility after a timing/cue change.
    Not a scored third constructor or an evaluation intervention arm.
    """
    fixed=copy.deepcopy(art);arrival={'P':0};todo=['P'];paths={}
    while todo:
        new=[]
        for k in sorted(set(todo)):
            for e in fixed['edges']:
                if e['u']==k:
                    val=arrival[k]+e['delay'];paths.setdefault(e['v'],{})[e['channel']]=val
                    arrival[e['v']]=val;new.append(e['v'])
        todo=new
    for k,v in fixed['cells'].items():
        if v['role']=='interface':
            top=max(paths[k].values());v['holds']={ch:top-paths[k][ch] for ch in ['C','M']}
    return fixed

def challenge(art,group,namespace):
    results=[]
    for pair in range(4):
        import random
        rng=random.Random(task.seed(namespace,'challenge',group,pair));bank=[1,-1,1,-1];rng.shuffle(bank)
        if art is None:
            for side,a in enumerate([bank.index(1),bank.index(-1)]):
                e=dict(id=200+pair,t=4*58,kind='Q',a=a,v=0)
                results.append(dict(pair=pair,side=side,bank=bank,event=e,target=bank[a],success=0,correct=0,missed_deadline=1,failure='construction_invalid',latency='',responses=[]))
            continue
        rt=c.Runtime(art)
        for a,v in enumerate(bank):rt.transact(dict(id=100+a,t=a*58,kind='W',a=a,v=v),a*58+47)
        state=rt.snapshot(); encoded=json.loads(c.canonical(rt.compiled()))
        assert encoded['initial']==state
        addrs=[bank.index(1),bank.index(-1)];pair_rows=[]
        for side,a in enumerate(addrs):
            clone=c.Runtime(encoded)
            if clone.snapshot()!=state:raise AssertionError('complete clone defect')
            e=dict(id=200+pair,t=4*58,kind='Q',a=a,v=0)
            out=clone.transact(e,e['t']+47)
            pair_rows.append(dict(pair=pair,side=side,bank=bank,event=e,target=bank[a],**task.score(e,bank[a],out)))
        results.extend(pair_rows)
    return results

def fixtures():
    if (ROOT/'freeze.json').exists():raise SystemExit('fixtures forbidden after freeze')
    cases=[]
    for j,s in enumerate([(1,1,1),(1,2,1),(2,1,2),(2,2,1)]):
        env=c.substrate(*s,task.seed('fixture',j));stream=task.events(task.seed('fixture-events',j))
        reference=task.reference(env,stream)
        one={}
        for arm,constructor in [('D',c.develop),('B',c.build)]:
            base,cost=constructor(c.cue_view(env));art,_=c.execute_in(env,base)
            rt,rows,stats=simulate(art,stream)
            assert all(x['success'] for x in rows)
            assert all(x['latency']==reference[x['event']['id']]['latency'] for x in rows)
            replay(art,stream,rt)
            assert c.Runtime(art).s['stored']==art['initial']['stored']
            altered,swaps=c.execute_in(env,base,True)
            fixed=retime_for_fixture(altered)
            # Relevant arrival difference changes at every address; multiset same.
            assert all(fixed['cells'][f'G{a}']['holds']!=art['cells'][f'G{a}']['holds'] for a in range(4))
            _,broken,_=simulate(altered,stream);assert any(not x['success'] for x in broken)
            _,restored,_=simulate(fixed,stream);assert all(x['success'] for x in restored)
            mis,_=constructor(c.cue_view(env,True));mis,_=c.execute_in(env,mis)
            _,mismatch,_=simulate(mis,stream)
            _,repair,_=simulate(retime_for_fixture(mis),stream);assert all(x['success'] for x in repair)
            ch=challenge(art,0,'fixture');assert all(x['success'] for x in ch)
            # Memory removal after settled write; query should return +1.
            live=c.Runtime(art);live.transact(dict(id=0,t=0,kind='W',a=0,v=1),47)
            live.s['stored']={k:0 for k in live.s['stored']}
            q=dict(id=1,t=58,kind='Q',a=0,v=0)
            assert not task.score(q,1,live.transact(q,105))['success']
            # Query-blind control always chooses address zero, against an opposite answer.
            live=c.Runtime(art)
            for a,v in enumerate([1,-1,1,-1]):live.transact(dict(id=a,t=a*58,kind='W',a=a,v=v),a*58+47)
            q=dict(id=9,t=232,kind='Q',a=1,v=0);blind=dict(q,a=0)
            assert not task.score(q,-1,live.transact(blind,279))['success']
            # Remove transport to selected interface, retaining all primitive roles.
            cut=copy.deepcopy(art);cut['edges']=[e for e in cut['edges'] if not(e['v']=='G0' and e['channel']=='C')]
            _,cutrows,_=simulate(cut,[dict(id=0,t=0,kind='W',a=0,v=1),dict(id=1,t=58,kind='Q',a=0,v=0)])
            assert not cutrows[0]['success']
            # The scorer rejects duplicates, wrong/late/missing responses.
            assert not task.score(q,-1,[[9,233,-1,'G1'],[9,234,-1,'G1']])['success']
            assert not task.score(q,-1,[[9,277,-1,'G1']])['success']
            one[arm]=dict(cost=cost,intact_queries=len(rows),intact_correct=sum(x['success'] for x in rows),
                context_correct=sum(x['success'] for x in mismatch),timing_correct=sum(x['success'] for x in broken),
                retimed_correct=sum(x['success'] for x in restored),swaps=swaps,exact_replay=True,
                memory_query_transport_shortcuts_detected=True)
        assert c.sha(c.develop(c.cue_view(env))[0])==c.sha(c.build(c.cue_view(env))[0])
        cases.append(dict(setting=s,results=one))
    destination=ROOT/'fixture-results.json'
    if destination.exists():destination=ROOT/'fixture-results-hardened.json'
    if destination.exists():destination=ROOT/'fixture-results-final.json'
    save(destination,dict(completed_utc=now(),cases=cases,scored_evaluation_generated=False))
    print('All four tiny fixtures passed for D and B; shortcuts, physical timing, reset, exact replay, clone and scoring checks passed.')

def freeze():
    if (ROOT/'freeze.json').exists():raise SystemExit('freeze already exists')
    gate=load('fixture-results.json');cfg=load('config.json')
    stamp=now();manifest=dict(frozen_utc=stamp,label='locally frozen exploratory experiment; not public preregistration',
        files={name:digest(ROOT/name) for name in SOURCES+['fixture-results.json','fixture-results-hardened.json','fixture-results-final.json']},
        rule=c.RULE,primitive_version=c.VERSION,schedule=c.SCHEDULE,config=cfg,
        fixture_results_sha256=digest(ROOT/'fixture-results.json'),
        evaluated_seeds_before_freeze=False,outer_groups=12,scheduled_model_conditions=12*8*3*2,
        episodes=12*8*3*2*4,ordinary_external_events=12*8*3*2*4*64,
        challenge_dependent_pairs=12*8*3*2*4,
        seeds=[dict(group=g,layout=task.seed(cfg['namespace'],'layout',g),
             episodes=[task.seed(cfg['namespace'],'ordinary',g,i) for i in range(4)],
             challenge=[task.seed(cfg['namespace'],'challenge',g,i) for i in range(4)]) for g in range(12)])
    save(ROOT/'freeze.json',manifest)
    save(ROOT/'status.json',dict(execution_status='frozen',publication_status='not_prepared',active_run=False,frozen_utc=stamp))
    print('Local freeze sealed:',digest(ROOT/'freeze.json'),stamp)

def verify_freeze():
    f=load('freeze.json')
    for n,h in f['files'].items():
        if digest(ROOT/n)!=h:raise AssertionError('source changed after freeze: '+n)
    return f['config']

QUERY_FIELDS=['group','setting','arm','condition','episode','event_id','input_tick','address','target','retention_ticks','success','correct','missed_deadline','failure','latency','responses','validity']
CHALLENGE_FIELDS=['group','setting','arm','condition','pair','side','bank','input_tick','address','target','success','correct','failure','latency','responses','validity']

def run_matrix(out,small=False,rerun=False):
    cfg=verify_freeze();out.mkdir(exist_ok=False)
    start=now();cpu=time.process_time();wall=time.perf_counter()
    construction=[];summaries=[];circuits={};event_count=0;completed=0;total_queries=0;trace_written=False
    def writer(name,fields):
        f=(out/name).open('x',newline='');w=csv.DictWriter(f,fields);w.writeheader();return f,w
    qf,qw=writer('queries.csv',QUERY_FIELDS);hf,hw=writer('challenge.csv',CHALLENGE_FIELDS)
    ef,ew=writer('events.csv',['group','episode','event_id','t','kind','a','v','target','retention_ticks'])
    resources=[]
    try:
        groups=[0] if small or rerun else range(12)
        combos=settings(cfg)
        for group in groups:
            streams=[task.events(task.seed(cfg['namespace'],'ordinary',group,ep)) for ep in range(4)]
            for ep,stream in enumerate(streams):
                target,age=task.oracle(stream)
                for e in stream:ew.writerow(dict(group=group,episode=ep,event_id=e['id'],**{k:e[k] for k in ['t','kind','a','v']},target=target.get(e['id'],''),retention_ticks=age.get(e['id'],'')))
            order=list(range(8));rotate=group%8;order=order[rotate:]+order[:rotate]
            if small or rerun:order=[7]
            for si in order:
                s=combos[si];setting=key(s);env=c.substrate(*s,task.seed(cfg['namespace'],'layout',group))
                for stream in streams:task.reference(env,stream)
                conditions=cfg['conditions'];rotate=(group+si)%3;conditions=conditions[rotate:]+conditions[:rotate]
                for ci,condition in enumerate(conditions):
                    arms=['D','B'] if (group+si+ci)%2==0 else ['B','D']
                    for arm in arms:
                        constructor=c.develop if arm=='D' else c.build
                        validity='valid';art=None;cost={};swaps=[]
                        try:
                            generated,cost=constructor(c.cue_view(env,condition=='context'))
                            art,swaps=c.execute_in(env,generated,condition=='timing')
                            c.validate(art)
                            # Relevant path arrival holds must differ, not merely edge labels.
                            if condition=='timing':
                                fixed=retime_for_fixture(art)
                                if any(fixed['cells'][f'G{a}']['holds']==art['cells'][f'G{a}']['holds'] for a in range(4)):
                                    raise AssertionError('causal timing path no-op')
                        except (ValueError,TimeoutError) as error:
                            validity='construction_invalid:'+str(error)
                        record=dict(group=group,setting=setting,arm=arm,condition=condition,validity=validity,
                            circuit_sha256=c.sha(art) if art else '',swaps=swaps,extents=c.module_extents(env),**cost)
                        construction.append(record)
                        if art is not None and group==0 and si in [0,7]:
                            circuits[f'{setting}/{arm}/{condition}']=dict(environment=env,circuit=art)
                        all_rows=[];metrics=[];hashes=[]
                        for ep,stream in enumerate(streams):
                            if art is not None:
                                try:
                                    rt,rows,stats=simulate(art,stream,trace=group==0 and si==7 and condition=='intact' and arm=='D' and ep==0)
                                except ValueError as error:
                                    targets,ages=task.oracle(stream)
                                    rows=[dict(event=e,target=targets[e['id']],age=ages[e['id']],success=0,correct=0,missed_deadline=1,failure='model_budget:'+str(error),latency='',responses=[]) for e in stream if e['kind']=='Q']
                                    for x in rows:
                                        e=x['event'];qw.writerow(dict(group=group,setting=setting,arm=arm,condition=condition,episode=ep,event_id=e['id'],input_tick=e['t'],address=e['a'],target=x['target'],retention_ticks=x['age'],validity='model_budget',**{k:x[k] for k in ['success','correct','missed_deadline','failure','latency']},responses='[]'))
                                    total_queries+=len(rows);event_count+=len(stream);all_rows.extend(rows)
                                    continue
                                if arm=='D':hashes.append(replay(art,stream,rt))
                                if rt.trace_enabled:
                                    save(out/'trace.json',dict(group=group,setting=setting,arm=arm,condition=condition,episode=ep,inputs=stream,deliveries=rt.trace,outputs=rt.s['outputs']))
                                    trace_written=True
                                metrics.append(stats)
                            else:
                                targets,ages=task.oracle(stream)
                                rows=[dict(event=e,target=targets[e['id']],age=ages[e['id']],success=0,correct=0,missed_deadline=1,failure=validity,latency='',responses=[]) for e in stream if e['kind']=='Q']
                            for x in rows:
                                e=x['event'];qw.writerow(dict(group=group,setting=setting,arm=arm,condition=condition,episode=ep,event_id=e['id'],input_tick=e['t'],address=e['a'],target=x['target'],retention_ticks=x['age'],validity=validity,**{k:x[k] for k in ['success','correct','missed_deadline','failure','latency']},responses=c.canonical(x['responses'])))
                            total_queries+=len(rows);event_count+=len(stream);all_rows.extend(rows)
                        challenge_cpu=time.process_time();challenge_wall=time.perf_counter()
                        ch=challenge(art,group,cfg['namespace'])
                        challenge_cost=dict(execution_cpu_seconds=time.process_time()-challenge_cpu,execution_wall_seconds=time.perf_counter()-challenge_wall,write_exposures=16,query_exposures=8)
                        for x in ch:
                            e=x['event'];hw.writerow(dict(group=group,setting=setting,arm=arm,condition=condition,pair=x['pair'],side=x['side'],bank=c.canonical(x['bank']),input_tick=e['t'],address=e['a'],target=x['target'],validity=validity,**{k:x[k] for k in ['success','correct','failure','latency']},responses=c.canonical(x['responses'])))
                        latency=[x['latency'] for x in all_rows if x['latency']!='']
                        failures={kind:sum(x['failure']==kind for x in all_rows) for kind in sorted({x['failure'] for x in all_rows})}
                        row=dict(group=group,setting=setting,arm=arm,condition=condition,validity=validity,
                            queries=len(all_rows),successes=sum(x['success'] for x in all_rows),
                            fraction=statistics.mean(x['success'] for x in all_rows),correct_fraction=statistics.mean(x['correct'] for x in all_rows),
                            missed_deadline_fraction=statistics.mean(x['missed_deadline'] for x in all_rows),failures=failures,
                            mean_latency=statistics.mean(latency) if latency else None,max_latency=max(latency) if latency else None,
                            challenge_queries=8,challenge_fraction=sum(x['success'] for x in ch)/8,
                            challenge_both_fraction=sum(all(y['success'] for y in ch if y['pair']==p) and len([y for y in ch if y['pair']==p])==2 for p in range(4))/4,
                            replay_agreement=(True if arm=='D' and len(hashes)==4 else None),trajectory_hashes=hashes,
                            execution=metrics,challenge_cost=challenge_cost)
                        summaries.append(row);completed+=1
                        # Checkpoint after every model/condition so interruption is retained.
                        save(out/'progress.json',dict(completed_model_conditions=completed,execution_status='running',last=[group,setting,arm,condition]))
                print(f'completed group={group} setting={setting}',flush=True)
        if not small and not rerun and completed!=576:raise AssertionError('incomplete matrix')
        save(out/'construction.json',construction);save(out/'representative-circuits.json',circuits)
        save(out/'summary.json',dict(started_utc=start,ended_utc=now(),model_conditions=completed,
            ordinary_event_exposures=event_count,ordinary_query_exposures=total_queries,
            rows=summaries,execution_status='completed',study_cpu_seconds=time.process_time()-cpu,
            study_wall_seconds=time.perf_counter()-wall,trace_written=trace_written,
            stored_source_bytes={n:(ROOT/n).stat().st_size for n in SOURCES},freeze_sha256=digest(ROOT/'freeze.json')))
        save(out/'progress.json',dict(completed_model_conditions=completed,execution_status='completed'))
    except BaseException as error:
        save(out/'failure.json',dict(error_type=type(error).__name__,message=str(error),completed_model_conditions=completed,execution_status='stopped',utc=now()))
        save(out/'partial-summary.json',summaries);save(out/'partial-construction.json',construction)
        raise
    finally:qf.close();hf.close();ef.close()

def analysis():
    s=load('results/summary.json');rows=s['rows'];cfg=load('config.json');base=key((2,2,1))
    def get(g,sett,a,cond):return next(x['fraction'] for x in rows if (x['group'],x['setting'],x['arm'],x['condition'])==(g,sett,a,cond))
    held=[key(x) for x in settings(cfg) if key(x)!=base]
    diffs=[statistics.mean(get(g,t,'D','intact')-get(g,t,'B','intact') for t in held) for g in range(12)]
    av={a:statistics.mean(statistics.mean(get(g,t,a,'intact') for t in held) for g in range(12)) for a in ['D','B']}
    effects={}
    for condition in ['context','timing']:
        d={a:[statistics.mean(get(g,t,a,'intact')-get(g,t,a,condition) for t in held) for g in range(12)] for a in ['D','B']}
        effects[condition]=dict(seed_effects=d,mean_effect={a:statistics.mean(v) for a,v in d.items()},paired_difference=[x-y for x,y in zip(d['D'],d['B'])])
    sd=statistics.stdev(diffs);se=sd/math.sqrt(12);mean=statistics.mean(diffs)
    save(ROOT/'analysis.json',dict(primary_D=av['D'],primary_B=av['B'],paired_D_minus_B=mean,
        seed_differences=diffs,sample_sd=sd,se=se,descriptive_t11_interval=[mean-2.200985*se,mean+2.200985*se],screening_margin=.02,
        effects=effects,complete_model_conditions=s['model_conditions'],held_out_settings=held,
        uncertainty_limit='Seed-cluster descriptive interval; all-zero differences and ceiling do not bound untested environments or constructors.'))
    fields=['group','setting','arm','condition','validity','queries','successes','fraction','correct_fraction','missed_deadline_fraction','mean_latency','max_latency','challenge_fraction','challenge_both_fraction','replay_agreement','failures']
    with (ROOT/'aggregate.csv').open('x',newline='') as f:
        w=csv.DictWriter(f,fields);w.writeheader()
        for x in rows:w.writerow({k:c.canonical(x[k]) if isinstance(x[k],dict) else x[k] for k in fields})

def main():
    p=argparse.ArgumentParser();p.add_argument('mode',choices=['fixtures','freeze','smoke','full','rerun']);a=p.parse_args()
    if a.mode=='fixtures':fixtures()
    elif a.mode=='freeze':freeze()
    elif a.mode=='smoke':
        with tempfile.TemporaryDirectory(prefix='context-dynamics-smoke-') as temp:
            out=pathlib.Path(temp)/'out';run_matrix(out,small=True)
            print('Smoke completed in disposable directory; six model/conditions, four episodes each.')
    elif a.mode=='full':
        cfg=verify_freeze();status=load('status.json')
        if status['execution_status']!='frozen' or (ROOT/'results').exists():raise SystemExit('duplicate/replacement matrix forbidden')
        save(ROOT/'status.json',dict(execution_status='running',publication_status='not_prepared',active_run=True,started_utc=now()))
        try:
            run_matrix(ROOT/'results');analysis()
            save(ROOT/'status.json',dict(execution_status='completed',publication_status='not_prepared',active_run=False,ended_utc=now()))
        except BaseException:
            save(ROOT/'status.json',dict(execution_status='stopped',publication_status='not_prepared',active_run=False,ended_utc=now()));raise
    else:
        if load('status.json')['execution_status']!='completed':raise SystemExit('rerun requires completed attempt')
        with tempfile.TemporaryDirectory(prefix='context-dynamics-author-') as temp:
            out=pathlib.Path(temp)/'out';run_matrix(out,rerun=True)
            rer=json.loads((out/'summary.json').read_text());old=load('results/summary.json')
            keys=['queries','successes','fraction','correct_fraction','failures','challenge_fraction','challenge_both_fraction','trajectory_hashes']
            for x in rer['rows']:
                y=next(y for y in old['rows'] if all(y[k]==x[k] for k in ['group','setting','arm','condition']))
                assert all(x[k]==y[k] for k in keys),'author rerun mismatch'
            newcons=json.loads((out/'construction.json').read_text());oldcons=load('results/construction.json')
            assert all(x['circuit_sha256']==next(y['circuit_sha256'] for y in oldcons if all(y[k]==x[k] for k in ['group','setting','arm','condition'])) for x in newcons)
            save(ROOT/'author-rerun.json',dict(utc=now(),label='author rerun; not independent reproduction',group=0,setting='m4-r4-d5',model_conditions=6,episodes=24,score_and_trajectory_fields_match=True,circuit_hashes_match=True,disposable_directory_removed=True))
            print('Author rerun matched all six paired model/intervention cases, exact D trajectories and circuit hashes.')
if __name__=='__main__':main()