src/ui/app.py

"""
Gradio UI for satellite image retrieval.

Vaporwave/Outrun interface: neon grids, pink-cyan-purple palette, retro-futurism.
"""

import gradio as gr
import time
import traceback
import numpy as np
from PIL import Image
from pathlib import Path
from typing import Optional

from ..retrieval.cross_modal_retrieval import CrossModalRetrieval
from ..features.extractor import FeatureExtractor


_retrieval: Optional[CrossModalRetrieval] = None
_feature_extractor: Optional[FeatureExtractor] = None
_gallery_dir: Optional[Path] = None
_gallery_metadata: Optional[list] = None


def initialize(
    retrieval: CrossModalRetrieval,
    feature_extractor: Optional[FeatureExtractor],
    gallery_dir: Optional[Path] = None,
    gallery_metadata: Optional[list] = None,
) -> None:
    global _retrieval, _feature_extractor, _gallery_dir, _gallery_metadata
    _retrieval = retrieval
    _feature_extractor = feature_extractor
    _gallery_dir = Path(gallery_dir) if gallery_dir else None
    _gallery_metadata = gallery_metadata


def _gallery_image_path(idx: int, modality: str) -> Optional[str]:
    if _gallery_metadata is not None and idx < len(_gallery_metadata):
        entry = _gallery_metadata[idx]
        path = Path(entry["gallery_path"]).resolve()
        if path.exists():
            return str(path)
    if _gallery_dir is not None:
        path = (_gallery_dir / f"{modality}_{idx}.png").resolve()
        if path.exists():
            return str(path)
    return None


def _load_image_tensor(path, modality):
    """Load an image and return (PIL preview, torch tensor with proper channels)."""
    import torch
    ext = Path(path).suffix.lower()
    # Try multi-channel TIFF first
    if ext in (".tif", ".tiff"):
        try:
            import tifffile
            arr = tifffile.imread(str(path))
            # Handle different channel arrangements
            if arr.ndim == 2:
                # Grayscale → make 3-channel for preview, keep 1ch for features
                preview = Image.fromarray(arr).convert("RGB")
                tensor = torch.from_numpy(arr).float().unsqueeze(0)  # (1, H, W)
                tensor = tensor.unsqueeze(0)  # (1, 1, H, W)
                return preview, tensor
            elif arr.ndim == 3:
                if arr.shape[-1] in (2, 3, 4, 13):
                    # Channels-last: (H, W, C)
                    tensor = torch.from_numpy(arr).float()
                    tensor = tensor.permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
                    # RGB preview
                    if arr.shape[-1] >= 3:
                        preview = Image.fromarray(arr[:, :, :3].astype(np.uint8))
                    else:
                        preview = Image.fromarray(arr[:, :, 0].astype(np.uint8)).convert("RGB")
                    return preview, tensor
                elif arr.shape[0] in (2, 3, 4, 13):
                    # Channels-first: (C, H, W)
                    tensor = torch.from_numpy(arr).float().unsqueeze(0)
                    # For preview: use first 3 channels or repeat
                    if arr.shape[0] >= 3:
                        preview_arr = np.transpose(arr[:3], (1, 2, 0))
                    else:
                        preview_arr = np.stack([arr[0]] * 3, axis=-1)
                    if arr.dtype == np.uint16:
                        preview_arr = (preview_arr / 65535.0 * 255).astype(np.uint8)
                    preview = Image.fromarray(preview_arr)
                    return preview, tensor
        except ImportError:
            pass
    # Fallback: PIL
    img = Image.open(path).convert("RGB")
    return img, None


