# "From Two Views to 3D": experiments on synthetic scenes with exact ground truth (CPU, numpy, OpenCV, SciPy).
#   1. the 8-point algorithm with and without Hartley's normalisation, vs pixel noise and number of points
#   2. relative pose (R, t) from the 8-point and the 5-point algorithm, vs noise and baseline
#   3. triangulation: depth error vs depth and baseline, against the formula Z^2 * sigma_d / (f b)
#   4. bundle adjustment of 6 cameras and 300 points from a perturbed start (snapshots for the lab)
#   5. checkerboard calibration (Zhang's method, cv2.calibrateCamera) vs number of views, and a degenerate set of views
#   python twoview_experiments.py OUT.json
import sys, json, time, numpy as np, cv2
from scipy.optimize import least_squares
from scipy.sparse import lil_matrix
rng = np.random.default_rng(0); t0 = time.time()
K = np.array([[800., 0, 320], [0, 800, 240], [0, 0, 1]]); W, H = 640, 480
def rot(yaw, pitch, roll):   # degrees, about the camera's y, x, z axes
    y, p, r = np.radians([yaw, pitch, roll])
    Ry = np.array([[np.cos(y), 0, np.sin(y)], [0, 1, 0], [-np.sin(y), 0, np.cos(y)]]); Rx = np.array([[1, 0, 0], [0, np.cos(p), -np.sin(p)], [0, np.sin(p), np.cos(p)]])
    Rz = np.array([[np.cos(r), -np.sin(r), 0], [np.sin(r), np.cos(r), 0], [0, 0, 1]]); return Rz @ Rx @ Ry
def proj(R, t, X, Kc=K): x = (Kc @ (R @ X.T + t[:, None])).T; return x[:, :2] / x[:, 2:]
def scene(n, zmin=6, zmax=14):
    return np.column_stack([rng.uniform(-3, 3, n), rng.uniform(-2, 2, n), rng.uniform(zmin, zmax, n)])
def visible(R, t, X, Kc=K):
    Xc = (R @ X.T + t[:, None]).T; x = proj(R, t, X, Kc); return (Xc[:, 2] > 0.1) & (x[:, 0] >= 0) & (x[:, 0] < W) & (x[:, 1] >= 0) & (x[:, 1] < H)
def skew(v): return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
R1 = rot(-8, 3, 2); C1 = np.array([1.0, 0.1, 0.2]); t1 = -R1 @ C1          # camera 0 = world frame; camera 1 moved 1 m and turned
Ftrue = np.linalg.inv(K).T @ skew(t1) @ R1 @ np.linalg.inv(K)
res = {"K": K.tolist(), "R1": R1.tolist(), "t1": t1.tolist(), "F": (Ftrue / Ftrue[2, 2]).tolist()}

def normT(x):
    c = x.mean(0); s = np.sqrt(2) / np.mean(np.linalg.norm(x - c, axis=1)); return np.array([[s, 0, -s * c[0]], [0, s, -s * c[1]], [0, 0, 1]])
def eight_point(x0, x1, normalise=True):
    T0, T1 = (normT(x0), normT(x1)) if normalise else (np.eye(3), np.eye(3))
    a = (T0 @ np.column_stack([x0, np.ones(len(x0))]).T).T; b = (T1 @ np.column_stack([x1, np.ones(len(x1))]).T).T
    A = np.column_stack([b[:, 0] * a[:, 0], b[:, 0] * a[:, 1], b[:, 0], b[:, 1] * a[:, 0], b[:, 1] * a[:, 1], b[:, 1], a[:, 0], a[:, 1], np.ones(len(a))])
    F = np.linalg.svd(A)[2][-1].reshape(3, 3)
    U, S, Vt = np.linalg.svd(F); F = U @ np.diag([S[0], S[1], 0]) @ Vt      # enforce rank 2
    return T1.T @ F @ T0
def epi_dist(F, x0, x1):   # symmetric point-to-epipolar-line distance in pixels
    a = np.column_stack([x0, np.ones(len(x0))]); b = np.column_stack([x1, np.ones(len(x1))])
    l1 = (F @ a.T).T; l0 = (F.T @ b.T).T; e = np.abs(np.sum(b * l1, 1))
    return 0.5 * (e / np.hypot(l1[:, 0], l1[:, 1]) + e / np.hypot(l0[:, 0], l0[:, 1]))

