import os
import sys
import json
import tempfile
import subprocess
import shutil
from pathlib import Path

# Root of the repo — wherever app.py lives
REPO_ROOT = Path(__file__).parent.resolve()

import gradio as gr
import spaces
import numpy as np
import torch


# ──────────────────────────────────────────────
# Lazy model loading (done inside GPU decorator)
# ──────────────────────────────────────────────
_model_loaded = False


def _ensure_models():
    global _model_loaded
    if _model_loaded:
        return
    # Models are expected to be pre-downloaded to ./checkpoints/
    # on the Space via the HF repo or a setup script.
    _model_loaded = True


# ──────────────────────────────────────────────
# Core inference helpers
# ──────────────────────────────────────────────

def _write_caption_txt(image_path: str, caption: str) -> str:
    """Write a .txt caption file beside the image and return the directory."""
    img_path = Path(image_path)
    txt_path = img_path.with_suffix(".txt")
    txt_path.write_text(caption)
    return str(img_path.parent)


def _run(cmd: list[str], desc: str = "") -> tuple[bool, str]:
    """Run a Python -m command using the current interpreter with PYTHONPATH set."""
    # Replace bare "python" / "python3" with the real interpreter path
    if cmd[0] in ("python", "python3"):
        cmd = [sys.executable] + cmd[1:]

    run_env = {
        **os.environ,
        "PYTHONPATH": str(REPO_ROOT),          # makes lyra_2._src.* importable
        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
    }
    print(f"[Lyra] {desc}: {' '.join(cmd)}")
    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        env=run_env,
        cwd=str(REPO_ROOT),                    # run from repo root
    )
    log = result.stdout + "\n" + result.stderr
    return result.returncode == 0, log


# ──────────────────────────────────────────────
# Zoom-in / Zoom-out trajectory (Option 1)
# ──────────────────────────────────────────────

@spaces.GPU(duration=190)
def run_zoomgs(
    image,
    caption: str,
    sample_id: int,
    zoom_in_strength: float,
    zoom_out_strength: float,
    num_frames_in: int,
    num_frames_out: int,
    use_dmd: bool,
    run_reconstruction: bool,
):
    _ensure_models()

    with tempfile.TemporaryDirectory() as tmp:
        # Save uploaded image + caption
        img_path = Path(tmp) / "input.png"
        caption_path = Path(tmp) / "input.txt"

        from PIL import Image
        Image.fromarray(image).save(img_path)
        caption_path.write_text(caption.strip() or "A scenic outdoor environment.")

        output_dir = Path(tmp) / "outputs" / "zoomgs"
        output_dir.mkdir(parents=True, exist_ok=True)

        cmd = [
            "python", "-m", "lyra_2._src.inference.lyra2_zoomgs_inference",
            "--input_image_path", str(tmp),
            "--sample_id", "0",           # we always name it input.png → id 0 equivalent
            "--experiment", "lyra2",
            "--checkpoint_dir", "checkpoints/model",
            "--prompt_dir", str(tmp),
            "--output_path", str(output_dir),
            "--num_frames_zoom_in", str(num_frames_in),
            "--num_frames_zoom_out", str(num_frames_out),
            "--zoom_in_strength", str(zoom_in_strength),
            "--zoom_out_strength", str(zoom_out_strength),
        ]
        if use_dmd:
            cmd.append("--use_dmd")

        ok, log = _run(cmd, "ZoomGS video generation")

        # Locate output video
        video_path = output_dir / "0" / "videos" / "0.mp4"
        if not video_path.exists():
            # Fallback: search recursively
            candidates = list(output_dir.rglob("*.mp4"))
            video_path = candidates[0] if candidates else None

        gs_video = None
        if run_reconstruction and video_path and video_path.exists():
            ok2, log2 = _run(
                ["python", "-m", "lyra_2._src.inference.vipe_da3_gs_recon",
                 "--input_video_path", str(video_path)],
                "GS reconstruction",
            )
            log += "\n" + log2
            ply_candidates = list(output_dir.rglob("gs_trajectory.mp4"))
            if ply_candidates:
                gs_video = str(ply_candidates[0])

        return (
            str(video_path) if video_path and video_path.exists() else None,
            gs_video,
            log[-4000:],
        )


# ──────────────────────────────────────────────
# Custom trajectory (Option 2)
# ──────────────────────────────────────────────

