"""Fit the picture's camera and the model's key points together (numpy, Levenberg-Marquardt).

python tg_fit.py   -> build/cam.json (camera params, lens, the key points as fitted)

One picture cannot give depth by itself, so every key point starts from the Mako sheet's reading of the same design
(tender-mako/build/mk_measure.py: SIDE heights, TOP and FRONT lateral offsets, same frame and scale) and may move from it
where the picture needs it, at a price (sigma). Points on the centre plane stay on it; the two big fins' tips are one
mirrored pair, which is what fixes the view's yaw. The nose and the tail fin's tip are pinned in length to keep the scale.
"""
import math
import json
import numpy as np
from tg_cam import project, look_at, lens_mm, HERE

S, X0, Z0, DZ = 0.0075, 324.0, 805.333, 0.10      # the Mako sheet's frame (v04)


def sheet(xs, ys, lat=0.0):
    return [lat, (xs - X0) * S, (Z0 - ys) * S + DZ]


# name: (prior [X, Y, Z], sigma [X, Y, Z] (0 = fixed), kind, picture pixel(s))
#   kind "c": on the centre plane (X = 0); "m": mirrored pair, obs = (near/port, far/starboard); "p": a free point
MARKS = {
    "nose":     (sheet(41, 788), (0, 0.0, 0.10), "c", [(15, 412)]),
    "dorsal":   (sheet(538, 547), (0, 0.6, 0.5), "c", [(427, 90)]),
    "tail":     (sheet(860, 657), (0, 0.0, 0.4), "c", [(711, 201)]),
    "tail_te":  (sheet(893, 748), (0, 0.4, 0.3), "c", [(750, 321)]),
    "ventral":  (sheet(868, 948), (0, 0.5, 0.5), "c", [(746, 523)]),
    "fin_tip":  ([1.42, 2.1075, -1.68], (0.5, 0.9, 0.6), "m", [(594, 659), (201, 679)]),
    "belly":    (sheet(632, 910, 0.30), (0.15, 0.5, 0.4), "p", [(604, 486)]),
    "prong_up": (sheet(312, 975, 0.195), (0.10, 0.4, 0.4), "p", [(333, 591)]),
    "prong_lo": (sheet(305, 982, 0.195), (0.10, 0.4, 0.4), "p", [(316, 619)]),
}
PRIOR_PX = 15.0          # one sigma off the sheet costs as much as 15 px off the picture


def unpack(q, names):
    cam = q[:7]
    pts, k = {}, 7
    for n in names:
        prior, sig, kind, _ = MARKS[n]
        v = list(prior)
        for i in range(3):
            if sig[i] > 0:
                v[i] = q[k]
                k += 1
        pts[n] = np.array(v)
    return cam, pts


def residuals(q, names, focal_prior):
    cam, pts = unpack(q, names)
    r = []
    for n in names:
        prior, sig, kind, obs = MARKS[n]
        P = pts[n]
        X = [P] if kind != "m" else [P, P * np.array([-1, 1, 1])]
        uv, z = project(cam, np.array(X))
        for (u, v), o in zip(uv, obs):
            r += [u - o[0], v - o[1]]
        for i in range(3):
            if sig[i] > 0:
                r.append((P[i] - prior[i]) / sig[i] * PRIOR_PX)
    r.append((cam[6] - math.log(focal_prior)) / 0.35 * PRIOR_PX)        # a weak pull toward a normal lens
    return np.array(r)


def solve(names, q):
    lam = 1e-2
    focal_prior = 1150.0
    r = residuals(q, names, focal_prior)
    for it in range(600):
        J = np.zeros((len(r), len(q)))
        for i in range(len(q)):
            h = 1e-6 if 3 <= i < 6 else 1e-5
            dq = q.copy()
            dq[i] += h
            J[:, i] = (residuals(dq, names, focal_prior) - r) / h
        A, g = J.T @ J, J.T @ r
        while True:
            step = np.linalg.solve(A + lam * np.diag(np.diag(A) + 1e-9), -g)
            q2 = q + step
            r2 = residuals(q2, names, focal_prior)
            if r2 @ r2 < r @ r:
                q, r = q2, r2
                lam = max(lam * 0.3, 1e-10)
                break
            lam *= 10
            if lam > 1e10:
                break
        if lam > 1e10 or np.abs(step).max() < 1e-9:
            break
    return q, r


def main():
    names = list(MARKS)
    # start: a camera ahead and to port, a little above, looking at the hull's middle
    tgt = np.array([0.0, 0.9, 0.2])
    yaw = math.radians(50)            # from abeam toward the nose
    dist = 14.0
    cam0 = tgt + dist * np.array([math.cos(yaw), -math.sin(yaw), 0.12])
    q = list(cam0) + look_at(cam0, tgt) + [math.log(1150.0)]
    for n in names:
        prior, sig, kind, _ = MARKS[n]
        q += [prior[i] for i in range(3) if sig[i] > 0]
    q, r = solve(names, np.array(q, float))
    cam, pts = unpack(q, names)
    out = {"params": [float(x) for x in cam], "lens_mm": lens_mm(cam), "points": {}, "err_px": {}}
    for n in names:
        prior, sig, kind, obs = MARKS[n]
        P = pts[n]
        X = [P] if kind != "m" else [P, P * np.array([-1, 1, 1])]
        uv, z = project(cam, np.array(X))
        out["points"][n] = [round(float(x), 4) for x in P]
        out["err_px"][n] = [round(float(np.hypot(*(a - np.array(o)))), 1) for a, o in zip(uv, obs)]
        print(f"{n:9s} fitted {np.round(P, 3)}  prior {np.round(prior, 3)}  err px {out['err_px'][n]}")
    C = np.array(cam[:3])
    print("camera", np.round(C, 3), "rot deg", np.round(np.degrees(cam[3:6]), 2), "lens %.1f mm" % lens_mm(cam),
          "dist to (0,1,0) %.2f m" % np.linalg.norm(C - np.array([0, 1, 0])))
    (HERE / "cam.json").write_text(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()