# ---------- 1. normalised vs unnormalised 8-point ----------
Xtest = scene(2000); vt = visible(np.eye(3), np.zeros(3), Xtest) & visible(R1, t1, Xtest); Xtest = Xtest[vt][:500]
x0t, x1t = proj(np.eye(3), np.zeros(3), Xtest), proj(R1, t1, Xtest)
res["eight"] = []
for n in (8, 15, 30, 100):
    for sig in (0.0, 0.25, 0.5, 1.0, 2.0):
        e_n, e_u = [], []
        for trial in range(200):
            X = scene(4 * n); v = visible(np.eye(3), np.zeros(3), X) & visible(R1, t1, X); X = X[v][:n]
            x0 = proj(np.eye(3), np.zeros(3), X) + rng.normal(0, sig, (n, 2)); x1 = proj(R1, t1, X) + rng.normal(0, sig, (n, 2))
            e_n.append(np.mean(epi_dist(eight_point(x0, x1, True), x0t, x1t))); e_u.append(np.mean(epi_dist(eight_point(x0, x1, False), x0t, x1t)))
        res["eight"].append({"n": n, "sigma": sig, "normalised": float(np.median(e_n)), "unnormalised": float(np.median(e_u))})
        print("eight", n, sig, round(np.median(e_n), 3), round(np.median(e_u), 3), flush=True)

# ---------- 2. relative pose: 8-point (normalised) vs 5-point, vs noise and baseline ----------
def ang_err(Ra, Rb): return float(np.degrees(np.arccos(np.clip((np.trace(Ra.T @ Rb) - 1) / 2, -1, 1))))
def dir_err(ta, tb): return float(np.degrees(np.arccos(np.clip(abs(ta @ tb) / np.linalg.norm(ta) / np.linalg.norm(tb), -1, 1))))
def pose_trial(base, sig, n=50, outl=0.0):
    C = C1 / np.linalg.norm(C1) * base; t = -R1 @ C
    X = scene(6 * n); v = visible(np.eye(3), np.zeros(3), X) & visible(R1, t, X); X = X[v][:n]
    x0 = proj(np.eye(3), np.zeros(3), X) + rng.normal(0, sig, (len(X), 2)); x1 = proj(R1, t, X) + rng.normal(0, sig, (len(X), 2))
    bad = rng.random(len(X)) < outl; x1[bad] = rng.uniform([0, 0], [W, H], (bad.sum(), 2))   # wrong matches
    out = {}
    F = eight_point(x0, x1); E8 = K.T @ F @ K                     # 8-point, least squares on all matches
    _, R8, t8, _ = cv2.recoverPose(E8, x0, x1, K); out["8pt"] = (ang_err(R8, R1), dir_err(t8.ravel(), t))
    for key, meth in (("5pt", cv2.RANSAC), ("5pt_lo", cv2.USAC_ACCURATE)):   # 5-point inside RANSAC; and with local optimisation
        E5, _ = cv2.findEssentialMat(x0, x1, K, method=meth, prob=0.999, threshold=3 * max(sig, 0.3))
        if E5 is not None and E5.shape == (3, 3):
            _, R5, t5, _ = cv2.recoverPose(E5, x0, x1, K); out[key] = (ang_err(R5, R1), dir_err(t5.ravel(), t))
    return out
res["pose"] = {"noise": [], "baseline": [], "outliers": []}
for key, sweep in [("noise", [(1.0, s, 0.0) for s in (0.25, 0.5, 1.0, 2.0, 4.0)]), ("baseline", [(b, 1.0, 0.0) for b in (0.05, 0.1, 0.25, 0.5, 1.0, 2.0)]),
                   ("outliers", [(1.0, 1.0, o) for o in (0.0, 0.05, 0.1, 0.2, 0.4)])]:
    for base, sig, outl in sweep:
        acc = {"8pt": [], "5pt": [], "5pt_lo": []}
        for trial in range(200):
            for k, v in pose_trial(base, sig, outl=outl).items(): acc[k].append(v)
        row = {"baseline": base, "sigma": sig, "outliers": outl}
        for k, v in acc.items(): v = np.array(v); row[k] = {"rot_deg": float(np.median(v[:, 0])), "dir_deg": float(np.median(v[:, 1]))}
        res["pose"][key].append(row); print("pose", key, base, sig, outl, {k: (round(v["rot_deg"], 2), round(v["dir_deg"], 2)) for k, v in row.items() if isinstance(v, dict)}, flush=True)

