# "When Attention Meets Geometry": learned matchers on the benchmark of "Finding the Same Point Twice", what LightGlue's
# cross-attention looks at, and monocular depth from a DINOv2-based network. CPU only (4 threads).
#   1. SIFT, ORB, SuperPoint + SuperGlue, SuperPoint + LightGlue and LoFTR (detector-free) on the same 4 photos and the same
#      known transformations as before (+ a "combined" family that stacks several changes), with time per image pair
#   2. LightGlue with early stopping and pruning switched off: per layer, how much of each keypoint's cross-attention
#      lands on its true match (and the top-20 attention weights of 100 keypoints, for the lab)
#   3. Depth Anything V2 (small) on the 4 photos
#   python bridge_experiments.py IMG_DIR OUT_DIR   (needs scripts/features/match_experiments.py next to this folder)
import sys, os, json, time, numpy as np, cv2, torch
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForKeypointMatching, AutoModelForDepthEstimation
torch.set_num_threads(4)
IMG_DIR, OUT = sys.argv[1], sys.argv[2]; os.makedirs(OUT, exist_ok=True)
HERE = os.path.dirname(os.path.abspath(__file__))
src = open(f"{HERE}/../features/match_experiments.py").read().split('res = {"suite": {}, "dets": DETS}')[0]
sys.argv = [sys.argv[0], IMG_DIR, OUT]; g = {"__name__": "shared"}; exec(src, g)   # the same photos, transformations and helpers
IMGS, SUITE, make, rot_scale, perspective, W, H, NAMES = (g[k] for k in ["IMGS", "SUITE", "make", "rot_scale", "perspective", "W", "H", "NAMES"])
SUITE = dict(SUITE)
SUITE["combined"] = [(f"level {k}", dict(Hm=perspective(0.06 * k, seed=11) @ rot_scale(20 * k, 1 / (1 + 0.2 * k)), gain=1 - 0.1 * k, blur=0.4 * k)) for k in range(6)]

# LoFTR comes from Kornia (outdoor weights): the EfficientLoFTR port in transformers 5.18 returned wrong matches as soon as
# the second image was rotated (pure shifts were fine), while Kornia's LoFTR matched rotated images correctly.
import kornia.feature as KF
LEARNED = {"SuperPoint + SuperGlue": "magic-leap-community/superglue_outdoor", "SuperPoint + LightGlue": "ETH-CVG/lightglue_superpoint"}
models = {}
for name, hub in LEARNED.items():
    proc = AutoImageProcessor.from_pretrained(hub); m = AutoModelForKeypointMatching.from_pretrained(hub).eval()
    if hasattr(m.config, "keypoint_detector_config"):
        try: m.keypoint_detector.config.max_keypoints = 2048     # the same budget as SIFT and ORB (2,000)
        except AttributeError: pass
    models[name] = (proc, m)
loftr = KF.LoFTR(pretrained="outdoor").eval()   # needs SSL_CERT_FILE=$(python -m certifi) on python.org builds for the download

def classical(name, a, b):
    det, norm = (cv2.SIFT_create(nfeatures=2000), cv2.NORM_L2) if name == "SIFT" else (cv2.ORB_create(nfeatures=2000), cv2.NORM_HAMMING)
    ka, da = det.detectAndCompute(a, None); kb, db = det.detectAndCompute(b, None)
    if da is None or db is None or len(ka) < 2 or len(kb) < 2: return np.zeros((0, 2)), np.zeros((0, 2))
    good = [m for m, n in (p for p in cv2.BFMatcher(norm).knnMatch(da, db, k=2) if len(p) == 2) if m.distance < 0.8 * n.distance]
    return np.float32([ka[m.queryIdx].pt for m in good]).reshape(-1, 2), np.float32([kb[m.trainIdx].pt for m in good]).reshape(-1, 2)
