"""Private repeated-refinement screen: native, verified input reuse, reuse plus CUDA."""
from __future__ import annotations

import argparse
from collections import Counter
from datetime import datetime, timezone
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import statistics
import subprocess
import sys
import time

ROOT=Path(__file__).resolve().parent


def sha(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def canonical(value):
    return hashlib.sha256(json.dumps(value,sort_keys=True,separators=(",",":"),allow_nan=False).encode()).hexdigest()


def worker(args):
    source=Path(args.source).resolve()
    sys.path.insert(0,str(source/"src"))
    from calibrex.core.io import read_mapping
    from calibrex.evaluation import probabilistic_refinement_ablation as ablation
    from calibrex.evaluation import probabilistic_camera_lidar_run as run
    counts=Counter()
    caches={}
    expected={Path(args.correspondence).resolve():sha(args.correspondence),
              Path(args.problem).resolve():sha(args.problem)}

    def loader(role, native):
        def load(path):
            path=Path(path).resolve()
            counts[role+"_requests"]+=1
            assert sha(path)==expected[path]
            if args.variant!="native" and role in caches:
                counts[role+"_reuse"]+=1
                return caches[role]
            value=native(path)
            counts[role+"_validated"]+=1
            assert sha(path)==expected[path]
            if args.variant!="native":
                caches[role]=value
            return value
        return load

    original_correspondence=run.load_probabilistic_correspondence
    original_problem=run.load_camera_lidar_problem
    correspondence_loader=loader("correspondence",original_correspondence)
    problem_loader=loader("problem",original_problem)
    run.load_probabilistic_correspondence=correspondence_loader
    ablation.load_probabilistic_correspondence=correspondence_loader
    run.load_camera_lidar_problem=problem_loader
    ablation.load_camera_lidar_problem=problem_loader
    cuda_counts={}
    if args.variant=="reuse-cuda":
        spec=importlib.util.spec_from_file_location("cuda_screen",ROOT/".tmp-probabilistic-cuda-screen-v1.py")
        extension=importlib.util.module_from_spec(spec)
        spec.loader.exec_module(extension)
        cuda_counts=extension.install_cuda()
    output=Path(args.output)
    output.mkdir()
    begin=time.perf_counter()
    definition,benchmark=ablation.run_probabilistic_refinement_ablation(
        args.correspondence,args.problem,result_directory=output/"results",
        command="complete-paired-uncertainty-ablation",split_seeds=(0,1),
    )
    definition.save(output/"definition.json")
    benchmark.save(output/"benchmark.json")
    api=time.perf_counter()-begin
    assert len(definition.trials)==8
    normalized_results={}
    for trial in definition.trials:
        seed=int(trial.split_id.removeprefix("seed-"))
        path=output/"results"/f"{trial.method_id}.seed-{seed:04d}.yaml"
        assert sha(path)==trial.provenance.output_sha256
        value=read_mapping(path)
        value["provenance"].pop("created_at")
        normalized_results[path.name]=canonical(value)
    def normalize(value):
        value=json.loads(json.dumps(value))
        value["provenance"].pop("generated_at")
        for trial in value["trials"]:
            trial.pop("runtime_seconds")
            trial["provenance"].pop("output_sha256")
        for summary in value.get("method_summaries",{}).values():
            summary.pop("runtime_seconds")
        return value
    manifest=json.loads((source/"manifest.json").read_bytes())
    assert manifest=={p.relative_to(source/"src").as_posix():sha(p) for p in (source/"src").rglob("*.py")}
    for name,mod in list(sys.modules.items()):
        if (name=="calibrex" or name.startswith("calibrex.")) and getattr(mod,"__file__",None):
            path=Path(mod.__file__).resolve()
            assert path.is_relative_to(source/"src") and sha(path)==manifest[path.relative_to(source/"src").as_posix()]
    result={
        "variant":args.variant,"api_seconds":api,"input_counts":dict(counts),"cuda_counts":cuda_counts,
        "complete_results":normalized_results,"definition_sha256":canonical(normalize(definition.model_dump(mode="json"))),
        "benchmark_sha256":canonical(normalize(benchmark.model_dump(mode="json"))),
        "trial_statuses":[(t.method_id,t.split_id,t.status) for t in definition.trials],
        "provenance":{"recorded_at":datetime.now(timezone.utc).isoformat(),"source_commit":args.commit,
                      "source_manifest_sha256":sha(source/"manifest.json"),"launcher_sha256":sha(Path(__file__)),
                      "cuda_extension_sha256":sha(ROOT/".tmp-probabilistic-cuda-screen-v1.py"),
                      "inputs":{str(p):digest for p,digest in expected.items()}},
    }
    (output/"measurement.json").write_text(json.dumps(result,indent=2)+"\n",encoding="utf-8")
    print(json.dumps({"variant":args.variant,"api_seconds":api,"input_counts":dict(counts),"cuda_counts":cuda_counts}),flush=True)


def main():
    parser=argparse.ArgumentParser()
    parser.add_argument("--variant",choices=["native","reuse-cpu","reuse-cuda"])
    parser.add_argument("--source")
    parser.add_argument("--commit")
    parser.add_argument("--output")
    parser.add_argument("--correspondence")
    parser.add_argument("--problem")
    args=parser.parse_args()
    if args.variant:
        worker(args)
        return
    output=ROOT/"outputs/probabilistic-ablation-screen-v5"
    output.mkdir()
    source=ROOT/"outputs/probabilistic-cuda-screen-v1/source"
    fixture=ROOT/"outputs/probabilistic-cuda-screen-inputs-v1"
    commit=subprocess.check_output(["git","rev-parse","HEAD"],text=True).strip()
    protocol={
        "scope":"Complete existing default uncertainty ablation: 2 split seeds x 4 uncertainty settings, 147456 public-geometry correspondences.",
        "order":["native","reuse-cpu","reuse-cuda","reuse-cpu","reuse-cuda","native"],
        "gates":{"native_cuda_whole_ratio_at_least":2.0,"reuse_cpu_cuda_whole_ratio_at_least":2.0,
                 "results":"exact complete 8 fits, definitions and quality aggregates; only bound output hashes and measured durations normalized"},
        "source_commit":commit,"source_manifest_sha256":sha(source/"manifest.json"),
        "launcher_sha256":sha(Path(__file__)),"cuda_extension_sha256":sha(ROOT/".tmp-probabilistic-cuda-screen-v1.py"),
        "inputs":{name:sha(fixture/name) for name in ["correspondence.json","problem.json","fixture-protocol.json"]},
        "provenance":{"recorded_at":datetime.now(timezone.utc).isoformat(),"activity":"unreserved host","blas_threads":1},
    }
    (output/"protocol.json").write_text(json.dumps(protocol,indent=2)+"\n",encoding="utf-8")
    measurements=[]
    for index,variant in enumerate(protocol["order"]):
        target=output/f"{index:02d}-{variant}"
        env=os.environ.copy()
        env["CUPY_CACHE_DIR"]=str(output/f"cache-{index:02d}")
        env["CUDA_CACHE_DISABLE"]="1"
        assert not Path(env["CUPY_CACHE_DIR"]).exists()
        command=[sys.executable,str(Path(__file__)),"--variant",variant,"--source",str(source),
                 "--commit",commit,"--output",str(target),"--correspondence",str(fixture/"correspondence.json"),
                 "--problem",str(fixture/"problem.json")]
        begin=time.perf_counter()
        with (output/f"{index:02d}-{variant}.stdout.log").open("w",encoding="utf-8") as stdout:
            with (output/f"{index:02d}-{variant}.stderr.log").open("w",encoding="utf-8") as stderr:
                completed=subprocess.run(command,env=env,stdout=stdout,stderr=stderr)
        if completed.returncode:
            raise RuntimeError(f"{variant} failed; inspect retained stderr.")
        value=json.loads((target/"measurement.json").read_bytes())
        value["whole_seconds"]=time.perf_counter()-begin
        measurements.append(value)
        (output/"attempts.json").write_text(json.dumps(measurements,indent=2)+"\n",encoding="utf-8")
        print(json.dumps({"index":index,"variant":variant,"whole_seconds":value["whole_seconds"],
                          "api_seconds":value["api_seconds"],"input_counts":value["input_counts"]}),flush=True)
    exact=len({canonical(m["complete_results"]) for m in measurements})==1
    exact=exact and len({m["definition_sha256"] for m in measurements})==1
    exact=exact and len({m["benchmark_sha256"] for m in measurements})==1
    medians={variant:statistics.median(m["whole_seconds"] for m in measurements if m["variant"]==variant)
             for variant in ["native","reuse-cpu","reuse-cuda"]}
    result={"exact_complete_results":exact,"median_whole_seconds":medians,
            "native_cuda_ratio":medians["native"]/medians["reuse-cuda"],
            "reuse_cpu_cuda_ratio":medians["reuse-cpu"]/medians["reuse-cuda"],
            "provenance":protocol["provenance"],"protocol_sha256":sha(output/"protocol.json")}
    (output/"summary.json").write_text(json.dumps(result,indent=2)+"\n",encoding="utf-8")
    print(json.dumps(result),flush=True)


if __name__=="__main__":
    main()
