"""Fixed paired analysis. No recipe selection, test-set tuning or gate revision."""
from __future__ import annotations

import csv
import gzip
import json
from collections import defaultdict

import torch

from .data import ROOT, PROTOCOL, write
from .study import verify_lock


def mean(xs): return sum(xs)/len(xs)


def run():
    verify_lock()
    protocol = json.loads(PROTOCOL.read_text())
    device = protocol['device']
    grouped = defaultdict(list)
    with gzip.open(ROOT/'evaluation/metrics.csv.gz','rt') as f:
        for row in csv.DictReader(f):
            if int(row['budget']) == 0: continue
            key = (row['dataset'],row['environment'],row['task'],row['regime'],
                   row['model'],row['representation'])
            grouped[key].append(row)
    fields = ['utility','queries','abstention','estimated_utility','fitted_compression_loss','fitted_dispersion',
              'empirical_compression_loss','empirical_model_loss','posterior_brier',
              'posterior_log_loss','joint_conditional_log_loss','observation_given_h_log_loss',
              'raw_posterior_brier','raw_posterior_log_loss','raw_joint_conditional_log_loss']
    task_results = {key:{f:mean([float(r[f]) for r in rows]) for f in fields}
                    for key,rows in grouped.items()}
    environments = defaultdict(list)
    for key, scores in task_results.items():
        d,e,t,regime,model,representation = key
        environments[(d,e,regime,model,representation)].append(scores)
    env_results = {key:{f:mean([s[f] for s in scores]) for f in fields}
                  for key,scores in environments.items()}
    datasets = defaultdict(list)
    for key,scores in env_results.items():
        d,e,regime,model,representation = key
        datasets[(d,regime,model,representation)].append(scores)
    dataset_results = {key:{f:mean([s[f] for s in scores]) for f in fields}
                       for key,scores in datasets.items()}
    all_results = defaultdict(list)
    for key,scores in dataset_results.items():
        d,regime,model,representation = key
        all_results[(regime,model,representation)].append(scores)
    overall = [{**dict(zip(('regime','model','representation'),key)),
                **{f:mean([s[f] for s in scores]) for f in fields}}
               for key,scores in all_results.items()]
    write(ROOT/'analysis/overall.json',overall)
    write(ROOT/'analysis/by_dataset.json',[
        {**dict(zip(('dataset','regime','model','representation'),key)),**scores}
        for key,scores in dataset_results.items()])
    write(ROOT/'analysis/by_environment.json',[
        {**dict(zip(('dataset','environment','regime','model','representation'),key)),**scores}
        for key,scores in env_results.items()])
    write(ROOT/'analysis/by_task.json',[
        {**dict(zip(('dataset','environment','task','regime','model','representation'),key)),**scores}
        for key,scores in task_results.items()])
    primary = protocol['primary']
    regime,model = primary['regime'],primary['model']
    comparisons = [(model,m) for m in primary['compact_controls']]+[
        (m,'identity') for m in protocol['estimation_models']]
    datasets_present = sorted(set(k[0] for k in env_results if k[2] == regime))
    generator = torch.Generator(device=device).manual_seed(primary['bootstrap_seed'])
    # One set of paired environment resamples reused for all nine comparisons.
    draws = primary['bootstrap_draws']
    bootstraps, points, per_dataset = [], [], {}
    nenv = {}
    for d in datasets_present:
        names = sorted(k[1] for k in env_results if k[0] == d and k[2:] == (regime,model,'query_joint'))
        nenv[d] = len(names)
        differences = []
        for e in names:
            q = env_results[(d,e,regime,model,'query_joint')]['utility']
            differences.append([q-env_results[(d,e,regime,m,r)]['utility'] for m,r in comparisons])
        delta = torch.tensor(differences,device=device,dtype=torch.float64)
        indices = torch.randint(len(names),(draws,len(names)),device=device,generator=generator)
        bootstraps.append(delta[indices].mean(1))
        points.append(delta.mean(0))
        per_dataset[d] = delta.mean(0).tolist()
    boot = torch.stack(bootstraps).mean(0)
    point = torch.stack(points).mean(0)
    # Registered conservative family includes all compact and full comparisons.
    alpha = .05/(2*len(comparisons))
    low,high = torch.quantile(boot,boot.new_tensor([alpha,1-alpha]),dim=0)
    paired = [dict(model=m,representation=r,gain=float(point[i]),
                   simultaneous_95_interval=[float(low[i]),float(high[i])],
                   dataset_gains={d:per_dataset[d][i] for d in datasets_present})
              for i,(m,r) in enumerate(comparisons)]
    write(ROOT/'analysis/paired_comparisons.json',dict(comparisons=paired,
        eligible_environments=nenv,bootstrap_unit='target environment, stratified by fixed dataset',
        family_size=len(comparisons),draws=draws,
        limitations='Few environmental clusters; conditional benchmark intervals, not universal transfer certification.'))
    compact_gate = len(datasets_present) == 2 and min(nenv.values()) >= 2 and all(
        r['gain'] >= primary['effect_minimum'] and r['simultaneous_95_interval'][0] > 0 and
        min(r['dataset_gains'].values()) > 0 for r in paired[:6])
    full_gate = all(r['simultaneous_95_interval'][0] > -.01 for r in paired[6:])
    write(ROOT/'analysis/decision.json',dict(
        external_compact_representation_gain=bool(compact_gate),
        strongest_full_statistical_noninferiority=bool(full_gate),
        complete_cost_saving_gate='pending benchmark',practical_algorithm_promotion=False,
        learned_controller_transfer_tested=False,representation_construction_refit=True,
        external_recorded_measurements=True,provider_authored_sequential_task=False,
        no_threshold_changes_after_opening=True,no_test_selected_recipe=True,
        empirical_hindsight_decomposition_is_descriptive_only=True))
    print(json.dumps({'compact_gate':bool(compact_gate),'full_reference_gate':bool(full_gate),
                      'environment_counts':nenv,'paired':paired}),flush=True)


if __name__ == '__main__': run()