import os
os.environ.setdefault("HF_HOME", "/tmp/hf_cache")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import spaces
import gradio as gr
import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
import io
import time

from terrain_diffusion.inference.world_pipeline import WorldPipeline
from terrain_diffusion.inference.relief_map import get_relief_map

MODEL_ID = "xandergos/terrain-diffusion-30m"

# Global pipeline — loaded once at startup
_world = None


def _get_world():
    global _world
    if _world is not None:
        return _world
    print("Loading terrain diffusion model...")
    t0 = time.time()
    _world = WorldPipeline.from_pretrained(
        MODEL_ID,
        seed=42,
        latents_batch_size=[1, 2, 4, 8, 16],
        log_mode="info",
        torch_compile=False,
        dtype=None,
        caching_strategy="direct",
        cache_limit=100 * 1024 * 1024,
    )
    _world.to("cpu")
    _world.bind(hdf5_file=None)
    print(f"Model loaded in {time.time() - t0:.1f}s, seed={_world.seed}")
    return _world


@spaces.GPU(duration=1)
def _zerogpu_probe():
    return "ready"


# ─────────────────────────────────────────────────────────────────────────────
# Preset conditioning maps
# Each preset is a small (cells × cells) elevation conditioning grid in meters.
# 1 cell = 256 pixels, so a 4×4 grid generates a 1024×1024 region.
# PADDING adds border cells so the model has context.
# ─────────────────────────────────────────────────────────────────────────────

PADDING = 4
GRID_SIZE = 4
TOTAL = GRID_SIZE + 2 * PADDING


def _empty_grid(default_val=-1000.0):
    return np.full((TOTAL, TOTAL), default_val, dtype=np.float32)


def _mountains_grid():
    g = _empty_grid()
    cx = TOTAL // 2
    for i in range(TOTAL):
        for j in range(TOTAL):
            di = (i - cx) / (GRID_SIZE / 2)
            dj = (j - cx) / (GRID_SIZE / 2)
            dist = np.sqrt(di**2 + dj**2)
            if dist < 1.0:
                g[i, j] = 2500.0 * (1.0 - dist)
            elif dist < 1.5:
                g[i, j] = 500.0 * (1.5 - dist) / 0.5
    return g


def _plains_grid():
    g = _empty_grid()
    cx = TOTAL // 2
    for i in range(PADDING, PADDING + GRID_SIZE):
        for j in range(PADDING, PADDING + GRID_SIZE):
            di = (i - cx) / (GRID_SIZE / 2)
            dj = (j - cx) / (GRID_SIZE / 2)
            g[i, j] = 200.0 + 150.0 * np.sin(di * 3) * np.cos(dj * 3)
    return g


def _islands_grid():
    g = _empty_grid(-2000.0)
    cx = TOTAL // 2
    for i in range(TOTAL):
        for j in range(TOTAL):
            di = (i - cx) / (GRID_SIZE / 2)
            dj = (j - cx) / (GRID_SIZE / 2)
            dist = np.sqrt(di**2 + dj**2)
            if dist < 0.5:
                g[i, j] = 800.0 * (1.0 - dist / 0.5)
            elif dist < 0.8:
                g[i, j] = 100.0 * (0.8 - dist) / 0.3
    return g


def _canyon_grid():
    g = _empty_grid()
    cx = TOTAL // 2
    for i in range(TOTAL):
        for j in range(PADDING, PADDING + GRID_SIZE):
            dist_from_center = abs(j - cx) / (GRID_SIZE / 2)
            if dist_from_center < 0.3:
                g[i, j] = -300.0
            elif dist_from_center < 0.5:
                g[i, j] = 1500.0
            else:
                g[i, j] = 800.0
    return g


PRESETS = {
    "Mountains": _mountains_grid,
    "Plains": _plains_grid,
    "Islands": _islands_grid,
    "Canyon": _canyon_grid,
}


def grid_to_display(grid):
    """Convert conditioning grid to display image."""
    display = np.kron(grid, np.ones((16, 16)))
    vmin, vmax = float(display.min()), float(display.max())
    if vmax == vmin:
        vmax = vmin + 1
    norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)
    rgba = plt.get_cmap("terrain")(norm(display))
    return (np.clip(rgba, 0, 1) * 255).astype(np.uint8)


