"""
DINO Hierarchies — upload an image, get its condensed tree (hover) + reveal animation.
Runs DINOv3 ViT-L/16, builds the same condensed tree we precomputed for the static site,
and injects it into the existing JS widgets (no server-side rendering of the video).
"""
import os, io, json, html, base64, sys
print("[app] starting import", flush=True)
import numpy as np, torch
try:
    import gradio as gr
    print("[app] gradio", gr.__version__, flush=True)
except ImportError:
    gr = None          # allows importing the pipeline without gradio (local testing)
from PIL import Image
print("[app] transformers import...", flush=True)
from scipy.cluster.hierarchy import linkage, to_tree
from scipy.spatial.distance import squareform
from transformers import AutoModel

REPO = "facebook/dinov3-vitl16-pretrain-lvd1689m"
RES, N_PREFIX, MS = 224, 5, 4
DEV = "cuda" if torch.cuda.is_available() else "cpu"
TOKEN = os.environ.get("HF_TOKEN")

MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
TREE_HTML = open("tree_widget.html").read()
REVEAL_HTML = open("reveal_widget.html").read()

_model = None
def get_model():
    """Lazy-load so the app binds its port immediately (avoids startup timeout)."""
    global _model
    if _model is None:
        _model = AutoModel.from_pretrained(REPO, dtype=torch.float32, token=TOKEN,
                                           attn_implementation="eager").eval().to(DEV)
    return _model


def build_tree(image):
    model = get_model()
    img = image.convert("RGB").resize((RES, RES), Image.BICUBIC)
    arr = np.asarray(img).astype(np.float32) / 255
    x = ((torch.from_numpy(arr).permute(2, 0, 1) - MEAN) / STD).unsqueeze(0).to(DEV)
    with torch.no_grad():
        o = model(x, output_attentions=True)
    feat = o.last_hidden_state[0, N_PREFIX:].float().cpu().numpy()
    att = o.attentions[-1][0][:, 0, N_PREFIX:].mean(0).float().cpu().numpy()
    att_pct = 100.0 * att / att.sum()
    P = feat.shape[0]; g = int(round(P ** 0.5))
    fn = feat / (np.linalg.norm(feat, axis=1, keepdims=True) + 1e-8)
    root = to_tree(linkage(squareform(1.0 - fn @ fn.T, checks=False), method="average"))

    nodes = []
    def add(cn, depth, parent):
        idx = len(nodes)
        nodes.append({"leaves": cn.pre_order(lambda v: v.id), "children": [], "depth": depth,
                      "split_tau": None, "parent": parent})
        cur, fo = cn, []
        while not cur.is_leaf():
            l, r = cur.left, cur.right
            if l.count >= MS and r.count >= MS: break
            small, big = (r, l) if l.count >= r.count else (l, r); fo.append(small); cur = big
        if not cur.is_leaf():
            nodes[idx]["split_tau"] = round(1.0 - cur.dist, 3)
            for ch in (cur.left, cur.right):
                nodes[idx]["children"].append(add(ch, depth + 1, idx))
            for f in fo:
                fi = len(nodes); nodes.append({"leaves": f.pre_order(lambda v: v.id), "children": [],
                                               "depth": depth + 1, "split_tau": None, "parent": idx})
                nodes[idx]["children"].append(fi)
        return idx
    add(root, 0, -1)

    xpos = {}; cnt = [0]
    def setx(i):
        ch = nodes[i]["children"]
        if not ch: xpos[i] = cnt[0]; cnt[0] += 1
        else:
            for c in ch: setx(c)
            xpos[i] = float(np.mean([xpos[c] for c in ch]))
    setx(0)
    nleaves = cnt[0]; maxdepth = max(n["depth"] for n in nodes)
    tree = {"g": g, "nleaves": nleaves, "maxdepth": maxdepth,
            "nodes": [{"id": i, "x": round(xpos[i], 3), "depth": n["depth"], "parent": n["parent"],
                       "leaf": len(n["children"]) == 0, "n": len(n["leaves"]),
                       "att": round(float(att_pct[n["leaves"]].sum()), 1),
                       "attm": round(float(att_pct[n["leaves"]].mean()), 4),
                       "tau": n["split_tau"], "patches": n["leaves"]} for i, n in enumerate(nodes)]}
    buf = io.BytesIO(); img.save(buf, "PNG")
    return tree, base64.b64encode(buf.getvalue()).decode()


def iframe(widget_html, tree, img_b64, height):
    inject = (f'<script>window.TREE={json.dumps(tree)};'
              f'window.IMG_SRC="data:image/png;base64,{img_b64}";</script>')
    doc = widget_html.replace("<body>", "<body>" + inject, 1)
    return (f'<iframe srcdoc="{html.escape(doc, quote=True)}" '
            f'style="width:100%;height:{height}px;border:none;background:#0f1117;border-radius:8px"></iframe>')


def process(image):
    if image is None:
        return "<p style='color:#999'>Upload an image first.</p>", ""
    tree, img_b64 = build_tree(image)
    return iframe(TREE_HTML, tree, img_b64, 660), iframe(REVEAL_HTML, tree, img_b64, 700)


if gr is not None:
    with gr.Blocks(title="Visual Tree", theme=gr.themes.Base()) as demo:
        gr.Markdown("# 🌳 Visual Tree\nUpload an image → get its **condensed tree** (hover a node) "
                    "and the **reveal animation** (best-first by CLS attention). Runs DINOv3 ViT-L/16.")
        with gr.Row():
            inp = gr.Image(type="pil", label="Upload an image", height=300)
            btn = gr.Button("Build tree + reveal", variant="primary", scale=0)
        with gr.Tab("Condensed tree (hover)"):
            out_tree = gr.HTML()
        with gr.Tab("Reveal animation"):
            out_rev = gr.HTML()
        btn.click(process, inp, [out_tree, out_rev])
        inp.upload(process, inp, [out_tree, out_rev])

    print("[app] launching gradio", flush=True)
    demo.queue().launch()   # module-level: HF Spaces imports this file, so no __main__ guard