import os

os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import base64
import hashlib
import io
import json
import time
import logging

import numpy as np
import torch
from einops import rearrange
from PIL import Image
from sklearn.decomposition import PCA

import gradio as gr
from gradio.routes import HTMLResponse, JSONResponse

from olmoearth_pretrain_minimal import load_model_from_id, ModelID, Normalizer
from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import Modality
from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import (
    MaskedOlmoEarthSample,
    MaskValue,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Model loading (at startup, CPU only)
# ---------------------------------------------------------------------------
logger.info("Loading OlmoEarth-v1-Base model...")
_t0 = time.perf_counter()
model = load_model_from_id(ModelID.OLMOEARTH_V1_BASE)
model.eval()
encoder = model.encoder
num_params = sum(p.numel() for p in encoder.parameters())
logger.info(f"Model loaded in {time.perf_counter() - _t0:.1f}s — {num_params:,} params ({num_params / 1e6:.1f}M)")

normalizer = Normalizer(std_multiplier=2.0)
S2_SPEC = Modality.SENTINEL2_L2A
S2_BAND_ORDER = S2_SPEC.band_order  # 12 bands
NORM_CONFIG = normalizer.norm_config["sentinel2_l2a"]

PATCH_SIZE = 4
EMBED_DIM = 768
PAPER_URL = "https://allenai.org/papers/olmoearth"
MODEL_URL = "https://huggingface.co/allenai/OlmoEarth-v1-Base"
GITHUB_URL = "https://github.com/allenai/olmoearth_pretrain"
MAX_EMBED_CACHE = 12
EMBEDDING_CACHE: dict[str, np.ndarray] = {}
EMBEDDING_CACHE_ORDER: list[str] = []


# ---------------------------------------------------------------------------
# Image processing helpers
# ---------------------------------------------------------------------------
def normalize_sentinel2(data: np.ndarray) -> np.ndarray:
    """Normalize data using OlmoEarth's computed Sentinel-2 stats.

    data shape: (C, H, W) or (T, C, H, W) — last-matched band dim is normalized.
    """
    data = data.astype(np.float32).copy()
    band_means = np.array([NORM_CONFIG[b]["mean"] for b in S2_BAND_ORDER], dtype=np.float32)
    band_stds = np.array([NORM_CONFIG[b]["std"] for b in S2_BAND_ORDER], dtype=np.float32)
    min_vals = band_means - 2.0 * band_stds
    max_vals = band_means + 2.0 * band_stds
    ranges = max_vals - min_vals
    ranges[ranges == 0] = 1.0
    # Broadcast: data shape (..., C, H, W), normalize along C axis
    if data.ndim == 3:
        data = (data - min_vals[:, None, None]) / ranges[:, None, None]
    elif data.ndim == 4:
        data = (data - min_vals[None, :, None, None]) / ranges[None, :, None, None]
    return np.clip(data, 0, 1)


def rgb_to_12band(rgb: np.ndarray) -> np.ndarray:
    """Expand a 3-channel RGB array to 12-channel Sentinel-2 band order.

    rgb: (H, W, 3) uint8 or float
    returns: (12, H, W) float32 normalized
    """
    if rgb.dtype == np.uint8:
        rgb = rgb.astype(np.float32)
    h, w, c = rgb.shape
    # Map RGB -> approximate Sentinel-2 reflectance values
    # Scale 0-255 to approximate 0-3000 reflectance range
    scaled = rgb * (3000.0 / 255.0)
    bands_12 = np.zeros((12, h, w), dtype=np.float32)
    # B02=blue, B03=green, B04=red, B08=NIR ~ near red
    bands_12[0] = scaled[:, :, 2]   # B02 (blue) <- R channel (reversed for visual)
    bands_12[1] = scaled[:, :, 1]   # B03 (green)
    bands_12[2] = scaled[:, :, 0]   # B04 (red)
    bands_12[3] = scaled[:, :, 1] * 1.2  # B08 (NIR) - approx from green
    bands_12[4] = scaled[:, :, 0] * 0.9   # B05
    bands_12[5] = scaled[:, :, 1] * 0.9   # B06
    bands_12[6] = scaled[:, :, 1] * 0.8   # B07
    bands_12[7] = scaled[:, :, 1] * 1.1   # B8A
    bands_12[8] = scaled[:, :, 2] * 0.7   # B11 (SWIR)
    bands_12[9] = scaled[:, :, 2] * 0.5   # B12 (SWIR)
    bands_12[10] = scaled[:, :, 2] * 0.3  # B01 (coastal aerosol)
    bands_12[11] = scaled[:, :, 0] * 0.4  # B09 (water vapor)
    return bands_12


def extract_embeddings(img: Image.Image) -> tuple[np.ndarray, int, int]:
    """Extract per-patch embeddings from a PIL image.

    Returns (patch_embeddings, n_patch_h, n_patch_w).
    patch_embeddings shape: (n_patch_h, n_patch_w, embed_dim)
    """
    # Convert to RGB and resize to a multiple of patch_size
    img = img.convert("RGB")
    w, h = img.size

    # Target size: nearest multiple of PATCH_SIZE, capped at 224
    max_dim = 224
    scale = min(max_dim / w, max_dim / h, 1.0)
    new_w = int(round(w * scale / PATCH_SIZE)) * PATCH_SIZE
    new_h = int(round(h * scale / PATCH_SIZE)) * PATCH_SIZE
    new_w = max(new_w, PATCH_SIZE)
    new_h = max(new_h, PATCH_SIZE)
    img = img.resize((new_w, new_h), Image.BILINEAR)

    rgb = np.array(img, dtype=np.float32)  # (H, W, 3)
    bands_12 = rgb_to_12band(rgb)  # (12, H, W)
    bands_12 = normalize_sentinel2(bands_12)  # normalize per band

    # Build model input: (B, H, W, T=1, BandSets=3, Bands)
    # Sentinel-2 bandsets: [B02,B03,B04,B08], [B05,B06,B07,B8A,B11,B12], [B01,B09]
    bandset_indices = S2_SPEC.bandsets_as_indices()  # [[0,1,2,3], [4,5,6,7,8,9], [10,11]]
    bandsets_data = []
    for indices in bandset_indices:
        bandsets_data.append(bands_12[indices])  # (n_bands, H, W)

    # Stack: (BandSets, Bands, H, W) -> (H, W, T=1, BandSets, Bands)
    # But the model expects (B, H, W, T, BandSets, Bands) via the sample
    # Actually the sample format is (B, H, W, T, BandSets, Bands)
    # Wait - looking at the tutorial, the data is (B, H, W, T, BandSets) where each entry has D dim after embedding
    # Let me re-check: the input tensor is (B, H, W, T, BandSets) where each entry is the raw data
    # No - looking more carefully at the tutorial code:
    # s2_tensor = rearrange(s2_tensor, 't c h w -> 1 h w t c')
    # This gives (1, H, W, T, C=12) - all 12 bands in one dimension
    # And mask is (1, H, W, T, 3) - 3 bandsets
    # So the input is (B, H, W, T, Bands=12) and the encoder internally splits into bandsets

    s2_data = np.stack(bands_12, axis=0)  # (12, H, W)
    s2_tensor = torch.from_numpy(s2_data.astype(np.float32))  # (12, H, W)
    s2_tensor = rearrange(s2_tensor, 'c h w -> 1 h w 1 c')  # (1, H, W, T=1, 12)

    h, w = new_h, new_w
    mask = torch.full((1, h, w, 1, 3), MaskValue.ONLINE_ENCODER.value, dtype=torch.int32)
    timestamps = torch.zeros((1, 1, 3), dtype=torch.int32)
    timestamps[0, :, 1] = 0  # month = 0 (January)

    sample = MaskedOlmoEarthSample(
        sentinel2_l2a=s2_tensor,
        sentinel2_l2a_mask=mask,
        timestamps=timestamps,
    )

    with torch.no_grad():
        output = encoder(sample, fast_pass=True, patch_size=PATCH_SIZE)
        tokens = output["tokens_and_masks"].sentinel2_l2a
        # Shape: (B, P_H, P_W, T, BandSets, D)
        # Pool across T and BandSets to get (B, P_H, P_W, D)
        patch_emb = tokens.mean(dim=[3, 4])  # (1, P_H, P_W, D)
        patch_emb = patch_emb[0].float().cpu().numpy()  # (P_H, P_W, D)

    n_ph, n_pw, d = patch_emb.shape
    return patch_emb, n_ph, n_pw


def embeddings_to_rgb_pca(patch_emb: np.ndarray) -> np.ndarray:
    """PCA project embeddings to 3 dims, normalize to 0-255 RGB.

    patch_emb: (P_H, P_W, D)
    returns: (P_H, P_W, 3) uint8
    """
    ph, pw, d = patch_emb.shape
    flat = patch_emb.reshape(-1, d)  # (P_H*P_W, D)

    pca = PCA(n_components=3)
    pca.fit(flat)
    projected = pca.transform(flat)  # (P_H*P_W, 3)

    # Normalize each component independently to [0, 1]
    for i in range(3):
        col = projected[:, i]
        cmin, cmax = col.min(), col.max()
        if cmax - cmin > 1e-8:
            projected[:, i] = (col - cmin) / (cmax - cmin)
        else:
            projected[:, i] = 0.5

    rgb = (projected * 255).clip(0, 255).astype(np.uint8)
    rgb = rgb.reshape(ph, pw, 3)
    return rgb


def embeddings_to_heatmap(patch_emb: np.ndarray) -> np.ndarray:
    """Create a similarity heatmap from embeddings.

    Shows cosine similarity of each patch to the mean embedding.
    returns: (P_H, P_W, 3) uint8
    """
    ph, pw, d = patch_emb.shape
    flat = patch_emb.reshape(-1, d)
    mean_emb = flat.mean(axis=0)
    # Cosine similarity
    norms = np.linalg.norm(flat, axis=1) * np.linalg.norm(mean_emb)
    norms[norms == 0] = 1.0
    sim = (flat @ mean_emb) / norms
    sim = sim.reshape(ph, pw)

    # Map to color: low sim = dark blue, high sim = cyan/green
    sim_norm = (sim - sim.min()) / (sim.max() - sim.min() + 1e-8)
    r = (sim_norm * 50).astype(np.uint8)
    g = (sim_norm * 200 + 55).astype(np.uint8)
    b = (sim_norm * 200 + 55).astype(np.uint8)
    return np.stack([r, g, b], axis=-1)


def embeddings_to_query_heatmap(
    patch_emb: np.ndarray,
    x_norm: float,
    y_norm: float,
) -> tuple[np.ndarray, dict]:
    """Create a brightness heatmap from cosine similarity to one clicked patch."""
    ph, pw, d = patch_emb.shape
    q_col = int(np.clip(np.floor(float(x_norm) * pw), 0, pw - 1))
    q_row = int(np.clip(np.floor(float(y_norm) * ph), 0, ph - 1))

    flat = patch_emb.reshape(-1, d).astype(np.float32)
    flat_norms = np.linalg.norm(flat, axis=1, keepdims=True)
    flat_norms[flat_norms == 0] = 1.0
    flat_unit = flat / flat_norms

    query_idx = q_row * pw + q_col
    query = flat_unit[query_idx]
    sim = (flat_unit @ query).reshape(ph, pw)

    lo, hi = np.percentile(sim, [2, 98])
    if hi - lo < 1e-8:
        lo, hi = float(sim.min()), float(sim.max())
    sim_norm = np.clip((sim - lo) / (hi - lo + 1e-8), 0, 1)
    sim_norm = sim_norm ** 0.85

    # Brightness encodes similarity; color keeps the map legible on a dark UI.
    r = (12 + sim_norm * 243).astype(np.uint8)
    g = (18 + sim_norm * 226).astype(np.uint8)
    b = (34 + sim_norm * 166).astype(np.uint8)
    rgb = np.stack([r, g, b], axis=-1)

    # Mark the clicked query patch.
    for rr in range(max(0, q_row - 1), min(ph, q_row + 2)):
        rgb[rr, q_col] = [255, 205, 64]
    for cc in range(max(0, q_col - 1), min(pw, q_col + 2)):
        rgb[q_row, cc] = [255, 205, 64]

    return rgb, {
        "row": int(q_row),
        "col": int(q_col),
        "grid_h": int(ph),
        "grid_w": int(pw),
        "x_norm": float(np.clip(x_norm, 0, 1)),
        "y_norm": float(np.clip(y_norm, 0, 1)),
        "max_similarity": float(sim.max()),
        "mean_similarity": float(sim.mean()),
    }


def upsample_to_image(small: np.ndarray, target_h: int, target_w: int) -> np.ndarray:
    """Nearest-neighbor upsample a small array to target size."""
    pil = Image.fromarray(small)
    pil = pil.resize((target_w, target_h), Image.NEAREST)
    return np.array(pil)


def resize_for_display(img: Image.Image, max_display: int = 512) -> Image.Image:
    """Resize an input image to the size used by the frontend panels."""
    img_rgb = img.convert("RGB")
    w, h = img_rgb.size
    scale = min(max_display / w, max_display / h, 1.0)
    disp_w, disp_h = int(w * scale), int(h * scale)
    return img_rgb.resize((disp_w, disp_h), Image.BILINEAR)


def decode_image_b64(image_b64: str) -> tuple[bytes, Image.Image]:
    """Decode either a raw base64 payload or a data URL."""
    if "," in image_b64 and image_b64.strip().startswith("data:"):
        image_b64 = image_b64.split(",", 1)[1]
    img_data = base64.b64decode(image_b64)
    return img_data, Image.open(io.BytesIO(img_data))


def image_cache_key(img_data: bytes) -> str:
    return hashlib.sha256(img_data).hexdigest()[:20]


def remember_embeddings(cache_key: str, patch_emb: np.ndarray) -> None:
    EMBEDDING_CACHE[cache_key] = patch_emb
    if cache_key in EMBEDDING_CACHE_ORDER:
        EMBEDDING_CACHE_ORDER.remove(cache_key)
    EMBEDDING_CACHE_ORDER.append(cache_key)
    while len(EMBEDDING_CACHE_ORDER) > MAX_EMBED_CACHE:
        stale_key = EMBEDDING_CACHE_ORDER.pop(0)
        EMBEDDING_CACHE.pop(stale_key, None)


def img_to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def numpy_to_b64(arr: np.ndarray) -> str:
    img = Image.fromarray(arr)
    return img_to_b64(img)


# ---------------------------------------------------------------------------
# Real satellite imagery examples (pre-computed, cached)
# ---------------------------------------------------------------------------
EXAMPLES_DIR = os.path.join(os.path.dirname(__file__), "examples")


def _load_examples():
    """Load real satellite examples with pre-computed model outputs from examples/ dir."""
    manifest_path = os.path.join(EXAMPLES_DIR, "manifest.json")
    with open(manifest_path) as f:
        manifest = json.load(f)

    examples = []
    for entry in manifest:
        img_path = os.path.join(EXAMPLES_DIR, os.path.basename(entry["filename"]))
        pca_path = os.path.join(EXAMPLES_DIR, os.path.basename(entry["pca_file"]))
        heat_path = os.path.join(EXAMPLES_DIR, os.path.basename(entry["heatmap_file"]))
        img = Image.open(img_path)
        pca_img = Image.open(pca_path)
        heat_img = Image.open(heat_path)

        examples.append({
            "id": entry["id"],
            "label": entry["label"],
            "emoji": entry["emoji"],
            "description": entry["description"],
            "source": entry["source"],
            "image": img_to_b64(img),
            "pca_rgb": img_to_b64(pca_img),
            "heatmap": img_to_b64(heat_img),
            "stats": {
                "n_patches": entry["n_patches"],
                "patch_grid": entry["patch_grid"],
                "embed_dim": EMBED_DIM,
                "mean_norm": entry["mean_norm"],
                "inference_time": 0.0,
            },
            "patch_emb_shape": [56, 56, EMBED_DIM],
        })
        logger.info(f"Loaded cached example: {entry['id']} - {entry['label']}")
    return examples


CACHED_EXAMPLES = _load_examples()

# ---------------------------------------------------------------------------
# API endpoints
# ---------------------------------------------------------------------------
app = gr.Server()


@app.api(name="process_image", concurrency_limit=2, time_limit=120)
def process_image(image_b64: str) -> str:
    """Process a base64-encoded image and return visualization data as JSON string.

    Returns JSON with: input_image, pca_rgb, heatmap, embedding_stats, n_patches.
    All images are base64-encoded PNGs.
    """
    try:
        img_data, img = decode_image_b64(image_b64)
    except Exception as e:
        return json.dumps({"error": f"Failed to decode image: {e}"})

    t0 = time.perf_counter()
    patch_emb, ph, pw = extract_embeddings(img)
    inference_time = time.perf_counter() - t0
    cache_key = image_cache_key(img_data)
    remember_embeddings(cache_key, patch_emb)

    # Generate visualizations
    pca_rgb = embeddings_to_rgb_pca(patch_emb)
    heatmap = embeddings_to_heatmap(patch_emb)

    # Resize model outputs to match input image size
    img_rgb = resize_for_display(img)
    disp_w, disp_h = img_rgb.size

    pca_up = upsample_to_image(pca_rgb, disp_h, disp_w)
    heat_up = upsample_to_image(heatmap, disp_h, disp_w)

    # Compute embedding stats
    flat = patch_emb.reshape(-1, EMBED_DIM)
    stats = {
        "n_patches": int(ph * pw),
        "patch_grid": f"{ph} × {pw}",
        "embed_dim": EMBED_DIM,
        "mean_norm": float(np.linalg.norm(flat, axis=1).mean()),
        "inference_time": round(inference_time, 3),
    }

    result = {
        "input_image": numpy_to_b64(np.array(img_rgb)),
        "pca_rgb": numpy_to_b64(pca_up),
        "heatmap": numpy_to_b64(heat_up),
        "stats": stats,
        "patch_emb_shape": [ph, pw, EMBED_DIM],
        "cache_key": cache_key,
    }
    return json.dumps(result)


@app.api(name="query_similarity", concurrency_limit=2, time_limit=120)
def query_similarity(image_b64: str, x_norm: float, y_norm: float, cache_key: str = "") -> str:
    """Return a similarity heatmap using the clicked patch as the query."""
    try:
        img_data, img = decode_image_b64(image_b64)
    except Exception as e:
        return json.dumps({"error": f"Failed to decode image: {e}"})

    cache_key = cache_key or image_cache_key(img_data)
    t0 = time.perf_counter()
    patch_emb = EMBEDDING_CACHE.get(cache_key)
    cache_hit = patch_emb is not None
    if patch_emb is None:
        patch_emb, _, _ = extract_embeddings(img)
        remember_embeddings(cache_key, patch_emb)
    elapsed = time.perf_counter() - t0

    heatmap, query = embeddings_to_query_heatmap(patch_emb, x_norm, y_norm)
    img_rgb = resize_for_display(img)
    disp_w, disp_h = img_rgb.size
    heat_up = upsample_to_image(heatmap, disp_h, disp_w)

    flat = patch_emb.reshape(-1, EMBED_DIM)
    stats = {
        "n_patches": int(patch_emb.shape[0] * patch_emb.shape[1]),
        "patch_grid": f"{patch_emb.shape[0]} × {patch_emb.shape[1]}",
        "embed_dim": EMBED_DIM,
        "mean_norm": float(np.linalg.norm(flat, axis=1).mean()),
        "inference_time": round(elapsed, 3),
    }

    return json.dumps({
        "heatmap": numpy_to_b64(heat_up),
        "stats": stats,
        "query": query,
        "cache_key": cache_key,
        "cache_hit": cache_hit,
    })


@app.api(name="get_samples", concurrency_limit=1, time_limit=30)
def get_samples() -> str:
    """Return all cached sample images with pre-computed model outputs."""
    return json.dumps(CACHED_EXAMPLES)


@app.api(name="get_model_info", concurrency_limit=1, time_limit=10)
def get_model_info() -> str:
    """Return model information."""
    info = {
        "model_name": "OlmoEarth-v1-Base",
        "model_id": "allenai/OlmoEarth-v1-Base",
        "architecture": "ViT-Base (Flexible ViT)",
        "parameters": f"{num_params / 1e6:.1f}M",
        "embedding_dim": EMBED_DIM,
        "encoder_depth": 12,
        "encoder_heads": 12,
        "patch_size": PATCH_SIZE,
        "modalities": [
            "Sentinel-1",
            "Sentinel-2",
            "Landsat-8",
            "WorldCover",
            "OpenStreetMap",
            "SRTM",
            "Canopy Height",
        ],
        "description": "A spatio-temporal, multimodal Earth observation foundation model that turns aligned satellite observations into reusable patch embeddings.",
        "paper_url": PAPER_URL,
        "huggingface_url": MODEL_URL,
        "github_url": GITHUB_URL,
        "paper_context": {
            "earth_observation_foundation_models": (
                "Earth observation foundation models are reusable encoders for satellite and geospatial data. "
                "They learn patch-level representations that can transfer to land-cover mapping, crop and ecosystem classification, segmentation, and detection with limited labels."
            ),
            "why_special": [
                "Multimodal inputs: Sentinel-1 radar, Sentinel-2 optical, Landsat-8, and map layers including WorldCover, OpenStreetMap, SRTM elevation, and canopy height.",
                "Multi-temporal context: each pretraining sample covers a 2.56 km by 2.56 km region with up to 12 monthly timesteps over one year.",
                "Latent MIM Lite: a stable latent masked-image-modeling objective that predicts frozen random-projection targets instead of pixels.",
                "Modality-aware masking: the model learns to reconstruct missing bandsets or modalities rather than relying on easy nearby tokens.",
                "Strong transfer: the paper reports best performance on 15 of 24 frozen embedding tasks and 19 of 29 fine-tuning tasks against 12 other EO foundation models.",
            ],
            "pretraining_dataset": "285,288 globally sampled locations selected from OpenStreetMap categories, with observations resampled to 10 m per pixel.",
        },
    }
    return json.dumps(info)


# ---------------------------------------------------------------------------
# Custom HTML frontend
# ---------------------------------------------------------------------------
FRONTEND_HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no">
<title>OlmoEarth — Satellite Vision Explorer</title>
<style>
:root {
  --bg: #0a0a0a;
  --bg-elevated: #111111;
  --bg-glass: rgba(15, 15, 15, 0.85);
  --accent: #00e5a0;
  --accent-dim: rgba(0, 229, 160, 0.15);
  --accent-glow: rgba(0, 229, 160, 0.4);
  --text: #e8e8e8;
  --text-dim: #888888;
  --text-faint: #555555;
  --border: rgba(255, 255, 255, 0.06);
  --transition: 0.4s cubic-bezier(0.4, 0, 0.2, 1);
  --radius: 12px;
}

* { margin: 0; padding: 0; box-sizing: border-box; }

html, body {
  height: 100%;
  overflow: hidden;
  background: var(--bg);
  color: var(--text);
  font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif;
  font-size: 14px;
  line-height: 1.5;
  -webkit-font-smoothing: antialiased;
  -moz-osx-font-smoothing: grayscale;
}

#app { height: 100vh; height: 100dvh; display: flex; flex-direction: column; }

/* ===== Landing ===== */
#landing {
  position: fixed; inset: 0; z-index: 100;
  display: flex; align-items: center; justify-content: center;
  background: linear-gradient(135deg, #0a0a0a 0%, #0d1a14 50%, #0a0a0a 100%);
  transition: opacity var(--transition), visibility var(--transition);
}
#landing.hidden { opacity: 0; visibility: hidden; pointer-events: none; }

.landing-bg {
  position: absolute; inset: 0;
  background-image: url('data:image/svg+xml;utf8,<svg xmlns="http://www.w3.org/2000/svg" width="100" height="100"><rect width="100" height="100" fill="%230a0a0a"/><circle cx="20" cy="30" r="1" fill="%2300e5a0" opacity="0.3"/><circle cx="60" cy="70" r="1" fill="%2300e5a0" opacity="0.2"/><circle cx="80" cy="20" r="1" fill="%2300e5a0" opacity="0.15"/><circle cx="40" cy="80" r="0.5" fill="%23fff" opacity="0.1"/></svg>');
  background-size: 200px 200px;
  opacity: 0.5;
}

.landing-content { text-align: center; z-index: 1; padding: 20px; }

.landing-title {
  font-size: clamp(32px, 7vw, 64px);
  font-weight: 200;
  letter-spacing: -1px;
  margin-bottom: 12px;
  background: linear-gradient(135deg, #e8e8e8, var(--accent));
  -webkit-background-clip: text;
  -webkit-text-fill-color: transparent;
  background-clip: text;
}

.landing-subtitle {
  font-size: clamp(13px, 2.5vw, 16px);
  color: var(--text-dim);
  margin-bottom: 40px;
  font-weight: 300;
  letter-spacing: 2px;
  text-transform: uppercase;
}

.landing-enter {
  padding: 14px 40px;
  background: transparent;
  border: 1px solid var(--accent);
  border-radius: 100px;
  color: var(--accent);
  font-size: 15px;
  cursor: pointer;
  transition: all var(--transition);
  letter-spacing: 1px;
}
.landing-enter:hover {
  background: var(--accent);
  color: var(--bg);
  box-shadow: 0 0 30px var(--accent-glow);
}

.landing-stats {
  margin-top: 60px;
  display: flex; gap: 40px; justify-content: center;
  flex-wrap: wrap;
}
.landing-stat { text-align: center; }
.landing-stat-value { font-size: 24px; font-weight: 300; color: var(--accent); }
.landing-stat-label { font-size: 11px; color: var(--text-faint); text-transform: uppercase; letter-spacing: 1px; }

/* ===== Main App ===== */
#main-app {
  flex: 1; display: flex; flex-direction: column;
  opacity: 0; transition: opacity var(--transition);
  overflow: hidden;
}
#main-app.visible { opacity: 1; }

/* Header */
.app-header {
  display: flex; align-items: center; justify-content: space-between;
  padding: 12px 20px;
  border-bottom: 1px solid var(--border);
  background: var(--bg-glass);
  backdrop-filter: blur(20px);
  z-index: 10;
  flex-shrink: 0;
}
.header-left { display: flex; align-items: center; gap: 12px; }
.header-logo {
  width: 28px; height: 28px;
  border-radius: 50%;
  background: linear-gradient(135deg, var(--accent), #00b380);
  display: flex; align-items: center; justify-content: center;
  font-size: 14px; color: var(--bg); font-weight: 600;
}
.header-title { font-size: 16px; font-weight: 400; letter-spacing: 0.5px; }
.header-title .accent { color: var(--accent); }

.header-right { display: flex; gap: 8px; }
.icon-btn {
  width: 36px; height: 36px;
  border: 1px solid var(--border);
  border-radius: var(--radius);
  background: transparent;
  color: var(--text-dim);
  cursor: pointer;
  display: flex; align-items: center; justify-content: center;
  transition: all var(--transition);
  font-size: 16px;
}
.icon-btn:hover { color: var(--accent); border-color: var(--accent); }
.icon-btn.active { color: var(--accent); border-color: var(--accent); background: var(--accent-dim); }

/* Content area */
.content {
  flex: 1 1 auto;
  display: flex;
  overflow: hidden;
  width: 100%;
  min-width: 0;
  min-height: 0;
}

/* Sidebar */
.sidebar {
  width: 250px;
  flex: 0 0 250px;
  border-right: 1px solid var(--border);
  overflow-y: auto;
  padding: 16px;
  transition: width var(--transition), opacity var(--transition);
}
.sidebar.collapsed { width: 0; flex-basis: 0; opacity: 0; padding: 0; overflow: hidden; }

.sidebar-section { margin-bottom: 24px; }
.sidebar-label {
  font-size: 11px; text-transform: uppercase; letter-spacing: 1.5px;
  color: var(--text-faint); margin-bottom: 10px; font-weight: 500;
}

/* Sample gallery */
.sample-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 8px; }
.sample-card {
  border: 1px solid var(--border);
  border-radius: var(--radius);
  overflow: hidden;
  cursor: pointer;
  transition: all var(--transition);
  background: var(--bg-elevated);
}
.sample-card:hover { border-color: var(--accent); transform: translateY(-2px); }
.sample-card.active { border-color: var(--accent); box-shadow: 0 0 12px var(--accent-dim); }
.sample-card img { width: 100%; height: 80px; object-fit: cover; display: block; }
.sample-card-label {
  padding: 6px 8px; font-size: 11px; color: var(--text-dim);
  display: flex; align-items: center; gap: 4px;
  font-weight: 500;
}
.sample-card-desc {
  padding: 0 8px 6px; font-size: 9px; color: var(--text-faint);
  line-height: 1.3;
}

/* Upload */
.upload-zone {
  border: 2px dashed var(--border);
  border-radius: var(--radius);
  padding: 20px;
  text-align: center;
  cursor: pointer;
  transition: all var(--transition);
}
.upload-zone:hover { border-color: var(--accent); background: var(--accent-dim); }
.upload-zone.dragover { border-color: var(--accent); background: var(--accent-dim); transform: scale(1.02); }
.upload-zone-icon { font-size: 28px; margin-bottom: 8px; opacity: 0.5; }
.upload-zone-text { font-size: 12px; color: var(--text-dim); }
.upload-zone input { display: none; }

/* Main view */
.main-view {
  flex: 1 1 auto;
  display: flex;
  flex-direction: column;
  overflow: hidden;
  width: 100%;
  min-width: 0;
  min-height: 0;
}

/* Split view */
.split-view {
  flex: 1 1 auto;
  display: grid;
  grid-template-columns: minmax(0, 1fr) minmax(0, 1fr);
  gap: 1px;
  background: var(--border);
  overflow: hidden;
  width: 100%;
  min-width: 0;
  min-height: 0;
}
.panel {
  display: flex; flex-direction: column;
  background: var(--bg);
  overflow: hidden;
  position: relative;
  min-width: 0;
  min-height: 0;
}
.panel-header {
  padding: 10px 16px;
  font-size: 11px; text-transform: uppercase; letter-spacing: 1.5px;
  color: var(--text-faint);
  border-bottom: 1px solid var(--border);
  display: flex; align-items: center; justify-content: space-between;
  flex-shrink: 0;
  gap: 12px;
}
.panel-help {
  width: 18px; height: 18px;
  border: 1px solid var(--border);
  border-radius: 50%;
  display: inline-flex;
  align-items: center;
  justify-content: center;
  color: var(--accent);
  font-size: 11px;
  letter-spacing: 0;
  cursor: help;
  flex: 0 0 auto;
}
.input-hint {
  color: var(--accent);
  font-size: 10px;
  letter-spacing: 1px;
  text-align: right;
  text-transform: uppercase;
}
.panel-content {
  flex: 1 1 auto;
  display: flex; align-items: center; justify-content: center;
  overflow: hidden;
  padding: 24px;
  position: relative;
  width: 100%;
  min-width: 0;
  min-height: 0;
}
.image-frame {
  width: 100%;
  height: 100%;
  position: relative;
  display: flex;
  align-items: center;
  justify-content: center;
  min-width: 0;
  min-height: 0;
}
.panel-image {
  width: 100%;
  height: 100%;
  max-width: 100%; max-height: 100%;
  border-radius: 8px;
  object-fit: contain;
  display: none;
}
.panel-image.loaded { display: block; }
.panel-image.clickable { cursor: crosshair; }
.click-marker {
  position: absolute;
  width: 18px;
  height: 18px;
  transform: translate(-50%, -50%);
  border: 2px solid #ffd166;
  border-radius: 50%;
  box-shadow: 0 0 0 2px rgba(0, 0, 0, 0.65), 0 0 16px rgba(255, 209, 102, 0.85);
  pointer-events: none;
  display: none;
}
.click-marker.visible { display: block; }
.panel-empty {
  color: var(--text-faint); font-size: 13px; text-align: center;
  display: block;
}
.panel-empty.hidden { display: none; }

/* Loading */
.loading-overlay {
  position: absolute; inset: 0;
  background: rgba(10, 10, 10, 0.8);
  display: flex; align-items: center; justify-content: center;
  z-index: 5;
  display: none;
}
.loading-overlay.active { display: flex; }
.loading-spinner {
  width: 40px; height: 40px;
  border: 2px solid var(--border);
  border-top-color: var(--accent);
  border-radius: 50%;
  animation: spin 0.8s linear infinite;
}
@keyframes spin { to { transform: rotate(360deg); } }

/* View toggle */
.view-toggle {
  display: flex; gap: 4px;
  padding: 4px;
  background: var(--bg-elevated);
  border-radius: var(--radius);
  margin: 12px 16px;
  align-self: flex-start;
}
.toggle-btn {
  padding: 6px 16px;
  border: none;
  border-radius: 8px;
  background: transparent;
  color: var(--text-dim);
  font-size: 12px;
  cursor: pointer;
  transition: all var(--transition);
}
.toggle-btn.active { background: var(--accent); color: var(--bg); }
.toggle-btn:hover:not(.active) { color: var(--text); }

/* Stats bar */
.stats-bar {
  padding: 10px 16px;
  border-top: 1px solid var(--border);
  display: flex; gap: 20px;
  flex-shrink: 0;
  overflow-x: auto;
  background: var(--bg-elevated);
}
.stat-item { display: flex; flex-direction: column; gap: 2px; }
.stat-label { font-size: 10px; color: var(--text-faint); text-transform: uppercase; letter-spacing: 1px; }
.stat-value { font-size: 13px; color: var(--accent); font-weight: 400; }

/* Info overlay */
.info-overlay {
  position: fixed; top: 0; right: 0; bottom: 0;
  width: 360px; max-width: 100vw;
  background: var(--bg-glass);
  backdrop-filter: blur(20px);
  border-left: 1px solid var(--border);
  padding: 24px;
  overflow-y: auto;
  transform: translateX(100%);
  transition: transform var(--transition);
  z-index: 50;
}
.info-overlay.open { transform: translateX(0); }
.info-title { font-size: 20px; font-weight: 300; margin-bottom: 16px; }
.info-section { margin-bottom: 20px; }
.info-section-title { font-size: 11px; text-transform: uppercase; letter-spacing: 1.5px; color: var(--text-faint); margin-bottom: 8px; }
.info-text { font-size: 13px; color: var(--text-dim); line-height: 1.6; }
.info-text ul { margin-left: 18px; display: flex; flex-direction: column; gap: 8px; }
.info-text strong { color: var(--text); font-weight: 500; }
.info-link { color: var(--accent); text-decoration: none; font-size: 13px; }
.info-link:hover { text-decoration: underline; }
.info-close {
  position: absolute; top: 16px; right: 16px;
  width: 32px; height: 32px;
  border: 1px solid var(--border); border-radius: 50%;
  background: transparent; color: var(--text-dim); cursor: pointer;
  display: flex; align-items: center; justify-content: center;
}

.info-backdrop {
  position: fixed; inset: 0; background: rgba(0,0,0,0.5);
  z-index: 49; display: none;
}
.info-backdrop.active { display: block; }

/* Comparison slider */
.compare-stage {
  position: relative;
  display: none;
  width: 100%;
  height: 100%;
  max-width: 100%;
  max-height: 100%;
  overflow: hidden;
  border-radius: 8px;
  background: #050505;
}
.compare-stage.visible { display: block; }
.compare-stage-image {
  position: absolute;
  inset: 0;
  width: 100%;
  height: 100%;
  max-width: 100%;
  max-height: 100%;
  object-fit: contain;
  border-radius: 8px;
}
.compare-base { z-index: 1; }
.compare-overlay {
  z-index: 2;
  clip-path: inset(0 50% 0 0);
}
.compare-divider {
  position: absolute;
  top: 0;
  bottom: 0;
  left: 50%;
  width: 2px;
  transform: translateX(-1px);
  background: rgba(255, 255, 255, 0.9);
  box-shadow: 0 0 12px rgba(0, 0, 0, 0.65);
  z-index: 3;
  pointer-events: none;
}
.compare-divider span {
  position: absolute;
  top: 50%;
  left: 50%;
  width: 34px;
  height: 34px;
  transform: translate(-50%, -50%);
  border: 2px solid rgba(255, 255, 255, 0.95);
  border-radius: 50%;
  background: rgba(10, 10, 10, 0.75);
  box-shadow: 0 0 18px var(--accent-glow);
}
.compare-slider {
  -webkit-appearance: none;
  position: absolute;
  left: 18px;
  right: 18px;
  bottom: 18px;
  z-index: 5;
  width: calc(100% - 36px);
  height: 8px;
  margin: 0;
  display: none;
  border: 1px solid rgba(255, 255, 255, 0.22);
  border-radius: 999px;
  background: linear-gradient(90deg, var(--accent) var(--compare-position, 50%), rgba(255,255,255,0.26) var(--compare-position, 50%));
  box-shadow: 0 8px 22px rgba(0, 0, 0, 0.5);
  cursor: ew-resize;
  touch-action: none;
}
.compare-slider.visible { display: block; }
.compare-slider::-webkit-slider-thumb {
  -webkit-appearance: none;
  width: 28px; height: 28px;
  border: 2px solid var(--bg);
  border-radius: 50%;
  background: var(--accent);
  cursor: ew-resize;
  box-shadow: 0 0 12px var(--accent-glow);
}
.compare-slider::-moz-range-thumb {
  width: 28px; height: 28px;
  border: 2px solid var(--bg);
  border-radius: 50%;
  background: var(--accent);
  cursor: ew-resize;
  box-shadow: 0 0 12px var(--accent-glow);
}
.compare-slider::-moz-range-track {
  height: 8px;
  border-radius: 999px;
  background: transparent;
}

/* Mobile */
@media (max-width: 768px) {
  .sidebar {
    width: 100%;
    position: fixed; bottom: 0; left: 0; right: 0;
    border-right: none; border-top: 1px solid var(--border);
    max-height: 50vh;
    z-index: 20;
  }
  .sidebar.collapsed { max-height: 0; }
  .split-view {
    grid-template-columns: 1fr;
    grid-template-rows: minmax(0, 1fr) minmax(0, 1fr);
  }
  .panel-content { padding: 16px; }
  .info-overlay { width: 100vw; }
  .stats-bar { flex-wrap: wrap; }
  .landing-stats { gap: 20px; }
}

/* Scrollbar */
::-webkit-scrollbar { width: 4px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: var(--border); border-radius: 2px; }
::-webkit-scrollbar-thumb:hover { background: var(--text-faint); }
</style>
</head>
<body>
<div id="app">

  <!-- Landing -->
  <div id="landing">
    <div class="landing-bg"></div>
    <div class="landing-content">
      <h1 class="landing-title">OlmoEarth</h1>
      <p class="landing-subtitle">Satellite Vision Explorer</p>
      <button class="landing-enter" onclick="enterApp()">Enter</button>
      <div class="landing-stats">
        <div class="landing-stat"><div class="landing-stat-value">89M</div><div class="landing-stat-label">Params</div></div>
        <div class="landing-stat"><div class="landing-stat-value">768</div><div class="landing-stat-label">Embed Dim</div></div>
        <div class="landing-stat"><div class="landing-stat-value">ViT-B</div><div class="landing-stat-label">Architecture</div></div>
        <div class="landing-stat"><div class="landing-stat-value">S1·S2·LS</div><div class="landing-stat-label">Modalities</div></div>
      </div>
    </div>
  </div>

  <!-- Main App -->
  <div id="main-app">
    <header class="app-header">
      <div class="header-left">
        <div class="header-logo">E</div>
        <div class="header-title">Olmo<span class="accent">Earth</span></div>
      </div>
      <div class="header-right">
        <button class="icon-btn" id="view-toggle-btn" onclick="toggleView()" title="Toggle view">⇄</button>
        <button class="icon-btn" id="sidebar-toggle-btn" onclick="toggleSidebar()" title="Toggle panel">☰</button>
        <button class="icon-btn" id="info-btn" onclick="toggleInfo()" title="Info">ⓘ</button>
      </div>
    </header>

    <div class="content">
      <aside class="sidebar" id="sidebar">
        <div class="sidebar-section">
          <div class="sidebar-label">Sample Imagery</div>
          <div class="sample-grid" id="sample-grid"></div>
        </div>
        <div class="sidebar-section">
          <div class="sidebar-label">Upload Your Own</div>
          <div class="upload-zone" id="upload-zone" onclick="document.getElementById('file-input').click()">
            <div class="upload-zone-icon">📎</div>
            <div class="upload-zone-text">Drop image or click to upload</div>
            <input type="file" id="file-input" accept="image/*">
          </div>
        </div>
      </aside>

      <div class="main-view">
        <div class="view-toggle">
          <button class="toggle-btn active" data-mode="pca" onclick="setMode('pca')">PCA RGB</button>
          <button class="toggle-btn" data-mode="heatmap" onclick="setMode('heatmap')">Similarity</button>
          <button class="toggle-btn" data-mode="compare" onclick="setMode('compare')">Compare</button>
        </div>

        <div class="split-view" id="split-view">
          <div class="panel">
            <div class="panel-header">
              <span>Satellite Input</span>
              <span class="input-hint">Click to find similar regions</span>
            </div>
            <div class="panel-content">
              <div class="panel-empty" id="input-empty">Select a sample or upload an image</div>
              <div class="image-frame" id="input-frame">
                <img class="panel-image clickable" id="input-image" alt="Input">
                <div class="click-marker" id="click-marker"></div>
              </div>
            </div>
          </div>
          <div class="panel" id="output-panel">
            <div class="panel-header">
              <span id="output-label">Model Output</span>
              <span class="panel-help" id="output-tooltip" title="">?</span>
            </div>
            <div class="panel-content">
              <div class="panel-empty" id="output-empty">Embedding visualization will appear here</div>
              <img class="panel-image" id="output-image" alt="Output">
              <div class="compare-stage" id="compare-stage" aria-hidden="true">
                <img class="compare-stage-image compare-base" id="compare-base-image" alt="PCA output">
                <img class="compare-stage-image compare-overlay" id="compare-image" alt="Satellite input overlay">
                <div class="compare-divider" id="compare-divider"><span></span></div>
                <input type="range" class="compare-slider" id="compare-slider" min="0" max="100" value="50" aria-label="Compare input and PCA output">
              </div>
              <div class="loading-overlay" id="loading"><div class="loading-spinner"></div></div>
            </div>
          </div>
        </div>

        <div class="stats-bar" id="stats-bar" style="display:none;">
          <div class="stat-item"><div class="stat-label">Patches</div><div class="stat-value" id="stat-patches">—</div></div>
          <div class="stat-item"><div class="stat-label">Grid</div><div class="stat-value" id="stat-grid">—</div></div>
          <div class="stat-item"><div class="stat-label">Dim</div><div class="stat-value" id="stat-dim">—</div></div>
          <div class="stat-item"><div class="stat-label">Norm</div><div class="stat-value" id="stat-norm">—</div></div>
          <div class="stat-item"><div class="stat-label">Inference</div><div class="stat-value" id="stat-time">—</div></div>
        </div>
      </div>
    </div>
  </div>

  <!-- Info overlay -->
  <div class="info-backdrop" id="info-backdrop" onclick="toggleInfo()"></div>
  <div class="info-overlay" id="info-overlay">
    <button class="info-close" onclick="toggleInfo()">✕</button>
    <h2 class="info-title">About OlmoEarth</h2>
    <div id="info-content"></div>
  </div>
</div>

<script>
const TOOLTIP_PCA = 'Each pixel shows the first 3 principal components of the 768-dim embedding for that patch, mapped to RGB. Similar colors = similar embeddings = similar land cover.';
const TOOLTIP_SIMILARITY = "Each pixel's brightness shows how similar its embedding is to the clicked patch. Bright = similar, dark = different.";

let currentMode = 'pca';
let currentResult = null;
let currentImageB64 = null;
let currentSampleId = null;
let currentCacheKey = null;
let isProcessing = false;

// ---- Landing ----
function enterApp() {
  document.getElementById('landing').classList.add('hidden');
  document.getElementById('main-app').classList.add('visible');
  loadSamples();
  loadInfo();
}

// ---- API calls ----
async function callApi(apiName, data) {
  const resp = await fetch('/gradio_api/call/' + apiName, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({ data: data })
  });
  if (!resp.ok) throw new Error('API call failed: ' + resp.status);
  const result = await resp.json();
  if (!result.event_id) throw new Error('No event_id returned');
  return result.event_id;
}