def retrieve(image, modality: str, k: int, retrieval_type: str,
             use_sar_adapter: bool = False, use_multiscale: bool = False,
             lat: float = None, lon: float = None, radius_km: float = 50.0):
    if image is None:
        return [], "", "Please upload an image first."
    if _retrieval is None:
        return [], "", "System not initialized. Please restart the app."

    start = time.perf_counter()

    try:
        import torch

        if isinstance(image, str):
            pil_img, img_tensor = _load_image_tensor(image, modality)
        else:
            pil_img = image
            img_tensor = None

        if _feature_extractor is not None:
            if img_tensor is not None and img_tensor.shape[1] not in (3,):
                # Multi-channel TIFF → use tensor extractor
                query_embedding = _feature_extractor.extract_features_from_tensor(
                    img_tensor, modality=modality, normalize=True
                )
            elif use_sar_adapter and modality == "sar":
                from ..features.sar_adapter import SARAdapter
                adapter = SARAdapter()
                adapter.eval()
                img_t = torch.from_numpy(np.array(pil_img)).permute(2, 0, 1).float() / 255.0
                if img_t.shape[0] == 3:
                    img_t = img_t[:2]
                img_t = img_t.unsqueeze(0)
                with torch.no_grad():
                    adapted = adapter(img_t)
                adapted_pil = Image.fromarray(
                    (adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).astype(np.uint8))
                query_embedding = _feature_extractor.extract_features(
                    adapted_pil, modality=modality, normalize=True)
            else:
                query_embedding = _feature_extractor.extract_features(
                    pil_img, modality=modality, normalize=True)
        else:
            embed_dim = _retrieval.embed_dim
            query_embedding = torch.randn(embed_dim)
            query_embedding = torch.nn.functional.normalize(query_embedding, dim=0)

        query_np = query_embedding.unsqueeze(0).numpy().astype(np.float32)

        if lat is not None and lon is not None:
            result = _retrieval.search(query_np, modality, k=k, lat=lat, lon=lon, radius_km=radius_km)
        elif retrieval_type == "same-modal":
            result = _retrieval.search(query_np, modality, target_modality=modality, k=k)
        else:
            result = _retrieval.search(query_np, modality, k=k, strategy="multi")

        elapsed_ms = (time.perf_counter() - start) * 1000

        gallery_images = []
        for i, (idx, score) in enumerate(zip(result.indices, result.scores)):
            mod = result.modalities[i] if result.modalities else modality
            img_path = _gallery_image_path(idx, mod)
            if img_path:
                gallery_images.append(Image.open(img_path))

        if not gallery_images:
            for idx, _ in zip(result.indices, result.scores):
                np.random.seed(idx)
                arr = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
                gallery_images.append(Image.fromarray(arr))

        timing_text = f"{elapsed_ms:.0f}ms"
        n_results = len(result.indices)
        mod_str = ", ".join(set(result.modalities)) if result.modalities else modality
        status_text = f"{n_results} results | {mod_str} | {elapsed_ms:.0f}ms"

        return gallery_images, timing_text, status_text

    except Exception as exc:
        tb = traceback.format_exc()
        return [], "", f"Error: {exc}\n\n{tb}"


# ---------------------------------------------------------------------------
# Vaporwave Design System
# ---------------------------------------------------------------------------

