"""발표 아키텍처 그림용 패널 — 실제 데이터에서 만든다 (scene-0402, 구간 1 모델).

출력: <out>/*.png. 그 다음 compose.py 가 논문식 도식 두 장으로 합친다.
GPU 를 쓰지 않는다 (가우시안 투영 · 깊이는 numpy 로 근사).
"""

import json
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
from matplotlib.patches import Ellipse, Polygon, Rectangle  # noqa: E402
from PIL import Image  # noqa: E402

from nusim.data.index import load_scene  # noqa: E402
from nusim.data.rig import box_corners, camera_pose, quat_to_matrix  # noqa: E402
from nusim.io.ply import read_origin, read_ply  # noqa: E402
from nusim.paths import dataset, workspace_root  # noqa: E402
from nusim.render.export import rig_frames  # noqa: E402

OUT = Path(sys.argv[1])
OUT.mkdir(parents=True, exist_ok=True)
ROOT = workspace_root()
SC = "scene-0402"
FRAME = 90  # rig 프레임 (CAM_FRONT 약 7.5 초 — 구간 1 이 맡는다)
SEG = ROOT / "outputs/gs_day" / SC / "seg1"
SCRATCH = Path(sys.argv[2])  # 앞서 만든 렌더 (f6/, pts0402.npy)
BG = "#14171c"
plt.rcParams["font.family"] = "NanumBarunGothic"
plt.rcParams["axes.unicode_minus"] = False

ds = dataset()
scene = load_scene(SC)
fr = {cam: f for k, cam, f in rig_frames(scene) if k == FRAME}
pose = camera_pose(fr["CAM_FRONT"], scene.calib["CAM_FRONT"], "CAM_FRONT")
W, H = 800, 450
K = pose.K.copy()
K[:2] *= W / pose.width
w2c = pose.w2c


def save(fig, name, dpi=160):
    fig.savefig(OUT / name, dpi=dpi, facecolor=fig.get_facecolor(), bbox_inches="tight", pad_inches=0)
    plt.close(fig)


def img_axes(w=8.0, h=4.5, bg=BG):
    fig = plt.figure(figsize=(w, h), facecolor=bg)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_facecolor(bg)
    ax.set_xlim(0, W)
    ax.set_ylim(H, 0)
    ax.axis("off")
    return fig, ax


def project(pg):
    pc = pg @ w2c[:3, :3].T + w2c[:3, 3]
    z = pc[:, 2]
    uv = pc[:, :2] / np.maximum(z[:, None], 1e-6) @ np.diag([K[0, 0], K[1, 1]]) + K[:2, 2]
    return uv, z, pc


# ---------------------------------------------------------------- 입력 6캠 (실제 사진)
lay = [["CAM_FRONT_LEFT", "CAM_FRONT", "CAM_FRONT_RIGHT"], ["CAM_BACK_RIGHT", "CAM_BACK", "CAM_BACK_LEFT"]]
for cam in sum(lay, []):
    Image.open(ds.root / fr[cam].filename).convert("RGB").resize((480, 270)).save(OUT / f"gt_{cam}.jpg", quality=92)
gt = Image.open(ds.root / fr["CAM_FRONT"].filename).convert("RGB").resize((W, H))
gt.save(OUT / "gt_front.jpg", quality=92)

# ---------------------------------------------------------------- 가우시안 (구간 1 모델)
g = read_ply(SEG / "point_cloud.ply")
mg = g.means + g.origin
op = 1 / (1 + np.exp(-g.opacities))
col = np.clip(0.5 + 0.28209479 * g.sh0[:, 0, :], 0, 1)
uv, z, pc = project(mg)
vis = (z > 1.0) & (z < 60) & (uv[:, 0] > -50) & (uv[:, 0] < W + 50) & (uv[:, 1] > -50) & (uv[:, 1] < H + 50) & (op > 0.3)
idx = np.flatnonzero(vis)
# 2D 공분산 Σ' = J W Σ Wᵀ Jᵀ
S = np.exp(g.scales[idx])
R = np.stack([quat_to_matrix(tuple(q)) for q in g.quats[idx] / np.linalg.norm(g.quats[idx], axis=1, keepdims=True)])
Sig = R @ (S[:, :, None] ** 2 * np.transpose(R, (0, 2, 1)))
Wr = w2c[:3, :3]
Sc = Wr @ Sig @ Wr.T
x, y, zz = pc[idx, 0], pc[idx, 1], pc[idx, 2]
J = np.zeros((len(idx), 2, 3))
J[:, 0, 0] = K[0, 0] / zz
J[:, 0, 2] = -K[0, 0] * x / zz ** 2
J[:, 1, 1] = K[1, 1] / zz
J[:, 1, 2] = -K[1, 1] * y / zz ** 2
S2 = J @ Sc @ np.transpose(J, (0, 2, 1))
ev, evec = np.linalg.eigh(S2)
ax_len = 2.0 * np.sqrt(np.maximum(ev, 1e-6))  # 2σ
angle = np.degrees(np.arctan2(evec[:, 1, 1], evec[:, 0, 1]))
order = np.argsort(-zz)  # 먼 것부터 (블렌딩 순서)


