"""Full-ontology confirmation with fixed candidates, controls and held bindings."""
from __future__ import annotations
import argparse,json,time
from pathlib import Path
import torch
from . import source,factorial,classifiers,learned_controls

ROOT=Path('artifacts/factorial_transfer_01');PROTOCOL=Path('configs/factorial_transfer_protocol_01.json')
MODELS=['additive_full','additive_diag','factorized_neural','joint_neural','factorized_nb','factorized_lda','factorized_kernel','factorized_ridge','joint_gaussian','factorized_logistic','factorized_classifier']
WORLD=['additive_full','additive_diag','factorized_neural','joint_neural','joint_gaussian']


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


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


def lock():
    assert not (ROOT/'confirmation_lock.json').exists() and not (ROOT/'sources').exists();ROOT.mkdir(exist_ok=True)
    previous=Path('artifacts/continuous_decision_01')
    protected=json.loads((previous/'preservation_lock.json').read_text())['files']
    for p in previous.rglob('*'):
        if p.is_file() and '__pycache__' not in p.parts:protected[str(p)]=source.sha(p)
    write('preservation_lock.json',dict(files=protected))
    files=[*Path('unionmind/continuous_decision').glob('*.py'),PROTOCOL,Path('tests/test_continuous_decision.py')]
    write('confirmation_lock.json',dict(files={str(p):source.sha(p) for p in files},
        timestamp_utc=time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()),
        before_new_gas_archive_or_numerical_values=True,
        hydraulic_is_opened_development=True,
        candidate_and_all_thresholds_fixed=True))
    print('FULL ONTOLOGY CONFIRMATION LOCKED',flush=True)


def prepare():
    verify();config=json.loads(PROTOCOL.read_text());start=time.perf_counter();source.acquire(ROOT,config)
    acquisition=time.perf_counter()-start;torch.cuda.set_device(config['device']);torch.backends.cuda.matmul.allow_tf32=False
    source.prepare(ROOT,config,config['device'])
    write('preparation_execution.json',dict(acquisition_seconds=acquisition,complete_seconds=time.perf_counter()-start))
    verify()


def fit_one(x,h,sizes,name,seed,steps=400):
    if name in WORLD:return factorial.world_fit(x,h,sizes,name,seed,steps)
    if name in ('factorized_logistic','factorized_classifier'):return learned_controls.fit(x,h,sizes,name,seed,steps)
    return classifiers.fit(x,h,sizes,name)


def predict(model,x,observed):
    if model['kind'] in WORLD:return factorial.posterior(model,x,observed)
    if model['kind'] in ('factorized_logistic','factorized_classifier'):return learned_controls.posterior(model,x,observed)
    return classifiers.posterior(model,x,observed)


@torch.no_grad()
def metrics(prob,true,grid,h):
    return dict(top1=factorial.topk_credit(prob,true,1),top3=factorial.topk_credit(prob,true,3),
        deterministic_top1=(prob.argmax(-1)==true).double(),
        deterministic_top3=(prob.topk(min(3,prob.shape[-1]),-1).indices==true[:,None]).any(-1).double(),
        nll=-prob.gather(1,true[:,None]).squeeze(-1).clamp_min(1e-300).log(),
        confidence=prob.max(-1).values,top3_mass=prob.topk(min(3,prob.shape[-1]),-1).values.sum(-1))


