"""
calib_zhang.py: calibrate a virtual camera with Zhang's method.

EECE 5554 Robotics Sensing and Navigation, demo for Lecture 14.

A virtual camera, whose true K is a little different from its spec sheet (as
real cameras are), takes photos of a flat checkerboard. Zhang's method then:

    1. fits one homography per view, H_i = lambda K [r1 r2 t]_i   (the DLT, Lecture 12)
    2. uses r1 . r2 = 0 and |r1| = |r2| to write two equations per view in
       omega = K^-T K^-1, stacks them as V b = 0, and solves with the SVD
    3. takes K out of omega, and each view's R and t out of K^-1 H
    4. refines K and every R and t by nonlinear least squares (Levenberg-Marquardt)
       on the reprojection error

Real tools (OpenCV, the ROS 2 camera_calibration package) also model lens
distortion; this virtual camera has none, so the steps stay easy to read.

Camera: 1616 x 1240 pixels, spec-sheet f = 1111 px (5 mm lens, 4.5 um pixels).
Board: 8 x 6 inner corners, 108 mm squares, 1.2 to 3 m away.

Run it from a terminal:
    python3 calib_zhang.py                12 views tilted up to 30 degrees
    python3 calib_zhang.py 12 5           the same with only 5 degrees of tilt
"""
import json
import math
import sys

import numpy as np

W, H = 1616, 1240
TRUE_K = {"fx": 1124.0, "fy": 1122.0, "cx": 821.0, "cy": 611.0}     # the camera you have
SPEC_K = {"fx": 1111.0, "fy": 1111.0, "cx": 808.0, "cy": 620.0}     # what the spec sheet says
BOARD = np.array([[(i - 3.5) * 0.108, (j - 2.5) * 0.108] for j in range(6) for i in range(8)])


def kmat(k):
    return np.array([[k["fx"], 0, k["cx"]], [0, k["fy"], k["cy"]], [0, 0, 1.0]])


def rot(rx, ry, rz):
    cx, sx, cy, sy, cz, sz = math.cos(rx), math.sin(rx), math.cos(ry), math.sin(ry), math.cos(rz), math.sin(rz)
    Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
    Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
    Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])
    return Rz @ Ry @ Rx


def rodrigues(r):
    """Rotation vector (axis times angle) to a rotation matrix."""
    th = np.linalg.norm(r)
    if th < 1e-12:
        return np.eye(3)
    k = r / th
    Kx = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
    return np.eye(3) + math.sin(th) * Kx + (1 - math.cos(th)) * Kx @ Kx


def rot_to_vec(R):
    th = math.acos(max(-1.0, min(1.0, (np.trace(R) - 1) / 2)))
    if th < 1e-9:
        return np.zeros(3)
    return th / (2 * math.sin(th)) * np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])


def project(K, R, t, pts):
    """Board points (X, Y, 0) to pixels."""
    pc = pts[:, 0:1] * R[:, 0] + pts[:, 1:2] * R[:, 1] + t
    q = pc @ K.T
    return q[:, :2] / q[:, 2:3]


def make_views(n=12, tilt_deg=30.0, noise_px=0.3, seed=11):
    """n photos of the board: the detected corners of each, with pixel noise. All corners stay in the image."""
    rng = np.random.default_rng(seed)
    K, tilt, views, poses = kmat(TRUE_K), math.radians(tilt_deg), [], []
    for _ in range(4000):
        if len(views) == n:
            break
        R = rot(rng.uniform(-tilt, tilt), rng.uniform(-tilt, tilt), rng.uniform(-0.35, 0.35))
        d = rng.uniform(1.2, 3.0)
        u, v = rng.uniform(280, 1336), rng.uniform(230, 1010)
        t = np.array([d * (u - TRUE_K["cx"]) / TRUE_K["fx"], d * (v - TRUE_K["cy"]) / TRUE_K["fy"], d])
        img = project(K, R, t, BOARD)
        if img[:, 0].min() < 10 or img[:, 0].max() > W - 10 or img[:, 1].min() < 10 or img[:, 1].max() > H - 10:
            continue
        views.append(img + rng.normal(0, noise_px, img.shape))
        poses.append((R, t))
    return views