def splat_image(name, color_mode="rgb", n_max=9000, grid=False, bg=BG):
    sel = order
    big = ax_len[sel].max(1) > 2.5
    sel = sel[big]
    if len(sel) > n_max:
        sel = sel[np.linspace(0, len(sel) - 1, n_max).astype(int)]
    fig, ax = img_axes(bg=bg)
    rank = np.empty(len(sel))
    rank[np.argsort(zz[sel])] = np.linspace(0, 1, len(sel))
    for j, i in enumerate(sel):
        a, b = ax_len[i, 1], ax_len[i, 0]
        c = col[idx[i]] if color_mode == "rgb" else plt.cm.turbo(1 - rank[j])[:3]
        ax.add_patch(Ellipse(uv[idx[i]], 2 * a, 2 * b, angle=angle[i], facecolor=c, edgecolor="none", alpha=min(0.85, op[idx[i]])))
    if grid:
        for gx in range(0, W, 16 * 2):
            ax.plot([gx, gx], [0, H], color="#ffffff", lw=0.25, alpha=0.35)
        for gy in range(0, H, 16 * 2):
            ax.plot([0, W], [gy, gy], color="#ffffff", lw=0.25, alpha=0.35)
    save(fig, name)


splat_image("splats.png")
splat_image("splats_sorted.png", color_mode="depth")
splat_image("splats_tiles.png", grid=True)

def chase_view(pg, colors, name, s=1.2, cmap=None, back=14.0, up=9.0, ahead=18.0):
    """자차 뒤 위쪽 가상 카메라에서 본 점구름 (먼 점부터 그린다)."""
    fwd = w2c[2, :3]  # 카메라 광축 (전역)
    fwd = np.array([fwd[0], fwd[1], 0.0]) / np.linalg.norm(fwd[:2])
    eye = ego + np.array([0, 0, 1.6]) - fwd * back + np.array([0, 0, up])
    tgt = ego + fwd * ahead
    zc = (tgt - eye) / np.linalg.norm(tgt - eye)
    xc = np.cross(zc, [0, 0, 1.0])
    xc /= np.linalg.norm(xc)
    yc = np.cross(zc, xc)
    Rv = np.stack([xc, yc, zc])
    pc_ = (pg - eye) @ Rv.T
    ok = pc_[:, 2] > 1.0
    f = 520.0
    u = f * pc_[ok, 0] / pc_[ok, 2] + W / 2
    v = f * pc_[ok, 1] / pc_[ok, 2] + H / 2
    o = np.argsort(-pc_[ok, 2])
    fig, ax = img_axes()
    cc = colors[ok][o] if cmap is None else colors[ok][o]
    ax.scatter(u[o], v[o], c=cc, cmap=cmap, vmin=-1 if cmap else None, vmax=8 if cmap else None,
               s=s * np.clip(12 / pc_[ok, 2][o], 0.15, 3), linewidths=0)
    # 자차 위치 표시
    pe = (ego - eye) @ Rv.T
    ax.plot(f * pe[0] / pe[2] + W / 2, f * pe[1] / pe[2] + H / 2, marker="^", color="#ED7D31", ms=12)
    save(fig, name)


