"""Learn2Splat — interactive demo of the meta-learned 3DGS optimizer.

A ZeroGPU / Gradio port of ``demo.py --with-gui gradio`` from the official
Learn2Splat release (https://huggingface.co/spaces/autonomousvision/Learn2Splat).

The scene (Mip-NeRF 360 "garden", shipped by the authors in
``autonomousvision/learn2splat``) is SfM-initialized from its COLMAP points on
CPU at startup; each run streams the decoder's render of a chosen training view
while the picked optimizer refines the Gaussians, then hands the finished splats
to an interactive ``gr.Model3D`` viewer and scores them on held-out views.

The two optimizers run through exactly the same ``OptGS.optimize_iter`` path,
from an identical SfM initialization, with the same view minibatches and step
budget — only the update rule differs (learned vs. Adam), which is the
comparison the paper makes.
"""

import glob
import importlib
import importlib.util
import os
import site
import subprocess
import sys

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import warnings

warnings.filterwarnings("ignore")

# --- CUDA extensions -------------------------------------------------------
# gsplat / pointops / fused_ssim / nerfacc / fused_knn_attn have no upstream
# wheel for the ZeroGPU cell (py3.12 · torch 2.11 · cu130 · sm_120), so they are
# prebuilt and committed under wheels/. requirements.txt cannot reference them:
# the Space build mounts only requirements.txt (the repo is copied in later), so
# a relative wheel path there fails with "No such file or directory". Installing
# them here — from local files, before torch is imported — is the reliable spot.
_APP_DIR = os.path.dirname(os.path.abspath(__file__))
_WHEELS = sorted(glob.glob(os.path.join(_APP_DIR, "wheels", "*.whl")))
_EXT_MODULES = ("gsplat", "pointops", "fused_ssim", "nerfacc", "fused_knn_attn_cuda")


def _install_cuda_extensions() -> None:
    # find_spec() locates without importing, so nothing touches CUDA before
    # `import spaces` below.
    if all(importlib.util.find_spec(m) is not None for m in _EXT_MODULES):
        return
    if not _WHEELS:
        raise RuntimeError("wheels/ is empty — the CUDA extensions cannot be installed")
    cmd = [
        sys.executable, "-m", "pip", "install", "--no-deps", "--no-input",
        "--disable-pip-version-check", "-q", *_WHEELS,
    ]
    print(f"[setup] installing {len(_WHEELS)} prebuilt CUDA extension wheels …", flush=True)
    if subprocess.run(cmd).returncode != 0:
        subprocess.run(cmd + ["--user"], check=True)
        sys.path.append(site.getusersitepackages())
    importlib.invalidate_caches()


_install_cuda_extensions()

import spaces  # noqa: E402  — must precede torch / any CUDA-touching import

import gc  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402
from types import SimpleNamespace  # noqa: E402

import gradio as gr  # noqa: E402
import imageio.v2 as imageio  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from torch import Tensor  # noqa: E402

from optgs.dataset.colmap.utils import Dataset, Parser
from optgs.experimental.api import OptGS
from optgs.experimental.api.integration.config_bridge import build_adam_baseline
from optgs.experimental.initializers_utils import knn, points_to_gaussians
from optgs.misc.image_io import prep_image
from optgs.model.ply_export import save_gaussian_ply
from optgs.model.types import Gaussians
from optgs.scene_trainer.common.gaussian_adapter import build_covariance

# --------------------------------------------------------------------------
# Constants (mirrors demo.py)
# --------------------------------------------------------------------------
NEAR_PLANE = 0.01
FAR_PLANE = 100.0
SH_C0 = 0.28209479177387814

DEMO_REPO = "autonomousvision/learn2splat"
DATA_DIR = "datasets/mip360/garden"
DATA_FACTOR = 4
TEST_EVERY = 8
SEED = 42
GLOBAL_SCALE = 1.0
INIT_OPA = 0.1
INIT_SCALE = 1.0

