"""spp スイープ: サンプル数に対する計算時間・GPU 率・ノイズ・指標の安定性を測る.

TEST_18 までは spp=1024 固定で，大きいキャビティほど生の輝度が低く相対ノイズが 7〜26% に
達していた（指標はノイズ補正で救っている）．研究室サーバの GPU で 1 条件 5 秒になったので，
spp を上げたときに (1) 計算時間がどう増えるか（CPU 側の固定費と GPU 側の spp 比例部分），
(2) ノイズが理論どおり 1/sqrt(spp) で下がるか，(3) 各指標（構造 CoV・輪帯 min/max・
画像全体 min/max）がどの spp で安定するかを測り，今後の標準 spp を決める．

条件（TEST_18 の既存 YAML を流用し spp と seed だけ書き換える）:
    球 φ340（r=170）θ=0/90，立方体 L=234 θ=0/90，球 φ500（r=250）θ=90，立方体 L=340 θ=90
    × spp = 256〜16384（2 倍刻み 7 段階）× seed 2 つ

計測:
    - 段階別の時間（YAML 読込 / build_scene=PLY 書き出し / load_dict / mi.render / 保存 / 評価）．
      mi.render は dr.sync_thread で GPU の完了を待ってから測る．GPU 率 = render / 全体
    - 残差ノイズ CoV（生 − 平滑化）・生の画素 CoV・構造 CoV・輪帯 min/max・画像全体 min/max
      （ノイズ補正あり / なし）・ノイズ場の極値
    - 併走する nvidia-smi のサンプリング（sm-remote run の外で回す）は解析側で突き合わせる

--set shapes6（スライド用の追試）: LED を軸上（θ=0）にそろえ，形状だけを 6 種類にする
    球 φ270/340/500（r=135/170/250）+ 立方体 L=187/234/340．出力は output/spp_noise_shapes6/

実行方法:
    uv run python scripts/sweep_spp_noise.py --prepare-only
    uv run sm-remote run --name spp --gpu 0 --threads 16 -- python scripts/sweep_spp_noise.py
"""

from __future__ import annotations

import argparse
import copy
import csv
import json
import os
import shutil
import sys
import time
from pathlib import Path

import numpy as np
import yaml

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

from build_radial_comparison import radial_stats, ring_min_max_ratio  # noqa: E402
from build_radius_sweep_visual import decompose_cov  # noqa: E402
from sweep_led_placement import TMP_DIR as PLACEMENT_TMP_DIR  # noqa: E402
from sweep_sphere_radius import EVAL_CONFIG_YAML  # noqa: E402

from sm_tias_lightsim.evaluation.fullimage import calc_fullimage_min_max_ratio  # noqa: E402
from sm_tias_lightsim.evaluation.pipeline import evaluate_image  # noqa: E402
from sm_tias_lightsim.evaluation.roi_presets import SIM_ROI, extract_roi_xywh  # noqa: E402
from sm_tias_lightsim.runtime.overrides import resolve_render_variant  # noqa: E402
from sm_tias_lightsim.scene.schema import load_scene_yaml  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
TMP_DIR = ROOT / "configs" / "_tmp_spp_noise"
OUTPUT_BASE = ROOT / "output" / "spp_noise_sweep"

# (表示名, TEST_18 の条件名)
CONDITIONS = [
    ("sphere_d340_th000", "sphere_r170_th000_board_n1"),
    ("sphere_d340_th090", "sphere_r170_th090_board_n1"),
    ("cube_L234_th000", "cube_L234_th000_board_n1"),
    ("cube_L234_th090", "cube_L234_th090_board_n1"),
    ("sphere_d500_th090", "sphere_r250_th090_board_n1"),
    ("cube_L340_th090", "cube_L340_th090_board_n1"),
]
# --set shapes6: LED 位置を θ=0 にそろえ形状だけを変える 6 条件
CONDITIONS_SHAPES6 = [
    ("sphere_d270_th000", "sphere_r135_th000_board_n1"),
    ("sphere_d340_th000", "sphere_r170_th000_board_n1"),
    ("sphere_d500_th000", "sphere_r250_th000_board_n1"),
    ("cube_L187_th000", "cube_L187_th000_board_n1"),
    ("cube_L234_th000", "cube_L234_th000_board_n1"),
    ("cube_L340_th000", "cube_L340_th000_board_n1"),
]
SPPS = [256, 512, 1024, 2048, 4096, 8192, 16384]
SEEDS = [0, 1]

SUMMARY_COLUMNS = [
    "name", "condition", "spp", "seed", "render_variant",
    "t_total", "t_load_yaml", "t_build_scene", "t_load_dict", "t_render", "t_save", "t_evaluate",
    "gpu_fraction", "t_render_per_kspp",
    "raw_mean_roi", "pixel_cov", "noise_cov", "structure_cov", "ring_min_max",
    "fullimage_min_max", "fullimage_raw_min_max", "fullimage_noise_min", "fullimage_noise_max",
    "radial_min_max_ratio", "cov",
]