async function getApiResult(apiName, eventId) {
  // Poll SSE endpoint
  const resp = await fetch('/gradio_api/call/' + apiName + '/' + eventId, { headers: { 'Accept': 'text/event-stream' } });
  const text = await resp.text();
  // Parse SSE events
  const lines = text.split('\n');
  let completeData = null;
  for (let i = 0; i < lines.length; i++) {
    if (lines[i].startsWith('event: complete')) {
      // Next data: line
      for (let j = i + 1; j < lines.length; j++) {
        if (lines[j].startsWith('data: ')) {
          completeData = lines[j].slice(6);
          break;
        }
      }
    }
  }
  if (completeData) {
    try {
      const parsed = JSON.parse(completeData);
      if (parsed.length > 0) return parsed[0];
    } catch (e) {
      return completeData;
    }
  }
  // Fallback: look for any data line
  for (const line of lines) {
    if (line.startsWith('data: ')) {
      try {
        const parsed = JSON.parse(line.slice(6));
        if (Array.isArray(parsed) && parsed.length > 0) return parsed[0];
        return line.slice(6);
      } catch (e) {}
    }
  }
  throw new Error('No result in SSE stream');
}

async function processImage(imageB64) {
  if (isProcessing) return;
  isProcessing = true;
  document.getElementById('loading').classList.add('active');
  document.getElementById('output-empty').style.display = 'none';

  try {
    const eventId = await callApi('process_image', [imageB64]);
    const resultStr = await getApiResult('process_image', eventId);
    const result = JSON.parse(resultStr);

    if (result.error) {
      document.getElementById('output-empty').textContent = result.error;
      document.getElementById('output-empty').style.display = 'block';
      return;
    }

    currentResult = result;
    currentCacheKey = result.cache_key || currentCacheKey;
    displayResult();
  } catch (e) {
    document.getElementById('output-empty').textContent = 'Error: ' + e.message;
    document.getElementById('output-empty').style.display = 'block';
  } finally {
    isProcessing = false;
    document.getElementById('loading').classList.remove('active');
  }
}

