"""LatticeMemory HF Space — interactive demo."""
from __future__ import annotations
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

import gradio as gr
import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer

from e8_utils import nestquant_snap, embedding_to_e8_address, index_size_bytes
import json
import plotly.graph_objects as go

# ---------------------------------------------------------------------------
# Load model and index at startup
# ---------------------------------------------------------------------------
MODEL_ID   = "dfrokido/bge-large-e8-snap"
INDEX_PATH = Path(__file__).parent / "data" / "index.pt"

print("Loading model...")
_device = "cuda" if torch.cuda.is_available() else "cpu"
_model  = SentenceTransformer(MODEL_ID, device=_device)

print("Loading index...")
_idx        = torch.load(INDEX_PATH, weights_only=False)
_float_norm = _idx["float_norm"]   # [N, D]
_snap_norm  = _idx["snap_norm"]    # [N, D]
_doc_ids    = _idx["doc_ids"]
_doc_texts  = _idx["doc_texts"]
_d_model    = _idx["d_model"]
_n_docs     = _idx["n_docs"]
_sizes      = _idx["sizes"]

_sz = index_size_bytes(_n_docs, _d_model)
INDEX_STATS = (
    f"**Index:** {_n_docs} docs · {_d_model}-dim\n\n"
    f"| Method | Size | Compression |\n"
    f"|---|---|---|\n"
    f"| float32 | {_sz['float32']/1e6:.2f} MB | 1x |\n"
    f"| int4    | {_sz['int4']/1e6:.2f} MB | 8x |\n"
    f"| **RF-Snap** | **{_sz['rfsnap']/1e6:.2f} MB** | **10.7x** |"
)

# ---------------------------------------------------------------------------
# Retrieval logic
# ---------------------------------------------------------------------------

def _encode(text: str) -> torch.Tensor:
    """Encode a single query string to a normalized float32 tensor [1, D]."""
    emb = _model.encode([text], convert_to_tensor=True)
    return F.normalize(emb.float().cpu(), p=2, dim=1)


def retrieve(query: str) -> tuple[str, str, str]:
    """Encode query, retrieve top-3 with float32 and RF-Snap, return results."""
    if not query.strip():
        return "", "", INDEX_STATS

    q_emb = _encode(query)  # [1, D]

    # Float32 retrieval
    t0 = time.perf_counter()
    float_scores = (_float_norm @ q_emb.T).squeeze()
    float_top = torch.topk(float_scores, k=3).indices.tolist()
    float_ms = (time.perf_counter() - t0) * 1000

    # RF-Snap retrieval
    q_snap = F.normalize(nestquant_snap(q_emb), p=2, dim=1)
    t0 = time.perf_counter()
    snap_scores = (_snap_norm @ q_snap.T).squeeze()
    snap_top = torch.topk(snap_scores, k=3).indices.tolist()
    snap_ms = (time.perf_counter() - t0) * 1000

    # E8 address of query
    address_hex = embedding_to_e8_address(q_emb.squeeze(0))
    address_display = " ".join(
        address_hex[i:i+8] for i in range(0, len(address_hex), 8)
    )

    # Results table
    rows = []
    for rank, (fi, si) in enumerate(zip(float_top, snap_top), 1):
        f_text = _doc_texts[fi][:120].replace("|", "/")
        s_text = _doc_texts[si][:120].replace("|", "/")
        match = "yes" if fi == si else "~"
        rows.append(f"| {rank} | {f_text}... | {s_text}... | {match} |")

    results_md = (
        f"**float32:** {float_ms:.2f} ms    "
        f"**RF-Snap:** {snap_ms:.2f} ms    "
        f"**Speedup:** {float_ms/max(snap_ms,0.01):.1f}x\n\n"
        f"| Rank | Float32 result | RF-Snap result | Match |\n"
        f"|---|---|---|---|\n"
        + "\n".join(rows)
    )

    return address_display, results_md, INDEX_STATS