# ---------- 3. triangulation depth error ----------
res["depth"] = []
for b in (0.1, 0.3, 1.0):
    P0 = K @ np.hstack([np.eye(3), np.zeros((3, 1))]); P1 = K @ np.hstack([np.eye(3), np.array([[-b], [0], [0]])])   # rectified pair
    for Z in (2, 4, 8, 16, 32):
        X = np.column_stack([rng.uniform(-0.3, 0.3, 2000) * Z / 4, rng.uniform(-0.2, 0.2, 2000) * Z / 4, np.full(2000, float(Z))])
        x0 = proj(np.eye(3), np.zeros(3), X) + rng.normal(0, 0.5, (2000, 2)); x1 = proj(np.eye(3), np.array([-b, 0, 0]), X) + rng.normal(0, 0.5, (2000, 2))
        Xh = cv2.triangulatePoints(P0, P1, x0.T, x1.T); Zest = Xh[2] / Xh[3]
        res["depth"].append({"baseline": b, "Z": Z, "rms": float(np.sqrt(np.mean((Zest - Z) ** 2))), "median_abs": float(np.median(np.abs(Zest - Z))),
                             "formula": float(Z ** 2 * 0.5 * np.sqrt(2) / (800 * b))})
        print("depth", b, Z, res["depth"][-1], flush=True)

# ---------- 4. bundle adjustment ----------
ncam, npt = 6, 300
Xw = np.column_stack([rng.uniform(-3, 3, npt), rng.uniform(-1.5, 1.5, npt), rng.uniform(-2, 2, npt)])   # a scene around the origin
cams = []
for i in range(ncam):   # cameras on an arc 10 m from the origin, looking at it
    a = np.radians(-40 + 16 * i); C = np.array([10 * np.sin(a), rng.uniform(-0.5, 0.5), -10 * np.cos(a)])
    fw = -C / np.linalg.norm(C); r = np.cross([0, 1, 0], fw); r /= np.linalg.norm(r); d = np.cross(fw, r); R = np.stack([r, d, fw]); cams.append((R, -R @ C))
obs = []
for ci, (R, t) in enumerate(cams):
    v = visible(R, t, Xw); x = proj(R, t, Xw) + rng.normal(0, 1.0, (npt, 2))
    obs += [(ci, j, x[j, 0], x[j, 1]) for j in np.where(v)[0]]
obs = np.array(obs); ci_, pj_ = obs[:, 0].astype(int), obs[:, 1].astype(int); uv = obs[:, 2:]
def pack(cs, X): return np.concatenate([np.concatenate([cv2.Rodrigues(R)[0].ravel(), t]) for R, t in cs[1:]] + [X.ravel()])
def unpack(p):
    cs = [cams[0]] + [(cv2.Rodrigues(p[6 * i:6 * i + 3])[0], p[6 * i + 3:6 * i + 6]) for i in range(ncam - 1)]
    return cs, p[6 * (ncam - 1):].reshape(-1, 3)
start = [cams[0]] + [(cv2.Rodrigues(cv2.Rodrigues(R)[0] + rng.normal(0, np.radians(4), (3, 1)))[0], t + rng.normal(0, 0.8, 3)) for R, t in cams[1:]]
X0 = Xw + rng.normal(0, 0.8, Xw.shape)
snaps = []
def resid(p):
    cs, X = unpack(p); r = np.empty((len(obs), 2))
    for c in range(ncam):
        m = ci_ == c; r[m] = proj(cs[c][0], cs[c][1], X[pj_[m]]) - uv[m]
    # gauge: camera 0 is fixed; the scale is fixed by a soft prior on camera 1's distance from camera 0
    C1e = -cs[1][0].T @ cs[1][1]; C0 = -cams[0][0].T @ cams[0][1]; d_true = np.linalg.norm((-cams[1][0].T @ cams[1][1]) - C0)
    g = np.array([100 * (np.linalg.norm(C1e - C0) - d_true)])
    return np.concatenate([r.ravel(), g])
def snapshot(p):
    cs, X = unpack(p); r = resid(p)[:-1].reshape(-1, 2)
    return {"rms": float(np.sqrt(np.mean(np.sum(r ** 2, 1)))), "cam_err": float(np.mean([np.linalg.norm((-R.T @ t) - (-Rt.T @ tt)) for (R, t), (Rt, tt) in zip(cs, cams)])),
            "pt_err": float(np.mean(np.linalg.norm(X - Xw, axis=1))), "X": np.round(X[:, [0, 2]], 3).tolist(), "C": [np.round((-R.T @ t)[[0, 2]], 3).tolist() for R, t in cs]}
