"""面法線修正後の ρ × ω × α 再スイープ（TEST_12 後の条件引き直し）.

面法線バグ（TEST_12）修正前に実施した `sweep_reflectance.py` / `sweep_blend_bsdf.py`
の結論は，評価板のシェーディング法線アーティファクトが不均一性の 6 割を占めていた
状態で得たものであり無効．本スクリプトは同じ 3 軸を **修正後の条件で引き直し**，
実機（全灯）に最も近い組合せを探す．

スイープ軸:
    ρ  : diffuse reflectance（内壁 printer_white の拡散反射率）
    ω  : blendbsdf weight（roughconductor 成分の混合比．0 なら純 diffuse）
    α  : roughconductor の GGX 粗さ

固定条件（既存レンダと完全に同一に保つ）:
    spp=1024 / max_depth=16 / 512x512 / llvm_ad_rgb / path / orthographic
    → configs/_tmp_fn_full_base.yaml をテンプレートとし materials のみ書き換える

探索戦略（無人実行を想定）:
    ρ×ω×α の全数グリッドを，粗い ρ（COARSE_RHO）→細かい ρ（FINE_RHO）の 2 パスで流す．
    時間予算で打ち切られても ρ 軸が粗いなりに全域カバーされている状態を保つための順序．
    1 ジョブ完了ごとに summary.csv を書き直し，既にレンダ済みのジョブは再利用する
    （基準点 ρ=0.85/ω=0.05/α=0.1 は output/full_base_fix/ を再利用し再レンダしない）.

出力（1 ジョブ完了ごとに逐次追記するので途中終了しても結果は残る）:
    output/bsdf_grid_fix/<job>/luminance.npy, metrics.csv
    output/bsdf_grid_fix/summary.csv
    configs/_tmp_bsdf_grid/scene_<job>.yaml

実行方法:
    uv run python scripts/sweep_bsdf_grid_fix.py [--budget-hours 5.5]
"""

from __future__ import annotations

import argparse
import copy
import csv
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import yaml
from scipy.ndimage import gaussian_filter, median_filter
from skimage.transform import resize

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

from sm_tias_lightsim.evaluation.roi_presets import SIM_ROI, extract_roi_xywh  # noqa: E402
from sm_tias_lightsim.evaluation.shape_metrics import calc_ncc  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]

TEMPLATE_YAML = ROOT / "configs" / "_tmp_fn_full_base.yaml"
EVAL_CONFIG_YAML = ROOT / "configs" / "eval_baseline_dng_target.yaml"
TMP_DIR = ROOT / "configs" / "_tmp_bsdf_grid"
OUTPUT_BASE = ROOT / "output" / "bsdf_grid_fix"
SUMMARY_CSV = OUTPUT_BASE / "summary.csv"

REAL_FULL = ROOT / "lightsim" / "luminance_maps" / "full_MiniTIAS_QM_20260601_124500.npy"
REFERENCE_JOB_DIR = ROOT / "output" / "full_base_fix"

# 基準点（= output/full_base_fix/ を生成した条件）
REF_RHO, REF_OMEGA, REF_ALPHA = 0.85, 0.05, 0.1

# 全数グリッド．1 ジョブ約 200 秒（実測）なので 5.5h でおよそ 90 ジョブ入る．
# 粗パス（COARSE_RHO）を先に流し切ってから細パス（FINE_RHO）で ρ を埋める順序にし，
# 時間予算で打ち切られても ρ 軸が粗いなりに全域カバーされている状態を保つ．
COARSE_RHO = [0.75, 0.85, 0.90]
FINE_RHO = [0.78, 0.82, 0.88]
OMEGA_VALUES = [0.0, 0.05, 0.15, 0.30]
ALPHA_VALUES = [0.1, 0.3, 0.6]

# 実機（全灯）の指標: configs/eval_baseline_dng_target.yaml ヘッダの記録値
REAL_COV = 0.0695
REAL_CP = 1.0911
REAL_GRADIENT = 8.35
REAL_RADIAL = 0.860

SUMMARY_COLUMNS = [
    "job", "rho", "omega", "alpha", "stage", "elapsed_s",
    "cov", "center_periphery_ratio", "gradient_magnitude", "radial_min_max_ratio",
    "ncc_smooth", "peak_mean_smooth", "top_bot_smooth",
    "e_cov", "e_cp", "e_grad", "e_radial", "e_ncc", "score",
    "truncation_loss_pct",
]