# 가우시안 3D (위에서 비스듬히 본 점구름, 실제 색)
ego = np.asarray(fr["CAM_FRONT"].ego_translation)
near = np.flatnonzero((np.linalg.norm(mg[:, :2] - ego[:2], axis=1) < 60) & (op > 0.35))
near = near[np.random.default_rng(0).choice(len(near), min(400000, len(near)), replace=False)]
def chase_splats(ids, name, n_max=14000, back=14.0, up=9.0, ahead=18.0):
    """가상 추격 카메라에서 본 3D 가우시안 (투영한 타원 — 3DGS 의 '얼룩' 모양)."""
    fwd = w2c[2, :3]
    fwd = np.array([fwd[0], fwd[1], 0.0]) / np.linalg.norm(fwd[:2])
    eye = ego + np.array([0, 0, 1.6]) - fwd * back + np.array([0, 0, up])
    tgt = ego + fwd * ahead
    zc = (tgt - eye) / np.linalg.norm(tgt - eye)
    xc = np.cross(zc, [0, 0, 1.0]); xc /= np.linalg.norm(xc)
    yc = np.cross(zc, xc)
    Rv = np.stack([xc, yc, zc])
    f = 520.0
    pcv = (mg[ids] - eye) @ Rv.T
    ok = (pcv[:, 2] > 2.0)
    ids, pcv = ids[ok], pcv[ok]
    Sx = np.exp(g.scales[ids])
    qn_ = g.quats[ids] / np.linalg.norm(g.quats[ids], axis=1, keepdims=True)
    Rg = np.stack([quat_to_matrix(tuple(q)) for q in qn_])
    Sig_ = Rv @ (Rg @ (Sx[:, :, None] ** 2 * np.transpose(Rg, (0, 2, 1)))) @ Rv.T
    x_, y_, z_ = pcv[:, 0], pcv[:, 1], pcv[:, 2]
    J_ = np.zeros((len(ids), 2, 3))
    J_[:, 0, 0] = f / z_; J_[:, 0, 2] = -f * x_ / z_ ** 2
    J_[:, 1, 1] = f / z_; J_[:, 1, 2] = -f * y_ / z_ ** 2
    S2_ = J_ @ Sig_ @ np.transpose(J_, (0, 2, 1))
    ev_, evec_ = np.linalg.eigh(S2_)
    al = 2.0 * np.sqrt(np.maximum(ev_, 1e-6))
    ang = np.degrees(np.arctan2(evec_[:, 1, 1], evec_[:, 0, 1]))
    u = f * x_ / z_ + W / 2
    v = f * y_ / z_ + H / 2
    keep = (u > -40) & (u < W + 40) & (v > -40) & (v < H + 40) & (al.max(1) > 1.0) & (al.max(1) < 120)
    sel_ = np.flatnonzero(keep)
    sel_ = sel_[np.argsort(-al[sel_].max(1))][:n_max]
    sel_ = sel_[np.argsort(-z_[sel_])]
    fig, ax = img_axes()
    for i in sel_:
        ax.add_patch(Ellipse((u[i], v[i]), 2 * al[i, 1], 2 * al[i, 0], angle=ang[i], facecolor=col[ids[i]],
                             edgecolor="none", alpha=min(0.9, op[ids[i]])))
    pe = (ego - eye) @ Rv.T
    ax.plot(f * pe[0] / pe[2] + W / 2, f * pe[1] / pe[2] + H / 2, marker="^", color="#ED7D31", ms=12)
    save(fig, name)


chase_splats(near, "gaussians_3d.png")

# ---------------------------------------------------------------- 라이다 초기화 (정적 라이다 누적, 높이 색)
pts = np.load(SCRATCH / "pts0402.npy").astype(np.float64) + read_origin(ROOT / f"outputs/gs_day/{SC}/seg0/point_cloud.ply")
sel = np.flatnonzero(np.linalg.norm(pts[:, :2] - ego[:2], axis=1) < 60)
sel = sel[np.random.default_rng(1).choice(len(sel), min(500000, len(sel)), replace=False)]
chase_view(pts[sel], pts[sel, 2] - ego[2], "lidar_init.png", s=1.0, cmap="turbo")

# 라이다 깊이 지도 (그 카메라로 투영, 깊이 감독)
uvl, zl, _ = project(pts)
m = (zl > 0.5) & (zl < 70) & (uvl[:, 0] >= 0) & (uvl[:, 0] < W) & (uvl[:, 1] >= 0) & (uvl[:, 1] < H)
fig, ax = img_axes()
o = np.argsort(-zl[m])
ax.scatter(uvl[m][o, 0], uvl[m][o, 1], c=np.clip(zl[m][o], 0, 50), cmap="turbo_r", s=0.8, linewidths=0)
save(fig, "lidar_depth.png")

