"""
homography_dlt.py: four points to a bird's-eye view.

EECE 5554 Robotics Sensing and Navigation, demo for Lecture 11.

A camera on the roof sees a parking lot. Every point on the flat ground maps
to the image through one 3 x 3 matrix, the homography H:

    [u, v, 1]^T  ~  H [X, Y, 1]^T          (~ means equal up to a scale)

The direct linear transform (DLT) finds H from four or more point pairs:
each pair gives two linear equations in the nine entries of H, and the
solution is the singular vector of the stacked equations with the smallest
singular value. Normalizing the points first (Hartley) keeps it accurate.

Run it from a terminal:
    python3 homography_dlt.py          fit H from noisy clicks and print the errors
"""
import json
import math
import sys

import numpy as np

# the same camera as pinhole_camera.py: Blackfly S with a 5 mm lens, 1.7 m up, pitched 12 degrees down
W, H_IMG, F = 1616, 1240, 5.0 / 0.0045
HEIGHT_M, PITCH_DEG = 1.7, 12.0

# one parking stall, 2.6 m wide and 5.5 m deep, starting 4.5 m ahead (X forward, Y left, in meters)
STALL = np.array([[4.5, 1.3], [4.5, -1.3], [10.0, -1.3], [10.0, 1.3]])
CHECK = np.array([[10.0, 3.9], [4.5, -3.9], [15.5, 1.3], [15.5, -3.9]])   # corners NOT used for the fit


def camera_homography(h=HEIGHT_M, pitch_deg=PITCH_DEG):
    """The true ground-to-image homography of the roof camera: H = K [r1 r2 t]."""
    K = np.array([[F, 0, W / 2], [0, F, H_IMG / 2], [0, 0, 1.0]])
    p = math.radians(pitch_deg)
    fwd = np.array([math.cos(p), 0.0, -math.sin(p)]); right = np.array([0.0, -1.0, 0.0]); down = np.cross(fwd, right)
    R = np.vstack([right, down, fwd])                  # world to camera
    t = -R @ np.array([0.0, 0.0, h])
    return K @ np.column_stack([R[:, 0], R[:, 1], t])


def apply_h(H, pts):
    pts = np.asarray(pts, float)
    q = np.column_stack([pts, np.ones(len(pts))]) @ H.T
    return q[:, :2] / q[:, 2:3]


def normalizer(pts):
    """Similarity transform that moves the points to their centroid and scales them to an average distance of sqrt(2)."""
    c = pts.mean(axis=0)
    s = math.sqrt(2) / max(np.linalg.norm(pts - c, axis=1).mean(), 1e-12)
    return np.array([[s, 0, -s * c[0]], [0, s, -s * c[1]], [0, 0, 1.0]])


def dlt(src, dst, normalize=True):
    """Homography H with dst ~ H src, from 4 or more point pairs (N x 2 each)."""
    src, dst = np.asarray(src, float), np.asarray(dst, float)
    if len(src) < 4:
        raise ValueError("need at least 4 point pairs")
    Ts, Td = (normalizer(src), normalizer(dst)) if normalize else (np.eye(3), np.eye(3))
    s, d = apply_h(Ts, src), apply_h(Td, dst)
    rows = []
    for (x, y), (u, v) in zip(s, d):
        rows.append([-x, -y, -1, 0, 0, 0, u * x, u * y, u])
        rows.append([0, 0, 0, -x, -y, -1, v * x, v * y, v])
    _, _, Vt = np.linalg.svd(np.array(rows))
    Hn = Vt[-1].reshape(3, 3)                          # smallest singular value
    Hm = np.linalg.inv(Td) @ Hn @ Ts                   # undo the normalization
    return Hm / Hm[2, 2]


def check_error_cm(H_est, H_true=None):
    """Map the check corners from the image back to the ground with the estimated H. Mean error in cm."""
    H_true = camera_homography() if H_true is None else H_true
    img = apply_h(H_true, CHECK)
    ground = apply_h(np.linalg.inv(H_est), img)
    return float(np.linalg.norm(ground - CHECK, axis=1).mean() * 100)


def layout():
    """Lines of the parking lot on the ground: two rows of stalls. Returns a list of [X0, Y0, X1, Y1]."""
    lines = []
    for x0, x1 in ((4.5, 10.0), (15.5, 21.0)):
        for k in range(-3, 4):
            y = 1.3 + 2.6 * k
            lines.append([x0, y, x1, y])
        lines.append([x1, -6.5, x1, 9.1])
    lines.append([3.0, -6.5, 3.0, 9.1])
    return lines


def experiment(n_points=4, noise_px=1.0, normalize=True, trials=200, seed=0):
    """Fit H from n noisy clicks on stall corners many times; return the check-point errors in cm."""
    rng = np.random.default_rng(seed)
    Ht = camera_homography()
    grid = np.array([[x, y] for x in (4.5, 10.0, 15.5, 21.0) for y in (-3.9, -1.3, 1.3, 3.9, 6.5)])
    errs = []
    for _ in range(trials):
        idx = rng.choice(len(grid), size=n_points, replace=False) if n_points > 4 else [5, 6, 10, 11]
        g = grid[idx]
        img = apply_h(Ht, g) + rng.normal(0, noise_px, (len(g), 2))
        try:
            errs.append(check_error_cm(dlt(g, img, normalize), Ht))
        except np.linalg.LinAlgError:
            pass
    e = np.array(errs)
    return {"median_cm": float(np.median(e)), "p90_cm": float(np.percentile(e, 90)), "n": int(len(e))}


def demo_api(params_json):
    """Web page entry point.
    {"clicks": [[u, v] x 4]}  fit H from the clicked stall corners (in the STALL order)
    {"experiment": true}  check-point error against the number of points, for three click-noise levels."""
    p = json.loads(params_json or "{}")
    Ht = camera_homography()
    if p.get("experiment"):
        ns = [4, 5, 6, 8, 10, 14, 20]
        curves = {str(nz): [experiment(n, nz, True, trials=120)["median_cm"] for n in ns] for nz in (0.5, 1.0, 2.0)}
        return json.dumps({"ns": ns, "curves": curves})
    clicks = np.array(p.get("clicks") or apply_h(Ht, STALL).tolist(), float)
    Hs = dlt(STALL, clicks, normalize=bool(p.get("normalize", True)))
    Hinv = np.linalg.inv(Hs)
    lines = layout()
    img_lines = [apply_h(Ht, [[l[0], l[1]], [l[2], l[3]]]).ravel().round(1).tolist() for l in lines]
    top_lines = [apply_h(Hinv, apply_h(Ht, [[l[0], l[1]], [l[2], l[3]]])).ravel().round(3).tolist() for l in lines]
    return json.dumps({"H": (Hs / Hs[2, 2]).tolist(), "true_clicks": apply_h(Ht, STALL).round(1).tolist(),
                       "img_lines": img_lines, "top_lines": top_lines, "truth_lines": lines,
                       "check_cm": check_error_cm(Hs, Ht), "W": W, "Hh": H_IMG})


def main(argv):
    Ht = camera_homography()
    print("true H (ground meters to pixels):\n", np.round(Ht / Ht[2, 2], 3))
    for noise in (0.5, 1.0, 2.0):
        for n in (4, 8, 20):
            r = experiment(n, noise)
            print(f"click noise {noise} px, {n:2d} points: check-point error {r['median_cm']:5.1f} cm (90% under {r['p90_cm']:.1f} cm)")


if __name__ == "__main__":
    main(sys.argv)