@spaces.GPU(duration=900)
def run_custom_traj(
    image,
    trajectory_file,
    captions_json: str,
    num_frames: int,
    pose_scale: float,
    use_dmd: bool,
    run_reconstruction: bool,
):
    _ensure_models()

    with tempfile.TemporaryDirectory() as tmp:
        from PIL import Image
        img_path = Path(tmp) / "first_frame.png"
        Image.fromarray(image).save(img_path)

        traj_path = Path(tmp) / "trajectory.npz"
        shutil.copy(trajectory_file.name, traj_path)

        captions_path = Path(tmp) / "captions.json"
        try:
            json.loads(captions_json)   # validate
            captions_path.write_text(captions_json)
        except json.JSONDecodeError:
            captions_path.write_text(json.dumps({"0": captions_json}))

        output_dir = Path(tmp) / "outputs" / "custom"
        output_dir.mkdir(parents=True, exist_ok=True)

        cmd = [
            "python", "-m", "lyra_2._src.inference.lyra2_custom_traj_inference",
            "--input_image_path", str(img_path),
            "--trajectory_path", str(traj_path),
            "--experiment", "lyra2",
            "--checkpoint_dir", "checkpoints/model",
            "--captions_path", str(captions_path),
            "--num_frames", str(num_frames),
            "--output_path", str(output_dir),
            "--pose_scale", str(pose_scale),
        ]
        if use_dmd:
            cmd.append("--use_dmd")

        ok, log = _run(cmd, "Custom trajectory video generation")

        video_candidates = list(output_dir.rglob("*.mp4"))
        video_path = video_candidates[0] if video_candidates else None

        gs_video = None
        if run_reconstruction and video_path:
            ok2, log2 = _run(
                ["python", "-m", "lyra_2._src.inference.vipe_da3_gs_recon",
                 "--input_video_path", str(video_path)],
                "GS reconstruction",
            )
            log += "\n" + log2
            ply_candidates = list(output_dir.rglob("gs_trajectory.mp4"))
            if ply_candidates:
                gs_video = str(ply_candidates[0])

        return (
            str(video_path) if video_path else None,
            gs_video,
            log[-4000:],
        )


# ──────────────────────────────────────────────
# UI
# ──────────────────────────────────────────────

CSS = """
/* ── Global reset & fonts ── */
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;600;700;800&family=DM+Mono:wght@300;400;500&display=swap');

:root {
    --bg:       #0a0c10;
    --surface:  #111318;
    --border:   #1e2230;
    --accent:   #5affb0;
    --accent2:  #a78bfa;
    --text:     #e8eaf0;
    --muted:    #5a5f72;
    --radius:   12px;
    --font-head: 'Syne', sans-serif;
    --font-mono: 'DM Mono', monospace;
}

body, .gradio-container {
    background: var(--bg) !important;
    color: var(--text) !important;
    font-family: var(--font-head) !important;
}

/* Header banner */
#header {
    background: linear-gradient(135deg, #0d1117 0%, #161b27 60%, #0f1520 100%);
    border: 1px solid var(--border);
    border-radius: var(--radius);
    padding: 32px 40px 28px;
    margin-bottom: 24px;
    position: relative;
    overflow: hidden;
}
#header::before {
    content: '';
    position: absolute;
    inset: 0;
    background: radial-gradient(ellipse 70% 60% at 80% 50%, rgba(94,255,176,0.06) 0%, transparent 70%),
                radial-gradient(ellipse 50% 80% at 20% 80%, rgba(167,139,250,0.06) 0%, transparent 70%);
    pointer-events: none;
}
#header h1 {
    font-size: 2.4rem;
    font-weight: 800;
    letter-spacing: -0.02em;
    margin: 0 0 8px;
    background: linear-gradient(90deg, var(--accent) 0%, var(--accent2) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
#header p {
    color: var(--muted);
    font-family: var(--font-mono);
    font-size: 0.85rem;
    margin: 0;
    letter-spacing: 0.02em;
}
#header .badge {
    display: inline-block;
    margin-right: 8px;
    padding: 3px 10px;
    background: rgba(94,255,176,0.1);
    border: 1px solid rgba(94,255,176,0.25);
    border-radius: 20px;
    color: var(--accent);
    font-size: 0.75rem;
    font-family: var(--font-mono);
}

/* Tabs */
.tab-nav button {
    background: transparent !important;
    border: none !important;
    border-bottom: 2px solid transparent !important;
    color: var(--muted) !important;
    font-family: var(--font-head) !important;
    font-weight: 600 !important;
    font-size: 0.95rem !important;
    padding: 10px 20px !important;
    transition: all .2s !important;
}
.tab-nav button.selected, .tab-nav button:hover {
    color: var(--accent) !important;
    border-bottom-color: var(--accent) !important;
    background: transparent !important;
}

/* Panels / blocks */
.gr-panel, .gr-box, .gradio-group {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius) !important;
}

/* Inputs */
input, textarea, .gr-input, .gr-textbox textarea {
    background: #0d0f14 !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
    font-family: var(--font-mono) !important;
    border-radius: 8px !important;
}
input:focus, textarea:focus {
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 2px rgba(94,255,176,0.12) !important;
}

/* Sliders */
input[type=range] { accent-color: var(--accent) !important; }

/* Buttons */
button.primary, .gr-button-primary {
    background: linear-gradient(135deg, var(--accent) 0%, #38d9a9 100%) !important;
    color: #0a0c10 !important;
    font-family: var(--font-head) !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 12px 28px !important;
    font-size: 0.95rem !important;
    letter-spacing: 0.01em !important;
    transition: opacity .2s !important;
}
button.primary:hover { opacity: 0.85 !important; }

button.secondary, .gr-button-secondary {
    background: transparent !important;
    border: 1px solid var(--border) !important;
    color: var(--muted) !important;
    font-family: var(--font-head) !important;
    border-radius: 8px !important;
}

/* Labels */
label, .gr-form > label, .block > label span {
    color: var(--muted) !important;
    font-family: var(--font-mono) !important;
    font-size: 0.8rem !important;
    letter-spacing: 0.04em !important;
    text-transform: uppercase !important;
}

/* Log box */
#log-box textarea {
    font-size: 0.78rem !important;
    color: #7af0b0 !important;
    background: #060709 !important;
}

/* Accordion */
.gr-accordion { border-color: var(--border) !important; }

/* Info note */
.info-note {
    background: rgba(167,139,250,0.07);
    border: 1px solid rgba(167,139,250,0.2);
    border-radius: 8px;
    padding: 12px 16px;
    font-family: var(--font-mono);
    font-size: 0.8rem;
    color: #c4b5fd;
    line-height: 1.6;
}
"""