def preset_to_display(preset_name):
    """Return a visualization of the preset conditioning map for display."""
    grid = PRESETS[preset_name]()
    return grid_to_display(grid)


def sketch_to_conditioning(sketch_img):
    """Convert a user-drawn sketch image to an elevation conditioning grid."""
    if sketch_img is None:
        return None

    if isinstance(sketch_img, dict):
        img_data = sketch_img.get("background") or sketch_img.get("composite")
        if img_data is None:
            for key in ("path", "url"):
                if sketch_img.get(key):
                    img_data = sketch_img[key]
                    break
        if isinstance(img_data, str):
            arr = np.array(Image.open(img_data))
        else:
            arr = np.array(img_data, dtype=np.float32)
    elif isinstance(sketch_img, str):
        arr = np.array(Image.open(sketch_img))
    else:
        arr = np.array(sketch_img, dtype=np.float32)

    if arr.ndim == 3:
        if arr.shape[2] == 4:
            arr = arr[:, :, :3]
        arr = arr.mean(axis=2)

    normalized = arr.astype(np.float32) / 255.0
    elev = (normalized - 0.3) * 5000.0

    pil_img = Image.fromarray(elev.astype(np.float32), mode="F")
    pil_resized = pil_img.resize((TOTAL, TOTAL), Image.BILINEAR)
    return np.array(pil_resized, dtype=np.float32)


# ─────────────────────────────────────────────────────────────────────────────
# Terrain generation
# ─────────────────────────────────────────────────────────────────────────────

@spaces.GPU(duration=120)
def generate_terrain(preset_choice, sketch_img, seed, input_mode, progress=gr.Progress(track_tqdm=True)):
    """Generate terrain from preset or user sketch."""
    world = _get_world()
    world.to("cuda")

    try:
        if seed is not None and int(seed) != world.seed:
            world.change_seed(int(seed))

        use_sketch = input_mode == "Custom Sketch"

        if use_sketch and sketch_img is not None:
            cond_elev = sketch_to_conditioning(sketch_img)
            if cond_elev is not None:
                world.set_custom_conditioning_import(0, cond_elev, 0, 0, default_value=-1000.0)
                world.set_cond_snr([0.5, 0.5, 0.5, 0.5, 0.5])
                cond_display = grid_to_display(cond_elev)
            else:
                grid = PRESETS[preset_choice]()
                world.set_custom_conditioning_import(0, grid, 0, 0, default_value=-1000.0)
                world.set_cond_snr([0.5, 0.5, 0.5, 0.5, 0.5])
                cond_display = grid_to_display(grid)
        else:
            grid = PRESETS[preset_choice]()
            world.set_custom_conditioning_import(0, grid, 0, 0, default_value=-1000.0)
            world.set_cond_snr([0.5, 0.5, 0.5, 0.5, 0.5])
            cond_display = grid_to_display(grid)

        pi1 = PADDING * 256
        pi2 = (PADDING + GRID_SIZE) * 256
        with world:
            result = world.get(pi1, pi1, pi2, pi2, with_climate=False)
            elev = result["elev"].cpu().numpy()

        # 1. Shaded relief map
        relief = get_relief_map(elev, None, None, None, resolution=30)
        relief_img = (np.clip(relief, 0, 1) * 255).astype(np.uint8)

        # 2. Elevation colormap
        vmin, vmax = float(elev.min()), float(elev.max())
        if vmax == vmin:
            vmax = vmin + 1
        norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax)
        elev_rgba = plt.get_cmap("terrain")(norm(elev))
        elev_img = (np.clip(elev_rgba, 0, 1) * 255).astype(np.uint8)

        # 3. 3D preview
        fig = plt.figure(figsize=(6, 5))
        ax = fig.add_subplot(111, projection="3d")
        h, w = elev.shape
        downsample = 4
        X, Y = np.meshgrid(
            np.arange(0, w, downsample),
            np.arange(0, h, downsample)
        )
        Z = elev[::downsample, ::downsample]
        ax.plot_surface(X, Y, Z, cmap="terrain", linewidth=0, antialiased=True)
        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.set_zlabel("Elevation (m)")
        ax.set_title(f"3D Terrain Preview\nRange: {elev.min():.0f}m to {elev.max():.0f}m")
        buf = io.BytesIO()
        fig.savefig(buf, format="png", dpi=120, bbox_inches="tight")
        plt.close(fig)
        buf.seek(0)
        preview_3d = np.array(Image.open(buf))

        # 4. Heightmap PNG (16-bit)
        elev_clipped = np.clip(elev, 0, 65535).astype(np.uint16)
        heightmap_path = "/tmp/heightmap.png"
        Image.fromarray(elev_clipped).save(heightmap_path)

        stats = f"Elevation range: {elev.min():.1f}m to {elev.max():.1f}m | Mean: {elev.mean():.1f}m"

        return (
            relief_img,
            elev_img,
            preview_3d,
            heightmap_path,
            cond_display,
            stats,
        )
    finally:
        world.to("cpu")