def job_name(condition: str, spp: int, seed: int) -> str:
    """出力ディレクトリ名."""
    return f"{condition}_spp{spp}_s{seed}"


def prepare(log) -> dict[str, Path]:
    """TEST_18 の YAML を spp・seed だけ書き換えて書き出す（メッシュは元の場所を相対参照）."""
    scenes: dict[str, Path] = {}
    TMP_DIR.mkdir(parents=True, exist_ok=True)
    for display, source_name in CONDITIONS:
        source = PLACEMENT_TMP_DIR / f"scene_{source_name}.yaml"
        if not source.exists():
            raise FileNotFoundError(
                f"{source} がありません．先に sweep_led_placement.py --prepare-only を実行"
            )
        config = yaml.safe_load(source.read_text(encoding="utf-8"))
        for mesh in config["geometry"]["meshes"]:
            path = Path(mesh["path"])
            if not path.is_absolute():
                # 元 YAML からの相対パスを，新しい YAML からの相対パスに付け替える
                target = (source.parent / path).resolve()
                mesh["path"] = Path(os.path.relpath(target, TMP_DIR)).as_posix()
        for spp in SPPS:
            for seed in SEEDS:
                cfg = copy.deepcopy(config)
                cfg["name"] = f"spp_noise_{job_name(display, spp, seed)}"
                cfg["render"]["spp"] = spp
                cfg["render"]["seed"] = seed
                dest = TMP_DIR / f"scene_{job_name(display, spp, seed)}.yaml"
                dest.write_text(yaml.safe_dump(cfg, allow_unicode=True, sort_keys=False),
                                encoding="utf-8")
                scenes[job_name(display, spp, seed)] = dest
    log(f"=== 幾何流用: {len(CONDITIONS)} 条件 × {len(SPPS)} spp × {len(SEEDS)} seed = "
        f"{len(scenes)} レンダ（YAML: {TMP_DIR}）===")
    return scenes


def timed_render(scene_yaml: Path, out_dir: Path, variant: str) -> tuple[dict, np.ndarray]:
    """段階別に時間を測りながらレンダし，(時間の dict, HDR 画像) を返す.

    render/pipeline.render_yaml と同じ手順を手でほどいたもの（段階ごとの時間を取るため）．
    mi.render は非同期に GPU へ投げられるので，dr.sync_thread() で完了を待ってから測る．
    """
    from sm_tias_lightsim.render.runner import ensure_variant
    from sm_tias_lightsim.scene.builder import build_scene

    times: dict[str, float] = {}
    t0 = time.perf_counter()
    config = load_scene_yaml(scene_yaml)
    render_cfg = config.render.model_copy(update={"variant": variant})
    config = config.model_copy(update={"render": render_cfg})
    t1 = time.perf_counter()
    times["t_load_yaml"] = t1 - t0

    ensure_variant(variant)
    import drjit as dr
    import mitsuba as mi

    scene_dict, tempdir = build_scene(config, project_root=scene_yaml.parent)
    t2 = time.perf_counter()
    times["t_build_scene"] = t2 - t1
    try:
        scene = mi.load_dict(scene_dict)
        dr.sync_thread()
        t3 = time.perf_counter()
        times["t_load_dict"] = t3 - t2
        image = mi.render(scene, spp=config.render.spp, seed=config.render.seed)
        dr.eval(image)
        dr.sync_thread()
        arr = np.array(image)
        t4 = time.perf_counter()
        times["t_render"] = t4 - t3
    finally:
        shutil.rmtree(tempdir, ignore_errors=True)
    from sm_tias_lightsim.render.output import save_luminance, save_png

    out_dir.mkdir(parents=True, exist_ok=True)
    save_png(arr, out_dir / "render.png")
    save_luminance(arr, out_dir / "luminance.npy")
    times["t_save"] = time.perf_counter() - t4
    return times, arr


def noise_metrics(luminance_npy: Path) -> dict[str, float]:
    """ノイズと均一性の指標を luminance.npy から計算する."""
    arr = np.load(luminance_npy).astype(np.float64)
    roi = extract_roi_xywh(arr, SIM_ROI)
    pixel_cov, structure, noise = decompose_cov(luminance_npy)
    centers, means, _ = radial_stats(luminance_npy, 20)
    full = calc_fullimage_min_max_ratio(roi)
    from scipy.ndimage import gaussian_filter

    smooth = gaussian_filter(roi, sigma=12.0)
    return {
        "raw_mean_roi": float(roi.mean()),
        "pixel_cov": pixel_cov,
        "noise_cov": noise,
        "structure_cov": structure,
        "ring_min_max": ring_min_max_ratio(centers, means),
        "fullimage_min_max": full["fullimage_min_max_ratio"],
        "fullimage_raw_min_max": float(smooth.min() / smooth.max()),
        "fullimage_noise_min": full["fullimage_noise_min"],
        "fullimage_noise_max": full["fullimage_noise_max"],
    }