# The released Learn2Splat checkpoint: the "joint" model, trained on the mixed
# dense+sparse regime with the COLMAP/SfM scene initializer — the one this
# SfM-initialized demo is in-distribution for. (``resplat_init`` in the same
# repo is the ReSplat feed-forward-initializer baseline and needs a different
# pipeline, so it is not offered here.)
CKPT_FILE = "joint/checkpoints/epoch_7-step_150000.ckpt"
CKPT_CONFIG = "joint/config.yaml"
CKPT_DIR = "checkpoints"  # same layout optgs.misc.hf_ckpt uses for hf:// refs

DEFAULT_STEPS = 100
DEFAULT_BATCH = 8
DEFAULT_STRATEGY = "random"
DEFAULT_ADC = "off"
DEFAULT_HEIGHT = 540

METHOD_LEARNED = "Learn2Splat (learned optimizer)"
METHOD_ADAM = "Adam (3DGS baseline)"
METHODS = [METHOD_LEARNED, METHOD_ADAM]

CPU = torch.device("cpu")
DTYPE = torch.float32


# --------------------------------------------------------------------------
# Scene preparation — all of this is CPU-only and runs once at startup
# --------------------------------------------------------------------------
def ensure_data() -> None:
    if os.path.isdir(DATA_DIR) and os.listdir(DATA_DIR):
        return
    from huggingface_hub import snapshot_download

    print(f"[setup] downloading {DATA_DIR} from {DEMO_REPO} …", flush=True)
    snapshot_download(repo_id=DEMO_REPO, allow_patterns=[f"{DATA_DIR}/**"], local_dir=".")
    print("[setup] scene ready", flush=True)


def fetch_checkpoint() -> str:
    """Download the released checkpoint + its training config, return the ckpt path.

    Mirrors ``optgs.misc.hf_ckpt``'s layout (``checkpoints/<tag>/…``) so
    ``_find_config_for_checkpoint`` picks the sibling ``config.yaml`` up, but
    fetches it ourselves so the config can be reconciled with the vendored
    code first (see ``_reconcile_released_config``).
    """
    from huggingface_hub import hf_hub_download

    ckpt = hf_hub_download(DEMO_REPO, CKPT_FILE, local_dir=CKPT_DIR)
    cfg = hf_hub_download(DEMO_REPO, CKPT_CONFIG, local_dir=CKPT_DIR)
    _reconcile_released_config(cfg)
    return ckpt


def _reconcile_released_config(path: str) -> None:
    """Make the released training config parse against the vendored optgs code.

    The published `joint` config was written by a slightly newer revision than
    the released code, and two values do not round-trip:

    * ``refiner.prune_scale2d: null`` — the typed ADC config still requires a
      float. The refiner is ``none`` here (no densification, no pruning), so the
      value is never read; 0.0 is what ``refiner/none.yaml`` ships.
    * ``init_state_wo_features: false`` — this checkpoint was trained with
      encoder features, which an external SfM/COLMAP scene does not carry.
      ``OptGS`` already handles that case (it drops the feature-conditioned
      ``state_proj`` weights and seeds the optimizer state randomly), but it
      flips the flag *after* the config is parsed, and the COLMAP branch of
      ``KnnBasedOptimizerCfg.update`` asserts on it first. Setting it up front
      takes exactly the path OptGS intends, without its warning firing late.
    """
    import re as _re

    txt = open(path).read()
    out = _re.sub(r"(?m)^(\s*)prune_scale2d:\s*null\s*$", r"\g<1>prune_scale2d: 0.0", txt)
    out = _re.sub(
        r"(?m)^(\s*)init_state_wo_features:\s*false\s*$",
        r"\g<1>init_state_wo_features: true",
        out,
    )
    if out != txt:
        with open(path, "w") as f:
            f.write(out)
        print(f"[setup] reconciled released config {path} with the vendored optgs", flush=True)