def normalizer(pts):
    c = pts.mean(axis=0)
    s = math.sqrt(2) / max(np.sqrt(((pts - c) ** 2).sum(axis=1)).mean(), 1e-12)
    return np.array([[s, 0, -s * c[0]], [0, s, -s * c[1]], [0, 0, 1.0]])


def homography(src, dst):
    """The DLT with Hartley normalization: dst ~ H src."""
    Ts, Td = normalizer(src), normalizer(dst)
    hs = np.column_stack([src, np.ones(len(src))]) @ Ts.T
    hd = np.column_stack([dst, np.ones(len(dst))]) @ Td.T
    rows = []
    for (x, y, _), (u, v, _) in zip(hs, hd):
        rows.append([-x, -y, -1, 0, 0, 0, u * x, u * y, u])
        rows.append([0, 0, 0, -x, -y, -1, v * x, v * y, v])
    Hn = np.linalg.svd(np.array(rows))[2][-1].reshape(3, 3)
    Hm = np.linalg.inv(Td) @ Hn @ Ts
    return Hm / Hm[2, 2]


def _v(Hm, i, j):
    hi, hj = Hm[:, i], Hm[:, j]
    return np.array([hi[0] * hj[0], hi[0] * hj[1] + hi[1] * hj[0], hi[1] * hj[1],
                     hi[2] * hj[0] + hi[0] * hj[2], hi[2] * hj[1] + hi[1] * hj[2], hi[2] * hj[2]])


def zhang_closed_form(views):
    """K from the homographies. Pixels are scaled first (like Hartley's normalization) and K is scaled back."""
    N = np.array([[1e-3, 0, -0.808], [0, 1e-3, -0.620], [0, 0, 1.0]])
    Hs = [homography(BOARD, v) for v in views]
    V = []
    for Hm in Hs:
        Hn = N @ Hm
        Hn = Hn / np.linalg.norm(Hn)
        V.append(_v(Hn, 0, 1))
        V.append(_v(Hn, 0, 0) - _v(Hn, 1, 1))
    B11, B12, B22, B13, B23, B33 = np.linalg.svd(np.array(V))[2][-1]
    den = B11 * B22 - B12 ** 2
    v0 = (B12 * B13 - B11 * B23) / den
    lam = B33 - (B13 ** 2 + v0 * (B12 * B13 - B11 * B23)) / B11
    if lam / B11 <= 0 or lam * B11 / den <= 0:
        return None, Hs
    al, be = math.sqrt(lam / B11), math.sqrt(lam * B11 / den)
    ga = -B12 * al * al * be / lam
    u0 = ga * v0 / be - B13 * al * al / lam
    K = np.linalg.inv(N) @ np.array([[al, ga, u0], [0, be, v0], [0, 0, 1.0]])
    return K / K[2, 2], Hs


def poses_from_k(K, Hs):
    """Each view's R and t from K^-1 H, with R made a true rotation."""
    Ki, out = np.linalg.inv(K), []
    for Hm in Hs:
        a = Ki @ Hm
        lam = 1 / np.linalg.norm(a[:, 0])
        r1, r2, t = lam * a[:, 0], lam * a[:, 1], lam * a[:, 2]
        if t[2] < 0:
            r1, r2, t = -r1, -r2, -t
        U, _, Vt = np.linalg.svd(np.column_stack([r1, r2, np.cross(r1, r2)]))
        R = U @ Vt
        if np.linalg.det(R) < 0:
            R = U @ np.diag([1, 1, -1]) @ Vt
        out.append((R, t))
    return out


def reprojection_rms(K, poses, views):
    e = [project(K, R, t, BOARD) - v for (R, t), v in zip(poses, views)]
    return float(np.sqrt((np.concatenate(e) ** 2).sum(axis=1).mean()))


