"""All registered full-hypothesis comparisons; held recording groups are paired."""
from __future__ import annotations
import json,time
import torch
from .factor_study import ROOT,PROTOCOL,MODELS,verify,write

STRONG=['factorized_nb','factorized_lda','factorized_kernel','factorized_ridge','joint_gaussian','factorized_logistic','factorized_classifier']


@torch.no_grad()
def main(device='cuda:2'):
    verify();start=time.perf_counter();torch.cuda.set_device(device);torch.backends.cuda.matmul.allow_tf32=False
    config=json.loads(PROTOCOL.read_text());data=torch.load(ROOT/'confirmation.pt',map_location=device,weights_only=True)
    manifest=json.loads((ROOT/'prepared/manifest.json').read_text());rng=torch.Generator(device=device).manual_seed(config['uncertainty']['seed'])
    draws=config['uncertainty']['draws'];summary=[];comparisons=[];gates=[];byseries={};group_audits=[]
    fields=('top1','top3','deterministic_top1','deterministic_top3','nll','confidence','top3_mass')
    for n in config['sample_budgets']:
        available=[];series_points={};boot_values={}
        for series in manifest['series']:
            records=[r for r in data['results'] if r['series']==series['name'] and r['samples']==n]
            if not records:
                summary.append(dict(series=series['name'],samples=n,available=False,reason='preregistered_eligibility_or_source_budget'));continue
            available.append(series)
            template=records[0];unseen=template['unseen'];unique,inverse=torch.unique(template['groups'][unseen],sorted=True,return_inverse=True)
            counts=torch.bincount(inverse,minlength=len(unique)).double();values=counts.new_zeros((len(unique),len(MODELS)))
            point=counts.new_zeros(len(MODELS))
            for mi,name in enumerate(MODELS):
                rows=[r for r in records if r['model']==name];assert len(rows)==3
                assert all(torch.equal(r['held_rows'],template['held_rows']) for r in rows)
                scores={f:torch.stack([r['metrics'][f] for r in rows]).mean(0) for f in fields}
                for stratum,mask in (('unseen',unseen),('seen',~unseen)):
                    summary.append(dict(series=series['name'],system=series['system'],samples=n,model=name,stratum=stratum,
                        available=True,records=int(mask.sum()),**{f:float(v[mask].mean()) for f,v in scores.items()}))
                for source_support in (False,True):
                    total=weight=0.
                    for row in rows:
                        mask=unseen&(row['supported_levels']==source_support)
                        total+=float(row['metrics']['top3'][mask].sum());weight+=int(mask.sum())
                    summary.append(dict(series=series['name'],samples=n,model=name,
                        stratum='unseen_supported_levels' if source_support else 'unseen_novel_level',available=weight>0,
                        records_across_three_repetitions=weight,top3=total/weight if weight else None))
                # Selective accuracy is measured per repetition, without ensembling.
                for threshold in (.5,.7,.9):
                    selected=[]
                    for row in rows:
                        mask=unseen&(row['metrics']['confidence']>=threshold)
                        selected.append((int(mask.sum()),float(row['metrics']['top1'][mask].sum())))
                    accepted=sum(r[0] for r in selected);correct=sum(r[1] for r in selected)
                    summary.append(dict(series=series['name'],samples=n,model=name,stratum='selective_unseen',threshold=threshold,
                        coverage=accepted/(3*int(unseen.sum())),accuracy=correct/accepted if accepted else None))
                v=scores['top3'][unseen];values[:,mi].scatter_add_(0,inverse,v);point[mi]=v.mean()
            sampled=torch.randint(len(unique),(draws,len(unique)),device=device,generator=rng)
            weights=counts.new_zeros((draws,len(unique))).scatter_add_(1,sampled,torch.ones_like(sampled,dtype=counts.dtype))
            bootstrap=(weights@values)/(weights@counts)[:,None]
            series_points[series['name']]=point;boot_values[series['name']]=bootstrap
            group_audits.append(dict(series=series['name'],samples=n,groups=len(unique),unseen_records=int(unseen.sum())))
            byseries[f'{series["name"]}:{n}']=point.tolist()
        if not available:
            gates.append(dict(samples=n,complete_cohort=False,reason='no_eligible_series',promotion=False));continue
        # Hydraulic data were already opened. They remain a complete development
        # result, but cannot contribute to a fresh-confirmation effect estimate.
        fresh=[s for s in available if s['system']=='gas_array_rig']
        if not fresh:
            gates.append(dict(samples=n,complete_cohort=False,reason='no_fresh_eligible_series',promotion=False));continue
        systems=sorted({s['system'] for s in fresh});weights={}
        for system in systems:
            group=[s for s in fresh if s['system']==system]
            for series in group:weights[series['name']]=1/(len(systems)*len(group))
        point=sum(series_points[name]*w for name,w in weights.items())
        bootstrap=sum(boot_values[name]*w for name,w in weights.items())
        complete=len(available)==3
        for candidate in config['candidates']:
            ci=MODELS.index(candidate);refs=[i for i in range(len(MODELS)) if i!=ci]
            diff=bootstrap[:,ci,None]-bootstrap[:,refs];alpha=.05/(2*config['uncertainty']['family'])
            lo,hi=torch.quantile(diff,diff.new_tensor([alpha,1-alpha]),dim=0)
            rows=[]
            for j,index in enumerate(refs):
                system_gains={}
                for system in systems:
                    group=[s for s in fresh if s['system']==system]
                    system_gains[system]=sum(float(series_points[s['name']][ci]-series_points[s['name']][index]) for s in group)/len(group)
                row=dict(samples=n,candidate=candidate,control=MODELS[index],gain=float(point[ci]-point[index]),
                    simultaneous_95_interval=[float(lo[j]),float(hi[j])],system_gains=system_gains,
                    series_gains={name:float(p[ci]-p[index]) for name,p in series_points.items()},complete_cohort=complete)
                rows.append(row);comparisons.append(row)
            strong=[r for r in rows if r['control'] in STRONG];assert len(strong)==7
            large=complete and all(r['gain']>=.10 and r['simultaneous_95_interval'][0]>0 and all(r['series_gains'][s['name']]>0 for s in fresh) for r in strong)
            best=complete and all(r['simultaneous_95_interval'][0]>0 for r in rows)
            learning=next((r['simultaneous_95_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,utility_joint_top3=float(point[ci]),complete_cohort=complete,
                substantial_gain=large,strongest_available=best,learning_specific=learning,promotion=large and best,
                strongest_observed_model=MODELS[int(point.argmax())],strongest_observed_top3=float(point.max())))
    write('confirmation_summary.json',dict(rows=summary,comparisons=comparisons,gates=gates,models=MODELS,
        series_points=byseries,bootstrap_draws=draws,comparison_family=config['uncertainty']['family'],group_audits=group_audits,
        all_three_primary_eligible=manifest['all_three_primary_eligible'],
        conditional_on_fixed_rigs_and_fitted_models=True,no_unavailable_secondary_used_to_rescue_primary=True))
    torch.cuda.synchronize();write('analysis_execution.json',dict(seconds=time.perf_counter()-start,device=device,
        peak_gpu_bytes=torch.cuda.max_memory_allocated(device)))
    verify();print(json.dumps(dict(gates=gates)),flush=True)


if __name__=='__main__':main()