"""MJev-Spatial — one-forward-pass spatial decision demo.

Model: https://huggingface.co/Fr0zenCr4nE/MJev-Spatial-Qwen-3.5-4B
Paper: MJev-Spatial: System One Spatial Intelligence via Native Attention Matching
(arXiv 2610.22276). Inference follows the official scripts/predict.py path:
the model scores a finite candidate set with the language backbone's native
final-attention Q/K projections and returns the argmax decision — no text is generated.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  -- MUST precede torch / CUDA-touching imports
import json  # noqa: E402
import sys  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image, ImageDraw  # noqa: E402

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "src"))

# Vendored inference runtime (Apache-2.0) from github.com/Fr0zenCrane/MJev-Spatial
from spatial_jev.kev import KevBuilder, build_model, make_processor, predict_batch  # noqa: E402
from spatial_jev.numeric_tasks import ANGLE_CODEBOOK, COUNT_CODEBOOK  # noqa: E402

CHECKPOINT = "Fr0zencr4nE/MJev-Spatial-Qwen-3.5-4B"
MAX_SIDE = 512  # keep encode time and sequence length comfortable on ZeroGPU

# ---------------------------------------------------------------------------
# Model loading (module scope, eager .to("cuda") — intercepted by the spaces hijack)
# ---------------------------------------------------------------------------

def snapshot_repo(repo_id):
    """Fetch the checkpoint into a local directory path (build_model needs a *directory*)."""
    from huggingface_hub import snapshot_download
    return Path(snapshot_download(repo_id))


ckpt_dir = snapshot_repo(CHECKPOINT)  # local dir; build_model requires a checkpoint *path*
config = json.loads((ckpt_dir / "config.json").read_text())
config.update(
    preserve_option_order=True,
    max_image_tokens=1024,
    max_sequence_length=max(8192, int(config.get("max_sequence_length", 4096))),
    inference_language_dtype="bf16",
)

print("[mjev] loading processor…", flush=True)
processor = make_processor(config)
print("[mjev] loading model…", flush=True)
builder = KevBuilder(processor, config)
MODEL = build_model(config, ckpt_dir, trainable=False)
MODEL = MODEL.to("cuda").eval()
print("[mjev] model ready", flush=True)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def load_image(image, max_side=MAX_SIDE):
    """PIL RGB image scaled so neon-style huge uploads don't blow the token budget."""
    if image is None:
        raise gr.Error("Please provide an image.")
    img = image if isinstance(image, Image.Image) else Image.open(image)
    img = img.convert("RGB")
    w, h = img.size
    scale = max_side / max(w, h)
    if scale < 1:
        img = img.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.Resampling.BILINEAR)
    return img


def parse_options(text):
    """One option per line; 'id | text' optional — falls back to the text as id."""
    options = []
    for line in (text or "").strip().splitlines():
        line = line.strip()
        if not line:
            continue
        if "|" in line:
            oid, otext = line.split("|", 1)
            oid, otext = oid.strip(), otext.strip() or oid.strip()
        else:
            oid = otext = line
        options.append({"id": oid, "text": otext})
    if len(options) < 2:
        raise gr.Error("Choice tasks need at least 2 options (one per line).")
    if len(options) > config["max_choices"]:
        raise gr.Error(f"Too many options — this checkpoint supports at most {config['max_choices']}.")
    return options


def draw_point(image, xy):
    """Overlay the predicted normalized point + crosshair on a copy of the image."""
    img = image.copy()
    w, h = img.size
    x, y = xy[0] * w, xy[1] * h
    draw = ImageDraw.Draw(img)
    r = max(6, int(min(w, h) * 0.02))
    for color, width in (((255, 255, 255), 5), ((255, 70, 0), 3)):
        draw.ellipse([x - r, y - r, x + r, y + r], outline=color, width=width)
        draw.line([x - 1.6 * r, y, x + 1.6 * r, y], fill=color, width=width)
        draw.line([x, y - 1.6 * r, x, y + 1.6 * r], fill=color, width=width)
    return img


def overlay_probability_grid(image, result):
    """Pointing: warm heatmap of the final-stage 9-way scores over the selected crop."""
    if "logits" not in result or result.get("pointing_method") != "visual_tokens":
        return None
    probs = np.asarray(result.get("probabilities", []), dtype=np.float64)
    if probs.size == 0:
        return None
    height, width = result.get("grid_hw", [1, 1])
    if probs.shape != (height, width):
        return None
    heat = np.zeros((height, width, 4), dtype=np.uint8)
    heat[..., 0] = 255
    heat[..., 3] = (probs / probs.max() * 255).astype(np.uint8) if probs.max() > 0 else 0
    heat_img = Image.fromarray(heat, "RGBA").resize(image.size, Image.Resampling.BILINEAR)
    out = image.copy()
    out.paste(heat_img, (0, 0), heat_img)
    return out