def build_app():
    with gr.Blocks() as demo:

        # ── Header ──────────────────────────────────
        gr.HTML("""
        <div id="header">
          <h1>✦ Lyra 2.0</h1>
          <p>
            <span class="badge">NVIDIA Research</span>
            <span class="badge">3D Gaussian Splatting</span>
            <span class="badge">arXiv 2604.13036</span>
          </p>
          <p style="margin-top:14px; color:#8892a4; font-size:0.9rem; font-family:'Syne',sans-serif;">
            Generate persistent, explorable 3D worlds from a single image.
            Walk through scenes, revisit areas — no spatial forgetting, no temporal drift.
          </p>
        </div>
        """)

        # ── Tabs ─────────────────────────────────────
        with gr.Tabs():

            # ════════════════════════════════════════
            # TAB 1 — Zoom-in / Zoom-out
            # ════════════════════════════════════════
            with gr.Tab("🔭 Zoom Trajectory"):
                gr.HTML('<div class="info-note">Generate a zoom-in → zoom-out exploration video from a single image, then optionally lift it to a 3D Gaussian Splatting scene.</div>')

                with gr.Row():
                    with gr.Column(scale=1):
                        z_image = gr.Image(label="Input Image", type="numpy", height=280)
                        z_caption = gr.Textbox(
                            label="Scene Caption",
                            placeholder="A sunlit forest clearing with tall pine trees…",
                            lines=2,
                        )
                        with gr.Accordion("Advanced Options", open=False):
                            with gr.Row():
                                z_in_str   = gr.Slider(0.1, 3.0, value=0.5,  step=0.1, label="Zoom-in Strength")
                                z_out_str  = gr.Slider(0.1, 3.0, value=1.5,  step=0.1, label="Zoom-out Strength")
                            with gr.Row():
                                z_frames_in  = gr.Slider(81, 401, value=81,  step=80, label="Frames Zoom-in  (1+80k)")
                                z_frames_out = gr.Slider(81, 401, value=241, step=80, label="Frames Zoom-out (1+80k)")
                            with gr.Row():
                                z_dmd    = gr.Checkbox(label="⚡ Fast Mode (DMD ×15 speedup, lower quality)", value=False)
                                z_recon  = gr.Checkbox(label="🧊 Run 3DGS Reconstruction after video", value=True)
                        z_btn = gr.Button("Generate World", variant="primary")

                    with gr.Column(scale=1):
                        z_video  = gr.Video(label="Generated Exploration Video", height=280)
                        z_gs_vid = gr.Video(label="3DGS Flythrough (if reconstruction enabled)", height=280)
                        z_log    = gr.Textbox(label="Log", lines=6, interactive=False, elem_id="log-box")

                z_btn.click(
                    fn=run_zoomgs,
                    inputs=[z_image, z_caption, gr.State(0),
                            z_in_str, z_out_str, z_frames_in, z_frames_out,
                            z_dmd, z_recon],
                    outputs=[z_video, z_gs_vid, z_log],
                )

            # ════════════════════════════════════════
            # TAB 2 — Custom Trajectory
            # ════════════════════════════════════════
            with gr.Tab("🎮 Custom Trajectory"):
                gr.HTML('<div class="info-note">Provide your own camera trajectory (.npz with <code>w2c</code>, <code>intrinsics</code>, <code>image_height</code>, <code>image_width</code>) and per-chunk captions (JSON keyed by frame index, e.g. <code>{"0": "…", "81": "…"}</code>).</div>')

                with gr.Row():
                    with gr.Column(scale=1):
                        c_image   = gr.Image(label="First Frame", type="numpy", height=240)
                        c_traj    = gr.File(label="Trajectory (.npz)", file_types=[".npz"])
                        c_captions = gr.Textbox(
                            label='Per-chunk Captions (JSON or single string)',
                            placeholder='{"0": "A grand hall interior", "81": "Corridor leading outside"}',
                            lines=3,
                        )
                        with gr.Accordion("Advanced Options", open=False):
                            with gr.Row():
                                c_frames     = gr.Slider(81, 961, value=481, step=80, label="Num Frames (1+80k)")
                                c_pose_scale = gr.Slider(0.1, 5.0, value=1.0, step=0.1, label="Pose Scale")
                            with gr.Row():
                                c_dmd   = gr.Checkbox(label="⚡ Fast Mode (DMD)", value=False)
                                c_recon = gr.Checkbox(label="🧊 Run 3DGS Reconstruction", value=True)
                        c_btn = gr.Button("Generate World", variant="primary")

                    with gr.Column(scale=1):
                        c_video  = gr.Video(label="Generated Video", height=260)
                        c_gs_vid = gr.Video(label="3DGS Flythrough", height=260)
                        c_log    = gr.Textbox(label="Log", lines=6, interactive=False, elem_id="log-box")

                c_btn.click(
                    fn=run_custom_traj,
                    inputs=[c_image, c_traj, c_captions,
                            c_frames, c_pose_scale, c_dmd, c_recon],
                    outputs=[c_video, c_gs_vid, c_log],
                )

            # ════════════════════════════════════════
            # TAB 3 — Model Info
            # ════════════════════════════════════════
            with gr.Tab("ℹ️ About"):
                gr.Markdown("""
## Lyra 2.0 — Explorable Generative 3D Worlds

**NVIDIA Research** · [Paper](https://arxiv.org/abs/2604.13036) · [Project Page](https://research.nvidia.com/labs/sil/projects/lyra2/) · [HuggingFace](https://huggingface.co/nvidia/Lyra-2.0)

### How it works

Lyra 2.0 solves two fundamental failure modes of long-horizon 3D world generation:

| Problem | Solution |
|---|---|
| **Spatial Forgetting** — previously seen regions fall out of context and are hallucinated on revisit | Per-frame 3D geometry used for information routing — retrieve past frames and establish dense correspondences |
| **Temporal Drifting** — autoregressive errors accumulate and distort appearance/geometry | Self-augmented training histories expose the model to its own degraded outputs, teaching correction not propagation |

The generated video is then lifted to a **3D Gaussian Splatting** scene via VIPE pose estimation + Depth Anything 3 depth.

### GPU Requirements

- Recommended: **H100 80 GB** (or A100 80 GB)  
- ~9 min per 80 frames at full quality · ~35 s with `--use_dmd` (DMD fast mode)
- GS reconstruction adds ~1 min on top

### Checkpoint Setup

Checkpoints are expected at `./checkpoints/model/`.  
Download from HuggingFace:

```bash
huggingface-cli download nvidia/Lyra-2.0 \\
  --include "checkpoints/*" \\
  --local-dir .
```

### Citation

```bibtex
@article{shen2026lyra2,
    title={Lyra 2.0: Explorable Generative 3D Worlds},
    author={Shen, Tianchang and Bahmani, Sherwin and He, Kai and ...},
    journal={arXiv preprint arXiv:2604.13036},
    year={2026}
}
```

*Model weights released under [NVIDIA Internal Scientific Research and Development Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-internal-scientific-research-and-development-model-license/).*
                """)

    return demo


if __name__ == "__main__":
    demo = build_app()
    demo.launch(css=CSS)