"""Explore remaining CPU work in all three immutable-source TUM processes."""
import cProfile
import gzip
import hashlib
import json
import os
from pathlib import Path
import pstats
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parent
OUT = ROOT / "outputs/tum-small-process-profile-v1"
SOURCE = OUT / "source"
sha = lambda p: hashlib.sha256(p.read_bytes()).hexdigest()

def write(path, value):
    path.write_bytes((json.dumps(value, indent=2, allow_nan=False) + "\n").encode())

def save_profile(profile, role, seconds):
    folder = OUT / role
    folder.mkdir(exist_ok=False)
    profile.dump_stats(str(folder / "profile.pstats"))
    stats = pstats.Stats(profile)
    rows = []
    for (filename, line, name), (primitive, calls, own, cumulative, _) in stats.stats.items():
        path = Path(filename)
        if path.is_relative_to(SOURCE):
            filename = path.relative_to(SOURCE).as_posix()
        rows.append({"file": filename, "line": line, "function": name,
                     "primitive_calls": primitive, "calls": calls,
                     "self_seconds": own, "cumulative_seconds": cumulative})
    manifest = json.loads((SOURCE / "manifest.json").read_bytes())
    actual = {p.relative_to(SOURCE / "src").as_posix(): sha(p)
              for p in (SOURCE / "src/calibrex").rglob("*.py")}
    assert actual == manifest
    imports = {}
    for name, module in sys.modules.copy().items():
        if name == "calibrex" or name.startswith("calibrex."):
            path = Path(module.__file__).resolve()
            assert path.is_relative_to(SOURCE / "src"), path
            relative = path.relative_to(SOURCE / "src").as_posix()
            assert sha(path) == manifest[relative]
            imports[name] = relative
    write(folder / "profile.json", {
        "role": role, "elapsed_profiled_seconds": seconds,
        "total_calls": stats.total_calls, "total_self_seconds": stats.total_tt,
        "top_self": sorted(rows, key=lambda v: v["self_seconds"], reverse=True)[:100],
        "top_cumulative": sorted(rows, key=lambda v: v["cumulative_seconds"], reverse=True)[:100],
        "provenance": {"profile_sha256": sha(folder / "profile.pstats"),
                       "launcher_sha256": sha(Path(__file__)),
                       "source_manifest_sha256": sha(SOURCE / "manifest.json"),
                       "run_provenance_sha256": sha(OUT / "provenance.json"),
                       "verified_inputs_sha256":sha(OUT/"verified-inputs.json"),
                       "pid": os.getpid(), "parent_pid": os.getppid(), "imports": imports,
                       "scope": "Exploratory per-process CPU profile with profiler overhead; timers overlap across processes and are not latency speedups."},
    })