VAPORWAVE_CSS = """
<style>
@import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;700;900&family=Outfit:wght@300;400;600&display=swap');

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

:root {
    --neon-pink: #ff6bcd;
    --neon-cyan: #00f0ff;
    --neon-purple: #b300ff;
    --neon-blue: #0044ff;
    --dark-bg: #0a0015;
    --card-bg: #12002a;
    --card-border: #2a0050;
    --text-primary: #e0c0ff;
    --text-secondary: #9a6fb0;
    --glow-pink: 0 0 20px rgba(255, 107, 205, 0.5);
    --glow-cyan: 0 0 20px rgba(0, 240, 255, 0.5);
    --glow-purple: 0 0 20px rgba(179, 0, 255, 0.5);
}

/* Grid background */
body, .gradio-container {
    font-family: 'Outfit', sans-serif !important;
    max-width: 1200px !important;
    margin: 0 auto !important;
    background: var(--dark-bg) !important;
    color: var(--text-primary) !important;
    position: relative;
    overflow-x: hidden;
}

body::before {
    content: '';
    position: fixed;
    top: 0; left: 0; right: 0; bottom: 0;
    background:
        linear-gradient(transparent 0%, rgba(179, 0, 255, 0.03) 50%, transparent 100%),
        repeating-linear-gradient(
            0deg,
            transparent,
            transparent 40px,
            rgba(0, 240, 255, 0.04) 40px,
            rgba(0, 240, 255, 0.04) 41px
        ),
        repeating-linear-gradient(
            90deg,
            transparent,
            transparent 40px,
            rgba(255, 107, 205, 0.04) 40px,
            rgba(255, 107, 205, 0.04) 41px
        );
    pointer-events: none;
    z-index: 0;
}

/* Scanline overlay */
body::after {
    content: '';
    position: fixed;
    top: 0; left: 0; right: 0; bottom: 0;
    background: repeating-linear-gradient(
        0deg,
        transparent,
        transparent 2px,
        rgba(0, 0, 0, 0.15) 2px,
        rgba(0, 0, 0, 0.15) 4px
    );
    pointer-events: none;
    z-index: 1;
}

.gradio-container {
    position: relative;
    z-index: 2;
    background: transparent !important;
}

/* Header */
.vapor-header {
    text-align: center;
    padding: 2rem 1rem 1.5rem;
    margin: -1rem -1rem 0 -1rem;
    position: relative;
    background: linear-gradient(180deg, rgba(179, 0, 255, 0.15) 0%, transparent 100%);
    border-bottom: 2px solid var(--neon-purple);
    box-shadow: var(--glow-purple);
}

.vapor-header::after {
    content: '';
    position: absolute;
    bottom: -2px;
    left: 10%; right: 10%;
    height: 1px;
    background: linear-gradient(90deg, transparent, var(--neon-cyan), transparent);
}

.vapor-title {
    font-family: 'Orbitron', monospace;
    font-size: 1.6rem;
    font-weight: 900;
    text-transform: uppercase;
    letter-spacing: 4px;
    background: linear-gradient(90deg, var(--neon-cyan), var(--neon-pink), var(--neon-purple));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    text-shadow: none;
    filter: drop-shadow(0 0 10px rgba(255, 107, 205, 0.3));
}

.vapor-subtitle {
    font-size: 0.85rem;
    color: var(--text-secondary);
    letter-spacing: 3px;
    text-transform: uppercase;
    margin-top: 0.3rem;
}

.vapor-sun {
    display: inline-block;
    width: 60px; height: 60px;
    border-radius: 50%;
    background: linear-gradient(135deg, var(--neon-pink), var(--neon-purple));
    box-shadow: 0 0 40px rgba(255, 107, 205, 0.4), 0 0 80px rgba(179, 0, 255, 0.2);
    margin-bottom: 0.5rem;
    animation: pulse-glow 3s ease-in-out infinite;
}

@keyframes pulse-glow {
    0%, 100% { box-shadow: 0 0 40px rgba(255, 107, 205, 0.4), 0 0 80px rgba(179, 0, 255, 0.2); }
    50% { box-shadow: 0 0 60px rgba(255, 107, 205, 0.6), 0 0 100px rgba(179, 0, 255, 0.3); }
}

/* Cards */
.vapor-card-left, .vapor-card-right {
    background: var(--card-bg) !important;
    border: 1px solid var(--card-border) !important;
    box-shadow: 0 0 15px rgba(179, 0, 255, 0.1), inset 0 0 30px rgba(0, 0, 0, 0.3) !important;
    padding: 1rem !important;
    border-radius: 4px !important;
    position: relative;
    backdrop-filter: blur(10px);
}

.vapor-card-left::before, .vapor-card-right::before {
    content: '';
    position: absolute;
    top: 0; left: 0; right: 0;
    height: 1px;
    background: linear-gradient(90deg, transparent, var(--neon-cyan), transparent);
}

.vapor-card-left::after, .vapor-card-right::after {
    content: '';
    position: absolute;
    bottom: 0; left: 0; right: 0;
    height: 1px;
    background: linear-gradient(90deg, transparent, var(--neon-pink), transparent);
}

.section-label {
    font-family: 'Orbitron', monospace;
    font-size: 0.7rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 3px;
    color: var(--neon-cyan);
    display: block;
    margin-bottom: 0.75rem;
    text-shadow: var(--glow-cyan);
}

/* Search button */
.vapor-btn {
    background: linear-gradient(135deg, var(--neon-purple), var(--neon-pink)) !important;
    color: #fff !important;
    border: none !important;
    border-radius: 4px !important;
    font-family: 'Orbitron', monospace !important;
    font-weight: 700 !important;
    font-size: 0.85rem !important;
    text-transform: uppercase !important;
    letter-spacing: 3px !important;
    padding: 0.8rem 1.2rem !important;
    width: 100% !important;
    cursor: pointer !important;
    box-shadow: 0 0 20px rgba(179, 0, 255, 0.3) !important;
    transition: all 0.3s ease !important;
    position: relative;
    overflow: hidden;
}

.vapor-btn::before {
    content: '';
    position: absolute;
    top: -50%; left: -50%;
    width: 200%; height: 200%;
    background: linear-gradient(45deg, transparent, rgba(255,255,255,0.1), transparent);
    transform: rotate(45deg);
    transition: all 0.5s ease;
}

.vapor-btn:hover {
    box-shadow: 0 0 40px rgba(179, 0, 255, 0.5), 0 0 60px rgba(255, 107, 205, 0.3) !important;
    transform: translateY(-2px) !important;
}

.vapor-btn:hover::before {
    left: 100%;
}

.vapor-btn:active {
    transform: translateY(1px) !important;
}

/* Status */
.vapor-status textarea {
    background: rgba(0, 0, 0, 0.4) !important;
    color: var(--neon-cyan) !important;
    border: 1px solid var(--card-border) !important;
    font-family: 'Orbitron', monospace !important;
    font-weight: 400 !important;
    font-size: 0.75rem !important;
    letter-spacing: 2px !important;
    border-radius: 4px !important;
    box-shadow: inset 0 0 10px rgba(0, 0, 0, 0.3) !important;
}

/* Gallery */
.vapor-gallery {
    border: 1px solid var(--card-border) !important;
    border-radius: 4px !important;
    background: rgba(0, 0, 0, 0.2) !important;
    box-shadow: inset 0 0 20px rgba(0, 0, 0, 0.3) !important;
}

.vapor-gallery img {
    transition: all 0.3s ease !important;
    border: 1px solid transparent !important;
}

.vapor-gallery img:hover {
    transform: scale(1.05) !important;
    border-color: var(--neon-cyan) !important;
    box-shadow: 0 0 15px rgba(0, 240, 255, 0.3) !important;
    z-index: 10;
}

/* Modality cards */
.modality-cards {
    display: flex;
    gap: 2px;
    margin-top: 0.75rem;
}

.modality-card {
    flex: 1;
    padding: 0.6rem 0.4rem;
    text-align: center;
    font-size: 0.65rem;
    font-weight: 600;
    letter-spacing: 1px;
    text-transform: uppercase;
    color: var(--text-secondary);
    background: rgba(0, 0, 0, 0.3);
    border: 1px solid var(--card-border);
    transition: all 0.3s ease;
    cursor: pointer;
}

.modality-card strong {
    display: block;
    font-size: 0.75rem;
    margin-bottom: 0.1rem;
}

.modality-card:nth-child(1) { border-color: rgba(0, 240, 255, 0.3); }
.modality-card:nth-child(2) { border-color: rgba(255, 107, 205, 0.3); }
.modality-card:nth-child(3) { border-color: rgba(179, 0, 255, 0.3); }

.modality-card:nth-child(1):hover { background: rgba(0, 240, 255, 0.1); box-shadow: 0 0 10px rgba(0, 240, 255, 0.2); }
.modality-card:nth-child(2):hover { background: rgba(255, 107, 205, 0.1); box-shadow: 0 0 10px rgba(255, 107, 205, 0.2); }
.modality-card:nth-child(3):hover { background: rgba(179, 0, 255, 0.1); box-shadow: 0 0 10px rgba(179, 0, 255, 0.2); }

/* Tags */
.tag {
    display: inline-block;
    padding: 0.2rem 0.5rem;
    font-size: 0.6rem;
    font-family: 'Orbitron', monospace;
    letter-spacing: 1px;
    text-transform: uppercase;
    margin: 0 0.1rem;
    transition: all 0.2s ease;
}

.tag:hover {
    transform: translateY(-1px);
    filter: brightness(1.3);
}

.tag-optical { background: rgba(0, 240, 255, 0.2); color: var(--neon-cyan); border: 1px solid rgba(0, 240, 255, 0.3); }
.tag-sar { background: rgba(255, 107, 205, 0.2); color: var(--neon-pink); border: 1px solid rgba(255, 107, 205, 0.3); }
.tag-ms { background: rgba(179, 0, 255, 0.2); color: var(--neon-purple); border: 1px solid rgba(179, 0, 255, 0.3); }

/* Footer */
.vapor-footer {
    text-align: center;
    font-size: 0.7rem;
    color: var(--text-secondary);
    padding: 1rem;
    margin: 1rem -1rem -1rem;
    border-top: 1px solid var(--card-border);
    letter-spacing: 2px;
    text-transform: uppercase;
    position: relative;
}

.vapor-footer::before {
    content: '';
    position: absolute;
    top: -1px;
    left: 20%; right: 20%;
    height: 1px;
    background: linear-gradient(90deg, transparent, var(--neon-pink), transparent);
}

/* Gradio overrides */
.gradio-container .wrap { border-radius: 0 !important; }

.gradio-container input, .gradio-container textarea, .gradio-container select {
    border-radius: 4px !important;
    border: 1px solid var(--card-border) !important;
    background: rgba(0, 0, 0, 0.4) !important;
    color: var(--text-primary) !important;
    font-family: 'Outfit', sans-serif !important;
    transition: all 0.2s ease !important;
}

.gradio-container input:focus, .gradio-container textarea:focus, .gradio-container select:focus {
    border-color: var(--neon-cyan) !important;
    box-shadow: 0 0 10px rgba(0, 240, 255, 0.2) !important;
}

.gradio-container .slider-container input[type="range"] {
    accent-color: var(--neon-pink) !important;
}

/* Labels and dropdowns */
.gradio-container label {
    color: var(--text-secondary) !important;
    font-family: 'Outfit', sans-serif !important;
    font-weight: 400 !important;
    letter-spacing: 1px !important;
    text-transform: uppercase !important;
    font-size: 0.7rem !important;
}

/* Scrollbar */
::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: var(--dark-bg); }
::-webkit-scrollbar-thumb { background: var(--neon-purple); border-radius: 3px; }
::-webkit-scrollbar-thumb:hover { background: var(--neon-pink); }

/* File upload */
.gradio-container input[type="file"]::file-selector-button {
    background: linear-gradient(135deg, var(--neon-purple), var(--neon-pink)) !important;
    color: #fff !important;
    border: none !important;
    border-radius: 4px !important;
    padding: 0.4rem 0.8rem !important;
    font-family: 'Orbitron', monospace !important;
    font-size: 0.65rem !important;
    text-transform: uppercase !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
}

.gradio-container input[type="file"]::file-selector-button:hover {
    box-shadow: 0 0 10px rgba(179, 0, 255, 0.4) !important;
}

/* Gallery caption text */
.gradio-container .gallery-item p {
    font-family: 'Outfit', sans-serif !important;
    font-size: 0.65rem !important;
    color: var(--text-secondary) !important;
}

/* Keep the neon terminal vibes */
@keyframes flicker {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.98; }
}

.vapor-header {
    animation: flicker 0.15s infinite;
}
</style>
"""


