"""スイープ 1 条件の所要時間の内訳を測る（高速化の当たりを付けるための計測）.

sm-render / sm-evaluate をサブプロセスで呼ぶ現行方式と，同じ処理をプロセス内で
直接呼ぶ方式の各段階（起動・import・PLY 書き出し・シーン読込・レンダ・評価・図・遮蔽検証）
を秒で出す．

実行方法:
    uv run python scripts/profile_sweep_step.py \
        configs/_tmp_led_placement/scene_sphere_r170_th090_board_n1.yaml
"""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

ROOT = Path(__file__).resolve().parents[1]
EVAL_YAML = ROOT / "configs" / "eval_baseline_dng_target_shape.yaml"


class Timer:
    """段階ごとの経過秒を記録する."""

    def __init__(self) -> None:
        self.rows: list[tuple[str, float]] = []
        self._t = time.perf_counter()

    def lap(self, label: str) -> None:
        now = time.perf_counter()
        self.rows.append((label, now - self._t))
        self._t = now

    def report(self, title: str) -> None:
        print(f"\n=== {title} ===")
        for label, sec in self.rows:
            print(f"  {label:<40} {sec:8.2f} s")
        print(f"  {'合計':<40} {sum(s for _, s in self.rows):8.2f} s")


def profile_subprocess(scene: Path, out_dir: Path) -> None:
    """現行方式: uv run sm-render → uv run sm-evaluate をサブプロセスで."""
    t = Timer()
    subprocess.run(["uv", "run", "python", "-c", "pass"], cwd=ROOT, check=True,
                   capture_output=True)
    t.lap("uv run python -c pass（起動のみ）")
    subprocess.run(["uv", "run", "python", "-c", "import mitsuba"], cwd=ROOT, check=True,
                   capture_output=True)
    t.lap("uv run python -c 'import mitsuba'")
    subprocess.run(["uv", "run", "sm-render", "--scene", str(scene), "--output", str(out_dir)],
                   cwd=ROOT, check=True, capture_output=True)
    t.lap("uv run sm-render（全体）")
    subprocess.run(["uv", "run", "sm-evaluate", "--image", str(out_dir / "luminance.npy"),
                    "--config", str(EVAL_YAML), "--output", str(out_dir)],
                   cwd=ROOT, check=True, capture_output=True)
    t.lap("uv run sm-evaluate（全体）")
    t.report("現行方式（サブプロセス）")


def profile_inprocess(scene: Path, out_dir: Path, variant: str) -> None:
    """プロセス内で各段階を直接呼ぶ."""
    t = Timer()
    import numpy as np

    from sm_tias_lightsim.runtime.llvm_init import ensure_drjit_libllvm_path

    if variant.startswith("llvm_"):
        ensure_drjit_libllvm_path()
    import mitsuba as mi

    t.lap("import mitsuba")
    mi.set_variant(variant)
    t.lap(f"set_variant({variant})")

    from sm_tias_lightsim.scene.builder import build_scene
    from sm_tias_lightsim.scene.schema import load_scene_yaml

    config = load_scene_yaml(scene)
    config = config.model_copy(
        update={"render": config.render.model_copy(update={"variant": variant})}
    )
    t.lap("load_scene_yaml")
    scene_dict, tempdir = build_scene(config, project_root=scene.parent)
    t.lap("build_scene（STL 読込 + PLY 書き出し）")
    try:
        mscene = mi.load_dict(scene_dict)
        t.lap("mi.load_dict（PLY 読込 + 加速構造）")
        for i in range(2):
            img = mi.render(mscene, spp=config.render.spp)
            arr = np.array(img)
            t.lap(f"mi.render spp={config.render.spp}（{i + 1} 回目）")
    finally:
        shutil.rmtree(tempdir, ignore_errors=True)
    from sm_tias_lightsim.render.output import save_luminance, save_png

    save_png(arr, out_dir / "render.png")
    save_luminance(arr, out_dir / "luminance.npy")
    t.lap("save_png + save_luminance")

    from sm_tias_lightsim.evaluation.pipeline import evaluate_image

    image = out_dir / "luminance.npy"
    evaluate_image(image, EVAL_YAML, out_dir, save_figures=True)
    t.lap("evaluate_image（図あり）")
    evaluate_image(image, EVAL_YAML, out_dir, save_figures=False)
    t.lap("evaluate_image（図なし）")

    from verify_baffle_occlusion import check_occlusion

    result = check_occlusion(scene, board_grid=41, led_count=32)
    t.lap(f"check_occlusion（variant={mi.variant()}・{result['blocked']:.0f}/{result['total']:.0f}）")
    result = check_occlusion(scene, board_grid=41, led_count=32)
    t.lap("check_occlusion（2 回目）")
    t.report(f"プロセス内方式（variant={variant}）")


def main() -> int:
    parser = argparse.ArgumentParser(description="スイープ 1 条件の所要時間の内訳")
    parser.add_argument("scene", type=Path)
    parser.add_argument("--variant", default=os.environ.get("SM_RENDER_VARIANT", "llvm_ad_rgb"))
    parser.add_argument("--skip-subprocess", action="store_true")
    args = parser.parse_args()
    out_dir = Path(tempfile.mkdtemp(prefix="sm_profile_"))
    try:
        if not args.skip_subprocess:
            profile_subprocess(args.scene.resolve(), out_dir)
        profile_inprocess(args.scene.resolve(), out_dir, args.variant)
    finally:
        shutil.rmtree(out_dir, ignore_errors=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())