#!/usr/bin/env python3
"""Decode before/after SLAT from bench ``*.npz`` and rasterize Gaussians at one eval camera.

Camera matches ``eval_3d_editing._eval_view_cameras`` (300 virtual steps, stride ``300//10``).
Requires GPU and local TRELLIS decode weights (``TrellisImageTo3DPipeline``).
"""
from __future__ import annotations

import argparse
import os
import sys

import numpy as np
import torch
from PIL import Image

_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)
os.chdir(_REPO_ROOT)

from trellis.modules import sparse as sp  # noqa: E402
from trellis.pipelines import TrellisImageTo3DPipeline  # noqa: E402
from trellis.utils import render_utils  # noqa: E402


def _eval_view_cameras(total_render_view_num: int = 300, eval_view_num: int = 10):
    yaws = torch.linspace(0, 2 * 3.1415, total_render_view_num)
    pitchs = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, total_render_view_num))
    step = total_render_view_num // eval_view_num
    yaws = yaws[::step].tolist()
    pitchs = pitchs[::step].tolist()
    radius = 2.0
    fov_deg = 40.0
    return yaws, pitchs, radius, fov_deg


def _load_slat_npz(path: str, device: torch.device) -> sp.SparseTensor:
    z = np.load(path, allow_pickle=True)
    feats = torch.from_numpy(z["feats"]).to(device=device, dtype=torch.float32)
    coords = torch.from_numpy(z["coords"]).to(device=device)
    if coords.shape[1] == 3:
        coords = torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=1)
    return sp.SparseTensor(feats=feats, coords=coords.int())


def _denorm_slat(slat: sp.SparseTensor, pipeline: TrellisImageTo3DPipeline) -> sp.SparseTensor:
    std = torch.tensor(pipeline.slat_normalization["std"])[None].to(device=slat.device, dtype=slat.feats.dtype)
    mean = torch.tensor(pipeline.slat_normalization["mean"])[None].to(device=slat.device, dtype=slat.feats.dtype)
    return slat * std + mean


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ori-npz", required=True, help="Before SLAT, e.g. ori_latents.npz (feats, coords)")
    ap.add_argument("--edit-npz", required=True, help="After SLAT, e.g. edit_latents.npz")
    ap.add_argument("--out-dir", required=True, help="Output directory for PNGs")
    ap.add_argument(
        "--trellis-pipeline-path",
        default="/mnt/zsn/ckpts/TRELLIS-image-large",
        help="Local dir or Hub id for TrellisImageTo3DPipeline",
    )
    ap.add_argument("--view-idx", type=int, default=0, help="Which of the 10 eval views (0..9)")
    ap.add_argument("--resolution", type=int, default=512)
    ap.add_argument("--total-views", type=int, default=300)
    ap.add_argument("--eval-views", type=int, default=10)
    args = ap.parse_args()

    if not (0 <= args.view_idx < args.eval_views):
        raise SystemExit(f"--view-idx must be in [0, {args.eval_views - 1}]")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type != "cuda":
        raise SystemExit("CUDA is required for Trellis decode + Gaussian rasterization")

    os.makedirs(args.out_dir, exist_ok=True)

    pipe = TrellisImageTo3DPipeline.from_pretrained(args.trellis_pipeline_path)
    pipe.cuda()

    slat_before = _load_slat_npz(args.ori_npz, device)
    slat_after = _load_slat_npz(args.edit_npz, device)
    slat_before = _denorm_slat(slat_before, pipe)
    slat_after = _denorm_slat(slat_after, pipe)

    dec_b = pipe.decode_slat(slat_before, ["gaussian"])
    dec_a = pipe.decode_slat(slat_after, ["gaussian"])
    g_b = dec_b["gaussian"][0]
    g_a = dec_a["gaussian"][0]

    yaws, pitchs, r, fov_deg = _eval_view_cameras(args.total_views, args.eval_views)
    yaw = yaws[args.view_idx]
    pitch = pitchs[args.view_idx]
    extrinsics, intrinsics = render_utils.yaw_pitch_r_fov_to_extrinsics_intrinsics(
        [yaw], [pitch], r, fov_deg
    )
    opts = {"resolution": args.resolution, "bg_color": (0, 0, 0)}

    meta_path = os.path.join(args.out_dir, "camera.txt")
    with open(meta_path, "w", encoding="utf-8") as f:
        f.write(
            f"view_idx={args.view_idx} (of {args.eval_views})\n"
            f"total_virtual={args.total_views} stride={args.total_views // args.eval_views}\n"
            f"yaw_rad={yaw}\n pitch_rad={pitch}\n radius={r}\n fov_deg={fov_deg}\n"
        )

    for tag, g in ("before", g_b), ("after", g_a):
        rets = render_utils.render_frames(g, extrinsics, intrinsics, options=opts, verbose=False)
        out_png = os.path.join(args.out_dir, f"gs_{tag}_view{args.view_idx:02d}.png")
        Image.fromarray(rets["color"][0]).save(out_png)
        print("wrote", out_png)

    row = np.concatenate(
        [
            np.array(Image.open(os.path.join(args.out_dir, f"gs_before_view{args.view_idx:02d}.png"))),
            np.array(Image.open(os.path.join(args.out_dir, f"gs_after_view{args.view_idx:02d}.png"))),
        ],
        axis=1,
    )
    pair_path = os.path.join(args.out_dir, f"gs_before_after_view{args.view_idx:02d}_sidebyside.png")
    Image.fromarray(row).save(pair_path)
    print("wrote", pair_path)


if __name__ == "__main__":
    main()