"""Opened-data sensitivity analyses; original protocols and decisions are unchanged."""
from pathlib import Path
import argparse,csv,gzip,hashlib,json,sys,time
PROCESS_START=time.perf_counter()
import numpy as np
import torch
from reproduce import ROOT,load,predict,verify
from unionmind.decision_sufficiency import core
from unionmind.decision_sufficiency.analysis import paired_interval

METHODS=['identity','query_joint','posterior_clusters','posterior_bins',
         'future_marginal','raw_information','raw_entropy','hash_code','action_values']


def sync():
    torch.cuda.synchronize()


def action_profiles(fitted):
    """63 finite action values, computed exclusively from the estimated law."""
    mass=fitted.sum((2,3))
    conditional=fitted/mass[...,None,None].clamp_min(1e-300)
    masks,children=core.context_tables(fitted.device)
    masks=masks.to(fitted.dtype)
    utility=fitted.new_tensor([[1.,-2.,0.],[-2.,1.,0.]])
    terminal=torch.einsum('wrhy,ty,ha->wrta',conditional,masks,utility)
    context_mass=torch.einsum('wrhy,ty->wrt',conditional,masks)
    raw_plan=core.plan(conditional)
    profile=[terminal.flatten(2)]
    w,r=fitted.shape[:2]
    costs=fitted.new_tensor([0.,.01,.05])
    for budget in [1,2]:
        previous=raw_plan['values'][budget-1].reshape(w,r,3,9)
        for q in [0,1]:
            for t,child in enumerate(children):
                if child[q] is not None:
                    profile.append(previous[...,child[q]].sum(-1)-context_mass[...,t,None]*costs)
    result=torch.cat(profile,-1)
    assert result.shape==(w,r,63) and bool(torch.isfinite(result).all())
    return result


def encode(fitted,name):
    if name=='action_values':
        return core.weighted_clusters(action_profiles(fitted),fitted.sum((2,3))),8
    return core.encode(fitted,name)


def csv_write(path,records):
    with gzip.open(path,'wt',newline='') as f:
        writer=csv.DictWriter(f,fieldnames=list(records[0]))
        writer.writeheader();writer.writerows(records)


@torch.no_grad()
def controlled(out,device):
    archive=load('data/controlled/worlds',device)
    manifest=json.loads((ROOT/'data/controlled/manifest.json').read_text())['worlds']
    families=np.array([r['family'] for r in manifest])
    worlds=[r['id'] for r in manifest]
    rows=[];frontiers=[];timings=[]
    for available in [6,3]:
        truth=archive['joint'] if available==6 else core.project(archive['joint'],3)
        for n in [64,256,1024,4096]:
            for estimator in ['joint_jeffreys','naive_bayes','support_mixture']:
                group={}
                for stream in range(3):
                    sync();start=time.perf_counter()
                    fitted=core.estimate(core.count_samples(archive['samples'][:,stream,:n],available),estimator)
                    sync();estimation=time.perf_counter()-start
                    for method in METHODS:
                        sync();start=time.perf_counter()
                        code,k=encode(fitted,method);compressed=core.push(fitted,code,k)
                        policy=core.plan(compressed)
                        sync();construction=time.perf_counter()-start
                        gold=core.push(truth,code,k)
                        score=core.realized(gold,policy)['utility']
                        if available==6 and n==1024 and estimator=='support_mixture':
                            oracle=core.plan(gold)['expected'][:,1:,:].mean((1,2))
                        else:oracle=torch.full_like(score[:,0,0],float('nan'))
                        mean=score[:,1:,:].mean((1,2)).tolist()
                        oracle_values=oracle.tolist()
                        for i,value in enumerate(mean):
                            rows.append(dict(available=available,labels=n,estimator=estimator,
                                stream=stream,method=method,world=worlds[i],family=families[i],
                                utility=value,encoded_oracle=oracle_values[i]))
                        group.setdefault(method,[]).append(score.mean(0))
                        timings.append(dict(available=available,labels=n,estimator=estimator,
                            stream=stream,method=method,estimation_seconds=estimation,
                            encoding_push_planning_seconds=construction))
                for method,records in group.items():
                    scores=torch.stack(records).mean(0)
                    for b in range(3):
                        for c,cost in enumerate([0.,.01,.05]):
                            frontiers.append(dict(available=available,labels=n,estimator=estimator,
                                method=method,budget=b,cost=cost,utility=float(scores[b,c])))
    csv_write(out/'controlled_worlds.csv.gz',rows)
    csv_write(out/'controlled_frontiers.csv.gz',frontiers)
    primary=[r for r in rows if r['available']==6 and r['labels']==1024
             and r['estimator']=='support_mixture']
    per={}
    for method in METHODS:
        per[method]=np.array([[r['utility'] for r in primary if r['method']==method
                              and r['stream']==s] for s in range(3)]).mean(0)
    old=json.loads((ROOT/'results/controlled/decision.json').read_text())
    for reference in old['primary_metrics']:
        assert abs(float(per[reference['method']].mean())-reference['utility'])<1e-10
    delta=per['query_joint']-per['action_values']
    comparisons=dict(joint_minus_action_values=float(delta.mean()),
        descriptive_paired_95=paired_interval(delta,families,18701),
        mean_utility={name:float(a.mean()) for name,a in per.items()},
        status='Post-confirmation sensitivity on opened worlds; no new confirmation or promotion')
    # Every declared encoder is timed, with construction and planner work charged.
    bench=[];warmup=[]
    for repetition in range(-1,11):
        names=METHODS if repetition%2==0 else list(reversed(METHODS))
        for method in names:
            sync();start=time.perf_counter()
            fitted=core.estimate(core.count_samples(archive['samples'][:,0,:1024],6),'support_mixture')
            code,k=encode(fitted,method);compressed=core.push(fitted,code,k);policy=core.plan(compressed)
            sync();elapsed=time.perf_counter()-start
            row=dict(repetition=repetition,method=method,complete_construction_seconds=elapsed,
                retained_encoder_and_law_bytes=code.numel()*code.element_size()+compressed.numel()*compressed.element_size())
            (warmup if repetition<0 else bench).append(row)
    comparisons['warm_batched_construction_medians']={
        name:float(np.median([r['complete_construction_seconds'] for r in bench if r['method']==name]))
        for name in METHODS}
    (out/'controlled_timing.json').write_text(json.dumps(dict(
        scope='192-world batch; estimation, profiles, clustering, pushforward and exact plans included',
        warmup=warmup,repetitions=bench,grid_timings=timings),indent=2)+'\n')
    (out/'controlled_summary.json').write_text(json.dumps(comparisons,indent=2)+'\n')
    return comparisons


