"""낮 맵 재구성 진행도 그림 — 지역 3곳의 주행가능영역 위에 학습된 구간(정적 PSNR 색)과 아직 안 된 낮 씬(회색) 궤적.

도로 커버리지는 recon.batch.coverage_order 와 같은 10 m 칸 기준 (낮 씬 자차 궤적이 지나간 칸 중 학습된 구간이 맡는 칸).
GPU 를 쓰지 않는다 — 학습이 도는 중에도 돌려도 된다.
실행: source scripts/local_env.sh && python scripts/map_coverage.py [출력.png]   (기본 outputs/coverage/map_coverage_<날짜>.png)
"""
import json
import sys
import time
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import font_manager
from matplotlib.collections import LineCollection
from matplotlib.patches import Polygon

from nusim.data.index import load_scene
from nusim.maps.expansion import MapExpansion
from nusim.paths import output_dir
from nusim.recon.batch import COVER_CELL_M, day_scenes, segment_dirs

ROOT = output_dir() / "gs_day"
OUT = Path(sys.argv[1]) if len(sys.argv) > 1 else output_dir() / "coverage" / f"map_coverage_{time.strftime('%Y%m%d')}.png"
LOCS = [("singapore-onenorth", "One-North"), ("singapore-queenstown", "Queenstown"),
        ("singapore-hollandvillage", "Holland Village")]
FONT = "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc"
if Path(FONT).exists():
    font_manager.fontManager.addfont(FONT)
    plt.rcParams["font.family"] = font_manager.FontProperties(fname=FONT).get_name()
PSNR_LO, PSNR_HI = 20.0, 28.0


def cells(xy: np.ndarray) -> set:
    return set(map(tuple, np.floor(xy / COVER_CELL_M).astype(int)))


def scene_track(name: str) -> tuple[np.ndarray, np.ndarray]:
    """정면 카메라 자차 위치 (n,2) 와 씬 시작부터의 초."""
    fr = load_scene(name).frames["CAM_FRONT"]
    xy = np.array([f.ego_translation[:2] for f in fr])
    t = (np.array([f.timestamp for f in fr]) - fr[0].timestamp) / 1e6
    return xy, t


def seg_psnr(d: Path) -> float | None:
    p = d / "eval.json"
    return json.loads(p.read_text()).get("psnr_static") if p.exists() else None


fig, axes = plt.subplots(1, 3, figsize=(24, 9.2), gridspec_kw={"width_ratios": [1.25, 1, 1]})
cmap = plt.get_cmap("viridis")
tot = {"scenes": 0, "done": 0, "segs": 0, "cells": 0, "cells_done": 0}
for ax, (loc, title) in zip(axes, LOCS):
    names = day_scenes(loc)
    all_xy, done_lines, done_c, todo_lines, done_cells, all_cells = [], [], [], [], set(), set()
    n_done, n_segs = 0, 0
    for s in names:
        xy, t = scene_track(s)
        all_xy.append(xy)
        all_cells |= cells(xy)
        segs = segment_dirs(ROOT, s)
        if not segs:
            todo_lines.append(xy)
            continue
        n_done += 1
        owned = np.zeros(len(t), bool)
        for info, d in segs:
            n_segs += 1
            o0 = info.get("own0") if info.get("own0") is not None else -1
            o1 = info.get("own1") if info.get("own1") is not None else 1e9
            m = (t >= o0) & (t <= o1)
            owned |= m
            if m.sum() >= 2:
                done_lines.append(xy[m])
                v = seg_psnr(d)
                done_c.append(cmap((np.clip(v, PSNR_LO, PSNR_HI) - PSNR_LO) / (PSNR_HI - PSNR_LO)) if v else (0.5, 0.5, 0.5, 1))
                done_cells |= cells(xy[m])
        if (~owned).sum() >= 2:  # 학습이 덜 끝난 씬의 나머지 구간
            todo_lines.append(xy[~owned])
    pts = np.concatenate(all_xy)
    lo, hi = pts.min(0) - 80, pts.max(0) + 80
    m = MapExpansion(loc)
    for rec in m._d["drivable_area"]:
        for tok in rec["polygon_tokens"]:
            rings = m.polygon_rings(tok)
            if not rings:
                continue
            r0 = np.asarray(rings[0])
            if r0[:, 0].max() < lo[0] or r0[:, 0].min() > hi[0] or r0[:, 1].max() < lo[1] or r0[:, 1].min() > hi[1]:
                continue
            ax.add_patch(Polygon(r0, closed=True, fc="#E9ECF1", ec="none", zorder=0))
            for h in rings[1:]:
                ax.add_patch(Polygon(np.asarray(h), closed=True, fc="white", ec="none", zorder=0))
    ax.add_collection(LineCollection(todo_lines, colors="#B9BEC7", linewidths=1.2, zorder=1))
    ax.add_collection(LineCollection(done_lines, colors=done_c, linewidths=3.2, zorder=2, capstyle="round"))
    ax.set_xlim(lo[0], hi[0]); ax.set_ylim(lo[1], hi[1]); ax.set_aspect("equal")
    ax.set_xticks([]); ax.set_yticks([])
    cov = len(done_cells & all_cells) / max(len(all_cells), 1)
    ax.set_title(f"{title}\n씬 {n_done}/{len(names)} · 구간 {n_segs}개 · 도로 커버리지 {cov:.0%}", fontsize=17)
    sb = 200.0
    ax.plot([lo[0] + 30, lo[0] + 30 + sb], [lo[1] + 30] * 2, color="k", lw=3)
    ax.text(lo[0] + 30 + sb / 2, lo[1] + 45, "200 m", ha="center", fontsize=12)
    for k, v in (("scenes", len(names)), ("done", n_done), ("segs", n_segs), ("cells", len(all_cells)),
                 ("cells_done", len(done_cells & all_cells))):
        tot[k] += v
    print(f"{loc}: 씬 {n_done}/{len(names)}, 구간 {n_segs}, 도로 칸 {len(done_cells & all_cells)}/{len(all_cells)} ({cov:.1%})")

sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(PSNR_LO, PSNR_HI))
cb = fig.colorbar(sm, ax=axes, orientation="horizontal", fraction=0.035, pad=0.03, aspect=60)
cb.set_label("학습된 구간의 정적 PSNR (dB) — 회색 선: 아직 학습 안 된 낮 씬 궤적 · 바탕: 주행가능영역", fontsize=13)
fig.suptitle(f"싱가포르 낮 맵 재구성 진행 — {time.strftime('%Y-%m-%d %H:%M')}  ·  씬 {tot['done']}/{tot['scenes']}  ·  "
             f"구간 {tot['segs']}개  ·  도로 커버리지 {tot['cells_done'] / tot['cells']:.0%} (10 m 칸 {tot['cells_done']}/{tot['cells']})",
             fontsize=20, fontweight="bold")
OUT.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(OUT, dpi=110, bbox_inches="tight")
print("saved", OUT)