def sfm_initialization(parser: Parser, sh_degree: int) -> Gaussians:
    """SfM point cloud -> an optgs ``Gaussians`` (batch=1), on CPU. From demo.py."""
    points = torch.from_numpy(parser.points).float()
    rgbs = torch.from_numpy(parser.points_rgb / 255.0).float()

    dist2_avg = (knn(points, 4)[:, 1:] ** 2).mean(dim=-1)
    scales = (torch.sqrt(dist2_avg) * INIT_SCALE).unsqueeze(-1).repeat(1, 3)
    opacities = torch.full((points.shape[0],), INIT_OPA)

    g = points_to_gaussians(
        {"xyz": points, "rgb": rgbs, "scales": scales, "opacities": opacities},
        sh_degree=sh_degree,
        device=CPU,
    )
    sh0, shN = g["sh0"], g["shN"]
    harmonics = torch.cat([sh0, shN], dim=1) if shN is not None else sh0
    harmonics = harmonics.permute(0, 2, 1)  # [N, 3, K]

    scales_act = torch.exp(g["scales_raw"])
    opacities_act = torch.sigmoid(g["opacities_raw"])
    rotations = F.normalize(g["rotations_unnorm"], dim=-1)
    covariances = build_covariance(scale=scales_act, rotation_xyzw=rotations)

    def _b(t: Tensor) -> Tensor:
        return t.unsqueeze(0).to(DTYPE)

    return Gaussians(
        means=_b(g["xyz"]),
        covariances=_b(covariances),
        harmonics=_b(harmonics),
        opacities=_b(opacities_act),
        scales=_b(scales_act),
        rotations=_b(rotations),
        rotations_unnorm=_b(g["rotations_unnorm"]),
    )


def collect_cameras(dataset: Dataset, indices):
    c2ws, ks, imgs = [], [], []
    for i in indices:
        data = dataset[i]
        c2ws.append(data["camtoworld"])
        ks.append(data["K"])
        imgs.append(data["image"] / 255.0)
    return torch.stack(c2ws), torch.stack(ks), torch.stack(imgs)


def build_batched_views(camtoworlds, Ks, images, scene_scale) -> dict:
    """COLMAP cameras -> an optgs ``BatchedViews`` dict (batch=1), on CPU."""
    v, h, w = images.shape[0], images.shape[1], images.shape[2]
    Ks_norm = Ks.clone()
    Ks_norm[:, 0, :] /= w
    Ks_norm[:, 1, :] /= h
    image = images.permute(0, 3, 1, 2)  # [V, 3, H, W]

    def _b(t: Tensor) -> Tensor:
        return t.unsqueeze(0).to(device=CPU, dtype=DTYPE)

    return {
        "extrinsics": _b(camtoworlds),
        "intrinsics": _b(Ks_norm),
        "image": _b(image),
        "near": torch.full((1, v), NEAR_PLANE, dtype=DTYPE),
        "far": torch.full((1, v), FAR_PLANE, dtype=DTYPE),
        "index": torch.arange(v).unsqueeze(0),
        "scene_scale": torch.tensor([scene_scale], dtype=DTYPE),
    }


print("[setup] fetching scene …", flush=True)
ensure_data()

print("[setup] fetching checkpoint …", flush=True)
CHECKPOINT = fetch_checkpoint()

print("[setup] building OptGS (learned optimizer) …", flush=True)
OPTGS = OptGS(
    checkpoint=CHECKPOINT,
    device="cuda",
    num_refine=DEFAULT_STEPS,
    opt_batch_size=DEFAULT_BATCH,
    opt_batch_strategy=DEFAULT_STRATEGY,
)

print("[setup] parsing COLMAP scene …", flush=True)
PARSER = Parser(data_dir=DATA_DIR, factor=DATA_FACTOR, normalize=True, verbose=False)
DATASET = Dataset(PARSER)
VAL_IDX = [i for i in range(len(DATASET)) if i % TEST_EVERY == 0]
TRAIN_IDX = [i for i in range(len(DATASET)) if i % TEST_EVERY != 0]
SCENE_SCALE = PARSER.scene_scale * 1.1 * GLOBAL_SCALE

_tc2w, _tK, _timg = collect_cameras(DATASET, TRAIN_IDX)
TRAIN_BV_CPU = build_batched_views(_tc2w, _tK, _timg, SCENE_SCALE)
H_FULL, W_FULL = _timg.shape[1], _timg.shape[2]
N_TRAIN_VIEWS = len(TRAIN_IDX)
del _tc2w, _tK, _timg

_vc2w, _vK, _vimg = collect_cameras(DATASET, VAL_IDX)
VAL_BV_CPU = build_batched_views(_vc2w, _vK, _vimg, SCENE_SCALE)
del _vc2w, _vK, _vimg
gc.collect()

