"""Grouper: the concept render's camera, fitted to the sub's landmarks (the single-picture way of the Tiger's tg_fit.py:
a pinhole in Blender's conventions, Levenberg-Marquardt in numpy). Plain Python 3.14 + numpy, no Blender:
    python build/gr_concept_fit.py   -> build/checks/concept_cam.json

The model's points come from the model itself (build/checks/landmarks.json, gr_landmarks.py); the picture's pixels were read
on a gridded 2.4x zoom of reference/user_concept.jpg (1408 x 768). The concept is an AI render of the same design, not a
photo of this model: the fit gives the camera that sees the model most like the picture, and the residuals (listed) say how
far the picture's own proportions differ from the sheet's. Unknowns: camera location, rotation, focal length (7)."""
import json
import math
import os
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
W_IMG, H_IMG = 1408, 768
PP = (W_IMG / 2.0, H_IMG / 2.0)
L = json.load(open(os.path.join(HERE, "checks", "landmarks.json")))
# name: (picture pixel, weight)
OBS = {
    "SOCKET_Lamp_R": ((531.5, 128.8), 1.0),       # the near emitter's lens (bright)
    "SOCKET_Lamp_L": ((630.4, 126.7), 1.0),       # the far emitter's lens
    "SOCKET_PodLamp_R": ((485.4, 395.4), 1.0),    # the near pod's lamp
    "SOCKET_DrumLamp_L": ((703.3, 357.9), 1.0),   # the lamp beside the drum (critique round 1 m4: it is not the far pod's lamp,
                                                  # which the drum hides; the model has a lamp there now, on the cradle's side)
    "SOCKET_Scanner": ((662.5, 333.0), 1.0),      # the drum's face, its hexagon's centre
    "porthole_R": ((398.3, 228.8), 1.0),          # the door's porthole
    "side_glass_R": ((526.0, 230.0), 0.7),        # the near side window's middle
    "front_glass": ((622.0, 220.4), 0.7),         # the front window's middle
    "lettering_R": ((418.8, 160.8), 0.7),         # "Subtub 1"
    "seal_R": ((347.0, 185.0), 0.7),
    "insignia_R": ((362.5, 155.8), 0.7),
    "duct_aft": ((274.0, 218.0), 0.2),            # the stern duct's aft opening (its lip's ellipse)
    "hatch_top": ((457.5, 110.0), 0.3),           # the collar disc's top (soft edge)
}


def euler_xyz(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 project(p, X3):
    X3 = np.atleast_2d(np.asarray(X3, float))
    R = euler_xyz(*p[3:6])
    pc = (X3 - np.asarray(p[:3])) @ R
    f = math.exp(p[6])
    z = -pc[:, 2]
    return np.stack([PP[0] + f * pc[:, 0] / z, PP[1] - f * pc[:, 1] / z], 1), z


def look_at(cam, tgt):
    f = np.asarray(tgt, float) - np.asarray(cam, float)
    f /= np.linalg.norm(f)
    r = np.cross(f, [0, 0, 1.0])
    r /= np.linalg.norm(r)
    u = np.cross(r, f)
    R = np.stack([r, u, -f], 1)
    return [math.atan2(R[2, 1], R[2, 2]), math.asin(-max(-1.0, min(1.0, R[2, 0]))), math.atan2(R[1, 0], R[0, 0])]


# the focal length: a picture of this kind is taken with a normal lens; free, the fit runs off to a telephoto 150 m away
# (the AI picture's own proportions are not the sheet's, and a far camera hides the difference). A pull toward a 35 mm
# lens (36 mm sensor) keeps the perspective the picture plainly has (the stern far smaller than the nose).
F0 = float(os.environ.get("GR_F0_MM", "35")) * W_IMG / 36.0
F_SIG = float(os.environ.get("GR_F_SIG", "0.1"))
names = list(OBS)
X = np.array([L[n] for n in names])
U = np.array([OBS[n][0] for n in names])
Wt = np.repeat(np.array([OBS[n][1] for n in names]), 2)


def res(q):
    uv, z = project(q, X)
    r = (uv - U).ravel() * Wt
    return np.concatenate([r, [(q[6] - math.log(F0)) / F_SIG * 12.0]])


best = None
for yaw in (25, 35, 45, 55):
    for el in (0.15, 0.30):
        for dist in (10.0, 14.0, 18.0):
            tgt = np.array([0.0, -0.6, 0.0])
            a = math.radians(yaw)
            cam0 = tgt + dist * np.array([-math.cos(a), -math.sin(a), el])
            q = np.array(list(cam0) + look_at(cam0, tgt) + [math.log(1400.0)], float)
            lam = 1e-2
            r = res(q)
            for it in range(400):
                J = np.zeros((len(r), 7))
                for i in range(7):
                    h = 1e-6 if 3 <= i < 6 else 1e-5
                    dq = q.copy()
                    dq[i] += h
                    J[:, i] = (res(dq) - r) / h
                A, g = J.T @ J, J.T @ r
                ok = False
                while lam < 1e10:
                    step = np.linalg.solve(A + lam * np.diag(np.diag(A) + 1e-9), -g)
                    r2 = res(q + step)
                    if r2 @ r2 < r @ r:
                        q, r, lam, ok = q + step, r2, max(lam * 0.3, 1e-10), True
                        break
                    lam *= 10
                if not ok or np.abs(step).max() < 1e-10:
                    break
            _, z = project(q, X)
            if np.all(z > 0) and (best is None or r @ r < best[1]):
                best = (q, float(r @ r))
q = best[0]
uv, z = project(q, X)
err = {n: round(float(np.hypot(*(uv[i] - U[i]))), 1) for i, n in enumerate(names)}
f = math.exp(q[6])
out = {"picture": "reference/user_concept.jpg", "resolution": [W_IMG, H_IMG],
       "location": [round(float(c), 4) for c in q[:3]], "rotation_euler_xyz_rad": [round(float(c), 6) for c in q[3:6]],
       "rotation_euler_xyz_deg": [round(math.degrees(c), 3) for c in q[3:6]], "focal_px": round(f, 2),
       "lens_mm_36mm_sensor_horizontal": round(f * 36.0 / W_IMG, 3),
       "err_px": err, "rms_px": round(float(np.sqrt(np.mean([e * e for e in err.values()]))), 1),
       "rms_px_weight_0.7_up": round(float(np.sqrt(np.mean([err[n] ** 2 for n in names if OBS[n][1] >= 0.7]))), 1),
       "picture_px": {n: list(OBS[n][0]) for n in names}, "model_m": {n: L[n] for n in names},
       "projected_px": {n: [round(float(c), 1) for c in uv[i]] for i, n in enumerate(names)}}
json.dump(out, open(os.path.join(HERE, "checks", "concept_cam.json"), "w"), indent=1)
print("CONCEPT CAM", json.dumps({k: out[k] for k in ("location", "rotation_euler_xyz_deg", "lens_mm_36mm_sensor_horizontal", "rms_px")}))
print("ERR", err)