def job_name(rho: float, omega: float, alpha: float) -> str:
    """ρ/ω/α からジョブ名（ディレクトリ名）を作る."""
    return f"r{rho:.2f}_w{omega:.2f}_a{alpha:.1f}"


def truncation_loss_pct(rho: float, omega: float) -> float:
    """max_depth=16 で取りこぼす相対エネルギー [%] を実効反射率から見積もる.

    blendbsdf の roughconductor 成分は material="none"（反射率 1.0 の完全反射体）なので
    実効的な壁の反射率は rho*(1-omega) + omega となり，rho 単体より高くなる．
    TEST_13 の check_truncation_bias と同じ rho_eff^(max_depth-1) を用いる．
    """
    rho_eff = rho * (1.0 - omega) + omega
    return round(100.0 * rho_eff ** 15, 2)


def build_scene_yaml(rho: float, omega: float, alpha: float, dest: Path) -> None:
    """テンプレートの printer_white だけを差し替えたシーン YAML を書き出す.

    ω=0 のときは blendbsdf を使わず純 diffuse にする（weight=0 の縮退を避ける）.
    他のセクション（geometry / lights / sensor / render）は一切変更しない.
    """
    config = yaml.safe_load(TEMPLATE_YAML.read_text(encoding="utf-8"))
    config = copy.deepcopy(config)
    if omega <= 0.0:
        material = {"type": "diffuse", "reflectance": [rho, rho, rho]}
    else:
        material = {
            "type": "blendbsdf",
            "weight": omega,
            "components": [
                {"type": "diffuse", "reflectance": [rho, rho, rho]},
                {"type": "roughconductor", "alpha": alpha, "distribution": "ggx"},
            ],
        }
    config["materials"]["printer_white"] = material
    config["name"] = f"bsdf_grid_{job_name(rho, omega, alpha)}"
    # メッシュの相対パスは YAML 自身のディレクトリ基準で解決されるため，
    # テンプレート（configs/ 直下）から別ディレクトリへ移す時点で絶対パス化する．
    for mesh in config["geometry"]["meshes"]:
        mesh["path"] = str((TEMPLATE_YAML.parent / mesh["path"]).resolve())
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_text(yaml.safe_dump(config, allow_unicode=True, sort_keys=False), encoding="utf-8")


def _smooth(arr: np.ndarray, is_sim: bool) -> np.ndarray:
    """build_full_vs_singles_panel.py と同一の平滑化（median5 → gauss σ4）."""
    if arr.ndim == 3:
        arr = 0.2126 * arr[..., 0] + 0.7152 * arr[..., 1] + 0.0722 * arr[..., 2]
    arr = arr.astype(np.float64)
    if is_sim:
        arr = extract_roi_xywh(arr, SIM_ROI)
    return gaussian_filter(median_filter(arr, size=5), sigma=4)


def _shape_stats(sim_path: Path, real_smooth: np.ndarray) -> dict[str, float]:
    """平滑化後の NCC 対実機・peak/mean・top/bot を返す（MC ノイズ支配を避けるため平滑後）."""
    sim = _smooth(np.load(sim_path), True)
    resized = resize(sim, real_smooth.shape, anti_aliasing=True)
    half = sim.shape[0] // 2
    return {
        "ncc_smooth": float(calc_ncc(resized, real_smooth)),
        "peak_mean_smooth": float(sim.max() / sim.mean()),
        "top_bot_smooth": float(sim[:half].mean() / sim[half:].mean()),
    }


def _score(metrics: dict[str, float]) -> dict[str, float]:
    """実機との相対乖離から総合スコアを組む（小さいほど実機に近い）.

    形（NCC）と振幅系 4 指標（CoV / C-P / 勾配 / 動径 min-max）を半々で重み付けする．
    C/P と動径 min/max は「1 からの超過分」で相対化する（1 近傍で比が鈍るため）.
    """
    e_cov = abs(metrics["cov"] - REAL_COV) / REAL_COV
    e_cp = abs((metrics["center_periphery_ratio"] - 1.0) - (REAL_CP - 1.0)) / (REAL_CP - 1.0)
    e_grad = abs(metrics["gradient_magnitude"] - REAL_GRADIENT) / REAL_GRADIENT
    e_radial = (abs((1.0 - metrics["radial_min_max_ratio"]) - (1.0 - REAL_RADIAL))
                / (1.0 - REAL_RADIAL))
    e_ncc = 1.0 - metrics["ncc_smooth"]
    amplitude = (e_cov + e_cp + e_grad + e_radial) / 4.0
    return {
        "e_cov": e_cov, "e_cp": e_cp, "e_grad": e_grad, "e_radial": e_radial,
        "e_ncc": e_ncc, "score": 0.5 * amplitude + 0.5 * e_ncc,
    }


