"""Locked independent-day retest of unchanged full-hypothesis models.

The preceding dynamic-mixture study failed its source-clock requirement.
It stays failed. This separately registered study reports every day and control.
"""
from __future__ import annotations
import argparse,hashlib,io,json,math,re,time,urllib.request,zipfile
from pathlib import Path
import numpy as np
import torch
from unionmind.continuous_decision import source,factorial
from unionmind.continuous_decision.factor_study import MODELS,fit_one,predict,metrics
from unionmind.continuous_decision.factor_analysis import STRONG
from unionmind.continuous_decision.factor_acquisition import episode

ROOT=Path('artifacts/day_transfer_01');PROTOCOL=Path('configs/day_transfer_protocol_01.json')


def write(name,value):
    path=ROOT/name;path.parent.mkdir(parents=True,exist_ok=True);path.write_text(json.dumps(value,indent=2)+'\n')


def verify():
    lock=json.loads((ROOT/'lock.json').read_text())
    assert all(source.sha(p)==h for p,h in lock['files'].items())


def lock():
    assert not ROOT.exists();ROOT.mkdir()
    files=[Path(__file__),PROTOCOL,*Path('unionmind/continuous_decision').glob('*.py'),Path('tests/test_continuous_decision.py')]
    write('lock.json',dict(files={str(p):source.sha(p) for p in files},before_archive_or_numerical_access=True,
        timestamp_utc=time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()),
        previous_strict_clock_study_not_amended=True,unchanged_algorithms_and_hyperparameters=True))
    print('INDEPENDENT DAY STUDY LOCKED',flush=True)


def text_array(payload):
    text=payload.decode('utf-8-sig');lines=text.strip().splitlines();header=[]
    number=re.compile(r'^[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?$')
    while lines and not number.match(re.split(r'[\s,;]+',lines[0].strip())[0]):header.append(lines.pop(0))
    body=re.sub('[,;]',' ','\n'.join(lines));array=np.fromstring(body,sep=' ',dtype=np.float64)
    assert array.size==len(lines)*20,(array.size,len(lines))
    return array.reshape(-1,20),header


