"""Chipoint — image geolocation demo (ZeroGPU).

Full 3-encoder model + density-weighted consensus, on a ZeroGPU (Blackwell) GPU. Weights are pulled from the
public chiikabu-labs/chipoint repo; DINOv3 comes ungated from ModelScope. Each guess takes a few seconds.
"""
import os
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
from pathlib import Path
from PIL import Image
import gradio as gr
import spaces                                              # ZeroGPU
from huggingface_hub import snapshot_download

DEV = "cuda"                                               # ZeroGPU: load on cuda at module level (emulated until @spaces.GPU)
M, K = 200, 64
Gv, LAM, MU, BW1, BW2 = 0.8, 1.7, 0.8, 1492.7, 750.0
ENC_D = {"so400m256": 1152, "dinov3": 1024, "gopt": 1536}

print("downloading weights ...", flush=True)
W = Path(snapshot_download("chiikabu-labs/chipoint", allow_patterns=["*.pt", "*.npy"]))

class Head(nn.Module):
    def __init__(s, D):
        super().__init__(); s.net = nn.Sequential(nn.Linear(D, 1024), nn.GELU(), nn.Dropout(0.1), nn.Linear(1024, 512))
    def forward(s, x): return F.normalize(s.net(x), dim=-1)

def load_head(tag, D):
    h = Head(D); sd = torch.load(W / f"proj_head_{tag}.pt", map_location="cpu")
    h.load_state_dict({k: v for k, v in sd.items() if k.startswith("net.")})
    return h.to(DEV).eval()

print("loading heads + galleries + encoders ...", flush=True)
HEADS = {t: load_head(t, D) for t, D in ENC_D.items()}
GAL = {t: np.load(W / f"proj_head_gallery_512_{t}.npy", mmap_mode="r") for t in ENC_D}   # mmap, streamed to GPU per query
GC = np.load(W / "gallery_all_C_so400m256.npy").astype("float32")
def cell_key(ll, C): return (np.round(ll[..., 0]/C).astype(np.int64) << 20) ^ (np.round(ll[..., 1]/C).astype(np.int64) & 0xFFFFF)
DENS = {C: dict(zip(*[a.tolist() for a in np.unique(cell_key(GC, C), return_counts=True)])) for C in (0.25, 0.5, 1.0, 2.0)}

import open_clip, torchvision.transforms as T
from transformers import AutoModel, AutoImageProcessor
from modelscope import snapshot_download as ms_snapshot
_OC = {}
for tag, name in [("so400m256", "ViT-SO400M-16-SigLIP2-256"), ("gopt", "ViT-gopt-16-SigLIP2-256")]:
    m, _, pre = open_clip.create_model_and_transforms(name, pretrained="webli"); _OC[tag] = (m.to(DEV).eval(), pre)
_DINO = ms_snapshot("facebook/dinov3-vitl16-pretrain-lvd1689m")
_DM = AutoModel.from_pretrained(_DINO).to(DEV).eval()
_DP = AutoImageProcessor.from_pretrained(_DINO)
_MEAN = torch.tensor(_DP.image_mean).view(1, 3, 1, 1).to(DEV); _STD = torch.tensor(_DP.image_std).view(1, 3, 1, 1).to(DEV)
_DPRE = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])

def zc(a): return (a - a.mean()) / (a.std() + 1e-6)

@torch.no_grad()
def _encode(tag, pil):
    if tag == "dinov3":
        x = ((_DPRE(pil).unsqueeze(0).to(DEV) - _MEAN) / _STD)
        o = _DM(x); e = o.pooler_output if getattr(o, "pooler_output", None) is not None else o.last_hidden_state[:, 0]
    else:
        m, pre = _OC[tag]; e = m.encode_image(pre(pil).unsqueeze(0).to(DEV))
    return F.normalize(e.float(), dim=-1)

@torch.no_grad()
def _retrieve(tag, q):                                     # q:[1,512] cuda -> top-M over the gallery
    g = GAL[tag]; step = 500000
    bs = torch.full((M,), -1e9, device=DEV); bi = torch.zeros(M, dtype=torch.long, device=DEV)
    for i in range(0, g.shape[0], step):
        gg = torch.from_numpy(np.asarray(g[i:i+step], dtype="float16")).to(DEV)
        s = (q.half() @ gg.T)[0].float(); cs, ci = s.topk(min(M, s.shape[0]))
        al = torch.cat([bs, cs]); ai = torch.cat([bi, ci + i]); o = al.topk(M).indices
        bs, bi = al[o], ai[o]; del gg, s
    return bi.cpu().numpy(), bs.cpu().numpy()

@spaces.GPU(duration=120)
def locate(pil):
    pil = pil.convert("RGB"); sc = {}
    for t in ENC_D:
        q = F.normalize(HEADS[t](_encode(t, pil)), dim=-1)
        idx, sim = _retrieve(t, q); z = zc(sim)
        for j, zz in zip(idx, z): sc[int(j)] = sc.get(int(j), 0.0) + float(zz)
    ci = np.array(sorted(sc, key=lambda k: -sc[k])[:K], np.int64); csim = np.array([sc[i] for i in ci], "float32")
    cc = GC[ci]
    DW = np.mean([[-np.log(DENS[C].get(int(k), 1)+1.0) for k in cell_key(cc, C)] for C in (0.25,0.5,1.0,2.0)], 0)
    la, lo = np.radians(cc[:, 0]), np.radians(cc[:, 1])
    D = 2*6371.0*np.arcsin(np.sqrt(np.clip(np.sin((la[:,None]-la[None,:])/2)**2 +
        np.cos(la)[:,None]*np.cos(la)[None,:]*np.sin((lo[:,None]-lo[None,:])/2)**2, 0, 1)))
    vote = np.maximum(csim, 0)*np.exp(Gv*zc(DW))
    KV1 = (vote[None,:]*np.exp(-D/BW1)).sum(1); KV2 = (vote[None,:]*np.exp(-D/BW2)).sum(1)
    score = zc(csim) + 0.3*zc(DW) + LAM*zc(KV1) + MU*zc(KV2)
    lat, lon = cc[int(score.argmax())]; return float(lat), float(lon)

def predict(img):
    if img is None: return "Upload a photo first.", ""
    lat, lon = locate(img)
    md = f"### 📍 {lat:.4f}, {lon:.4f}\n[Open in Google Maps](https://maps.google.com/?q={lat:.5f},{lon:.5f})"
    d = 1.5; bbox = f"{lon-d}%2C{lat-d}%2C{lon+d}%2C{lat+d}"
    iframe = (f'<iframe width="100%" height="360" frameborder="0" '
              f'src="https://www.openstreetmap.org/export/embed.html?bbox={bbox}&marker={lat}%2C{lon}"></iframe>')
    return md, iframe

with gr.Blocks(title="Chipoint") as demo:
    gr.Markdown("# Chipoint — where was this photo taken?\n"
                "Upload an outdoor photo; the model guesses its GPS by matching it against 4.9M geotagged images "
                "and beats the published state of the art on OSV-5M. It's a region/city model, not street-address.")
    with gr.Row():
        inp = gr.Image(type="pil", label="Photo")
        with gr.Column():
            out_md = gr.Markdown(); out_map = gr.HTML()
    gr.Button("Locate", variant="primary").click(predict, inp, [out_md, out_map])

demo.queue(max_size=8).launch()