# ---------------------------------------------------------------------------
# Benchmark figure
# ---------------------------------------------------------------------------

def _build_benchmark_figure() -> go.Figure:
    """Load benchmark.json and build a Plotly latency-scaling chart."""
    data_path = Path(__file__).parent / "data" / "benchmark.json"
    if not data_path.exists():
        fig = go.Figure()
        fig.add_annotation(text="Benchmark data not found", showarrow=False)
        return fig

    with open(data_path) as f:
        raw = json.load(f)["results"]

    sizes = sorted(int(k) for k in raw.keys())
    methods = {
        "float32":     ("Float32 scan",    "#ef4444", "dash"),
        "int4":        ("Int4 scan",       "#f97316", "dot"),
        "rfsnap_hit":  ("RF-Snap hit",     "#22c55e", "solid"),
        "rfsnap_miss": ("RF-Snap miss",    "#86efac", "dashdot"),
    }

    fig = go.Figure()
    for key, (label, color, dash) in methods.items():
        if key not in raw[str(sizes[0])]:
            continue
        y = [raw[str(n)][key]["p50"] for n in sizes]
        fig.add_trace(go.Scatter(
            x=sizes, y=y, name=label,
            line=dict(color=color, dash=dash, width=2),
            mode="lines+markers",
        ))

    fig.update_layout(
        title="Retrieval latency p50 vs corpus size (bge-large 1024-dim)",
        xaxis_title="Corpus size (docs)",
        yaxis_title="Latency p50 (ms)",
        xaxis_type="log",
        template="plotly_dark",
        legend=dict(orientation="h", yanchor="bottom", y=1.02),
    )
    return fig


_benchmark_fig = _build_benchmark_figure()

# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="LatticeMemory") as demo:
    gr.Markdown(
        "# LatticeMemory\n"
        "### 10.7x smaller index. Same retrieval quality. Every concept has a permanent address.\n"
        "Type anything — see its E8 lattice address and retrieve from a 500-doc corpus."
    )

    with gr.Tabs():
        with gr.Tab("Live Demo"):
            with gr.Row():
                query_box = gr.Textbox(
                    label="Query",
                    placeholder="What is the capital of France?",
                    lines=2,
                    scale=3,
                )
                submit_btn = gr.Button("Retrieve", variant="primary", scale=1)

            address_out = gr.Textbox(
                label="E8 Address (128 blocks - permanent coordinate for this concept)",
                lines=3,
                interactive=False,
            )
            results_out = gr.Markdown(label="Top-3 Results")
            stats_out   = gr.Markdown(value=INDEX_STATS)

            submit_btn.click(
                fn=retrieve,
                inputs=query_box,
                outputs=[address_out, results_out, stats_out],
            )
            query_box.submit(
                fn=retrieve,
                inputs=query_box,
                outputs=[address_out, results_out, stats_out],
            )

        with gr.Tab("Benchmark"):
            gr.Markdown(
                "## O(1) vs O(N) — Retrieval Latency Scaling\n"
                "RF-Snap hit latency stays **flat** regardless of corpus size. "
                "All scan methods (float32, int4) grow linearly with N. "
                "At 100K docs: **RF-Snap is 17x faster** than float32 scan.\n\n"
                "bge-large 1024-dim · 100 warmup queries per corpus size"
            )
            gr.Plot(value=_benchmark_fig)
            gr.Markdown(
                "| Method | Compression | 100K docs p50 | vs RF-Snap |\n"
                "|---|---|---|---|\n"
                "| Float32 | 1x | 20.8 ms | 17x slower |\n"
                "| Int8 | 4x | 17.8 ms | 14x slower |\n"
                "| Int4 | 8x | 18.5 ms | 15x slower |\n"
                "| **RF-Snap hit** | **10.7x** | **1.2 ms** | **baseline** |\n"
            )

if __name__ == "__main__":
    demo.launch()