# "From a 3D Point to a Pixel": reference values computed with OpenCV and numpy, used to check the browser lab
# (geo-lab.htm), which re-implements the same maths in JavaScript.
#   python geo_reference.py OUT.json
import sys, json, numpy as np, cv2
SRC = np.float32([[150, 120], [330, 120], [330, 240], [150, 240]]); DST = np.float32([[120, 70], [390, 110], [350, 300], [110, 260]])
ref = {"src": SRC.tolist(), "dst": DST.tolist()}
ref["projective"] = cv2.getPerspectiveTransform(SRC, DST).tolist()           # exact homography from 4 correspondences
A = np.hstack([SRC, np.ones((4, 1))]); ref["affine"] = np.vstack([np.linalg.lstsq(A, DST, rcond=None)[0].T, [0, 0, 1]]).tolist()   # least squares
rows, rhs = [], []                                                            # similarity: [a -b tx; b a ty], least squares
for (x, y), (u, v) in zip(SRC, DST): rows += [[x, -y, 1, 0], [y, x, 0, 1]]; rhs += [u, v]
a, b, tx, ty = np.linalg.lstsq(np.array(rows), np.array(rhs), rcond=None)[0]; ref["similarity"] = [[a, -b, tx], [b, a, ty], [0, 0, 1]]
cs, cd = SRC.mean(0), DST.mean(0); U, S, Vt = np.linalg.svd((SRC - cs).T @ (DST - cd))   # Euclidean: Kabsch / Procrustes without scale
R = (U @ np.diag([1, np.sign(np.linalg.det(U @ Vt))]) @ Vt).T; ref["euclidean"] = np.vstack([np.hstack([R, (cd - R @ cs)[:, None]]), [0, 0, 1]]).tolist()
ref["translation"] = [[1, 0, float(cd[0] - cs[0])], [0, 1, float(cd[1] - cs[1])], [0, 0, 1]]

# projection of 3D points with K, R, t, with and without lens distortion (OpenCV's model: k1, k2, p1, p2, k3)
K = np.array([[500., 0, 300], [0, 500, 240], [0, 0, 1]])
yaw, pitch, C = np.radians(10), np.radians(15), np.array([-1., -7, 2])
fw = np.array([np.sin(yaw) * np.cos(pitch), np.cos(yaw) * np.cos(pitch), -np.sin(pitch)]); r = np.array([np.cos(yaw), -np.sin(yaw), 0]); d = np.cross(fw, r)
Rc = np.stack([r, d, fw]); t = -Rc @ C
X = np.array([[0.5, -0.5, 1.6], [-0.5, 0.5, 0], [2.7, 3.7, 0.8], [-3, 6.2, 2.4]])
dist = np.array([-0.3, 0.08, 0.01, -0.005, 0.0])
ref["camera"] = {"K": K.tolist(), "yaw": 10, "pitch": 15, "C": C.tolist(), "R": Rc.tolist(), "t": t.tolist(), "X": X.tolist(), "dist": dist.tolist(),
                 "det_R": float(np.linalg.det(Rc)),
                 "px": cv2.projectPoints(X, cv2.Rodrigues(Rc)[0], t, K, None)[0].reshape(-1, 2).tolist(),
                 "px_dist": cv2.projectPoints(X, cv2.Rodrigues(Rc)[0], t, K, dist)[0].reshape(-1, 2).tolist()}
# field of view: 100 degrees on a 640 x 480 image, read as horizontal, vertical or diagonal
ref["fov"] = {k: float(e / 2 / np.tan(np.radians(50))) for k, e in [("horizontal", 640), ("vertical", 480), ("diagonal", np.hypot(640, 480))]}
json.dump(ref, open(sys.argv[1], "w"), indent=1)
print(json.dumps(ref["fov"]), np.round(ref["camera"]["px"], 2).tolist())