"""壁の反射特性（ω・α）と LED 配光（spot）のスイープ: 単灯の明部のずれを実機に近づける（TEST_20）.

実機と採用材質（ρ=0.85・ω=0.05・α=0.1）のシミュを全灯・単灯 3 灯で突き合わせると，
同心円方向（お椀）は合うが，単灯で明部が点灯基板の反対側へ寄る量がシミュは実機の半分以下
だった（compare_real_vs_sim_shapes.py）．原因の候補は (a) 壁の反射が拡散しすぎ（光沢成分が
一次照射を反対側の壁に集中させる），(b) LED の配光がランバートより狭い，の 2 つ．
本スイープは基準シーン（scene_baseline.yaml）から 4 構成（全灯・下・左・右）を作り，

    フェーズ A: 壁 = blendbsdf(diffuse ρ, roughconductor α) の重み ω と α を振る
    フェーズ B: 壁 = 採用材質のまま，LED を spot（半値全幅 FWHM を振る）にする

をサーバでレンダする（評価は手元で build_shape_fit_comparison.py）．

実行方法（サーバ）:
    uv run sm-remote run --name wall --gpu 0 --threads 16 -- \
        python scripts/sweep_wall_bsdf_spot.py --phase all --workers 2
"""

from __future__ import annotations

import argparse
import copy
import json
import multiprocessing
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import yaml

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

from sm_tias_lightsim.render.pipeline import render_yaml  # noqa: E402
from sm_tias_lightsim.runtime.overrides import resolve_render_variant  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
BASELINE_YAML = ROOT / "configs" / "scene_baseline.yaml"
TMP_DIR = ROOT / "configs" / "_tmp_wall_bsdf_spot"
OUTPUT_BASE = ROOT / "output" / "wall_bsdf_spot"

CONFIGS = ("full", "bottom", "left", "right")
BOARD_PREFIX = {"bottom": "led_bottom", "left": "led_left", "right": "led_right"}

# フェーズ A: 壁の反射特性
RHOS = [0.85, 0.81]
OMEGAS = [0.0, 0.05, 0.15, 0.30, 0.50, 0.70]
ALPHAS = [0.02, 0.05, 0.1, 0.3, 0.6]
SPP_A = 2048
# フェーズ B: LED 配光（半値全幅 FWHM）．採用材質で
SPOT_FWHMS = [30.0, 45.0, 60.0, 90.0, 120.0]
SPOT_MATERIAL = (0.85, 0.05, 0.1)
SPP_B = 4096
# 面光源（2.8×3.5 mm の矩形・radiance）と同じ全光束になる spot の軸上強度を作る
AREA_MM2 = 2.8 * 3.5


@dataclass(frozen=True)
class Condition:
    """1 レンダ条件."""

    phase: str
    config: str
    rho: float
    omega: float
    alpha: float
    fwhm: float | None = None

    @property
    def name(self) -> str:
        mat = f"r{self.rho:.2f}_w{self.omega:.2f}_a{self.alpha:.2f}"
        spot = f"_spot{self.fwhm:g}" if self.fwhm is not None else ""
        return f"{self.phase}_{self.config}_{mat}{spot}"


def conditions(phase: str) -> list[Condition]:
    """条件の一覧（フェーズ A → B の順）."""
    out: list[Condition] = []
    if phase in ("a", "all"):
        for rho in RHOS:
            for omega in OMEGAS:
                for alpha in ALPHAS if omega > 0 else [ALPHAS[2]]:
                    for cfg in CONFIGS:
                        out.append(Condition("a", cfg, rho, omega, alpha))
    if phase in ("b", "all"):
        rho, omega, alpha = SPOT_MATERIAL
        for fwhm in SPOT_FWHMS:
            for cfg in CONFIGS:
                out.append(Condition("b", cfg, rho, omega, alpha, fwhm))
    return out


def material_dict(rho: float, omega: float, alpha: float) -> dict:
    """内壁材質の YAML 断片（ω=0 なら純 diffuse）."""
    if omega <= 0.0:
        return {"type": "diffuse", "reflectance": [rho, rho, rho]}
    return {
        "type": "blendbsdf", "weight": omega,
        "components": [
            {"type": "diffuse", "reflectance": [rho, rho, rho]},
            {"type": "roughconductor", "alpha": alpha, "distribution": "ggx"},
        ],
    }


def spot_solid_angle(cutoff_deg: float, beam_deg: float) -> float:
    """Mitsuba spot（beam_width まで 1，cutoff まで smoothstep で 0）の実効立体角 [sr]."""
    theta = np.linspace(0.0, np.radians(cutoff_deg), 4001)
    width = max(np.radians(cutoff_deg - beam_deg), 1e-9)
    t = np.clip((np.radians(cutoff_deg) - theta) / width, 0, 1)
    falloff = np.where(theta <= np.radians(beam_deg), 1.0, t * t * (3 - 2 * t))
    return float(2 * np.pi * np.trapezoid(falloff * np.sin(theta), theta))


