"""All construction costs, shared preprocessing, and repeat-use inference."""
from __future__ import annotations

import json
import statistics
import time

import torch

from .data import ROOT, PROTOCOL, write
from .runtime import Timer, construct, replay, select_samples
from .study import verify_lock


def run():
    verify_lock()
    protocol = json.loads(PROTOCOL.read_text())
    device = protocol['device']
    start = time.perf_counter()
    torch.cuda.set_device(device)
    torch.cuda.init(); torch.cuda.synchronize(device)
    startup = time.perf_counter()-start
    torch.backends.cuda.matmul.allow_tf32 = False
    manifest = json.loads((ROOT/'prepared/manifest.json').read_text())
    sources = json.loads((ROOT/'sources/manifest.json').read_text())
    shared_preprocessing = manifest['raw_parse_and_aggregation_seconds']+sum(
        t['preprocessing_seconds'] for t in manifest['tasks'])
    shared_download = sum(a['download_seconds'] for a in sources['archives'])
    archives = []
    with Timer(device) as materialized:
        for task in manifest['tasks']:
            archive = torch.load(ROOT/'prepared'/task['file'],map_location=device,weights_only=True)
            samples = select_samples(archive,'target_64',protocol['adapter']['sample_seeds'][0])
            archives.append((samples,archive['evaluation']))
    # Equal common CUDA/library warmup is reported and charged to every recipe.
    alternatives = [('support_mixture','query_joint')]+[(m,'identity') for m in protocol['estimation_models']]
    with Timer(device) as common_warmup:
        for m,r in alternatives:
            _,code,_,policy = construct(archives[0][0],m,r)
            replay(policy,code,(archives[0][1][:32]//4)%64,archives[0][1][:32]%4,2,1)
    records = []
    for repetition in range(protocol['cost_gate']['repeats']):
        order = alternatives if repetition % 2 == 0 else list(reversed(alternatives))
        for model,method in order:
            with Timer(device) as construction:
                systems = []
                for samples,held in archives:
                    fitted,code,compact,policy = construct(samples,model,method)
                    systems.append((code,policy,held))
            for requests in protocol['cost_gate']['requests']:
                checksum = 0
                with Timer(device) as executed:
                    for code,policy,held in systems:
                        indices = torch.arange(requests,device=device) % len(held)
                        events = held[indices]
                        action,queries,mask = replay(policy,code,(events//4)%64,events%4,2,1)
                        checksum += int(action.sum()+queries.sum()+mask.sum())
                complete = (startup+shared_download+shared_preprocessing+materialized.seconds+
                            common_warmup.seconds+construction.seconds+executed.seconds)
                # All warm/replay results include reconstruction; common raw I/O
                # is additional and cannot hide a slower recipe's steady cost.
                records.append(dict(repetition=repetition,model=model,representation=method,
                    requests_per_task=requests,tasks=len(archives),
                    cuda_startup_seconds=startup,download_seconds=shared_download,
                    parse_aggregate_preprocess_seconds=shared_preprocessing,
                    sample_materialization_seconds=materialized.seconds,
                    common_warmup_seconds=common_warmup.seconds,
                    construction_seconds=construction.seconds,replay_seconds=executed.seconds,
                    steady_complete_seconds=construction.seconds+executed.seconds,
                    cold_complete_seconds=complete,checksum=checksum))
        print(json.dumps({'benchmark_repetition':repetition}),flush=True)
    medians = []
    cost_gate = True
    for requests in protocol['cost_gate']['requests']:
        selected = [r for r in records if r['requests_per_task'] == requests]
        summaries = []
        for model,method in alternatives:
            rs = [r for r in selected if (r['model'],r['representation']) == (model,method)]
            summaries.append(dict(model=model,representation=method,requests_per_task=requests,
                steady_complete_seconds=statistics.median(r['steady_complete_seconds'] for r in rs),
                cold_complete_seconds=statistics.median(r['cold_complete_seconds'] for r in rs),
                construction_seconds=statistics.median(r['construction_seconds'] for r in rs),
                replay_seconds=statistics.median(r['replay_seconds'] for r in rs)))
        candidate = summaries[0]
        best = min(summaries[1:],key=lambda r:r['cold_complete_seconds'])
        saving = 1-candidate['cold_complete_seconds']/best['cold_complete_seconds']
        steady_saving = 1-candidate['steady_complete_seconds']/min(
            r['steady_complete_seconds'] for r in summaries[1:])
        cost_gate &= saving >= protocol['cost_gate']['minimum_complete_time_reduction']
        medians.extend(summaries)
        candidate['cold_saving_against_fastest_full'] = saving
        candidate['steady_saving_against_fastest_full'] = steady_saving
    kcompact,kfull = 8,64
    def storage(k):
        return dict(probability_bytes=k*8*8,encoder_bytes=64,
                    cached_choices_bytes=3*k*3*9*8,cached_values_bytes=3*k*3*9*8,
                    expected_values_bytes=3*3*8,
                    retained_fit_counts_bytes=64*8*8,retained_fitted_model_bytes=64*8*8,
                    retained_64_adaptation_events_bytes=64*8)
    write(ROOT/'benchmark.json',dict(records=records,medians=medians,
        complete_cost_saving_gate=bool(cost_gate),task_count=len(archives),
        actual_tensor_payload_bytes_per_task=dict(compact=storage(kcompact),full=storage(kfull)),
        no_claim_of_physical_query_or_human_time_savings=True))
    decision = json.loads((ROOT/'analysis/decision.json').read_text())
    decision['complete_cost_saving_gate'] = bool(cost_gate)
    decision['practical_algorithm_promotion'] = bool(cost_gate and
        decision['external_compact_representation_gain'] and decision['strongest_full_statistical_noninferiority'])
    write(ROOT/'analysis/decision.json',decision)
    verify_lock()
    print(json.dumps({'cost_gate':bool(cost_gate),'medians':medians}),flush=True)


if __name__ == '__main__': run()