@torch.no_grad()
def sensors(out,device):
    accounting={};normalization=[];hydraulic=None
    for cohort in ['hydraulic','temperature']:
        p=load(f'data/{cohort}/prepared',device)
        if cohort=='hydraulic':p.update(p['masks']);hydraulic=p
        if cohort=='temperature':
            excluded=(p['days']<9)&p['unseen']
            assert int(excluded.sum())==128
            assert bool((p['source']|p['target']|excluded).all())
            assert not bool((p['source']&p['target']).any())
            accounting[cohort]=dict(retained=len(p['x']),source=int(p['source'].sum()),
                target=int(p['target'].sum()),excluded_source_bindings=int(excluded.sum()),
                per_day=[dict(day=int(d),retained=int((p['days']==d).sum()),
                    source=int(((p['days']==d)&p['source']).sum()),
                    excluded_source_bindings=int(((p['days']==d)&excluded).sum()),
                    target=int(((p['days']==d)&p['target']).sum())) for d in p['days'].unique()])
        for n in [64,128]:
            for directory in sorted((ROOT/'models'/cohort).glob(f'n{n}_s*')):
                common=None
                for path in sorted(directory.glob('*.json')):
                    c=load(str(path.relative_to(ROOT).with_suffix('')),device)
                    ids=c['fit_ids']
                    assert len(ids)==n and bool(p['source'][ids].all()) and not bool(p['unseen'][ids].any())
                    if common is None:common=ids
                    else:assert torch.equal(ids,common)
                    dm=float((c['mean']-p['x'][ids].mean(0)).abs().max())
                    ds=float((c['scale']-p['x'][ids].std(0).clamp_min(1e-6)).abs().max())
                    assert dm<1e-10 and ds<1e-10,(cohort,n,directory.name,path.name,dm,ds)
                    normalization.append(dict(cohort=cohort,prefix=n,fit=directory.name,
                        model=path.stem,maximum_mean_error=dm,maximum_scale_error=ds,
                        normalization_records=n,additional_unlabeled_records=0))
    p=hydraulic;sizes=[3,4,3,4];truth=load('results/hydraulic/predictions',device)['results'][0]
    held=truth['held_rows'][truth['unseen']]
    true=core.bits(1,device)  # no fitted module is instantiated or trained
    from unionmind.continuous_decision import factorial
    true=factorial.joint_index(p['h'][held],sizes)
    fits=[];cross=[];bindings=[]
    for seed in [20751,20752,20753]:
        ck={n:load(f'models/hydraulic/n{n}_s{seed}/factorized_neural',device) for n in [64,128]}
        assert torch.equal(ck[64]['fit_ids'],ck[128]['fit_ids'][:64])
        for n,c in ck.items():
            ids=c['fit_ids'];m=c['model'];x=(p['x'][ids]-c['mean'])/c['scale']
            hx=factorial.joint_index(p['h'][ids],sizes)
            evalx=(p['x'][held]-c['mean'])/c['scale']
            eig=torch.linalg.eigvalsh(m['cov'])
            fits.append(dict(seed=seed,labels=n,nested_prefix=True,
                observed_bindings=len(hx.unique()),
                factor_level_counts=[torch.bincount(p['h'][ids,f],minlength=k).tolist() for f,k in enumerate(sizes)],
                training_mean_squared_error=float((x-m['mean'][hx]).square().mean()),
                unseen_true_mean_squared_error=float((evalx-m['mean'][true]).square().mean()),
                covariance_trace=float(eig.sum()),covariance_minimum_eigenvalue=float(eig.min()),
                covariance_maximum_eigenvalue=float(eig.max()),covariance_condition_number=float(eig.max()/eig.min())))
        for mean_n in [64,128]:
            for cov_n in [64,128]:
                c=ck[mean_n];model=dict(c['model'])
                ratio=ck[cov_n]['scale']/c['scale']
                model['cov']=ck[cov_n]['model']['cov']*ratio[:,None]*ratio[None,:]
                x=(p['x'][held]-c['mean'])/c['scale']
                prob=factorial.posterior(model,x,range(56));score=factorial.topk_credit(prob,true,3)
                cross.append(dict(seed=seed,mean_labels=mean_n,covariance_labels=cov_n,
                    top3=float(score.mean()),hypothesis_log_loss=float(-prob.gather(1,true[:,None]).clamp_min(1e-300).log().mean())))
                for h in true.unique():
                    mask=true==h
                    bindings.append(dict(seed=seed,mean_labels=mean_n,covariance_labels=cov_n,
                        binding=int(h),records=int(mask.sum()),top3=float(score[mask].mean())))
    result=dict(status='Descriptive frozen-checkpoint diagnostics, no refitting',
        normalization_checkpoints=len(normalization),normalization=normalization,
        episode_accounting=accounting,frozen_hydraulic_fits=fits,crossed_mean_covariance=cross)
    (out/'sensor_diagnostics.json').write_text(json.dumps(result,indent=2)+'\n')
    csv_write(out/'hydraulic_bindings.csv.gz',bindings)
    return result


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--device',default='cuda:0');parser.add_argument('--output',required=True)
    a=parser.parse_args()
    assert a.device.startswith('cuda') and torch.cuda.is_available()
    out=Path(a.output).resolve();assert not out.exists(),'Use a new output directory.'
    protocol=json.loads((ROOT/'protocols/postconfirmation_analysis_01.json').read_text())
    assert protocol['status']=='Post-confirmation analysis of opened data'
    assert hashlib.sha256(Path(__file__).read_bytes()).hexdigest()==protocol['code_sha256']
    out.mkdir(parents=True);torch.cuda.set_device(a.device)
    torch.backends.cuda.matmul.allow_tf32=False
    sync();start=time.perf_counter();initialization=start-PROCESS_START
    with torch.no_grad():
        controlled_summary=controlled(out,a.device);sensor_summary=sensors(out,a.device)
    sync()
    ledger=dict(opened_data=True,new_neural_modules=0,refitted_models=0,backbone_executions=0,
        source_files_modified=False,initialization_seconds=initialization,
        complete_analysis_seconds=time.perf_counter()-start,
        complete_process_since_framework_import_seconds=time.perf_counter()-PROCESS_START,
        peak_cuda_allocated_bytes=torch.cuda.max_memory_allocated(),
        gpu=torch.cuda.get_device_name(),device=a.device,protocol_code_sha256=protocol['code_sha256'])
    (out/'compute.json').write_text(json.dumps(ledger,indent=2)+'\n')
    print(json.dumps(dict(controlled=controlled_summary,
        normalization_checkpoints=sensor_summary['normalization_checkpoints'],
        episode_accounting=sensor_summary['episode_accounting'],compute=ledger),indent=2))
if __name__=='__main__':main()