def to_spot(light: dict, fwhm: float) -> dict:
    """面光源 1 灯を同じ位置・向き・全光束の spot に置き換える（FWHM = cutoff + beam）."""
    cutoff = fwhm * 2.0 / 3.0
    beam = fwhm / 3.0
    radiance = light["radiance"]
    radiance = radiance[0] if isinstance(radiance, list) else radiance
    flux = np.pi * radiance * AREA_MM2 * 1e-6  # ランバート面光源の全光束 [W]
    intensity = flux / spot_solid_angle(cutoff, beam)
    return {
        "id": light["id"], "type": "spot", "position": light["center"],
        "direction": light["direction"], "cutoff_angle": round(cutoff, 4),
        "beam_width": round(beam, 4), "intensity": float(intensity),
    }


def build_yaml(cond: Condition, base: dict) -> Path:
    """条件の YAML を書き出す（メッシュは configs/meshes を相対参照）."""
    cfg = copy.deepcopy(base)
    for mesh in cfg["geometry"]["meshes"]:
        if not Path(mesh["path"]).is_absolute():
            mesh["path"] = "../" + mesh["path"]
    cfg["materials"]["printer_white"] = material_dict(cond.rho, cond.omega, cond.alpha)
    lights = cfg["lights"]
    if cond.config != "full":
        lights = [lt for lt in lights if lt["id"].startswith(BOARD_PREFIX[cond.config])]
    if cond.fwhm is not None:
        lights = [to_spot(lt, cond.fwhm) for lt in lights]
    cfg["lights"] = lights
    cfg["render"]["spp"] = SPP_B if cond.phase == "b" else SPP_A
    cfg["name"] = cond.name
    spot_note = f" / spot FWHM={cond.fwhm:g}°" if cond.fwhm else ""
    cfg["description"] = (f"TEST_20 {cond.phase.upper()}: {cond.config} / ρ={cond.rho} "
                          f"ω={cond.omega} α={cond.alpha}{spot_note}")
    TMP_DIR.mkdir(parents=True, exist_ok=True)
    dest = TMP_DIR / f"scene_{cond.name}.yaml"
    dest.write_text(yaml.safe_dump(cfg, allow_unicode=True, sort_keys=False), encoding="utf-8")
    return dest


def run_job(cond: Condition, scene_yaml: Path, variant: str) -> tuple[str, float, str]:
    """1 条件をレンダする（既存は再利用）．戻り値は (名前, 秒, 状態)."""
    out_dir = OUTPUT_BASE / cond.name
    if (out_dir / "row.json").exists():
        return cond.name, 0.0, "reuse"
    start = time.time()
    try:
        render_yaml(scene_yaml, out_dir, variant)
    except Exception as exc:  # noqa: BLE001 - 1 条件の失敗でスイープを止めない
        return cond.name, time.time() - start, f"fail: {exc}"
    row = {"name": cond.name, "phase": cond.phase, "config": cond.config, "rho": cond.rho,
           "omega": cond.omega, "alpha": cond.alpha, "fwhm": cond.fwhm,
           "spp": SPP_B if cond.phase == "b" else SPP_A, "variant": variant,
           "elapsed_s": round(time.time() - start, 1)}
    (out_dir / "row.json").write_text(json.dumps(row, ensure_ascii=False, indent=1),
                                      encoding="utf-8")
    return cond.name, time.time() - start, "ok"


def main() -> int:
    """スイープを実行する."""
    parser = argparse.ArgumentParser(description="壁の反射特性と spot 配光のスイープ")
    parser.add_argument("--phase", choices=("a", "b", "all"), default="all")
    parser.add_argument("--workers", type=int, default=1)
    parser.add_argument("--prepare-only", action="store_true")
    parser.add_argument("--list", action="store_true")
    args = parser.parse_args()
    conds = conditions(args.phase)
    if args.list:
        for c in conds:
            print(c.name)
        print(f"{len(conds)} 条件")
        return 0
    base = yaml.safe_load(BASELINE_YAML.read_text(encoding="utf-8"))
    scenes = {c.name: build_yaml(c, base) for c in conds}
    log_path = ROOT / "output" / "_sweep_wall_bsdf_spot.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()

    log(f"=== 壁反射特性 / spot スイープ: {len(conds)} 条件（phase={args.phase}）===")
    if args.prepare_only:
        return 0
    variant = resolve_render_variant(str(base["render"]["variant"]))
    done = 0
    if args.workers > 1:
        context = multiprocessing.get_context("spawn")
        with ProcessPoolExecutor(max_workers=args.workers, mp_context=context) as pool:
            futures = {pool.submit(run_job, c, scenes[c.name], variant): c for c in conds}
            for fut in as_completed(futures):
                name, sec, status = fut.result()
                log(f"  {name}: {status} ({sec:.0f}s)")
                done += status == "ok"
    else:
        for c in conds:
            name, sec, status = run_job(c, scenes[c.name], variant)
            log(f"  {name}: {status} ({sec:.0f}s)")
            done += status == "ok"
    log(f"=== 完了: 新規 {done} / {len(conds)} 条件 ===")
    log_file.close()
    return 0


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