function updateOutputTooltip() {
  const tooltip = document.getElementById('output-tooltip');
  if (!tooltip) return;
  if (currentMode === 'heatmap') {
    tooltip.title = TOOLTIP_SIMILARITY;
    tooltip.style.display = 'inline-flex';
  } else if (currentMode === 'pca') {
    tooltip.title = TOOLTIP_PCA;
    tooltip.style.display = 'inline-flex';
  } else {
    tooltip.title = 'Drag the slider to compare the satellite input with the PCA embedding visualization.';
    tooltip.style.display = 'inline-flex';
  }
}

function displayResult() {
  if (!currentResult) return;
  const inputImg = document.getElementById('input-image');
  const outputImg = document.getElementById('output-image');
  const inputEmpty = document.getElementById('input-empty');
  const outputEmpty = document.getElementById('output-empty');
  updateOutputTooltip();

  inputImg.src = 'data:image/png;base64,' + currentResult.input_image;
  inputImg.classList.add('loaded');
  inputEmpty.classList.add('hidden');
  outputEmpty.classList.add('hidden');
  outputEmpty.style.display = 'none';

  if (currentMode === 'compare') {
    outputImg.classList.remove('loaded');
    setupCompare();
    updateStats(currentResult.stats);
    return;
  }

  teardownCompare();

  let outputB64;
  if (currentMode === 'heatmap') {
    outputB64 = currentResult.heatmap;
    const query = currentResult.query;
    document.getElementById('output-label').textContent = query
      ? `Similarity Heatmap · Patch ${query.row + 1}, ${query.col + 1}`
      : 'Similarity Heatmap';
  } else {
    outputB64 = currentResult.pca_rgb;
    document.getElementById('output-label').textContent = 'PCA Embedding (RGB)';
  }

  outputImg.src = 'data:image/png;base64,' + outputB64;
  outputImg.classList.add('loaded');
  updateStats(currentResult.stats);
}