# 렌더 깊이 (가우시안 중심 z-버퍼 근사, 흐리게)
dimg = np.full((H // 2, W // 2), np.inf)
u2 = (uv[idx, 0] / 2).astype(int)
v2 = (uv[idx, 1] / 2).astype(int)
ok = (u2 >= 0) & (u2 < W // 2) & (v2 >= 0) & (v2 < H // 2)
for r_ in (0, 1, 2):
    for du in range(-r_, r_ + 1):
        for dv in range(-r_, r_ + 1):
            uu, vv = np.clip(u2 + du, 0, W // 2 - 1), np.clip(v2 + dv, 0, H // 2 - 1)
            np.minimum.at(dimg, (vv[ok], uu[ok]), zz[ok])
from scipy import ndimage  # noqa: E402

miss = np.isinf(dimg)
ii = ndimage.distance_transform_edt(miss, return_distances=False, return_indices=True)
dimg = dimg[tuple(ii)]

dimg = ndimage.grey_closing(dimg, size=5)
dimg = ndimage.gaussian_filter(dimg, 1.2)
dshow = plt.cm.turbo_r(np.clip(np.kron(dimg, np.ones((2, 2))), 0, 50) / 50)
skyp0 = ROOT / "cache/sky_masks" / SC / (Path(fr["CAM_FRONT"].filename).stem + ".png")
if skyp0.exists():
    skm = np.asarray(Image.open(skyp0).convert("L").resize((dshow.shape[1], dshow.shape[0]))) > 127
    dshow[skm] = (0.08, 0.09, 0.11, 1)
fig, ax = img_axes()
ax.imshow(dshow, extent=(0, W, H, 0), interpolation="bilinear")
save(fig, "render_depth.png")

# ---------------------------------------------------------------- 마스크: 하늘 (캐시) + 객체 박스
skyp = ROOT / "cache/sky_masks" / SC / (Path(fr["CAM_FRONT"].filename).stem + ".png")
from nusim.recon.dynamic import boxes_at, build_tracks, masked_instances, static_instances  # noqa: E402

tracks = build_tracks(scene)
boxes = boxes_at(tracks, masked_instances(scene, "agents"), fr["CAM_FRONT"].timestamp, static_instances(scene, "agents"))
fig, ax = img_axes()
ax.imshow(np.asarray(gt), extent=(0, W, H, 0))
if skyp.exists():
    sk = np.asarray(Image.open(skyp).convert("L").resize((W, H))) > 127
    ov = np.zeros((H, W, 4))
    ov[sk] = (0.25, 0.55, 1.0, 0.55)
    ax.imshow(ov, extent=(0, W, H, 0))
for tr, size, rot in boxes:
    cs = box_corners(tr, size, rot, pad=0.3)
    cu, cz, _ = project(cs)
    if (cz < 0.5).any():
        continue
    hull = cu[[0, 1, 5, 4, 7, 6, 2, 3]]
    from scipy.spatial import ConvexHull  # noqa: E402

    hh = ConvexHull(cu)
    ax.add_patch(Polygon(cu[hh.vertices], closed=True, facecolor=(1, 0.3, 0.2, 0.45), edgecolor=(1, 0.4, 0.3, 0.9), lw=1))
save(fig, "masks.png")

# ---------------------------------------------------------------- 렌더 결과 (Vulkan, 같은 프레임)
Image.open(SCRATCH / "deck/f6" / f"f{FRAME:04d}_CAM_FRONT.png").convert("RGB").resize((W, H)).save(OUT / "render_front.jpg", quality=92)
for cam in sum(lay, []):
    Image.open(SCRATCH / "deck/f6" / f"f{FRAME:04d}_{cam}.png").convert("RGB").resize((480, 270)).save(OUT / f"render_{cam}.jpg", quality=92)

# ---------------------------------------------------------------- 리그 변환 (위에서 본 6캠: 원본 · left 1 m)
cams = {"CAM_FRONT": (1.7, 0.0, 0), "CAM_FRONT_LEFT": (1.55, 0.5, 55), "CAM_FRONT_RIGHT": (1.55, -0.5, -55),
        "CAM_BACK": (0.2, 0.0, 180), "CAM_BACK_LEFT": (1.0, 0.55, 110), "CAM_BACK_RIGHT": (1.0, -0.55, -110)}
for name, shift in (("rig_transform.png", 1.0),):
    fig = plt.figure(figsize=(4.2, 4.6), facecolor="#ffffff")
    ax = fig.add_axes([0, 0, 1, 1])
    ax.set_aspect("equal")
    ax.axis("off")
    ax.add_patch(Rectangle((-0.87, -0.65), 1.74, 4.08, facecolor="#1F3B70", edgecolor="none", zorder=3))
    for (ex, ey, yaw), shf, colr, al in [((0, 0, 0), 0, "#4472C4", 0.55), ((0, 0, 0), shift, "#ED7D31", 0.45)]:
        for c, (cx, cy, cyaw) in cams.items():
            hx, hy = np.cos(np.radians(cyaw)), np.sin(np.radians(cyaw))
            # 자차 (x 앞, y 왼쪽) → 그림 (오른쪽 = −y, 위 = x)
            px, py = -(cy + shf), cx
            dx, dy = -hy, hx
            nx, ny = -dy, dx
            L, Wd = 2.2, 1.0
            tri = [(px, py), (px + dx * L + nx * Wd, py + dy * L + ny * Wd), (px + dx * L - nx * Wd, py + dy * L - ny * Wd)]
            ax.add_patch(Polygon(tri, closed=True, facecolor=colr, alpha=al, edgecolor="none", zorder=4))
            ax.plot(px, py, "o", color=colr, ms=4, zorder=5)
    ax.annotate("", xy=(-shift - 0.05, 4.4), xytext=(0, 4.4), arrowprops=dict(arrowstyle="->", color="#ED7D31", lw=2))
    ax.text(-shift / 2, 4.65, "T: 왼쪽 1 m", ha="center", fontsize=11, color="#ED7D31", fontweight="bold")
    ax.set_xlim(-4.2, 3.4)
    ax.set_ylim(-3.0, 5.2)
    save(fig, name)

# ---------------------------------------------------------------- 구간 담당 (자차 궤적을 구간별 색, BEV)
seg_txt = (ROOT / "outputs/gs_day" / SC / "segments.txt").read_text().splitlines()
P = {}
for ln in seg_txt:
    if ln.startswith("P "):
        _, k, px, py = ln.split()
        P.setdefault(int(k), []).append((float(px), float(py)))
fig = plt.figure(figsize=(4.6, 4.6), facecolor="#ffffff")
ax = fig.add_axes([0.02, 0.02, 0.96, 0.96])
ax.set_aspect("equal")
ax.axis("off")
cols = ["#1F3B70", "#4472C4", "#ED7D31", "#70AD47", "#A5A5A5"]
for k, pp in sorted(P.items()):
    pp = np.array(pp)
    ax.plot(pp[:, 0], pp[:, 1], color=cols[k % 5], lw=7, solid_capstyle="butt")
    mid = pp[len(pp) // 2]
    ax.text(mid[0] + 9, mid[1] - 3, f"seg{k}", fontsize=12, color=cols[k % 5], fontweight="bold")
save(fig, "segments_bev.png")

# ---------------------------------------------------------------- 엔진: 장애물 지도 조각, 주행 기록, 하늘 파노라마
b = open(ROOT / "outputs/gs_day" / SC / "obstacles.pgm", "rb").read()
hdr = b.split(b"\n", 4)
x0, y0, res = map(float, hdr[1].split()[3:6])
gw, gh = map(int, hdr[2].split())
grid = np.frombuffer(b[-gw * gh:], np.uint8).reshape(gh, gw)
cx_, cy_ = int((ego[0] - read_origin(ROOT / f"outputs/gs_day/{SC}/seg0/point_cloud.ply")[0] - x0) / res), \
    int((ego[1] - read_origin(ROOT / f"outputs/gs_day/{SC}/seg0/point_cloud.ply")[1] - y0) / res)
crop = grid[cy_ - 150:cy_ + 150, cx_ - 150:cx_ + 150][::-1]
fig = plt.figure(figsize=(4, 4), facecolor="#ffffff")
ax = fig.add_axes([0, 0, 1, 1])
ax.imshow(255 - crop, cmap="gray", vmin=0, vmax=255)
ax.add_patch(Rectangle((150 - 4, 150 - 10), 8, 20, facecolor="#ED7D31", edgecolor="none"))
ax.axis("off")
save(fig, "obstacles_crop.png")
log = np.loadtxt(ROOT / "outputs/gs_day/scene-0402/shots/drive_log.csv", delimiter=",", skiprows=1, usecols=(0, 4))
fig = plt.figure(figsize=(4, 2.4), facecolor="#ffffff")
ax = fig.add_axes([0.16, 0.22, 0.8, 0.7])
ax.plot(log[:, 0], log[:, 1] * 3.6, color="#4472C4", lw=1.6)
ax.set_xlabel("시간 (s)", fontsize=10)
ax.set_ylabel("속도 (km/h)", fontsize=10)
ax.spines[["top", "right"]].set_visible(False)
ax.tick_params(labelsize=9)
save(fig, "drive_log.png")
pano = ROOT / "outputs/ibl" / SC / "panorama.png"
if pano.exists():
    Image.open(pano).convert("RGB").resize((800, 400)).save(OUT / "sky_panorama.jpg", quality=92)
print("panels ->", OUT, sorted(p.name for p in OUT.iterdir()))