"""
pinhole_camera.py: aim a roof camera and see what it sees.

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

A pinhole camera at height h above flat ground, pitched down by an angle,
projects every ground point (X forward, Y left, Z = 0) into pixels with

    u = fx * xc / zc + cx,    v = fy * yc / zc + cy

where (xc, yc, zc) is the point in the camera optical frame (z forward,
x right, y down). The defaults are the course stereo rig: a Blackfly S
BFS-U3-19S4C (1616 x 1240 px, 4.5 um pixels) with a 5 mm lens.

Run it from a terminal:
    python3 pinhole_camera.py              print the coverage numbers
    python3 pinhole_camera.py 1.7 8 5      height (m), pitch (deg), focal length (mm)
"""
import json
import math
import sys

import numpy as np

WIDTH, HEIGHT, PIXEL_MM = 1616, 1240, 0.0045     # Blackfly S BFS-U3-19S4C


def intrinsics(f_mm=5.0):
    """K from the focal length and the pixel size, with the principal point at the image center."""
    f = f_mm / PIXEL_MM
    return np.array([[f, 0.0, WIDTH / 2], [0.0, f, HEIGHT / 2], [0.0, 0.0, 1.0]])


def rotation_world_to_camera(pitch_deg):
    """Rows are the camera axes (right, down, forward) written in world coordinates (x forward, y left, z up)."""
    p = math.radians(pitch_deg)
    forward = np.array([math.cos(p), 0.0, -math.sin(p)])
    right = np.array([0.0, -1.0, 0.0])
    down = np.cross(forward, right)
    return np.vstack([right, down, forward])


def project(points, h=1.7, pitch_deg=8.0, f_mm=5.0):
    """World points (N, 3) to pixels (N, 2). Points behind the camera come back as NaN."""
    K, R = intrinsics(f_mm), rotation_world_to_camera(pitch_deg)
    pc = (np.asarray(points, float) - np.array([0.0, 0.0, h])) @ R.T      # camera frame
    uv = np.full((len(pc), 2), np.nan)
    ok = pc[:, 2] > 1e-6
    uvw = pc[ok] @ K.T
    uv[ok] = uvw[:, :2] / uvw[:, 2:3]
    return uv


def coverage(h=1.7, pitch_deg=8.0, f_mm=5.0):
    """Field of view, nearest visible ground, and the size of one pixel on the ground."""
    K = intrinsics(f_mm)
    f = K[0, 0]
    hfov = 2 * math.degrees(math.atan(WIDTH / 2 / f))
    vfov = 2 * math.degrees(math.atan(HEIGHT / 2 / f))
    below = math.radians(pitch_deg) + math.atan((HEIGHT / 2) / f)        # ray through the bottom row
    nearest = h / math.tan(below) if below > 0 else float("inf")
    horizon_v = HEIGHT / 2 - f * math.tan(math.radians(pitch_deg))
    out = {"f_px": f, "hfov_deg": hfov, "vfov_deg": vfov, "nearest_ground_m": nearest, "horizon_row": horizon_v,
           "footprint": []}
    for X in (5.0, 10.0, 20.0, 40.0):
        zc = X * math.cos(math.radians(pitch_deg)) + h * math.sin(math.radians(pitch_deg))
        lateral = zc / f                                                  # meters per pixel across
        along = (X * X + h * h) / (h * f)                                 # meters per pixel row: dX/dangle times 1/f
        width = 2 * zc * (WIDTH / 2) / f
        out["footprint"].append({"X_m": X, "lateral_m_per_px": lateral, "along_m_per_px": along, "width_m": width})
    return out


def ground_lines(h=1.7, pitch_deg=8.0, f_mm=5.0, max_x=60.0):
    """Line segments of a ground grid and two lane lines, already clipped and projected to pixels."""
    segs = []
    inside = lambda q: -300 <= q[0] <= WIDTH + 300 and -300 <= q[1] <= HEIGHT + 300
    def add(p0, p1, kind, n=24):
        pts = np.linspace(p0, p1, n)
        uv = project(pts, h, pitch_deg, f_mm)
        for a, b in zip(uv[:-1], uv[1:]):
            if not (np.isnan(a).any() or np.isnan(b).any()) and (inside(a) or inside(b)):
                segs.append([round(a[0], 1), round(a[1], 1), round(b[0], 1), round(b[1], 1), kind])
    for x in [1, 2, 3, 4, 5, 6, 8, 10, 15, 20, 30, 40, 60]:
        if x <= max_x:
            add([x, -12, 0], [x, 12, 0], "row")
    for y in range(-12, 13, 2):
        add([0.5, y, 0], [max_x, y, 0], "col")
    for y in (-1.8, 1.8):
        add([0.5, y, 0], [max_x, y, 0], "lane")
    add([10, -2.5, 0], [10, -2.5, 1.7], "pole")      # a 1.7 m tall person 10 m ahead, 2.5 m to the right
    return segs


def demo_api(params_json):
    """Web page entry point. Input: {"h": m, "pitch": deg, "f_mm": mm}."""
    p = json.loads(params_json or "{}")
    h, pitch, f_mm = float(p.get("h", 1.7)), float(p.get("pitch", 8.0)), float(p.get("f_mm", 5.0))
    out = coverage(h, pitch, f_mm)
    out["segments"] = ground_lines(h, pitch, f_mm)
    out["K"] = intrinsics(f_mm).round(1).tolist()
    out["W"], out["H"] = WIDTH, HEIGHT
    return json.dumps(out)


def main(argv):
    h = float(argv[1]) if len(argv) > 1 else 1.7
    pitch = float(argv[2]) if len(argv) > 2 else 8.0
    f_mm = float(argv[3]) if len(argv) > 3 else 5.0
    c = coverage(h, pitch, f_mm)
    print(f"f = {c['f_px']:.0f} px, field of view {c['hfov_deg']:.1f} x {c['vfov_deg']:.1f} deg")
    print(f"camera {h} m up, pitched {pitch} deg down: nearest ground {c['nearest_ground_m']:.1f} m, horizon at row {c['horizon_row']:.0f}")
    for r in c["footprint"]:
        print(f"  {r['X_m']:4.0f} m ahead: {r['width_m']:5.1f} m wide, one pixel = {100 * r['lateral_m_per_px']:.1f} cm across, {100 * r['along_m_per_px']:.1f} cm along")


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