function updateStats(stats) {
  // Stats
  document.getElementById('stats-bar').style.display = 'flex';
  document.getElementById('stat-patches').textContent = stats.n_patches;
  document.getElementById('stat-grid').textContent = stats.patch_grid;
  document.getElementById('stat-dim').textContent = stats.embed_dim;
  document.getElementById('stat-norm').textContent = stats.mean_norm.toFixed(2);
  document.getElementById('stat-time').textContent = stats.inference_time + 's';
}

// ---- Click similarity ----
function getRenderedImagePoint(img, event) {
  if (!img.classList.contains('loaded') || !img.naturalWidth || !img.naturalHeight) return null;
  const rect = img.getBoundingClientRect();
  const renderedRatio = rect.width / rect.height;
  const naturalRatio = img.naturalWidth / img.naturalHeight;
  let drawW = rect.width;
  let drawH = rect.height;
  let offsetX = 0;
  let offsetY = 0;

  if (renderedRatio > naturalRatio) {
    drawW = rect.height * naturalRatio;
    offsetX = (rect.width - drawW) / 2;
  } else {
    drawH = rect.width / naturalRatio;
    offsetY = (rect.height - drawH) / 2;
  }

  const xPx = event.clientX - rect.left - offsetX;
  const yPx = event.clientY - rect.top - offsetY;
  if (xPx < 0 || yPx < 0 || xPx > drawW || yPx > drawH) return null;

  return {
    x: xPx / drawW,
    y: yPx / drawH,
    markerLeft: offsetX + xPx,
    markerTop: offsetY + yPx,
  };
}