torch.manual_seed(SEED)
np.random.seed(SEED)
print("[setup] SfM initialization …", flush=True)
INIT_GAUSSIANS_CPU = sfm_initialization(PARSER, OPTGS.sh_degree)
N_INIT_GS = int(INIT_GAUSSIANS_CPU.means.shape[1])
print(
    f"[setup] ready — {N_TRAIN_VIEWS} train / {len(VAL_IDX)} held-out views, "
    f"{N_INIT_GS:,} SfM Gaussians, {H_FULL}x{W_FULL} px",
    flush=True,
)

# Ground-truth thumbnails for the reference-view picker (CPU, no GPU needed).
GT_DIR = Path(tempfile.mkdtemp(prefix="l2s-gt-"))


def gt_view_path(view_idx: int) -> str:
    """Write (once) and return the ground-truth image for a training view."""
    i = int(view_idx)
    path = GT_DIR / f"gt_{i:04d}.jpg"
    if not path.exists():
        img = TRAIN_BV_CPU["image"][0, i].permute(1, 2, 0).numpy()
        imageio.imwrite(path, (np.clip(img, 0, 1) * 255).astype(np.uint8), quality=90)
    return str(path)


# --------------------------------------------------------------------------
# GPU-side helpers (run inside the @spaces.GPU worker)
# --------------------------------------------------------------------------
_GPU = {}  # warm-worker cache: CPU tensors are moved to VRAM once per worker


def _bv_to_cuda(bv: dict) -> dict:
    return {k: (v.to("cuda") if torch.is_tensor(v) else v) for k, v in bv.items()}


def _gpu_state():
    """Move the scene onto the GPU once per (warm) worker."""
    if "train_bv" not in _GPU:
        _GPU["train_bv"] = _bv_to_cuda(TRAIN_BV_CPU)
        _GPU["val_bv"] = _bv_to_cuda(VAL_BV_CPU)
        # Held-out ground truth, [V, H, W, 3] — a view of val_bv's image tensor.
        _GPU["val_images"] = _GPU["val_bv"]["image"][0].permute(0, 2, 3, 1)
        _GPU["init"] = INIT_GAUSSIANS_CPU.to(device=torch.device("cuda"), dtype=DTYPE)
    return _GPU


@torch.no_grad()
def _render(g: Gaussians, bv: dict, view_idx: int, height: int) -> np.ndarray:
    """Render ``g`` from one view of ``bv`` at ``height`` px (uint8 [H, W, 3])."""
    h = int(height)
    w = max(1, round(h * W_FULL / H_FULL))
    sl = slice(int(view_idx), int(view_idx) + 1)
    out = OPTGS.decoder.forward(
        g,
        bv["extrinsics"][:, sl],
        bv["intrinsics"][:, sl],
        bv["near"][:, sl],
        bv["far"][:, sl],
        image_shape=(h, w),
    )
    return prep_image(out.color[0, 0])


@torch.no_grad()
def _psnr_on_heldout(g: Gaussians, st: dict, chunk: int = 6) -> float:
    """Mean PSNR of ``g`` on the held-out views, at the dataset resolution."""
    val_bv, val_images = st["val_bv"], st["val_images"]
    n = val_bv["extrinsics"].shape[1]
    psnrs = []
    for s in range(0, n, chunk):
        sl = slice(s, min(s + chunk, n))
        out = OPTGS.decoder.forward(
            g,
            val_bv["extrinsics"][:, sl],
            val_bv["intrinsics"][:, sl],
            val_bv["near"][:, sl],
            val_bv["far"][:, sl],
            image_shape=(H_FULL, W_FULL),
        )
        colors = out.color[0].clamp(0.0, 1.0)  # [v, 3, H, W]
        mse = ((colors.permute(0, 2, 3, 1) - val_images[sl]) ** 2).mean(dim=(1, 2, 3))
        psnrs.append(-10.0 * torch.log10(mse))
    return float(torch.cat(psnrs).mean().item())