def refine(K, poses, views, iters=30):
    """Levenberg-Marquardt on fx, fy, cx, cy and every view's rotation vector and translation."""
    x = np.concatenate([[K[0, 0], K[1, 1], K[0, 2], K[1, 2]]] + [np.concatenate([rot_to_vec(R), t]) for R, t in poses])

    def unpack(x):
        Km = np.array([[x[0], 0, x[2]], [0, x[1], x[3]], [0, 0, 1.0]])
        return Km, [(rodrigues(x[4 + 6 * i:7 + 6 * i]), x[7 + 6 * i:10 + 6 * i]) for i in range(len(views))]

    def resid(x):
        Km, ps = unpack(x)
        return np.concatenate([(project(Km, R, t, BOARD) - v).ravel() for (R, t), v in zip(ps, views)])

    mu, r = 1e-3, resid(x)
    for _ in range(iters):
        J = np.empty((len(r), len(x)))
        for k in range(len(x)):                       # numerical Jacobian, one column per unknown
            h = 1e-6 * max(1.0, abs(x[k]))
            xp = x.copy(); xp[k] += h
            J[:, k] = (resid(xp) - r) / h
        A, g = J.T @ J, J.T @ r
        while True:
            dx = np.linalg.solve(A + mu * np.diag(np.diag(A) + 1e-12), -g)
            r_new = resid(x + dx)
            if (r_new ** 2).sum() < (r ** 2).sum():
                x, r, mu = x + dx, r_new, mu * 0.3
                break
            mu *= 10
            if mu > 1e8:
                break
        if np.abs(dx).max() < 1e-8 or mu > 1e8:
            break
    Km, ps = unpack(x)
    return Km, ps


def as_dict(K):
    return {"fx": float(K[0, 0]), "fy": float(K[1, 1]), "cx": float(K[0, 2]), "cy": float(K[1, 2]), "skew": float(K[0, 1])}


def calibrate(views, do_refine=True):
    K0, Hs = zhang_closed_form(views)
    if K0 is None or not np.all(np.isfinite(K0)):
        return {"ok": False, "reason": "The views do not pin K down: add tilt, or more views."}
    poses = poses_from_k(K0, Hs)
    out = {"ok": True, "closed": as_dict(K0), "rms_closed": reprojection_rms(K0, poses, views)}
    if do_refine:
        K1, p1 = refine(K0, poses, views)
        out["refined"], out["rms_refined"] = as_dict(K1), reprojection_rms(K1, p1, views)
    return out


def demo_api(params_json):
    """Web page entry point. Input: {"n": views, "tilt": deg, "noise": px, "seed": int, "refine": bool}."""
    p = json.loads(params_json or "{}")
    views = make_views(int(p.get("n", 12)), float(p.get("tilt", 30)), float(p.get("noise", 0.3)), int(p.get("seed", 11)))
    out = calibrate(views, bool(p.get("refine", True)))
    out.update({"W": W, "H": H, "true": TRUE_K, "spec": SPEC_K, "views": [v.round(1).tolist() for v in views]})
    return json.dumps(out)


def main(argv):
    n = int(argv[1]) if len(argv) > 1 else 12
    tilt = float(argv[2]) if len(argv) > 2 else 30.0
    r = json.loads(demo_api(json.dumps({"n": n, "tilt": tilt})))
    print("{} views, tilt up to {:.0f} deg".format(n, tilt))
    if not r["ok"]:
        print(r["reason"]); return
    for name in ("fx", "fy", "cx", "cy"):
        print("  {}: spec {:7.1f}   true {:7.1f}   closed form {:7.1f}   refined {:7.1f}".format(
            name, r["spec"][name], r["true"][name], r["closed"][name], r["refined"][name]))
    print("  reprojection RMS: closed form {:.2f} px, refined {:.2f} px".format(r["rms_closed"], r["rms_refined"]))


if __name__ == "__main__":
    main(sys.argv)