"""Fit the picture's camera to the cell's outer box, headless (numpy from Blender):

    blender.exe -b --factory-startup --python rc_fit.py        -> build/cam_fit.json

The cell's frame is a box: the six corners the picture shows (front four, back-right two) were read from the ruler
crops reference/z_frt.png, z_frb.png, z_tl.png, z_bl.png, z_brt.png, z_brb.png. Unknowns: the camera's place, its
yaw / pitch / roll, the focal length (principal point at the picture's middle), the box's width W and depth D; the
height H is 1 (the scale is set afterwards). World: x right along the front, y back (the front faces -y), z up.
"""
import math, os, json
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
RES = (1200, 960)
PIX = {"FLT": (76.5, 97.5), "FLB": (76.5, 829.0), "FRT": (1011.5, 45.5), "FRB": (1012.0, 922.0),
       "BRT": (1136.0, 172.0), "BRB": (1132.0, 763.0)}


def corners(W, D, H=1.0):
    return {"FLT": (-W / 2, -D / 2, H), "FLB": (-W / 2, -D / 2, 0.0), "FRT": (W / 2, -D / 2, H),
            "FRB": (W / 2, -D / 2, 0.0), "BRT": (W / 2, D / 2, H), "BRB": (W / 2, D / 2, 0.0)}


def rot(yaw, pitch, roll):
    """camera axes (right, up, forward) in world: forward = yaw about z from +y, pitch up, roll about forward"""
    f = np.array([math.sin(yaw) * math.cos(pitch), math.cos(yaw) * math.cos(pitch), math.sin(pitch)])
    r = np.cross(f, [0, 0, 1.0]); r /= np.linalg.norm(r)
    u = np.cross(r, f)
    cr, sr = math.cos(roll), math.sin(roll)
    return cr * r + sr * u, -sr * r + cr * u, f


def project(p, X):
    cx, cy, cz, yaw, pitch, roll, lf = p[:7]
    r, u, f = rot(yaw, pitch, roll)
    d = X - np.array([cx, cy, cz])
    zc = d @ f
    F = math.exp(lf)
    return np.stack([RES[0] / 2 + F * (d @ r) / zc, RES[1] / 2 - F * (d @ u) / zc], 1)


def resid(p, names, fixed):
    q = list(p)
    for k, v in fixed.items():
        q.insert(k, v)
    W, D = q[7], q[8]
    C = corners(W, D)
    X = np.array([C[n] for n in names])
    UV = np.array([PIX[n] for n in names])
    return (project(q, X) - UV).ravel()


def lm(p, names, fixed, it=300):
    p = np.array(p, float)
    lam = 1e-2
    r = resid(p, names, fixed)
    for _ in range(it):
        J = np.zeros((len(r), len(p)))
        for k in range(len(p)):
            dp = np.zeros(len(p)); dp[k] = 1e-6 * max(1.0, abs(p[k]))
            J[:, k] = (resid(p + dp, names, fixed) - r) / dp[k]
        A = J.T @ J
        g = J.T @ r
        while True:
            step = np.linalg.solve(A + lam * np.diag(np.diag(A) + 1e-9), -g)
            r2 = resid(p + step, names, fixed)
            if r2 @ r2 < r @ r:
                p, r = p + step, r2
                lam *= 0.3
                break
            lam *= 10
            if lam > 1e12:
                return p, r
    return p, r


names = list(PIX)
out = {}
# guess: the camera in front-right, at 0.6 H, about 3.5 H away
g = [1.2, -3.5, 0.6, math.radians(-20), 0.0, 0.0, math.log(2000.0), 1.3, 0.8]
for label, fixed in (("free", {}), ("level", {4: 0.0, 5: 0.0})):
    p0 = [x for i, x in enumerate(g) if i not in fixed]
    p, r = lm(p0, names, fixed)
    q = list(p)
    for k, v in fixed.items():
        q.insert(k, v)
    out[label] = dict(cam=[round(x, 4) for x in q[:3]], yaw_deg=round(math.degrees(q[3]), 3),
                      pitch_deg=round(math.degrees(q[4]), 3), roll_deg=round(math.degrees(q[5]), 3),
                      focal_px=round(math.exp(q[6]), 2), lens_mm_36=round(math.exp(q[6]) / RES[0] * 36, 2),
                      W=round(q[7], 4), D=round(q[8], 4), rms_px=round(float(np.sqrt(np.mean(r ** 2))), 3),
                      resid=[round(x, 2) for x in r])
json.dump(out, open(os.path.join(HERE, "cam_fit.json"), "w"), indent=1)
print("FIT", json.dumps(out))