def _init_psnr(st: dict) -> float:
    """PSNR of the untouched SfM initialization (computed once per worker)."""
    if "init_psnr" not in st:
        st["init_psnr"] = _psnr_on_heldout(st["init"], st)
    return st["init_psnr"]


def _reorient_for_viewer(g: Gaussians) -> Gaussians:
    """COLMAP Z-up -> the Y-up frame gr.Model3D/Babylon displays upright.

    Straight from demo.py: gradio's splat viewer flips Y, so exporting through
    the reflection E(p)=(x,-z,-y) makes the displayed scene (x,z,-y). E's
    point-reflection part leaves covariances unchanged; its proper-rotation
    part R_p=-E rotates the splat orientations to match.
    """
    from scipy.spatial.transform import Rotation as Rsp

    g = g.clone()
    m = g.means[0]
    g.means = torch.stack([m[:, 0], -m[:, 2], -m[:, 1]], dim=1)[None]
    q = F.normalize(g.rotations_unnorm[0], dim=-1).detach().cpu().numpy()  # xyzw
    R_p = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
    q = (Rsp.from_matrix(R_p) * Rsp.from_quat(q)).as_quat()
    g.rotations_unnorm = torch.from_numpy(q).to(g.means)[None]
    return g


def _export_ply(g: Gaussians) -> str:
    fd, path = tempfile.mkstemp(suffix=".ply", prefix="learn2splat-")
    os.close(fd)
    save_gaussian_ply(_reorient_for_viewer(g), save_path=Path(path))
    return path


# --------------------------------------------------------------------------
# Status panel
# --------------------------------------------------------------------------
def _msg(text: str, err: bool = False) -> str:
    return f"<div class='l2s-msg{' err' if err else ''}'>{text}</div>"


def _card(value: str, label: str) -> str:
    return f"<div class='l2s-stat'><b>{value}</b><span>{label}</span></div>"


def _status(head: str, cards: str, note: str = "") -> str:
    return (
        f"<div class='l2s-head'>{head}</div>"
        f"<div class='l2s-stats'>{cards}</div>{note}"
    )


def _running_status(step: int, n_steps: int, n_gs: int, t0: float) -> str:
    el = time.perf_counter() - t0
    ms = el / max(1, step) * 1000.0
    dg = n_gs - N_INIT_GS
    cards = (
        _card(f"{el:.1f}s", "elapsed")
        + _card(f"{ms:.0f}", "ms / step")
        + _card(f"{n_gs:,}", f"gaussians · +{dg:,}" if dg > 0 else "gaussians")
        + _card(f"{step}/{n_steps}", "step")
    )
    return _status(f"optimizing · {head_pct(step, n_steps)}", cards)


def head_pct(step: int, n_steps: int) -> str:
    return f"{100.0 * step / max(1, n_steps):.0f}%"


# --------------------------------------------------------------------------
# The run
# --------------------------------------------------------------------------
def _estimate_duration(method=METHOD_LEARNED, steps=DEFAULT_STEPS, *args, **kwargs):
    """GPU seconds to reserve. Measured on the live Space (ZeroGPU / Blackwell):
    ~0.11 s per step warm, plus ~13 s fixed (cold-worker warm-up, moving the
    scene to VRAM, the held-out PSNR pass and the PLY export). A cold 30-step
    run took 16.3 s end to end; this leaves ~1.4x headroom on top."""
    try:
        n = int(steps)
    except Exception:
        n = DEFAULT_STEPS
    return int(min(180, 20 + 0.18 * n))