function clearClickMarker() {
  const marker = document.getElementById('click-marker');
  marker.classList.remove('visible');
}

function setClickMarker(point) {
  const marker = document.getElementById('click-marker');
  marker.style.left = `${point.markerLeft}px`;
  marker.style.top = `${point.markerTop}px`;
  marker.classList.add('visible');
}

async function querySimilarityAt(point) {
  if (!currentResult || isProcessing) return;
  const sourceImage = currentImageB64 || currentResult.input_image;
  if (!sourceImage) return;

  isProcessing = true;
  document.getElementById('loading').classList.add('active');
  setMode('heatmap');

  try {
    const eventId = await callApi('query_similarity', [
      sourceImage,
      point.x,
      point.y,
      currentCacheKey || ''
    ]);
    const resultStr = await getApiResult('query_similarity', eventId);
    const result = JSON.parse(resultStr);
    if (result.error) throw new Error(result.error);

    currentResult.heatmap = result.heatmap;
    currentResult.query = result.query;
    currentResult.cache_key = result.cache_key;
    currentResult.stats = result.stats || currentResult.stats;
    currentCacheKey = result.cache_key || currentCacheKey;
    displayResult();
  } catch (e) {
    const outputEmpty = document.getElementById('output-empty');
    outputEmpty.textContent = 'Similarity query failed: ' + e.message;
    outputEmpty.style.display = 'block';
  } finally {
    isProcessing = false;
    document.getElementById('loading').classList.remove('active');
  }
}