def describe_path(result, kind, option_map=None):
    """Human-readable trace of the decision path."""
    lines = []
    logits = result.get("logits") or []
    path = result.get("path") or []
    names = {
        "top-center": "top-center", "top-left": "top-left", "top-right": "top-right",
        "center": "center", "middle-left": "middle-left", "middle-right": "middle-right",
        "bottom-center": "bottom-center", "bottom-left": "bottom-left",
        "bottom-right": "bottom-right",
    }
    from spatial_jev.hierarchy import POINT_NAMES, decode_point, decode_scalar  # local import
    if kind == "choice":
        if option_map:
            lines.append(f"selected option: **{option_map[path[0]]}**" if path else "")
        elif result.get("prediction") is not None:
            lines.append(f"selected option: **{result['prediction']}**")
    elif kind == "scalar":
        if path:
            edges = config["scalar_codebook"]["edges_m"]
            if path[0] == 0:
                lines.append("coarse stage: exactly zero")
            elif path[0] == len(edges):
                lines.append(f"coarse stage: greater than {edges[-1]:g} m (overflow)")
            else:
                from spatial_jev.hierarchy import scalar_interval, scalar_fine_bins
                lo, hi = scalar_interval(path[0], config["scalar_codebook"])
                lines.append(f"coarse stage: interval {lo:.6g} – {hi:.6g} m")
                if len(path) > 1:
                    fine = scalar_fine_bins(path[0], config["scalar_codebook"])
                    stepw = (hi - lo) / fine
                    lines.append(f"fine stage: refined bin {path[1] + 1}/{fine} "
                                 f"≈ {lo + path[1] * stepw:.6g} – {lo + (path[1] + 1) * stepw:.6g} m")
            lines.append(f"decoded estimate: **{result['prediction']:.6g} m**")
    elif kind == "point":
        steps = ["whole image", *(f"zoom #{i}" for i in range(1, len(path)))]
        for stage, (choice, scores) in enumerate(zip(path, logits)):
            label = names.get(POINT_NAMES[choice], POINT_NAMES[choice])
            best = scores[choice] if stage < len(scores) and choice < len(scores) else float("nan")
            lines.append(f"stage {stage + 1} ({steps[stage] if stage < len(steps) else ''}): "
                         f"region **{label}** (score {best:.3f})")
        if kind == "point" and result.get("prediction"):
            x, y = result["prediction"]
            lines.append(f"decoded point: **({x:.3f}, {y:.3f})** normalized "
                         f"(origin top-left, x→right, y→down)")
    return "\n".join(f"- {l}" for l in lines if l)


# ---------------------------------------------------------------------------
# Inference endpoint
# ---------------------------------------------------------------------------