@spaces.GPU(duration=_estimate_duration)
def optimize(
    method: str = METHOD_LEARNED,
    steps: int = DEFAULT_STEPS,
    batch_size: int = DEFAULT_BATCH,
    strategy: str = DEFAULT_STRATEGY,
    adc: str = DEFAULT_ADC,
    view_idx: int = 0,
    height: int = DEFAULT_HEIGHT,
):
    """Refine the demo scene's Gaussians and stream the live render.

    Both methods start from the same SfM initialization and run through the same
    optimize loop; only the update rule differs.

    Args:
        method: which update rule to run — the Learn2Splat learned optimizer or
            the 3DGS Adam baseline.
        steps: number of refinement steps.
        batch_size: training views fed to the optimizer per step (view minibatch).
        strategy: how the per-step view minibatch is sampled.
        adc: adaptive density control (densification) strategy.
        view_idx: which training view the live render is taken from.
        height: render height in pixels for the live preview.

    Yields:
        (live render, status HTML, refined splats as a .ply for the 3D viewer).
    """
    st = _gpu_state()
    train_bv = st["train_bv"]
    init_g = st["init"]

    n_steps = max(1, int(steps))
    bs = max(1, min(int(batch_size), N_TRAIN_VIEWS))
    adc_name = "none" if adc == "off" else adc
    use_adam = method == METHOD_ADAM
    opt = None

    try:
        yield (
            _render(init_g, train_bv, view_idx, height),
            _status(
                "initializing",
                _card(f"{N_INIT_GS:,}", "gaussians")
                + _card(f"{n_steps}", "steps")
                + _card(f"{bs}", "views / step")
                + _card("SfM", "init"),
                _msg("Building the optimizer state…"),
            ),
            None,
        )

        OPTGS.initialize_from_tensors(init_g.clone(), train_bv)
        OPTGS.num_refine = n_steps
        OPTGS.opt_batch_size = bs
        OPTGS.opt_batch_strategy = strategy
        OPTGS.configure_adc(adc_name)
        if use_adam:
            opt = build_adam_baseline(n_steps, adc=adc_name).to("cuda")

        torch.cuda.reset_peak_memory_stats()
        t0 = time.perf_counter()
        # Stream ~24 frames regardless of the step budget (rendering every step
        # would spend most of the GPU reservation on preview renders).
        every = max(1, n_steps // 24)
        g = init_g
        for step, g in OPTGS.optimize_iter(optimizer=opt):
            if (step + 1) % every == 0 or step + 1 == n_steps:
                yield (
                    _render(g, train_bv, view_idx, height),
                    _running_status(step + 1, n_steps, int(g.means.shape[1]), t0),
                    None,
                )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0

        psnr = _psnr_on_heldout(g, st)
        psnr0 = _init_psnr(st)
        peak = torch.cuda.max_memory_allocated() / 1e9
        cards = (
            _card(f"{psnr:.2f} dB", f"held-out psnr · {psnr - psnr0:+.2f}")
            + _card(f"{elapsed:.1f}s", f"{n_steps} steps")
            + _card(f"{int(g.means.shape[1]):,}", "gaussians")
            + _card(f"{peak:.1f} GB", "peak vram")
        )
        note = _msg(
            f"<b>{method}</b> · {elapsed / n_steps * 1000:.0f} ms/step · "
            f"SfM init scored {psnr0:.2f} dB · PSNR on {len(VAL_IDX)} held-out "
            f"views at {W_FULL}&times;{H_FULL}."
        )
        yield (
            _render(g, train_bv, view_idx, height),
            _status("done", cards, note),
            _export_ply(g),
        )
    except Exception as e:  # surface the failure in the panel instead of a toast
        yield (
            gr.update(),
            _msg(f"error — {type(e).__name__}: {str(e)[:300]}", err=True),
            None,
        )
        raise
    finally:
        opt = None
        gc.collect()
        torch.cuda.empty_cache()


def show_gt(view_idx: int) -> str:
    """Ground-truth image for a training view (CPU only)."""
    return gt_view_path(view_idx)


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1500px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
#l2s-status { border: 1px solid var(--block-border-color); border-radius: 10px;
  padding: 10px 12px; background: var(--block-background-fill); }
#l2s-status .l2s-head { font-size: 11px; letter-spacing: .12em; text-transform: uppercase;
  color: var(--body-text-color-subdued); font-weight: 600; margin: 2px 1px 10px; }
#l2s-status .l2s-stats { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; }
#l2s-status .l2s-stat { border: 1px solid var(--block-border-color); border-radius: 8px;
  padding: 7px 10px; }
#l2s-status .l2s-stat b { display: block; font-size: 18px; font-weight: 650; line-height: 1.2; }
#l2s-status .l2s-stat span { display: block; font-size: 9.5px; letter-spacing: .05em;
  text-transform: uppercase; color: var(--body-text-color-subdued); margin-top: 3px; }
