import numpy as np

pt2d = np.loadtxt('pt2d.txt').astype(np.float64)
pt3d = np.loadtxt('pt3d.txt').astype(np.float64)

print(pt2d.max(axis=0))

N = pt2d.shape[0]
x, y = pt2d.T
X, Y, Z = pt3d.T

A = np.zeros((2 * N, 12))
val1 = np.full((N,), 1.0)
val0 = np.full((N,), 0.0)
A[0::2] = np.stack(
    [X, Y, Z, val1, val0, val0, val0, val0, -x * X, -x * Y, -x * Z, -x], axis=1)
A[1::2] = np.stack(
    [val0, val0, val0, val0, X, Y, Z, val1, -y * X, -y * Y, -y * Z, -y], axis=1)

W, V = np.linalg.eig(A.T @ A)
P = V[:, W.argmin()].reshape(3, 4)
print(P)

M, P4 = P[:, :3], P[:, 3]
C = - np.linalg.inv(M) @ P4
print(C)
q, r = np.linalg.qr(np.linalg.inv(M))
K = np.linalg.inv(r)
R = np.linalg.inv(q)
T = np.linalg.inv(K) @ P4
Rt = np.concatenate([R, T.reshape(3, 1)], axis=1)
print(K)
assert np.allclose(P, K @ Rt)

homo3d = np.stack([X, Y, Z, val1], axis=1).T
homo2d = K @ Rt @ homo3d
coor2d = homo2d[:2, :] / homo2d[2, :]
coor2d = coor2d.T
error = np.linalg.norm(coor2d - pt2d, axis=1).mean()

print(error)