@spaces.GPU(duration=45)  # measured: 3.3s cold attach + 3-stage point, 0.1–0.8s warm
def run_spatial(image, task, question, options_text, quantity, unit, progress=gr.Progress()):
    """Answer a spatial question about an image in one forward pass per decision stage.

    Args:
        image: input image (PIL).
        task: one of "choice" | "scalar" | "point".
        question: the spatial question about the image.
        options_text: for choice tasks, one option per line ("id | text" accepted).
        quantity: for scalar tasks, what is being estimated (e.g. "distance").
        unit: scalar unit — "m", "count" or "deg".
        progress: Gradio progress bar.
    Returns:
        (answer markdown, decision-path markdown, annotated image, raw JSON result)
    """
    t0 = time.perf_counter()
    img = load_image(image)
    if not (question or "").strip():
        raise gr.Error("Please provide a question.")

    question = question.strip()
    with tempfile.TemporaryDirectory() as tmp:
        path = Path(tmp) / "input.png"
        img.save(path)
        if task == "choice":
            space = {"kind": "choice", "options": parse_options(options_text)}
        elif task == "scalar":
            space = {"kind": "scalar", "quantity": (quantity or "value").strip() or "value",
                     "unit": unit}
        else:
            space = {"kind": "point", "coordinate_system": "normalized_xy", "num_points": 1}
        request = {"id": "demo-request", "input": {
            "media": [{"kind": "image", "uri": str(path)}],
            "question": question,
            "answer_space": space,
        }}
        result = predict_batch(MODEL, builder, [request], torch.device("cuda"))[0]

    prediction = result["prediction"]
    if task == "point":
        # Accept both [[x, y]] (list of points) and [x, y] (single point)
        if isinstance(prediction, (list, tuple)) and len(prediction) >= 1:
            inner = prediction if not isinstance(prediction[0], (list, tuple)) else prediction[0]
            if len(inner) != 2 or any(isinstance(v, (list, tuple)) for v in inner):
                inner = [v for v in inner if isinstance(v, (int, float))][:2]
            point = [float(inner[0]), float(inner[1])]
        else:
            raise gr.Error(f"Unexpected point prediction: {prediction!r}")
        overlay = draw_point(img, point)
    else:
        overlay = img.copy()

    if task == "choice":
        answer = str(prediction)
    elif task == "scalar":
        label = {"m": "m", "count": "", "deg": "°"}[unit]
        answer = f"{prediction:g}{label}" if unit != "count" else f"{int(prediction)} {quantity or 'objects'}"
    else:
        answer = f"point (x, y) = ({point[0]:.3f}, {point[1]:.3f})"

    trace = describe_path(result, task)
    seconds = time.perf_counter() - t0
    summary = f"## {answer}\n\n*(one forward pass per decision stage, {seconds:.1f}s)*"
    raw = json.dumps({"prediction": result.get("prediction"), "path": result.get("path"),
                      "logits": result.get("logits")}, ensure_ascii=False)
    return summary, trace, overlay, raw


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

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="MJev-Spatial") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🧭 MJev-Spatial\n"
            "Spatial intelligence via **native attention matching** — a 4B VLM that "
            "answers spatial questions as a single candidate-scoring decision, with no "
            "generated answer tokens. [paper](https://arxiv.org/abs/2610.22276) · "
            "[code](https://github.com/Fr0zenCrane/MJev-Spatial) · "
            "[checkpoint](https://huggingface.co/Fr0zencr4nE/MJev-Spatial-Qwen-3.5-4B)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(type="pil", label="Image", height=300)
                task = gr.Radio(
                    ["choice", "scalar", "point"], value="choice", label="Answer space",
                    info="choice: pick an option · scalar: estimate a value · point: locate a target",
                )
                question = gr.Textbox(
                    label="Question", placeholder="e.g. Where is the red square relative to the blue circle?",
                    value="Where is the red square relative to the blue circle?",
                )
                options_text = gr.Textbox(
                    label="Options (choice only) — one per line, 'id | text' optional",
                    value="left | left\nright | right", lines=3,
                )
                quantity = gr.Textbox(
                    label="Quantity (scalar only)", value="distance",
                    info="What is being estimated, e.g. distance, width, height.",
                )
                unit = gr.Radio(["m", "count", "deg"], value="m", label="Unit (scalar only)")
                run_btn = gr.Button("Decide", variant="primary")

            with gr.Column(scale=1):
                answer_md = gr.Markdown(label="Answer")
                annotated = gr.Image(label="Predicted point (pointing tasks)", height=300)
                trace_md = gr.Markdown("Decision path will appear here.")
                with gr.Accordion("Raw result (JSON)", open=False):
                    raw_json = gr.JSON(label="scores / path")

        gr.Examples(
            examples=[
                ["examples/images/scene.png", "choice",
                 "Where is the red square relative to the blue circle?",
                 "left | left\nright | right", "distance", "m"],
                ["examples/images/scene.png", "scalar",
                 "What is the approximate distance between the two objects?",
                 "", "distance", "m"],
                ["examples/images/scene.png", "point",
                 "Point to the blue circle.", "", "distance", "m"],
            ],
            inputs=[image_in, task, question, options_text, quantity, unit],
            outputs=[answer_md, trace_md, annotated, raw_json],
            fn=run_spatial,
            label="Official examples (scene.png from the MJev-Spatial repo)",
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            "Pointing decodes three 3×3 region decisions into a normalized point; the "
            "score of each stage's picked region is shown in the decision path. "
            "Scalar questions are answered through an interval hierarchy over meters "
            "(or an integer count / degree codebook)."
        )

    run_btn.click(
        run_spatial,
        inputs=[image_in, task, question, options_text, quantity, unit],
        outputs=[answer_md, trace_md, annotated, raw_json],
        api_name="decide",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)