def prepare(device='cuda:1'):
    verify();config=json.loads(PROTOCOL.read_text());torch.cuda.set_device(device);start=time.perf_counter()
    path=ROOT/'source.zip';assert not path.exists();urllib.request.urlretrieve(config['source']['url'],path)
    write('source.json',dict(**config['source'],sha256=source.sha(path),bytes=path.stat().st_size,download_seconds=time.perf_counter()-start))
    sources=[];outer=zipfile.ZipFile(path)
    for entry in outer.infolist():
        if entry.filename.endswith('.zip'):
            inner=zipfile.ZipFile(io.BytesIO(outer.read(entry)));sources.append(inner)
    if not sources:sources=[outer]
    members=sorted([(z,a) for z in sources for a in z.infolist() if re.search(r'\d{8}_\d{6}',a.filename) and not a.is_dir()],key=lambda item:item[1].filename)
    assert len(members)==13,[(a.filename,a.file_size) for z,a in members]
    rows=[];labels=[];days=[];records=[];audits=[];excluded=[];headers={}
    for day,(archive,entry) in enumerate(members):
        values,header=text_array(archive.read(entry));headers[entry.filename]=header
        raw=torch.as_tensor(values,device=device,dtype=torch.float64)
        valid=bool(torch.isfinite(raw).all()) and bool((raw[1:,0]>=raw[:-1,0]).all())
        audit=dict(day=day,file=entry.filename,raw_rows=len(raw),finite=bool(torch.isfinite(raw).all()),
            duplicate_timestamps=int((raw[1:,0]==raw[:-1,0]).sum()),nondecreasing_clock=bool((raw[1:,0]>=raw[:-1,0]).all()))
        if not valid:audits.append({**audit,'valid':False});continue
        time_zero=raw[0,0];blocks=((raw[:,0]-time_zero)/900).floor().long()
        for block in blocks.unique(sorted=True).tolist():
            ids=torch.where(blocks==block)[0];a=int(ids[0]);b=int(ids[-1])+1
            duration=float(raw[b-1,0]-raw[a,0])
            if duration<850:
                excluded.append(dict(day=day,file=entry.filename,block=block,duration=duration,reason='less_than_850_seconds'));continue
            quarter=((raw[a:b,0]-time_zero-block*900)/225).floor().long().clamp(0,3)
            means=torch.stack([raw[a:b,6:][quarter==q].mean(0) for q in range(4)],-1)
            last=quarter==3;co=raw[a:b,1][last].median();rh=raw[a:b,2][last].mean()
            h0=(co*9/20).round().long().clamp(0,9);h1=(rh>rh.new_tensor([35.,55.])).sum()
            rows.append(means.flatten());labels.append(torch.stack((h0,h1)));days.append(day)
            records.append(dict(id=f'{entry.filename}:block{block}',day=day,file=entry.filename,block=block,
                first_row=a,last_row_exclusive=b,start_seconds=float(raw[a,0]),end_seconds=float(raw[b-1,0]),
                reference_co_ppm=float(co),reference_humidity_percent=float(rh)))
        audit.update(valid=True,episodes=sum(d==day for d in days));audits.append(audit)
        print(json.dumps(audit),flush=True);del raw,values
    x=torch.stack(rows);h=torch.stack(labels);day_tensor=torch.tensor(days,device=device)
    grid=factorial.hypotheses([10,3],device)
    ordering=sorted(grid.tolist(),key=lambda row:hashlib.sha256(('20821:'+','.join(map(str,row))).encode()).hexdigest())
    held_bindings=ordering[:6];unseen=torch.zeros(len(x),device=device,dtype=torch.bool)
    for binding in held_bindings:unseen|=(h==h.new_tensor(binding)).all(-1)
    source_mask=(day_tensor<9)&~unseen;target=day_tensor>=9
    eligibility=len(audits)==13 and all(a.get('valid') and a.get('episodes',0)>=90 for a in audits)
    eligibility &= int(source_mask.sum())>=128 and all(int((target&unseen&(day_tensor==d)).sum())>=8 for d in range(9,13))
    torch.save(dict(x=x,h=h,days=day_tensor,source=source_mask,target=target,unseen=unseen),ROOT/'prepared.pt')
    torch.cuda.synchronize();write('preparation.json',dict(eligible=eligibility,records=records,audits=audits,headers=headers,
        held_bindings=held_bindings,source_records=int(source_mask.sum()),held_records=int(target.sum()),
        unseen_held_records=int((target&unseen).sum()),exclusions=excluded,seconds=time.perf_counter()-start,
        measurements_are_episodes_not_independent_raw_rows=True,reference_columns_are_never_features=True))
    verify();print(json.dumps(dict(eligible=eligibility,source_records=int(source_mask.sum()),held_records=int(target.sum()))),flush=True)


