"""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 seed_worker(task):
    source,commit,paths,kwargs=task
    sys.path.insert(0,str(Path(source)/"src"))
    from calibrex.evaluation.probabilistic_refinement_ablation import run_probabilistic_refinement_ablation
    expected={p:sha(p) for p in paths}
    definition,_=run_probabilistic_refinement_ablation(*paths,**kwargs)
    assert all(sha(p)==digest for p,digest in expected.items())
    manifest=json.loads((Path(source)/"manifest.json").read_bytes())
    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(Path(source)/"src") and sha(path)==manifest[path.relative_to(Path(source)/"src").as_posix()]
    return definition


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=="reuse-cpu" 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=="reuse-cpu":
                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=="parallel":
        native_ablation=ablation.run_probabilistic_refinement_ablation
        def parallel(*paths,**kwargs):
            from concurrent.futures import ProcessPoolExecutor
            import multiprocessing
            from calibrex.core.benchmark import aggregate_benchmark_definition
            tasks=[(str(source),args.commit,tuple(str(p) for p in paths),dict(kwargs,split_seeds=(seed,))) for seed in (0,1,2,3,4)]
            with ProcessPoolExecutor(max_workers=4,mp_context=multiprocessing.get_context("spawn")) as pool:
                definitions=list(pool.map(seed_worker,tasks))
            first=definitions[0]
            definition=first.model_copy(update={"trials":[t for d in definitions for t in d.trials],
                "protocol":first.protocol.model_copy(update={"splits":[s for d in definitions for s in d.protocol.splits]})})
            return definition,aggregate_benchmark_definition(definition)
        ablation.run_probabilistic_refinement_ablation=parallel
    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",
    )
    definition.save(output/"definition.json")
    benchmark.save(output/"benchmark.json")
    api=time.perf_counter()-begin
    assert len(definition.trials)==20
    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__)),
                      "workers":4,
                      "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","parallel"])
    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-v3"
    output.mkdir()
    source=ROOT/"outputs/probabilistic-cuda-screen-v1/source"
    fixture=ROOT/"outputs/kitti-probabilistic-control-v2"
    commit=subprocess.check_output(["git","rev-parse","HEAD"],text=True).strip()
    protocol={
        "scope":"Complete existing default uncertainty ablation: 5 split seeds x 4 uncertainty settings.",
        "order":["native","parallel","parallel","native"],
        "gates":{"native_cuda_whole_ratio_at_least":2.0,"reuse_cpu_cuda_whole_ratio_at_least":2.0,
                 "results":"exact complete 20 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__)),"workers":4,
        "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","parallel"]}
    result={"exact_complete_results":exact,"median_whole_seconds":medians,
            "native_cuda_ratio":medians["native"]/medians["parallel"],
            
            "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()