document.getElementById('input-image').addEventListener('click', (event) => {
  const point = getRenderedImagePoint(event.currentTarget, event);
  if (!point || !currentResult) return;
  setClickMarker(point);
  querySimilarityAt(point);
});

// ---- Compare mode ----
function teardownCompare() {
  const stage = document.getElementById('compare-stage');
  const slider = document.getElementById('compare-slider');
  stage.classList.remove('visible');
  stage.setAttribute('aria-hidden', 'true');
  slider.classList.remove('visible');
}

function updateComparePosition(value) {
  const val = Math.max(0, Math.min(100, Number(value)));
  const compareImg = document.getElementById('compare-image');
  const divider = document.getElementById('compare-divider');
  const slider = document.getElementById('compare-slider');

  compareImg.style.clipPath = `inset(0 ${100 - val}% 0 0)`;
  divider.style.left = `${val}%`;
  slider.value = val;
  slider.style.setProperty('--compare-position', `${val}%`);
}

function setupCompare() {
  const slider = document.getElementById('compare-slider');
  const stage = document.getElementById('compare-stage');
  const baseImg = document.getElementById('compare-base-image');
  const compareImg = document.getElementById('compare-image');
  const outputImg = document.getElementById('output-image');

  if (currentMode === 'compare' && currentResult) {
    outputImg.classList.remove('loaded');
    stage.classList.add('visible');
    stage.setAttribute('aria-hidden', 'false');
    slider.classList.add('visible');

    baseImg.src = 'data:image/png;base64,' + currentResult.pca_rgb;
    compareImg.src = 'data:image/png;base64,' + currentResult.input_image;
    document.getElementById('output-label').textContent = 'Compare: Input ↔ PCA';
    updateComparePosition(slider.value || 50);
  } else {
    teardownCompare();
    if (currentResult) displayResult();
  }
}