A = lil_matrix((2 * len(obs) + 1, 6 * (ncam - 1) + 3 * npt), dtype=int)
for k, (c, j) in enumerate(zip(ci_, pj_)):
    if c > 0: A[2 * k:2 * k + 2, 6 * (c - 1):6 * c] = 1
    A[2 * k:2 * k + 2, 6 * (ncam - 1) + 3 * j:6 * (ncam - 1) + 3 * j + 3] = 1
A[2 * len(obs), 0:6] = 1
p0 = pack(start, X0); snaps.append(snapshot(p0))
for budget in range(2, 60):   # the same optimisation stopped after 2, 3, 4, ... evaluations (deterministic), to show its progress
    sol = least_squares(resid, p0, jac_sparsity=A, x_scale="jac", method="trf", xtol=1e-10, ftol=1e-10, max_nfev=budget)
    sn = snapshot(sol.x); sn["iter"] = budget - 1
    if abs(snaps[-1]["rms"] - sn["rms"]) > 1e-4: snaps.append(sn)
    if sol.status != 0: break
res["ba"] = {"n_obs": int(len(obs)), "n_params": int(6 * (ncam - 1) + 3 * npt), "noise_px": 1.0, "true_X": np.round(Xw[:, [0, 2]], 3).tolist(),
             "true_C": [np.round((-R.T @ t)[[0, 2]], 3).tolist() for R, t in cams], "snaps": snaps, "nfev": int(sol.nfev)}
print("ba", len(obs), "obs;", [(round(s["rms"], 3), round(s["cam_err"], 3), round(s["pt_err"], 3)) for s in snaps], flush=True)

# ---------- 5. calibration ----------
Kc = np.array([[800., 0, 330], [0, 800, 235], [0, 0, 1]]); Dc = np.array([-0.25, 0.08, 0, 0, 0])
board = np.array([[i * 0.03, j * 0.03, 0] for j in range(6) for i in range(9)], np.float32); bc = board.mean(0)
def view(tilt_max, rs):
    for _ in range(100):
        R = cv2.Rodrigues(np.array([rs.uniform(-tilt_max, tilt_max), rs.uniform(-tilt_max, tilt_max), rs.uniform(-np.pi, np.pi)]) * np.array([1, 1, 0.3]))[0]
        z = rs.uniform(0.4, 0.9); t = np.array([rs.uniform(-0.1, 0.1), rs.uniform(-0.08, 0.08), z]) - R @ bc
        x = cv2.projectPoints(board, cv2.Rodrigues(R)[0], t, Kc, Dc)[0].reshape(-1, 2)
        if (R @ board.T + t[:, None])[2].min() > 0.1 and x[:, 0].min() > 5 and x[:, 0].max() < W - 5 and x[:, 1].min() > 5 and x[:, 1].max() < H - 5: return x
    return None
def calib_trials(nv, tilt, trials=50):
    out = []
    for tr in range(trials):
        rs = np.random.default_rng(1000 * nv + tr + int(100 * tilt))
        imgs = [v for v in (view(tilt, rs) for _ in range(nv)) if v is not None]
        ip = [(v + rs.normal(0, 0.3, v.shape)).astype(np.float32) for v in imgs]
        try:
            rms, Ke, De, _, _ = cv2.calibrateCamera([board] * len(ip), ip, (W, H), None, None, flags=cv2.CALIB_FIX_K3)
            out.append([abs(Ke[0, 0] - 800), abs(Ke[0, 2] - 330), abs(De.ravel()[0] + 0.25), rms])
        except cv2.error: out.append([np.nan] * 4)
    o = np.array(out); return {"fx_err": float(np.nanmedian(o[:, 0])), "cx_err": float(np.nanmedian(o[:, 1])), "k1_err": float(np.nanmedian(o[:, 2])), "rms": float(np.nanmedian(o[:, 3])), "failed": int(np.isnan(o[:, 0]).sum())}
res["calib"] = {"views": [], "frontal": None}
for nv in (2, 3, 5, 10, 20, 40):
    r = {"n": nv, **calib_trials(nv, np.radians(40))}; res["calib"]["views"].append(r); print("calib", r, flush=True)
res["calib"]["frontal"] = {"n": 10, **calib_trials(10, np.radians(2))}; print("calib frontal", res["calib"]["frontal"], flush=True)
json.dump(res, open(sys.argv[1], "w"), separators=(",", ":"))
print("done", f"{time.time() - t0:.0f}s")