def main():
    if len(sys.argv) == 1:
        import yaml
        from tools.benchmark_bspline_lidar_imu import _freeze
        from tools.benchmark_lazy_runtime_loading import _git_identity
        OUT.mkdir(exist_ok=False)
        commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()
        _freeze(SOURCE, commit)
        _git_identity(SOURCE, commit)
        config = yaml.safe_load((ROOT / "examples/public_datasets/tum_rgbd_freiburg1_xyz/native_joint_slac_config.yaml").read_bytes())
        config["dataset"]["path"] = str(ROOT / "data/public/rgbd_dataset_freiburg1_xyz")
        config["project"]["output_dir"] = str(OUT / "run")
        config["pipeline"]["factors"]["tum_rgbd_joint_point_to_plane"]["options"]["window_workers"] = 3
        (OUT / "config.yaml").write_text(yaml.safe_dump(config, sort_keys=False), encoding="utf-8")
        write(OUT / "provenance.json", {
            "commit": commit, "launcher_sha256": sha(Path(__file__)),
            "source_manifest_sha256": sha(SOURCE / "manifest.json"),
            "config_sha256": sha(OUT / "config.yaml"),
            "input_inventory_sha256": sha(ROOT / "docs/assets/tum_window_processes/inputs.json"),
            "python": sys.version, "scope": "Post-process-acceleration development profile; no new speed or accuracy claim.",
        })
        env = dict(os.environ, PYTHONPATH=str(SOURCE / "src"), OPENBLAS_NUM_THREADS="1", MKL_NUM_THREADS="1", OMP_NUM_THREADS="1")
        with (OUT / "stdout.log").open("wb") as stdout, (OUT / "stderr.log").open("wb") as stderr:
            attempt = subprocess.run([sys.executable, "-u", str(Path(__file__)), "--parent"], cwd=SOURCE, env=env, stdout=stdout, stderr=stderr)
        if attempt.returncode:
            raise SystemExit(attempt.returncode)
        from tools.benchmark_tum_window_processes import comparable
        prior = json.loads(gzip.decompress((ROOT / "outputs/tum-factor-preparation-suite-v1/01-current/result.json.gz").read_bytes()))
        current = json.loads((OUT / "result.json").read_bytes())
        import yaml
        reference_config=ROOT/'outputs/tum-factor-preparation-suite-v1/current/config.yaml'
        left=yaml.safe_load(reference_config.read_bytes())
        right=yaml.safe_load((OUT/'config.yaml').read_bytes())
        left['project']['output_dir']=right['project']['output_dir']='<output>'
        assert left==right
        def normalize(value, config_path):
            result=comparable(value)
            run=result['run']
            assert run['config_sha256']==run['provenance']['config_sha256']==sha(config_path)
            assert run['provenance']['command']==['calibrex','calibrate',str(config_path)]
            run['config_sha256']=run['provenance']['config_sha256']='<verified config with output path normalized>'
            run['provenance']['command'][2]='<config>'
            return result
        assert normalize(prior,reference_config)==normalize(current,OUT/'config.yaml'), 'Complete profiled result differs'
        write(OUT/'comparison.json', {'complete_result_exact':True,'provenance':{
            'profile_run_provenance_sha256':sha(OUT/'provenance.json'),
            'profile_result_sha256':sha(OUT/'result.json'),
            'reference_result_sha256':sha(ROOT/'outputs/tum-factor-preparation-suite-v1/01-current/result.json.gz'),
            'reference_config_sha256':sha(reference_config)}})
        print("All three process profiles captured; complete result matches accepted evidence.")
        return
    sys.dont_write_bytecode = True
    sys.path.insert(0, str(SOURCE / "src"))
    from calibrex.solvers import _tum_window_processes as processes
    from calibrex.pipelines.calibrate import CalibrationRunOptions, run_calibration
    parent = sys.argv[1] == "--parent"
    if parent:
        from calibrex.core.config import load_config
        from calibrex.solvers import native_tum_joint_slac_solver as adapter
        from calibrex.data.tum_rgbd import associate_depth_groundtruth,read_image_index,read_groundtruth
        config=load_config(OUT/'config.yaml')
        options=adapter._options(config)
        data=Path(config.dataset.path)
        paired=associate_depth_groundtruth(read_image_index(data/'depth.txt'),read_groundtruth(data/'groundtruth.txt'),max_difference_sec=options.max_pose_time_delta_sec)
        paths={'depth.txt','groundtruth.txt'}
        from dataclasses import replace
        for start in options.replication_start_indices:
            frames=adapter._select_frames(paired,replace(options,frame_start_index=start))
            assert frames is not None
            paths.update(frame.depth.path for frame in frames)
        pins={x['name']:x for x in json.loads((ROOT/'docs/assets/tum_window_processes/inputs.json').read_bytes())}
        verified={}
        for name in sorted(paths):
            pin=pins[name]
            assert sha(data/name)==pin['sha256'] and (data/name).stat().st_size==pin['size_bytes']
            verified[name]=pin
        assert len(verified)==35
        write(OUT/'verified-inputs.json',{'verified':verified,'provenance':{'run_provenance_sha256':sha(OUT/'provenance.json'),'input_inventory_sha256':sha(ROOT/'docs/assets/tum_window_processes/inputs.json')}})
    if parent:
        processes._worker_command = lambda request, response: [sys.executable, "-u", str(Path(__file__)), "--child", str(request), str(response)]
    profile = cProfile.Profile()
    begin = time.perf_counter()
    profile.enable()
    try:
        if parent:
            result = run_calibration(OUT / "config.yaml", CalibrationRunOptions())
        else:
            processes._execute_request(Path(sys.argv[2]), Path(sys.argv[3]))
    finally:
        profile.disable()
        elapsed = time.perf_counter() - begin
        role = "parent" if parent else "window-" + Path(sys.argv[2]).stem.split("-")[0]
        save_profile(profile, role, elapsed)
    if parent:
        write(OUT / "result.json", result.model_dump(mode="json"))

if __name__ == "__main__":
    main()