document.getElementById('compare-slider').addEventListener('input', (e) => {
  updateComparePosition(e.target.value);
});

function setMode(mode) {
  currentMode = mode;
  document.querySelectorAll('.toggle-btn').forEach(btn => {
    btn.classList.toggle('active', btn.dataset.mode === mode);
  });
  if (currentResult) {
    displayResult();
  }
}

// ---- Samples ----
async function loadSamples() {
  try {
    const eventId = await callApi('get_samples', []);
    const resultStr = await getApiResult('get_samples', eventId);
    const samples = JSON.parse(resultStr);
    const grid = document.getElementById('sample-grid');
    grid.innerHTML = '';
    samples.forEach(s => {
      const card = document.createElement('div');
      card.className = 'sample-card';
      card.dataset.id = s.id;
      card.innerHTML = `
        <img src="data:image/png;base64,${s.image}" alt="${s.label}">
        <div class="sample-card-label">${s.emoji} ${s.label}</div>
        <div class="sample-card-desc">${s.description || ''}</div>
      `;
      card.onclick = () => selectSample(s, card);
      grid.appendChild(card);
    });
  } catch (e) {
    console.error('Failed to load samples:', e);
  }
}

function selectSample(sample, card) {
  document.querySelectorAll('.sample-card').forEach(c => c.classList.remove('active'));
  card.classList.add('active');
  currentSampleId = sample.id;
  currentCacheKey = sample.cache_key || null;
  currentImageB64 = sample.image;
  clearClickMarker();
  // Show input immediately
  document.getElementById('input-image').src = 'data:image/png;base64,' + sample.image;
  document.getElementById('input-image').classList.add('loaded');
  document.getElementById('input-empty').classList.add('hidden');

  // Use pre-computed results (instant display)
  if (sample.pca_rgb && sample.heatmap && sample.stats) {
    currentResult = {
      input_image: sample.image,
      pca_rgb: sample.pca_rgb,
      heatmap: sample.heatmap,
      stats: sample.stats,
      patch_emb_shape: sample.patch_emb_shape,
      cache_key: sample.cache_key || null,
      query: null,
    };
    document.getElementById('output-empty').style.display = 'none';
    document.getElementById('stats-bar').style.display = 'none';
    displayResult();
  } else {
    // Fallback: live inference for samples without pre-computed outputs
    document.getElementById('output-image').classList.remove('loaded');
    document.getElementById('output-empty').style.display = 'block';
    document.getElementById('output-empty').textContent = 'Processing...';
    document.getElementById('stats-bar').style.display = 'none';
    processImage(sample.image);
  }
}