def read_metrics_csv(path: Path) -> dict[str, float]:
    """sm-evaluate が出す metrics.csv の 1 行目を辞書で返す."""
    with path.open(encoding="utf-8") as fh:
        row = next(csv.DictReader(fh))
    out = {}
    for key in ("cov", "center_periphery_ratio", "gradient_magnitude", "radial_min_max_ratio"):
        out[key] = float(row[key])
    return out


def run_job(rho: float, omega: float, alpha: float, stage: str,
            real_smooth: np.ndarray, log) -> dict[str, float] | None:
    """1 条件をレンダ→評価→指標算出まで実行して集約行を返す（失敗時 None）.

    既に metrics.csv と luminance.npy が揃っているジョブは再レンダせず再利用する
    （中断・再開や，グリッドを広げて流し直すときに無駄な計算をしないため）.
    """
    name = job_name(rho, omega, alpha)
    out_dir = OUTPUT_BASE / name
    scene_yaml = TMP_DIR / f"scene_{name}.yaml"
    if (out_dir / "metrics.csv").exists() and (out_dir / "luminance.npy").exists():
        row = {"job": name, "rho": rho, "omega": omega, "alpha": alpha,
               "stage": stage + "(再利用)", "elapsed_s": 0.0}
        row.update(read_metrics_csv(out_dir / "metrics.csv"))
        row.update(_shape_stats(out_dir / "luminance.npy", real_smooth))
        row.update(_score(row))
        row["truncation_loss_pct"] = truncation_loss_pct(rho, omega)
        log(f"--- {stage}: {name} → 既存結果を再利用 (score={row['score']:.4f})")
        return row
    log(f"--- {stage}: {name} (rho={rho}, omega={omega}, alpha={alpha}) ---")
    build_scene_yaml(rho, omega, alpha, scene_yaml)
    out_dir.mkdir(parents=True, exist_ok=True)
    start = time.time()
    for cmd in (
        ["uv", "run", "sm-render", "--scene", str(scene_yaml), "--output", str(out_dir)],
        ["uv", "run", "sm-evaluate", "--image", str(out_dir / "luminance.npy"),
         "--config", str(EVAL_CONFIG_YAML), "--output", str(out_dir)],
    ):
        result = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True)
        if result.returncode != 0:
            log(f"  !! 失敗 ({cmd[2]}): {result.stderr[-600:]}")
            return None
    elapsed = time.time() - start
    row = {"job": name, "rho": rho, "omega": omega, "alpha": alpha,
           "stage": stage, "elapsed_s": round(elapsed, 1)}
    row.update(read_metrics_csv(out_dir / "metrics.csv"))
    row.update(_shape_stats(out_dir / "luminance.npy", real_smooth))
    row.update(_score(row))
    row["truncation_loss_pct"] = truncation_loss_pct(rho, omega)
    log(f"  -> {elapsed:.0f}s  CoV={row['cov']:.4f} (実機 {REAL_COV})  "
        f"NCC={row['ncc_smooth']:.4f}  勾配={row['gradient_magnitude']:.2f}%  "
        f"score={row['score']:.4f}")
    return row


def append_summary(rows: list[dict[str, float]]) -> None:
    """集約 CSV を毎回まるごと書き直す（途中終了でも最新状態が残る）."""
    OUTPUT_BASE.mkdir(parents=True, exist_ok=True)
    ordered = sorted(rows, key=lambda r: r["score"])
    with SUMMARY_CSV.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=SUMMARY_COLUMNS)
        writer.writeheader()
        for row in ordered:
            writer.writerow({k: row.get(k, "") for k in SUMMARY_COLUMNS})


