#!/usr/bin/env python3
"""
Cross-video exhaustive matching for two-flight drone footage.

Matches every v1_* frame against every v2_* frame. The within-video
matches from match_features.py are already in the database and are not
touched. After this script, re-run colmap mapper to stitch the scene.
"""
import argparse
import sqlite3
import numpy as np
import cv2
from pathlib import Path

MIN_INLIERS = 15
RATIO_TEST  = 0.75
RANSAC_ERROR = 4.0
KMAX = 2_147_483_647


def pair_id(id1: int, id2: int) -> int:
    lo, hi = (id1, id2) if id1 < id2 else (id2, id1)
    return KMAX * lo + hi


def load_desc_kpts(cur, image_id):
    cur.execute("SELECT rows, cols, data FROM descriptors WHERE image_id=?", (image_id,))
    r = cur.fetchone()
    desc = np.frombuffer(r[2], dtype=np.uint8).reshape(r[0], r[1]) if r else np.zeros((0,128), dtype=np.uint8)

    cur.execute("SELECT rows, cols, data FROM keypoints WHERE image_id=?", (image_id,))
    r = cur.fetchone()
    if r:
        kp = np.frombuffer(r[2], dtype=np.float32).reshape(r[0], r[1])
        kpts = kp[:, :2]
    else:
        kpts = np.zeros((0, 2), dtype=np.float32)
    return desc, kpts


def match_pair(desc1, desc2, kp1, kp2):
    if len(desc1) < 8 or len(desc2) < 8:
        return None, None
    bf  = cv2.BFMatcher(cv2.NORM_L2)
    raw = bf.knnMatch(desc1.astype(np.float32), desc2.astype(np.float32), k=2)
    good = []
    for m_pair in raw:
        if len(m_pair) == 2:
            m, n = m_pair
            if m.distance < RATIO_TEST * n.distance:
                good.append((m.queryIdx, m.trainIdx))
    if len(good) < MIN_INLIERS:
        return None, None
    arr  = np.array(good, dtype=np.uint32)
    pts1 = kp1[arr[:, 0]]
    pts2 = kp2[arr[:, 1]]
    F, mask = cv2.findFundamentalMat(
        pts1, pts2, cv2.FM_RANSAC,
        ransacReprojThreshold=RANSAC_ERROR,
        confidence=0.9999, maxIters=2000,
    )
    if F is None or mask is None:
        return None, None
    inliers = arr[mask.ravel().astype(bool)]
    return (inliers, F) if len(inliers) >= MIN_INLIERS else (None, None)


def write_pair(cur, pid, inliers, F):
    blob   = inliers.astype(np.uint32).tobytes()
    z9     = np.zeros(9, dtype=np.float64).tobytes()
    z4     = np.zeros(4, dtype=np.float64).tobytes()
    z3     = np.zeros(3, dtype=np.float64).tobytes()
    F_blob = F.flatten().astype(np.float64).tobytes()
    cur.execute(
        "INSERT OR REPLACE INTO matches (pair_id, rows, cols, data) VALUES (?,?,?,?)",
        (pid, len(inliers), 2, blob),
    )
    cur.execute(
        "INSERT OR REPLACE INTO two_view_geometries "
        "(pair_id, rows, cols, data, config, F, E, H, qvec, tvec) VALUES (?,?,?,?,?,?,?,?,?,?)",
        (pid, len(inliers), 2, blob, 3, F_blob, z9, z9, z4, z3),
    )


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--db", default="my_scene/database.db")
    args = p.parse_args()

    db  = sqlite3.connect(args.db)
    db.execute("PRAGMA journal_mode=WAL")
    cur = db.cursor()

    cur.execute("SELECT image_id, name FROM images ORDER BY name")
    rows = cur.fetchall()
    v1 = [(id, name) for id, name in rows if name.startswith("v1_")]
    v2 = [(id, name) for id, name in rows if name.startswith("v2_")]
    total_pairs = len(v1) * len(v2)
    print(f"v1={len(v1)} frames  v2={len(v2)} frames  cross-pairs={total_pairs}")

    # Preload all v1 and v2 descriptors into RAM
    print("Loading v1 descriptors…")
    v1_data = {id: load_desc_kpts(cur, id) for id, _ in v1}
    print("Loading v2 descriptors…")
    v2_data = {id: load_desc_kpts(cur, id) for id, _ in v2}

    matched = skipped = i = 0
    for id1, _ in v1:
        desc1, kp1 = v1_data[id1]
        for id2, _ in v2:
            desc2, kp2 = v2_data[id2]
            inliers, F = match_pair(desc1, desc2, kp1, kp2)
            if inliers is not None:
                write_pair(cur, pair_id(id1, id2), inliers, F)
                matched += 1
            else:
                skipped += 1
            i += 1
            if i % 500 == 0:
                pct = 100 * i / total_pairs
                print(f"  [{i}/{total_pairs} {pct:.0f}%] cross-matched={matched}", flush=True)
                db.commit()

    db.commit()
    db.close()
    print(f"\nDone. {matched} cross-video pairs matched, {skipped} below threshold.")
    print("Now delete my_scene/sparse/* and re-run colmap mapper.")


if __name__ == "__main__":
    main()