// ---- Upload ----
const uploadZone = document.getElementById('upload-zone');
const fileInput = document.getElementById('file-input');

fileInput.addEventListener('change', (e) => {
  const file = e.target.files[0];
  if (file) handleFile(file);
});

uploadZone.addEventListener('dragover', (e) => {
  e.preventDefault();
  uploadZone.classList.add('dragover');
});
uploadZone.addEventListener('dragleave', () => uploadZone.classList.remove('dragover'));
uploadZone.addEventListener('drop', (e) => {
  e.preventDefault();
  uploadZone.classList.remove('dragover');
  const file = e.dataTransfer.files[0];
  if (file && file.type.startsWith('image/')) handleFile(file);
});

function handleFile(file) {
  const reader = new FileReader();
  reader.onload = (e) => {
    // Convert to base64 (strip data URL prefix)
    const dataUrl = e.target.result;
    const base64 = dataUrl.split(',')[1];
    currentImageB64 = base64;
    currentSampleId = null;
    currentCacheKey = null;
    clearClickMarker();
    // Show input
    document.getElementById('input-image').src = dataUrl;
    document.getElementById('input-image').classList.add('loaded');
    document.getElementById('input-empty').classList.add('hidden');
    document.querySelectorAll('.sample-card').forEach(c => c.classList.remove('active'));
    // Reset output
    document.getElementById('output-image').classList.remove('loaded');
    document.getElementById('output-empty').style.display = 'block';
    document.getElementById('output-empty').textContent = 'Processing...';
    document.getElementById('stats-bar').style.display = 'none';
    // Process
    processImage(base64);
  };
  reader.readAsDataURL(file);
}

// ---- Info ----
async function loadInfo() {
  try {
    const eventId = await callApi('get_model_info', []);
    const resultStr = await getApiResult('get_model_info', eventId);
    const info = JSON.parse(resultStr);
    const paper = info.paper_context || {};
    document.getElementById('info-content').innerHTML = `
      <div class="info-section">
        <div class="info-section-title">Model</div>
        <div class="info-text">${info.model_name} — ${info.architecture}</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">Earth Observation Foundation Models</div>
        <div class="info-text">${paper.earth_observation_foundation_models || info.description}</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">Why OlmoEarth Is Interesting</div>
        <div class="info-text">
          <ul>
            ${(paper.why_special || []).map(item => `<li>${item}</li>`).join('')}
          </ul>
        </div>
      </div>
      <div class="info-section">
        <div class="info-section-title">Pretraining Data</div>
        <div class="info-text">${paper.pretraining_dataset || ''}</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">This Demo</div>
        <div class="info-text">${info.description}</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">Parameters</div>
        <div class="info-text">${info.parameters} · ${info.embedding_dim}-dim embeddings · ${info.encoder_depth} layers · ${info.encoder_heads} heads</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">Modalities</div>
        <div class="info-text">${info.modalities.join(' · ')}</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">What am I seeing?</div>
        <div class="info-text"><strong>PCA image:</strong> ${TOOLTIP_PCA}</div>
        <div class="info-text" style="margin-top:8px;"><strong>Similarity heatmap:</strong> ${TOOLTIP_SIMILARITY}</div>
      </div>
      <div class="info-section">
        <div class="info-section-title">Links</div>
        <div><a class="info-link" href="${info.paper_url}" target="_blank">Paper →</a></div>
        <div><a class="info-link" href="${info.huggingface_url}" target="_blank">Hugging Face Model →</a></div>
        <div><a class="info-link" href="${info.github_url}" target="_blank">GitHub Repository →</a></div>
      </div>
    `;
  } catch (e) {
    console.error('Failed to load info:', e);
  }
}

function toggleInfo() {
  const overlay = document.getElementById('info-overlay');
  const backdrop = document.getElementById('info-backdrop');
  overlay.classList.toggle('open');
  backdrop.classList.toggle('active');
}

function toggleSidebar() {
  document.getElementById('sidebar').classList.toggle('collapsed');
}

function toggleView() {
  document.getElementById('sidebar').classList.toggle('collapsed');
}
</script>
</body>
</html>"""


@app.get("/")
async def index():
    return HTMLResponse(FRONTEND_HTML)


@app.get("/health")
async def health():
    return JSONResponse({"status": "ok", "model_loaded": True})


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
demo = app

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
        _frontend=False,
        prevent_thread_lock=False,
        show_error=True,
    )