def _open_image(file):
    if file is None:
        return None
    if isinstance(file, dict):
        return Image.open(file.get('path') or file.get('url'))
    if hasattr(file, 'path'):
        return Image.open(file.path)
    if hasattr(file, 'name'):
        return Image.open(file.name)
    if isinstance(file, str):
        return Image.open(file)
    return Image.open(file)


def create_app() -> gr.Blocks:
    def on_upload(file):
        if file is None:
            return None
        path = None
        if isinstance(file, dict):
            path = file.get('path') or file.get('url')
        elif hasattr(file, 'path'):
            path = file.path
        elif hasattr(file, 'name'):
            path = file.name
        elif isinstance(file, str):
            path = file
        if path:
            pil_img, _ = _load_image_tensor(path, "optical")
            return pil_img
        return _open_image(file)

    def on_retrieve(file, modality, k, retrieval_type):
        if file is None:
            return [], "", "Upload an image first."
        return retrieve(file, modality, int(float(k)), retrieval_type)

    with gr.Blocks(title="SATCOM // Cross-Modal Retrieval") as app:
        gr.HTML(VAPORWAVE_CSS)

        gr.HTML("""
        <div class="vapor-header">
            <div class="vapor-sun"></div>
            <div class="vapor-title">SATCOM // RETRIEVAL</div>
            <div class="vapor-subtitle">Cross-Modal Satellite Image Search // Optical · SAR · Multispectral</div>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=1, elem_classes=["vapor-card-left"]):
                gr.HTML('<span class="section-label">// INPUT</span>')

                file_input = gr.File(
                    label="Upload Satellite Image",
                    file_types=[".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"],
                )

                preview = gr.Image(
                    label="Preview",
                    interactive=False,
                    height=160,
                )

                gr.HTML('<span class="section-label">// SETTINGS</span>')

                modality = gr.Dropdown(
                    ["optical", "sar", "multispectral"],
                    value="optical",
                    label="Query Modality",
                )

                retrieval_type = gr.Radio(
                    ["same-modal", "cross-modal"],
                    value="same-modal",
                    label="Retrieval Type",
                )

                k_slider = gr.Slider(
                    1, 10, value=5, step=1,
                    label="Results (K)",
                )

                btn = gr.Button(
                    "▶ EXECUTE SEARCH",
                    variant="primary",
                    elem_classes=["vapor-btn"],
                )

                gr.HTML("""
                <div class="modality-cards">
                    <div class="modality-card">
                        <strong>OPTICAL</strong>
                        RGB · 3ch
                    </div>
                    <div class="modality-card">
                        <strong>SAR</strong>
                        Radar · 2ch
                    </div>
                    <div class="modality-card">
                        <strong>MULTI</strong>
                        All · 13ch
                    </div>
                </div>
                """)

            with gr.Column(scale=2, elem_classes=["vapor-card-right"]):
                gr.HTML('<span class="section-label">// RESULTS</span>')

                status = gr.Textbox(
                    label="Status",
                    interactive=False,
                    lines=1,
                    elem_classes=["vapor-status"],
                )

                gallery = gr.Gallery(
                    label="Retrieved Images",
                    columns=5,
                    rows=2,
                    height=360,
                    elem_classes=["vapor-gallery"],
                )

                timing = gr.Textbox(
                    label="Query Time",
                    interactive=False,
                    lines=1,
                )

        gr.HTML("""
        <div class="vapor-footer">
            <span class="tag tag-optical">OPTICAL</span>
            <span class="tag tag-sar">SAR</span>
            <span class="tag tag-ms">MULTISPECTRAL</span>
            &nbsp;&nbsp;//&nbsp;&nbsp;
            SatCLIP + FAISS + Multi-Index
            &nbsp;&nbsp;//&nbsp;&nbsp;
            Ayush · Karan · Anurag
        </div>
        """)

        file_input.change(fn=on_upload, inputs=[file_input], outputs=[preview])
        btn.click(
            fn=on_retrieve,
            inputs=[file_input, modality, k_slider, retrieval_type],
            outputs=[gallery, timing, status],
        )

    return app


if __name__ == "__main__":
    app = create_app()
    print("Gradio app created. Run with: app.launch()")