def run(device='cuda:1'):
    verify();torch.cuda.set_device(device);torch.backends.cuda.matmul.allow_tf32=False
    config=json.loads(PROTOCOL.read_text());start=time.perf_counter();prep=json.loads((ROOT/'preparation.json').read_text())
    if not prep['eligible']:write('decision.json',dict(promotion=False,reason='locked_cohort_eligibility_failure'));return
    p=torch.load(ROOT/'prepared.pt',map_location=device,weights_only=True);pool=torch.where(p['source'])[0];held=torch.where(p['target'])[0]
    sizes=[10,3];grid=factorial.hypotheses(sizes,device);true=factorial.joint_index(p['h'][held],sizes);results=[]
    for n in config['sample_budgets']:
        for seed in config['seeds']:
            rng=torch.Generator(device=device).manual_seed(seed);ids=pool[torch.randperm(len(pool),device=device,generator=rng)[:n]]
            t=time.perf_counter();mean=p['x'][ids].mean(0);scale=p['x'][ids].std(0).clamp_min(1e-6)
            x=(p['x'][ids]-mean)/scale;target=(p['x'][held]-mean)/scale;torch.cuda.synchronize();normalization=time.perf_counter()-t
            supported=torch.ones(len(held),device=device,dtype=torch.bool)
            for f,k in enumerate(sizes):supported &= torch.bincount(p['h'][ids,f],minlength=k)[p['h'][held,f]]>0
            directory=ROOT/'models'/f'n{n}_s{seed}';directory.mkdir(parents=True,exist_ok=True)
            for name in MODELS:
                torch.cuda.synchronize();t=time.perf_counter();model=fit_one(x,p['h'][ids],sizes,name,seed)
                torch.cuda.synchronize();fit_seconds=time.perf_counter()-t;t=time.perf_counter()
                with torch.no_grad():prob=predict(model,target,range(x.shape[-1]));v=metrics(prob,true,grid,p['h'][held])
                torch.cuda.synchronize();prediction=time.perf_counter()-t
                checkpoint=directory/f'{name}.pt';torch.save(dict(model=model,mean=mean,scale=scale,fit_ids=ids),checkpoint)
                row=dict(samples=n,seed=seed,model=name,held_rows=held,days=p['days'][held],unseen=p['unseen'][held],
                    supported_levels=supported,metrics=v,normalization_seconds=normalization,fit_seconds=fit_seconds,
                    prediction_seconds=prediction,neural_parameters=model.get('neural_parameters',0),
                    checkpoint_bytes=checkpoint.stat().st_size,adaptation_fits=model.get('adaptation_fits',1))
                results.append(row)
                print(json.dumps(dict(samples=n,seed=seed,model=name,top3_unseen=float(v['top3'][row['unseen']].mean()))),flush=True)
    torch.save(dict(results=results),ROOT/'results.pt');write('execution.json',dict(seconds=time.perf_counter()-start,
        peak_gpu_bytes=torch.cuda.max_memory_allocated(device),device=device,cells=len(results),backbone_executions=0,
        timings=[{k:r[k] for k in ('samples','seed','model','normalization_seconds','fit_seconds','prediction_seconds','neural_parameters','checkpoint_bytes','adaptation_fits')} for r in results]))
    verify()


@torch.no_grad()
def analyze(device='cuda:2'):
    verify();torch.cuda.set_device(device);config=json.loads(PROTOCOL.read_text());start=time.perf_counter()
    results=torch.load(ROOT/'results.pt',map_location=device,weights_only=True)['results'];rng=torch.Generator(device=device).manual_seed(20891)
    summary=[];comparisons=[];gates=[]
    for n in config['sample_budgets']:
        records=[r for r in results if r['samples']==n];template=records[0];mask=template['unseen'];day=template['days'][mask]
        matrix=torch.stack([torch.stack([r['metrics']['top3'] for r in records if r['model']==name]).mean(0)[mask] for name in MODELS],-1)
        points=[];boot=[]
        for d in range(9,13):
            a=matrix[day==d];points.append(a.mean(0));sample=torch.randint(len(a),(20000,len(a)),device=device,generator=rng);boot.append(a[sample].mean(1))
        day_points=torch.stack(points);point=day_points.mean(0);conditional_boot=torch.stack(boot).mean(0)
        sampled_days=torch.randint(4,(20000,4),device=device,generator=rng)
        day_boot=day_points[sampled_days].mean(1)
        for name in MODELS:
            rows=[r for r in records if r['model']==name];assert len(rows)==3
            for stratum,include in (('unseen',mask),('seen',~mask)):
                for d in (None,9,10,11,12):
                    m=include if d is None else include&(template['days']==d)
                    summary.append(dict(samples=n,model=name,stratum=stratum,day=d,records=int(m.sum()),
                        **{f:float(torch.stack([r['metrics'][f][m].mean() for r in rows]).mean()) for f in ('top1','top3','nll','confidence','top3_mass')}))
            for support in (False,True):
                count=correct=0
                for r in rows:
                    m=r['unseen']&(r['supported_levels']==support);count+=int(m.sum());correct+=float(r['metrics']['top3'][m].sum())
                summary.append(dict(samples=n,model=name,stratum='supported' if support else 'novel_level',records_across_seeds=count,top3=correct/count if count else None))
            for threshold in (.5,.7,.9):
                accepted=correct=0
                for r in rows:
                    m=r['unseen']&(r['metrics']['confidence']>=threshold);accepted+=int(m.sum());correct+=float(r['metrics']['top1'][m].sum())
                summary.append(dict(samples=n,model=name,stratum='selective_unseen',threshold=threshold,
                    coverage=accepted/(3*int(mask.sum())),accuracy=correct/accepted if accepted else None))
        for candidate in config['candidates']:
            ci=MODELS.index(candidate);rows=[];alpha=.05/(2*40)
            for ri,reference in enumerate(MODELS):
                if ci==ri:continue
                diff=conditional_boot[:,ci]-conditional_boot[:,ri];lo,hi=torch.quantile(diff,diff.new_tensor([alpha,1-alpha]))
                dd=day_boot[:,ci]-day_boot[:,ri];dlo,dhi=torch.quantile(dd,dd.new_tensor([alpha,1-alpha]))
                row=dict(samples=n,candidate=candidate,control=reference,gain=float(point[ci]-point[ri]),
                    fixed_day_simultaneous_interval=[float(lo),float(hi)],day_cluster_simultaneous_interval=[float(dlo),float(dhi)],
                    day_gains=(day_points[:,ci]-day_points[:,ri]).tolist())
                rows.append(row);comparisons.append(row)
            strong=[r for r in rows if r['control'] in STRONG]
            substantial=all(r['gain']>=.10 and r['fixed_day_simultaneous_interval'][0]>0 and r['day_cluster_simultaneous_interval'][0]>0 and min(r['day_gains'])>0 for r in strong)
            best=all(r['fixed_day_simultaneous_interval'][0]>0 and r['day_cluster_simultaneous_interval'][0]>0 for r in rows)
            learned=next((r['fixed_day_simultaneous_interval'][0]>0 and r['day_cluster_simultaneous_interval'][0]>0 for r in rows if r['control']=='additive_full'),False) if candidate=='factorized_neural' else None
            gates.append(dict(samples=n,candidate=candidate,top3=float(point[ci]),substantial_gain=substantial,strongest_available=best,
                learning_specific=learned,promotion=n==64 and substantial and best,
                strongest_observed_model=MODELS[int(point.argmax())],strongest_top3=float(point.max())))
    write('summary.json',dict(rows=summary,comparisons=comparisons,gates=gates,all_thirteen_days_retained=True,
        confirmation_days=[9,10,11,12],one_physical_platform=True,draws=20000,comparison_family=40,
        original_dynamic_mixture_source_gate_remains_failed=True))
    write('analysis_execution.json',dict(seconds=time.perf_counter()-start,device=device));verify();print(json.dumps(gates),flush=True)


