"""Fit the robot's joint angles to the picture, headless (numpy from Blender):

    blender.exe -b --factory-startup --python rc_fitpose.py        -> build/pose_fit.json

Landmarks (rc_robot.py MARKS): the A2 and A3 gear covers' centres, the A5 hub, the flange, the tool's tip, as the
picture shows them. Unknowns A1, A2, A3, A5 (A4 and A6 at 0: the picture's wrist is not rolled). The robot's place and
size are fixed (it is drawn larger than its place in the cell allows, so the fit cannot be exact).
"""
import bpy, math, os, json
import numpy as np
from mathutils import Vector

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

KEYS = ["A1", "A2", "A3", "A5"]


def pix(q):
    th = dict(zip(KEYS, q))
    T = fk(th)
    out = []
    for nm, (fr, loc, px) in MARKS.items():
        w = ROB + (T[fr] @ Vector(loc))
        u, v = CAM.project(w)
        out += [u - px[0], v - px[1]]
    return np.array(out)


def lm(q, it=200):
    q = np.array(q, float)
    lam = 1e-2
    r = pix(q)
    for _ in range(it):
        J = np.zeros((len(r), len(q)))
        for k in range(len(q)):
            dq = np.zeros(len(q)); dq[k] = 1e-5
            J[:, k] = (pix(q + dq) - r) / 1e-5
        A, g = J.T @ J, J.T @ r
        while True:
            st = np.linalg.solve(A + lam * np.diag(np.diag(A) + 1e-9), -g)
            r2 = pix(q + st)
            if r2 @ r2 < r @ r:
                q, r = q + st, r2
                lam *= 0.3
                break
            lam *= 10
            if lam > 1e12:
                return q, r
    return q, r


q0 = [math.radians(a) for a in (28, 22, 2, 50)]
q, r = lm(q0)
res = {k: round(math.degrees(v), 2) for k, v in zip(KEYS, q)}
res.update(A4=0.0, A6=0.0)
names = list(MARKS)
rep = dict(pose_deg=res, rms_px=round(float(np.sqrt(np.mean(r ** 2))), 2),
           per_mark_px={n: [round(r[2 * i], 1), round(r[2 * i + 1], 1)] for i, n in enumerate(names)},
           start_px={n: [round(x, 1) for x in (pix(q0)[2 * i], pix(q0)[2 * i + 1])] for i, n in enumerate(names)})
json.dump(rep, open(os.path.join(HERE, "pose_fit.json"), "w"), indent=1)
print("FITPOSE", json.dumps(rep))