"""Independently re-run COLMAP SfM on each part to verify the reported failure.

Uses the colmap CLI (3.8, CPU SIFT). Records feature counts, match counts and
how many images the incremental mapper manages to register.
"""
import os, sys, json, shutil, subprocess, sqlite3

SRC = "repo/input_images"
OUT = "work/colmap_verify"
PARTS = ["bracket", "connector_module", "flange_housing", "timing_pulley"]


def sh(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    return p.returncode, p.stdout + p.stderr


def db_stats(db):
    con = sqlite3.connect(db); cur = con.cursor()
    kp = cur.execute("SELECT image_id, rows FROM keypoints").fetchall()
    nm = cur.execute("SELECT COUNT(*), SUM(rows) FROM matches WHERE rows>0").fetchone()
    ntv = cur.execute("SELECT COUNT(*), SUM(rows) FROM two_view_geometries WHERE rows>0").fetchone()
    con.close()
    return {"keypoints_per_image": {str(i): r for i, r in kp},
            "match_pairs": nm[0], "total_matches": nm[1] or 0,
            "verified_pairs": ntv[0], "verified_matches": ntv[1] or 0}


results = {}
for part in PARTS:
    ws = os.path.join(OUT, part)
    if os.path.exists(ws): shutil.rmtree(ws)
    os.makedirs(os.path.join(ws, "sparse"), exist_ok=True)
    imdir = os.path.join(ws, "images"); os.makedirs(imdir)
    src = os.path.join(SRC, part, "input")
    for f in sorted(os.listdir(src)):
        shutil.copy(os.path.join(src, f), imdir)
    db = os.path.join(ws, "database.db")

    rc, o = sh(["colmap", "feature_extractor", "--database_path", db,
                "--image_path", imdir,
                "--SiftExtraction.use_gpu", "0",
                "--SiftExtraction.estimate_affine_shape", "1",
                "--SiftExtraction.domain_size_pooling", "1"])
    rc2, o2 = sh(["colmap", "exhaustive_matcher", "--database_path", db,
                  "--SiftMatching.use_gpu", "0",
                  "--SiftMatching.guided_matching", "1"])
    st = db_stats(db)

    rc3, o3 = sh(["colmap", "mapper", "--database_path", db, "--image_path", imdir,
                  "--output_path", os.path.join(ws, "sparse"),
                  "--Mapper.init_min_num_inliers", "15",
                  "--Mapper.abs_pose_min_num_inliers", "10",
                  "--Mapper.min_num_matches", "8"])
    models = [d for d in sorted(os.listdir(os.path.join(ws, "sparse")))
              if os.path.isdir(os.path.join(ws, "sparse", d))]
    reg = 0; npts = 0
    if models:
        m = os.path.join(ws, "sparse", models[0])
        rc4, _ = sh(["colmap", "model_converter", "--input_path", m,
                     "--output_path", m, "--output_type", "TXT"])
        ip = os.path.join(m, "images.txt")
        if os.path.exists(ip):
            reg = sum(1 for l in open(ip) if l.strip() and not l.startswith("#")) // 2
        pp = os.path.join(m, "points3D.txt")
        if os.path.exists(pp):
            npts = sum(1 for l in open(pp) if l.strip() and not l.startswith("#"))
    n_imgs = len(os.listdir(imdir))
    results[part] = {"n_images": n_imgs, **st, "n_models": len(models),
                     "registered": reg, "points3D": npts,
                     "mapper_tail": o3.strip().splitlines()[-6:]}
    print(f"{part}: imgs={n_imgs} kpts={list(st['keypoints_per_image'].values())} "
          f"verified_pairs={st['verified_pairs']} verified_matches={st['verified_matches']} "
          f"-> models={len(models)} registered={reg} points={npts}")

os.makedirs(OUT, exist_ok=True)
json.dump(results, open(os.path.join(OUT, "colmap_verify.json"), "w"), indent=2)
print("\nwrote", os.path.join(OUT, "colmap_verify.json"))