def acquisition(device='cuda:1'):
    verify();torch.cuda.set_device(device);torch.backends.cuda.matmul.allow_tf32=False
    config=json.loads(PROTOCOL.read_text());p=torch.load(ROOT/'prepared.pt',map_location=device,weights_only=True);prep=json.loads((ROOT/'preparation.json').read_text())
    rows=[]
    for day in range(9,13):
        eligible=torch.where(p['target']&p['unseen']&(p['days']==day))[0].tolist()
        rows+=sorted(eligible,key=lambda i:hashlib.sha256(('20901:'+prep['records'][i]['id']).encode()).hexdigest())[:8]
    ids=torch.tensor(rows,device=device);true=factorial.joint_index(p['h'][ids],[10,3]);results=[];start=time.perf_counter()
    for name in config['candidates']:
        checkpoint=torch.load(ROOT/'models'/f'n64_s{config["seeds"][0]}'/f'{name}.pt',map_location=device,weights_only=True)
        x=(p['x'][ids]-checkpoint['mean'])/checkpoint['scale'];model=checkpoint['model']
        for policy in config['acquisition']['policies']:
            snapshots,trace=episode(model,x,policy,20911,(2,4,8),.01)
            for budget,out in snapshots.items():
                credit=factorial.topk_credit(out['probability'],true,3)
                results.append(dict(model=name,policy=policy,budget=budget,rows=ids,days=p['days'][ids],top3=credit,queries=out['queries'],trace=trace))
            print(json.dumps(dict(model=name,policy=policy,seconds=trace['seconds'])),flush=True)
    torch.save(dict(results=results),ROOT/'acquisition.pt');write('acquisition_execution.json',dict(seconds=time.perf_counter()-start,
        peak_gpu_bytes=torch.cuda.max_memory_allocated(device),secondary_only=True));verify()


if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('command',choices=['lock','prepare','run','analyze','acquisition']);a=parser.parse_args()
    {'lock':lock,'prepare':prepare,'run':run,'analyze':analyze,'acquisition':acquisition}[a.command]()