#l2s-status .l2s-msg { font-size: 13.5px; margin: 9px 1px 2px;
  color: var(--body-text-color-subdued); }
#l2s-status .l2s-msg.err { color: #c0392b; word-break: break-word; }
"""

INTRO = """# 🫧 Learn2Splat — learned 3D Gaussian Splatting optimizer

Learn2Splat replaces Adam with a **meta-learned update rule** for 3DGS. Pick an
optimizer, press **Run** and watch the render converge from the same SfM
initialization — then orbit the finished splats and compare held-out PSNR.

Scene: Mip-NeRF 360 *garden*, SfM-initialized from its COLMAP points.
[Paper](https://huggingface.co/papers/2605.15760) ·
[Project page](https://naamapearl.github.io/learn2splat/) ·
[Weights](https://huggingface.co/autonomousvision/learn2splat) ·
[Code](https://github.com/autonomousvision/learn2splat)
"""

with gr.Blocks(title="Learn2Splat") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)

        with gr.Row(equal_height=False):
            with gr.Column(scale=3, min_width=260):
                method_dd = gr.Dropdown(METHODS, value=METHOD_LEARNED, label="Optimizer")
                steps_sl = gr.Slider(
                    10, 200, value=DEFAULT_STEPS, step=10, label="Refinement steps"
                )
                run_btn = gr.Button("Run optimization", variant="primary")
                status_html = gr.HTML(
                    _msg("SfM-initialized — pick an optimizer, then <b>Run</b>."),
                    label="Run status", show_label=False, elem_id="l2s-status",
                )
                with gr.Accordion("Advanced settings", open=False):
                    view_sl = gr.Slider(
                        0, N_TRAIN_VIEWS - 1, value=0, step=1, label="Preview camera"
                    )
                    height_sl = gr.Slider(
                        240, 1080, value=DEFAULT_HEIGHT, step=60, label="Render height"
                    )
                    batch_sl = gr.Slider(
                        1, min(32, N_TRAIN_VIEWS), value=DEFAULT_BATCH, step=1,
                        label="Views per step (minibatch)",
                    )
                    strategy_dd = gr.Dropdown(
                        ["random", "sequential", "fps"], value=DEFAULT_STRATEGY,
                        label="Minibatch sampling",
                    )
                    adc_dd = gr.Dropdown(
                        ["off", "default", "mcmc"], value=DEFAULT_ADC,
                        label="Densification (ADC)",
                    )

            with gr.Column(scale=5, min_width=340):
                image_out = gr.Image(
                    label="Optimizer · live render", height=460, format="jpeg",
                    interactive=False,
                )
                gt_img = gr.Image(
                    value=gt_view_path(0), label="Ground truth · preview camera",
                    height=190, format="jpeg", interactive=False,
                )

            with gr.Column(scale=5, min_width=340):
                model3d_out = gr.Model3D(
                    label="Refined splats · drag to orbit", height=660, clear_color=[0, 0, 0, 0],
                )

        # Registered before gr.Examples so this handler owns the /optimize
        # endpoint (Gradio suffixes whichever event is registered second).
        run_btn.click(
            optimize,
            inputs=[
                method_dd, steps_sl, batch_sl, strategy_dd, adc_dd, view_sl, height_sl,
            ],
            outputs=[image_out, status_html, model3d_out],
            concurrency_id="gpu",
            api_name="optimize",
        )

        gr.Examples(
            examples=[
                [METHOD_LEARNED, 100],
                [METHOD_ADAM, 100],
                [METHOD_LEARNED, 30],
            ],
            inputs=[method_dd, steps_sl],
            outputs=[image_out, status_html, model3d_out],
            fn=optimize,
            # The generator yields None for the 3D viewer while the run streams,
            # which current Gradio cannot cache — and a cached static result
            # would defeat the point of watching the optimization converge.
            cache_examples=False,
            run_on_click=True,
            label="Try",
        )

    view_sl.change(show_gt, inputs=view_sl, outputs=gt_img, api_name="ground_truth")

if __name__ == "__main__":
    # Gradio 6 moved `theme` / `css` from the Blocks constructor to launch().
    demo.queue(default_concurrency_limit=1).launch(
        theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True
    )