# ─────────────────────────────────────────────────────────────────────────────
# Gradio UI
# ─────────────────────────────────────────────────────────────────────────────

CSS = """
.gradio-container {max-width: 1200px !important;}
"""

with gr.Blocks(title="Terrain Diffusion Demo") as demo:
    gr.Markdown("# 🏔️ Terrain Diffusion Demo")
    gr.Markdown(
        "Generate realistic terrain heightmaps using the "
        "[terrain-diffusion-30m](https://huggingface.co/xandergos/terrain-diffusion-30m) model. "
        "Based on the paper *Terrain Diffusion: A Diffusion-Based Successor to Perlin Noise "
        "in Infinite, Real-Time Terrain Generation* ([arxiv:2512.08309](https://arxiv.org/abs/2512.08309))."
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Choose Terrain Type")
            input_mode = gr.Radio(
                ["Presets", "Custom Sketch"],
                value="Presets",
                label="Input Mode",
            )

            with gr.Group(visible=True) as preset_group:
                preset = gr.Radio(
                    list(PRESETS.keys()),
                    value="Mountains",
                    label="Preset Terrain",
                )
                preset_display = gr.Image(
                    label="Preset Preview",
                    value=preset_to_display("Mountains"),
                    interactive=False,
                    height=200,
                )

            with gr.Group(visible=False) as sketch_group:
                sketch = gr.Sketchpad(
                    label="Draw Terrain Sketch",
                    brush=gr.Brush(default_size=40, colors=["#ffffff", "#000000"]),
                    layers=False,
                    height=200,
                    width=200,
                )
                gr.Markdown(
                    "💡 Draw white for high elevation, black for low/ocean. "
                    "The sketch guides the AI-generated terrain."
                )

            seed = gr.Slider(
                0, 2**31 - 1, value=42, step=1,
                label="Random Seed",
            )

            generate_btn = gr.Button("🚀 Generate Terrain", variant="primary", size="lg")

            gr.Markdown(
                "⚠️ Generation takes ~30-60 seconds on GPU. "
                "The model uses a 3-stage diffusion pipeline (coarse → base → decoder)."
            )

        with gr.Column(scale=2):
            gr.Markdown("### 2. Results")
            with gr.Tab("Shaded Relief"):
                relief_output = gr.Image(label="Shaded Relief Map", height=400)
            with gr.Tab("Elevation"):
                elev_output = gr.Image(label="Elevation Colormap", height=400)
            with gr.Tab("3D Preview"):
                preview_3d_output = gr.Image(label="3D Terrain Preview", height=400)
            with gr.Tab("Conditioning"):
                cond_output = gr.Image(label="Conditioning Map", height=300)

            stats_output = gr.Textbox(label="Statistics", interactive=False)
            download_output = gr.File(label="Download Heightmap (16-bit PNG)")

    def toggle_input(mode):
        if mode == "Presets":
            return gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=False), gr.update(visible=True)

    def update_preset_display(preset_name):
        return preset_to_display(preset_name)

    input_mode.change(
        toggle_input,
        inputs=[input_mode],
        outputs=[preset_group, sketch_group],
    )

    preset.change(
        update_preset_display,
        inputs=[preset],
        outputs=[preset_display],
    )

    generate_btn.click(
        generate_terrain,
        inputs=[preset, sketch, seed, input_mode],
        outputs=[relief_output, elev_output, preview_3d_output, download_output, cond_output, stats_output],
    )

demo.launch(css=CSS)