"""Fit a pinhole camera to each of the picture's three views from landmarks, headless (numpy from Blender):

    blender.exe -b --factory-startup --python nv_fit.py        -> build/cams.json

Landmarks are model points (middle-view pixels, see nv_lib.py) and where the picture shows them (picture pixels,
1000 x 579, read from the ruler crops). Each view is a panel of its own; its principal point is fixed at the panel's
middle. Solved by Levenberg-Marquardt over position, rotation and focal length.
"""
import bpy, math, os, json
import numpy as np
from mathutils import Vector, Matrix, Euler

HERE = os.path.dirname(os.path.abspath(__file__))
exec(open(os.path.join(HERE, "nv_lib.py")).read())

W_IMG, H_IMG = 1000, 579
PANELS = {"left": (172.0, 289.5), "mid": (520.0, 289.5), "right": (835.0, 289.5)}

# name: (model point x, u, v) and picture points per view
MARKS = {
    "knob1":    ((57, 540, 373), {"left": (250, 367), "mid": (540, 373), "right": (861, 416)}),
    "knob2":    ((57, 522, 393), {"left": (238, 388), "mid": (522, 393), "right": (841, 433)}),
    "can_top":  ((55, 578, 189), {"left": (285, 201), "mid": (578, 191), "right": (903, 247)}),
    "grille":   ((55, 495, 84), {"left": (221, 90), "mid": (495, 84), "right": (802, 132)}),
    "label":    ((55, 540, 132), {"left": (252, 145), "mid": (540, 132), "right": (854, 191)}),
    "decal":    ((55, 547, 348), {"left": (256, 337), "mid": (547, 348), "right": (861, 394)}),
    "head_fl":  ((-53, 409, 71), {"left": (93, 82), "mid": (409, 71)}),
    "head_fr":  ((53, 409, 71), {"left": (185, 84), "mid": (409, 71), "right": (723, 106)}),
    "head_rr":  ((53, 562, 73), {"mid": (562, 73), "right": (895, 132)}),
    "band":     ((0, 420, 226), {"left": (144, 225), "mid": (420, 226)}),
    "low_fl":   ((-53, 411, 394), {"left": (92, 402), "mid": (411, 394)}),
    "keypad":   ((-25, 409, 433), {"left": (124, 435)}),
    "glow":     ((0, 471, 152), {"left": (143, 155)}),
    "pill1_top": ((56, 551, 156), {"mid": (551, 156), "right": (872, 210)}),
    "plinth_fb": ((0, 400, 500), {"mid": (401, 500)}),
    "plinth_fbr": ((62, 401, 503), {"right": (737, 503)}),
    "plinth_rbr": ((62, 636, 503), {"right": (957, 526)}),
}


def look_rot(cam, tgt):
    d = (tgt - cam).normalized()
    return d.to_track_quat('-Z', 'Y').to_euler()


def project(p, X3):
    """p: [cx, cy, cz, rx, ry, rz, logf, ppx, ppy] -> picture pixels of points X3 (n x 3)"""
    C = np.array(p[:3])
    R = np.array(Euler((p[3], p[4], p[5]), 'XYZ').to_matrix())       # camera to world
    pc = (X3 - C) @ R                                                   # world to camera (R^T x)
    f = math.exp(p[6])
    x = p[7] + f * pc[:, 0] / -pc[:, 2]
    y = p[8] - f * pc[:, 1] / -pc[:, 2]
    return np.stack([x, y], 1)


def fit(view, guess_dir, elev, dist, focal_px):
    names = [n for n, (_, im) in MARKS.items() if view in im]
    X3 = np.array([list(P(*MARKS[n][0])) for n in names])
    UV = np.array([MARKS[n][1][view] for n in names], float)
    tgt = P(0, 486, 290)
    d = Vector(guess_dir).normalized()
    cam = tgt + Vector((d.x * math.cos(elev), d.y * math.cos(elev), math.sin(elev))) * dist
    e = look_rot(cam, tgt)
    cx, cy = PANELS[view]
    p = np.array([cam.x, cam.y, cam.z, e.x, e.y, e.z, math.log(focal_px), cx, cy])
    free = list(range(7))
    lam = 1e-2

    def res(q):
        return (project(q, X3) - UV).ravel()
    r = res(p)
    for it in range(400):
        J = np.zeros((len(r), len(free)))
        for k, i in enumerate(free):
            h = 1e-5 if i >= 3 else 1e-4
            q = p.copy()
            q[i] += h
            J[:, k] = (res(q) - r) / h
        A = J.T @ J
        g = J.T @ r
        while True:
            step = np.linalg.solve(A + lam * np.diag(np.diag(A) + 1e-9), -g)
            q = p.copy()
            q[free] += step
            r2 = res(q)
            if r2 @ r2 < r @ r:
                p, r = q, r2
                lam = max(lam * 0.3, 1e-9)
                break
            lam *= 10
            if lam > 1e9:
                break
        if lam > 1e9 or np.abs(step).max() < 1e-10:
            break
    err = np.sqrt((r.reshape(-1, 2) ** 2).sum(1))
    return dict(params=p.tolist(), names=names, err=err.round(2).tolist(), rms=float(np.sqrt((err ** 2).mean())),
                lens_mm=math.exp(p[6]) * 36.0 / W_IMG,
                shift=((W_IMG / 2 - p[7]) / W_IMG, (p[8] - H_IMG / 2) / W_IMG))   # Blender's shift moves the frame


def main():
    out = {}
    out["left"] = fit("left", (0.55, -1.0, 0), math.radians(5), 7.0, 2500)
    out["mid"] = fit("mid", (1.0, 0.0, 0), math.radians(2), 9.0, 3000)
    out["right"] = fit("right", (1.0, -0.35, 0), math.radians(18), 7.0, 2500)
    for k, v in out.items():
        c = v["params"]
        print("FIT", k, "rms %.2f" % v["rms"], "lens %.0f mm" % v["lens_mm"], "cam", [round(x, 2) for x in c[:3]],
              "rot", [round(math.degrees(x), 1) for x in c[3:6]], dict(zip(v["names"], v["err"])))
    with open(os.path.join(HERE, "cams.json"), "w") as fh:
        json.dump(out, fh, indent=1)


main()