def reference_row(real_smooth: np.ndarray, log) -> dict[str, float] | None:
    """基準点は既存の output/full_base_fix/ を再利用する（再レンダしない）."""
    metrics_csv = REFERENCE_JOB_DIR / "metrics.csv"
    luminance = REFERENCE_JOB_DIR / "luminance.npy"
    if not (metrics_csv.exists() and luminance.exists()):
        log("  基準点 output/full_base_fix/ が無いため再レンダします")
        return run_job(REF_RHO, REF_OMEGA, REF_ALPHA, "stage1", real_smooth, log)
    row = {"job": job_name(REF_RHO, REF_OMEGA, REF_ALPHA), "rho": REF_RHO,
           "omega": REF_OMEGA, "alpha": REF_ALPHA, "stage": "reference", "elapsed_s": 0.0}
    row.update(read_metrics_csv(metrics_csv))
    row.update(_shape_stats(luminance, real_smooth))
    row.update(_score(row))
    row["truncation_loss_pct"] = truncation_loss_pct(REF_RHO, REF_OMEGA)
    log(f"  基準点（既存 full_base_fix 再利用）: CoV={row['cov']:.4f} "
        f"NCC={row['ncc_smooth']:.4f} score={row['score']:.4f}")
    return row


def build_grid(rho_values: list[float]) -> list[tuple[float, float, float]]:
    """指定した ρ 値集合について ω×α の全組合せを返す.

    ω=0（純 diffuse）は α に依存しないので α は 1 値だけ回す.
    """
    combos = []
    for rho in rho_values:
        for omega in OMEGA_VALUES:
            for alpha in (ALPHA_VALUES[:1] if omega <= 0.0 else ALPHA_VALUES):
                combos.append((rho, omega, alpha))
    return combos


def main() -> int:
    """段階スイープを時間予算内で実行するメインエントリポイント."""
    parser = argparse.ArgumentParser(description="面法線修正後の ρ×ω×α 再スイープ")
    parser.add_argument("--budget-hours", type=float, default=5.5,
                        help="全体の時間予算（時間）．超過が見込まれる時点で打ち切る")
    args = parser.parse_args()

    deadline = time.time() + args.budget_hours * 3600.0
    log_path = ROOT / "output" / "_sweep_bsdf_grid_fix.log"
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_file = log_path.open("w", encoding="utf-8")

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

    if not REAL_FULL.exists():
        log(f"エラー: 実機マップ {REAL_FULL} が見つかりません")
        return 1
    real_smooth = _smooth(np.load(REAL_FULL), False)

    log(f"=== ρ×ω×α 再スイープ（面法線修正後）予算 {args.budget_hours}h ===")
    log(f"  固定条件: spp=1024 / max_depth=16 / 512^2（{TEMPLATE_YAML.name} 準拠）")
    log(f"  基準点: rho={REF_RHO}, omega={REF_OMEGA}, alpha={REF_ALPHA}")

    rows: list[dict[str, float]] = []
    ref = reference_row(real_smooth, log)
    if ref is not None:
        rows.append(ref)
        append_summary(rows)

    durations: list[float] = []

    def budget_ok() -> bool:
        """残り時間が想定 1 ジョブ分を確保できるか（実測の最大所要で判定）."""
        expected = max(durations) if durations else 900.0
        return time.time() + expected < deadline

    done = {(r["rho"], r["omega"], r["alpha"]) for r in rows}
    passes = (("粗パス", build_grid(COARSE_RHO)), ("細パス", build_grid(FINE_RHO)))
    for pass_name, combos in passes:
        pending = [c for c in combos if c not in done]
        log(f"\n=== {pass_name}: {len(pending)} ジョブ（ρ×ω×α 全組合せ） ===")
        for rho, omega, alpha in pending:
            if not budget_ok():
                log(f"  時間予算のため {pass_name} を打ち切ります")
                break
            row = run_job(rho, omega, alpha, pass_name, real_smooth, log)
            if row is not None:
                rows.append(row)
                done.add((rho, omega, alpha))
                if row["elapsed_s"] > 0:
                    durations.append(row["elapsed_s"])
                append_summary(rows)
        else:
            continue
        break

    log(f"\n=== 完了: {len(rows)} 条件 ===")
    log(f"  実機: CoV={REAL_COV:.4f}, C/P={REAL_CP:.4f}, "
        f"勾配={REAL_GRADIENT:.2f}%, 動径 min/max={REAL_RADIAL:.3f}")
    log("  スコア上位 5:")
    for row in sorted(rows, key=lambda r: r["score"])[:5]:
        log(f"    {row['job']:22s} score={row['score']:.4f}  CoV={row['cov']:.4f}  "
            f"NCC={row['ncc_smooth']:.4f}  勾配={row['gradient_magnitude']:.2f}%")
    log(f"  集約 CSV: {SUMMARY_CSV}")
    log_file.close()
    return 0


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