def run(device='cuda:1'):
    verify();start=time.perf_counter();torch.cuda.set_device(device);torch.backends.cuda.matmul.allow_tf32=False
    config=json.loads(PROTOCOL.read_text());manifest=json.loads((ROOT/'prepared/manifest.json').read_text())
    results=[];timings=[];executions=[]
    for si,series in enumerate(manifest['series']):
        p=torch.load(ROOT/'prepared'/series['file'],map_location=device,weights_only=True)
        if not series['primary_eligible']:
            executions.append(dict(series=series['name'],eligible=False,reasons=series['exclusions']));continue
        sizes=[len(v) for v in series['levels']];pool=torch.where(p['masks']['source'])[0]
        held=torch.where(p['masks']['unseen']|p['masks']['seen'])[0]
        true=factorial.joint_index(p['h'][held],sizes);grid=factorial.hypotheses(sizes,device)
        for n in config['sample_budgets']:
            if len(pool)<n:
                executions.append(dict(series=series['name'],samples=n,eligible=False,reasons=['insufficient_source_prefix']));continue
            for seed in config['sample_seeds']:
                rng=torch.Generator(device=device).manual_seed(seed+si*101)
                ids=pool[torch.randperm(len(pool),device=device,generator=rng)[:n]]
                mean=p['x'][ids].mean(0);scale=p['x'][ids].std(0).clamp_min(1e-6)
                torch.cuda.synchronize();prep_start=time.perf_counter()
                x=(p['x'][ids]-mean)/scale;h=p['h'][ids];target=(p['x'][held]-mean)/scale
                torch.cuda.synchronize();normalization_seconds=time.perf_counter()-prep_start
                support=torch.ones(len(held),device=device,dtype=torch.bool)
                for f,size in enumerate(sizes):support &= torch.bincount(h[:,f],minlength=size)[p['h'][held,f]]>0
                for name in MODELS:
                    torch.cuda.synchronize();t=time.perf_counter()
                    model=fit_one(x,h,sizes,name,seed+si*101)
                    torch.cuda.synchronize();fit_seconds=time.perf_counter()-t
                    t=time.perf_counter()
                    with torch.no_grad():prob=predict(model,target,range(target.shape[-1]));values=metrics(prob,true,grid,p['h'][held])
                    torch.cuda.synchronize();prediction_seconds=time.perf_counter()-t
                    directory=ROOT/'models'/series['name']/f'n{n}_s{seed}';directory.mkdir(parents=True,exist_ok=True)
                    torch.save(dict(model=model,mean=mean,scale=scale,fit_ids=ids,sizes=sizes),directory/f'{name}.pt')
                    result=dict(series=series['name'],system=series['system'],samples=n,seed=seed,model=name,
                        held_rows=held,groups=p['groups'][held],unseen=p['masks']['unseen'][held],supported_levels=support,
                        metrics=values,fit_seconds=fit_seconds,prediction_seconds=prediction_seconds,
                        neural_parameters=model.get('neural_parameters',0),fit_ids=ids,
                        checkpoint_path=str(directory/f'{name}.pt'))
                    result.update(normalization_seconds=normalization_seconds,
                        checkpoint_bytes=(directory/f'{name}.pt').stat().st_size,
                        adaptation_fits=model.get('adaptation_fits',1))
                    results.append(result);timings.append({k:result[k] for k in ('series','samples','seed','model','fit_seconds','prediction_seconds','neural_parameters','normalization_seconds','checkpoint_bytes','adaptation_fits')})
                print(json.dumps(dict(series=series['name'],samples=n,seed=seed,models=len(MODELS))),flush=True)
    torch.save(dict(results=results,models=MODELS,manifest_eligibility=manifest['all_three_primary_eligible']),ROOT/'confirmation.pt')
    write('confirmation_execution.json',dict(total_seconds=time.perf_counter()-start,device=device,
        peak_gpu_allocated_bytes=torch.cuda.max_memory_allocated(device),cells=len(results),timings=timings,
        unavailable=executions,all_three_primary_eligible=manifest['all_three_primary_eligible'],
        backbone_executions=0,cpu_model_offload=False))
    verify();print('FULL HYPOTHESIS DIAGNOSIS COMPLETE',flush=True)


if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('command',choices=('lock','prepare','run'));p.add_argument('--device',default='cuda:1');a=p.parse_args()
    if a.command=='lock':lock()
    elif a.command=='prepare':prepare()
    else:run(a.device)