def learned(name, a, b):
    if name == "LoFTR":
        t = lambda im: torch.from_numpy(im).float()[None, None] / 255
        with torch.no_grad(): r = loftr({"image0": t(a), "image1": t(b)})
        keep = r["confidence"].numpy() > 0.2
        return r["keypoints0"].numpy()[keep].astype(np.float32), r["keypoints1"].numpy()[keep].astype(np.float32)
    proc, m = models[name]; ia, ib = Image.fromarray(a).convert("RGB"), Image.fromarray(b).convert("RGB")
    inp = proc([[ia, ib]], return_tensors="pt")
    with torch.no_grad(): out = m(**inp)
    r = proc.post_process_keypoint_matching(out, [[(H, W), (H, W)]], threshold=0.2)[0]
    return r["keypoints0"].numpy().astype(np.float32), r["keypoints1"].numpy().astype(np.float32)
METHODS = ["SIFT", "ORB"] + list(LEARNED) + ["LoFTR"]
def evaluate(method, a, b, Hm, tol=3.0):
    t0 = time.time(); pa, pb = (classical if method in ("SIFT", "ORB") else learned)(method, a, b); dt = time.time() - t0
    if len(pa) == 0: return {"matches": 0, "correct": 0, "precision": 0.0, "h_ok": 0.0, "time": dt}
    gt = cv2.perspectiveTransform(pa[None].astype(np.float32), Hm)[0]; ok = np.linalg.norm(gt - pb, axis=1) < tol
    h_ok = 0.0
    if len(pa) >= 4:
        Hest, _ = cv2.findHomography(pa, pb, cv2.RANSAC, 3.0)
        if Hest is not None:
            c = np.float32([[0, 0], [W, 0], [W, H], [0, H]])[None]
            h_ok = float(np.linalg.norm(cv2.perspectiveTransform(c, Hest)[0] - cv2.perspectiveTransform(c, Hm)[0], axis=1).mean() < 5)
    return {"matches": int(len(pa)), "correct": int(ok.sum()), "precision": float(ok.mean()), "h_ok": h_ok, "time": dt}

res = {"methods": METHODS, "suite": {}}; t0 = time.time()
if os.environ.get("REUSE_SUITE") and os.path.exists(f"{OUT}/bridge_suite.json"): res = json.load(open(f"{OUT}/bridge_suite.json")); SUITE = {}
for fam, cases in SUITE.items():
    res["suite"][fam] = []
    for label, kw in cases:
        row = {"label": label}
        for meth in METHODS:
            rs = []
            for i, n in enumerate(NAMES):
                b, _ = make(IMGS[n], seed=i, **kw); rs.append(evaluate(meth, IMGS[n], b, kw["Hm"]))
            row[meth] = {k: float(np.mean([r[k] for r in rs])) for k in rs[0]}
        res["suite"][fam].append(row)
        print(fam, label, {m: (int(row[m]["correct"]), round(row[m]["precision"], 3), row[m]["h_ok"], round(row[m]["time"], 2)) for m in METHODS}, f"{time.time() - t0:.0f}s", flush=True)

json.dump(res, open(f"{OUT}/bridge_suite.json", "w"), separators=(",", ":"))   # saved before the slower analyses below
# ---------- 2. LightGlue's cross-attention, layer by layer ----------
proc, _ = models["SuperPoint + LightGlue"]
lg = AutoModelForKeypointMatching.from_pretrained(LEARNED["SuperPoint + LightGlue"], attn_implementation="eager").eval()
lg.config.depth_confidence = -1.0; lg.config.width_confidence = -1.0      # all 9 layers, no pruning
for attr in ("depth_confidence", "width_confidence"):
    if hasattr(lg, attr): setattr(lg, attr, -1.0)
