# -*- coding: utf-8 -*-
"""Verify COLMAP->c2w parsing against provided GT, and report trajectory radii."""
import numpy as np, json, sys
sys.path.insert(0, r"scripts")
from export_cameras import quat_to_rot, parse_images_txt, umeyama


def c2w_from_gt_sparse(images_txt):
    """Parse GT sparse/images.txt -> {name: c2w}."""
    out = {}
    with open(images_txt) as f:
        lines = f.readlines()
    i = 0
    while i < len(lines):
        line = lines[i].strip()
        if not line or line.startswith('#'):
            i += 1
            continue
        t = line.split()
        q = np.array([float(x) for x in t[1:5]])
        tvec = np.array([float(x) for x in t[5:8]])
        name = t[9]
        i += 1
        if i < len(lines):
            i += 1
        R = quat_to_rot(q)
        c2w = np.eye(4)
        c2w[:3, :3] = R.T
        c2w[:3, 3] = -R.T @ tvec
        out[name] = c2w
    return out


def main():
    # 1) sanity: GT sparse/images.txt vs transforms.json
    gt_sparse = c2w_from_gt_sparse(r"mvs_project\sparse\images.txt")
    gt_tr = json.load(open(r"mvs_project\poses\transforms.json"))
    gt_tr_map = {f["file_path"].split('/')[-1]: np.array(f["transform_matrix"]) for f in gt_tr["frames"]}
    common = sorted(set(gt_sparse) & set(gt_tr_map))
    errs = []
    for n in common:
        errs.append(np.abs(gt_sparse[n] - gt_tr_map[n]).max())
    print(f"[GT self-check] {len(common)} frames, max |sparse-c2w minus transforms-c2w| = {max(errs):.2e}")
    print("  (should be ~0 -> my quat_to_rot/c2w construction is correct)\n")

    # 2) GT trajectory radius
    C_gt = np.array([gt_tr_map[n][:3, 3] for n in common])
    print(f"GT centers: n={len(C_gt)}, centroid={C_gt.mean(0).round(3)}, "
          f"mean radius={np.linalg.norm(C_gt-C_gt.mean(0), axis=1).mean():.3f} m")

    # 3) COLMAP-estimated trajectory radius
    d = np.load(r"output\colmap\cameras_export\poses_c2w.npz", allow_pickle=True)
    names, c2w = d["names"], d["c2w"]
    C_est = c2w[:, :3, 3]
    print(f"COLMAP centers: n={len(C_est)}, centroid={C_est.mean(0).round(3)}, "
          f"mean radius={np.linalg.norm(C_est-C_est.mean(0), axis=1).mean():.3f} m")

    # 4) proper Umeyama + per-frame residual
    A, B, anames = [], [], []
    for n, c in zip(names, C_est):
        if n in gt_tr_map:
            A.append(c); B.append(gt_tr_map[n][:3, 3]); anames.append(n)
    A, B = np.array(A), np.array(B)
    s, R, t, rms, aligned = umeyama(A, B)
    res = np.linalg.norm(aligned - B, axis=1)
    print(f"\nUmeyama: scale={s:.4f}, RMS={rms:.3f} m, median={np.median(res):.3f} m, max={res.max():.3f} m")
    idx = np.argsort(-res)[:6]
    for i in idx:
        print(f"   worst: {anames[i]}  {res[i]:.3f} m")
    # per-ring check
    for n in anames:
        pass
    print("\n(COLMAP world is a similarity gauge; RMS measures shape accuracy after optimal alignment)")


if __name__ == "__main__":
    main()