"""Profile exact committed TUM preparation and full optimizer; not latency evidence."""
import cProfile
import gzip
import hashlib
import json
import os
from pathlib import Path
import pstats
import shutil
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parent
OUT = ROOT / "outputs/rgbd-current-profile-v1"
SOURCE = OUT / "source"

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

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

def main():
    if len(sys.argv) == 1:
        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)
        shutil.copyfile(ROOT / "docs/assets/cuda_rgbd_joint/config-cuda.yaml", OUT / "config.yaml")
        shutil.copyfile(Path(__file__), OUT / "launcher.py.txt")
        write(OUT / "provenance.json", {
            "commit": commit, "launcher_sha256": sha(Path(__file__)),
            "source_manifest_sha256": sha(SOURCE / "manifest.json"),
            "config_sha256": sha(OUT / "config.yaml"),
            "inputs_inventory_sha256": sha(ROOT / "docs/assets/tum_window_processes/inputs.json"),
            "python": sys.version, "blas_threads": 1,
            "scope": "Exploratory profile of current committed dense public TUM preparation and cold CUDA optimizer; fresh compiler cache and disabled driver cache. Profiler overhead; no speed or accuracy claim.",
        })
        env = dict(os.environ, CUPY_CACHE_DIR=str(OUT / 'cupy-compiler-cache'), CUDA_CACHE_DISABLE='1', 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:
            completed = subprocess.run([sys.executable, str(Path(__file__)), "--worker"], cwd=SOURCE, env=env, stdout=stdout, stderr=stderr)
        raise SystemExit(completed.returncode)
    sys.dont_write_bytecode = True
    sys.path.insert(0, str(SOURCE / "src"))
    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
    from calibrex.graph._point_plane_batches import _OwnedPointFactors
    from tools.benchmark_bspline_lidar_imu import _verify_sources
    manifest = _verify_sources(SOURCE)
    assert sha(Path(__file__)) == sha(OUT / "launcher.py.txt")
    config = load_config(OUT / "config.yaml")
    options = adapter._options(config)
    intrinsics = adapter._intrinsics(config, sorted(name for name, sensor in config.sensors.items() if sensor.type == "camera")[0])
    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)
    selected = adapter._select_frames(paired, options)
    pins = {item["name"]: item for item in json.loads((ROOT / "docs/assets/tum_window_processes/inputs.json").read_bytes())}
    verified = {}
    for name in ["depth.txt", "groundtruth.txt", *[frame.depth.path for frame in selected]]:
        assert (data / name).stat().st_size == pins[name]["size_bytes"]
        assert sha(data / name) == pins[name]["sha256"]
        verified[name] = pins[name]
    reports = {}
    def profile(label, function):
        profiler = cProfile.Profile()
        begin = time.perf_counter()
        profiler.enable()
        value = function()
        profiler.disable()
        elapsed = time.perf_counter() - begin
        profiler.dump_stats(OUT / (label + ".pstats"))
        stats = pstats.Stats(profiler)
        rows = [{"file": str(Path(file).relative_to(SOURCE)) if Path(file).is_relative_to(SOURCE) else file,
                 "line": line, "function": name, "calls": info[1], "self_seconds": info[2], "cumulative_seconds": info[3]}
                for (file, line, name), info in stats.stats.items()]
        reports[label] = {"profiled_seconds": elapsed, "profile_sha256": sha(OUT / (label + ".pstats")),
                          "top_self": sorted(rows, key=lambda row: row["self_seconds"], reverse=True)[:100],
                          "top_cumulative": sorted(rows, key=lambda row: row["cumulative_seconds"], reverse=True)[:100]}
        print(label, elapsed, flush=True)
        return value
    def prepare():
        problem = adapter._build_problem(data, selected, intrinsics, options)
        blocks = adapter._xyz_parameter_blocks(problem, options)
        controls = adapter.regular_lattice_control_points(minimum=adapter._SPATIAL_LATTICE_MINIMUM_M, maximum=adapter._SPATIAL_LATTICE_MAXIMUM_M, shape=adapter._SPATIAL_LATTICE_SHAPE)
        factors = _OwnedPointFactors((*adapter._xyz_data_factors(problem), *adapter._xyz_train_only_factors(problem, controls, ((0.0, 0.0, 0.0, 1.0),) * len(controls), options)))
        return problem, blocks, factors
    problem, blocks, factors = profile("preparation", prepare)
    result = profile("optimizer", lambda: adapter._optimize_joint(config, blocks, factors, options))
    imports = {}
    for name, module in sys.modules.copy().items():
        if name == "calibrex" or name.startswith("calibrex."):
            path = Path(module.__file__).resolve()
            relative = path.relative_to(SOURCE / "src").as_posix()
            assert sha(path) == manifest[relative]
            imports[name] = relative
    _verify_sources(SOURCE)
    provenance = {"run_provenance_sha256": sha(OUT / "provenance.json"), "imports": imports, "verified_inputs": verified}
    payload = result.as_dict()
    (OUT / "result.json.gz").write_bytes(gzip.compress((json.dumps({"payload": payload, "provenance": provenance}, indent=2) + "\n").encode(), mtime=0))
    baseline = json.loads(gzip.decompress((ROOT / "docs/assets/tum_initial_correspondences/01-current/result.json.gz").read_bytes()))["payload"]
    assert json.dumps(payload, sort_keys=True) == json.dumps(baseline, sort_keys=True)
    write(OUT / "profile.json", {"reports": reports, "data_factors": len(problem.measurements), "blocks": len(blocks), "result_exact": True, "result_sha256": sha(OUT / "result.json.gz"), "provenance": provenance})

if __name__ == "__main__":
    main()