lg.keypoint_detector.config.max_keypoints = 512
pair = json.load(open(f"{OUT}/match_pair.json")) if os.path.exists(f"{OUT}/match_pair.json") else None
Hlab = np.array(pair["Hgt"]) if pair else (perspective(0.1, seed=7) @ rot_scale(25, 0.85))
a = IMGS["street"]; b, _ = make(a, Hm=Hlab, gain=0.75, seed=1)
inp = proc([[Image.fromarray(a).convert("RGB"), Image.fromarray(b).convert("RGB")]], return_tensors="pt")
with torch.no_grad(): out = lg(**inp, output_attentions=True)
kp = out.keypoints[0].numpy() * np.array([W, H]) if out.keypoints.max() <= 1.5 else out.keypoints[0].numpy()
mask = out.mask[0].numpy().astype(bool); n0, n1 = int(mask[0].sum()), int(mask[1].sum())
k0, k1 = kp[0][:n0], kp[1][:n1]
att = out.attentions; print("attention tensors:", len(att), [tuple(t.shape) for t in att[:2]], flush=True)
gt = cv2.perspectiveTransform(k0[None].astype(np.float32), Hlab)[0]
dist = np.linalg.norm(gt[:, None] - k1[None], axis=2); true_j = dist.argmin(1); has = dist.min(1) < 3.0
layers = []
cross = [att[i] for i in range(1, len(att), 2)]   # LightGlue alternates self- and cross-attention
for li, A in enumerate(cross):
    Am = A[0].mean(0).numpy() if A.shape[0] == 2 else A.mean(1)[0].numpy()   # mean over heads; rows: queries of image 0
    A01 = Am[:n0, :n1] / Am[:n0, :n1].sum(1, keepdims=True)
    on_true = float(A01[np.where(has)[0], true_j[has]].mean()); ent = float((-A01 * np.log(A01 + 1e-12)).sum(1).mean())
    layers.append({"layer": li + 1, "mass_on_true_match": on_true, "entropy": ent, "uniform_mass": 1.0 / n1})
    print("cross layer", li + 1, round(on_true, 4), round(ent, 3), flush=True)
rs = np.random.default_rng(0); pick = rs.choice(np.where(has)[0], min(100, has.sum()), replace=False)
viz = []
for i in pick:
    per = []
    for A in cross:
        Am = A[0].mean(0).numpy() if A.shape[0] == 2 else A.mean(1)[0].numpy(); row = Am[i, :n1] / Am[i, :n1].sum(); top = np.argsort(-row)[:20]
        per.append([[int(j), round(float(row[j]), 4)] for j in top])
    viz.append({"i": int(i), "true": int(true_j[i]), "top": per})
res["attention"] = {"layers": layers, "n0": n0, "n1": n1, "k0": np.round(k0, 1).tolist(), "k1": np.round(k1, 1).tolist(), "viz": viz, "H": Hlab.tolist()}
cv2.imwrite(f"{OUT}/bridge-a.jpg", a, [cv2.IMWRITE_JPEG_QUALITY, 82]); cv2.imwrite(f"{OUT}/bridge-b.jpg", b, [cv2.IMWRITE_JPEG_QUALITY, 82])

# ---------- 3. monocular depth ----------
dproc = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf"); dm = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf").eval()
res["depth"] = []
for n in NAMES:
    im = Image.open(f"{IMG_DIR}/{n}.jpg").convert("RGB"); w, h = im.size; s = 448 / min(w, h)
    im = im.resize((round(w * s), round(h * s)), Image.BICUBIC); w, h = im.size; im = im.crop(((w - 448) // 2, (h - 448) // 2, (w - 448) // 2 + 448, (h - 448) // 2 + 448))
    t1 = time.time()
    with torch.no_grad(): d = dm(**dproc(images=im, return_tensors="pt")).predicted_depth[0].numpy()
    d = cv2.resize(d, (448, 448), interpolation=cv2.INTER_CUBIC); q = (255 * (d - d.min()) / (d.max() - d.min())).astype(np.uint8)
    cv2.imwrite(f"{OUT}/depth-{n}.png", q); res["depth"].append({"name": n, "time": round(time.time() - t1, 2), "min": float(d.min()), "max": float(d.max())})
    print("depth", n, res["depth"][-1], flush=True)
json.dump(res, open(f"{OUT}/bridge.json", "w"), separators=(",", ":"))
print("done", f"{time.time() - t0:.0f}s")