def run_job(name: str, condition: str, spp: int, seed: int, scene_yaml: Path,
            variant: str, log) -> dict | None:
    """1 レンダを計測付きで実行し row.json を書く（既存は再利用）."""
    out_dir = OUTPUT_BASE / name
    row_json = out_dir / "row.json"
    if row_json.exists():
        log(f"--- {name}: 既存結果を再利用")
        return json.loads(row_json.read_text(encoding="utf-8"))
    log(f"--- {name}: レンダ開始")
    t_start = time.perf_counter()
    try:
        times, _ = timed_render(scene_yaml, out_dir, variant)
        t_eval0 = time.perf_counter()
        import io
        from contextlib import redirect_stdout

        with redirect_stdout(io.StringIO()):
            metrics = evaluate_image(out_dir / "luminance.npy", EVAL_CONFIG_YAML, out_dir,
                                     save_figures=False)
        times["t_evaluate"] = time.perf_counter() - t_eval0
    except Exception as exc:  # noqa: BLE001 - 1 条件の失敗でスイープを止めない
        log(f"  !! 失敗: {exc}")
        return None
    times["t_total"] = time.perf_counter() - t_start
    row: dict = {"name": name, "condition": condition, "spp": spp, "seed": seed,
                 "render_variant": variant}
    row.update({k: round(v, 4) for k, v in times.items()})
    row["gpu_fraction"] = round(times["t_render"] / times["t_total"], 4)
    row["t_render_per_kspp"] = round(times["t_render"] / (spp / 1024.0), 4)
    row.update(noise_metrics(out_dir / "luminance.npy"))
    row["radial_min_max_ratio"] = float(metrics.get("radial_min_max_ratio", float("nan")))
    row["cov"] = float(metrics["cov"])
    row_json.write_text(json.dumps(row, ensure_ascii=False, indent=1), encoding="utf-8")
    log(f"  -> 全体 {row['t_total']:.1f}s（render {row['t_render']:.1f}s = "
        f"{row['gpu_fraction'] * 100:.0f}%）noise={row['noise_cov']:.4f} "
        f"structure={row['structure_cov']:.4f} min={row['fullimage_min_max']:.4f}")
    return row


def write_summary() -> Path:
    """全 row.json から集約 CSV を書く."""
    rows = [json.loads(p.read_text(encoding="utf-8"))
            for p in sorted(OUTPUT_BASE.glob("*/row.json"))]
    order = {job_name(d, s, k): i for i, (d, s, k)
             in enumerate((d, s, k) for d, _ in CONDITIONS for s in SPPS for k in SEEDS)}
    rows.sort(key=lambda r: order.get(r["name"], len(order)))
    path = OUTPUT_BASE / "summary.csv"
    with path.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=SUMMARY_COLUMNS)
        writer.writeheader()
        for r in rows:
            writer.writerow({k: r.get(k, "") for k in SUMMARY_COLUMNS})
    return path


def main() -> int:
    """spp スイープを実行する."""
    parser = argparse.ArgumentParser(description="spp スイープ（計算時間・GPU 率・ノイズ）")
    parser.add_argument("--prepare-only", action="store_true", help="YAML の生成だけ行う")
    parser.add_argument("--spps", type=int, nargs="*", default=None, help="実行する spp を絞る")
    parser.add_argument("--conditions", nargs="*", default=None, help="実行する条件名を絞る")
    parser.add_argument("--set", choices=["default", "shapes6"], default="default",
                        help="条件セット（shapes6: θ=0 で形状 6 種類．出力先も別）")
    args = parser.parse_args()
    if args.set == "shapes6":
        global CONDITIONS, TMP_DIR, OUTPUT_BASE
        CONDITIONS = CONDITIONS_SHAPES6
        TMP_DIR = ROOT / "configs" / "_tmp_spp_noise_shapes6"
        OUTPUT_BASE = ROOT / "output" / "spp_noise_shapes6"

    log_path = ROOT / "output" / "_sweep_spp_noise.log"
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_file = log_path.open("a", encoding="utf-8")

    def log(message: str) -> None:
        print(message, flush=True)
        log_file.write(message + "\n")
        log_file.flush()

    scenes = prepare(log)
    if args.prepare_only:
        log_file.close()
        return 0
    template = yaml.safe_load(next(iter(scenes.values())).read_text(encoding="utf-8"))
    variant = resolve_render_variant(str(template["render"]["variant"]))
    log(f"=== variant={variant} ===")
    done = 0
    # spp の小さい順に全条件を回す（途中で止めても低 spp 側が揃う）
    for spp in SPPS:
        if args.spps and spp not in args.spps:
            continue
        for display, _ in CONDITIONS:
            if args.conditions and display not in args.conditions:
                continue
            for seed in SEEDS:
                name = job_name(display, spp, seed)
                row = run_job(name, display, spp, seed, scenes[name], variant, log)
                if row is not None:
                    done += 1
                    write_summary()
    log(f"=== 完了: {done} レンダ．集約 CSV: {write_summary()} ===")
    log_file.close()
    return 0


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