"""
ground_mosaic.py: build a map of a parking lot from a roof camera and RTK poses.

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

A camera on the car's roof looks at the ground ahead. Each image is turned into
a bird's-eye view with the ground homography (Lecture 12), then placed on the
map with the car's pose at the moment the image was taken:

    p_map = p_base(t) + R(heading(t)) p_ground_in_base_link

RTK measures the antenna, not base_link (the center of the rear axle), so the
lever arm between them is subtracted with the heading. Every image is blended
into one grid, so errors show up as blur and doubled lines. You can add:

    heading error   a fixed bias, in degrees
    time offset     the image stamps are late: the pose used is from later on the track
    pitch error     the camera pitch in the calibration is wrong, so ground distances stretch
    no lever arm    the antenna position is used as if it were base_link

The car drives one loop of a 70 m by 46 m lot, counterclockwise, 10 images a
second. The lot, the path, and all the numbers are simulated.

Run it from a terminal:
    python3 ground_mosaic.py              misplacement for each kind of error
    python3 ground_mosaic.py --save       also save the maps as PNG files (needs matplotlib)
"""
import base64
import json
import math
import sys

import numpy as np

LOT_W, LOT_H = 70.0, 46.0                   # meters
CAM_FWD, CAM_H = 1.9, 1.7                   # camera 1.9 m ahead of base_link and 1.7 m up
ANT_FWD = 1.0                               # antenna 1.0 m ahead of base_link
FPS = 10.0
P0, P1, P2, P3, RC = (8, 12.25), (62, 12.25), (62, 33.75), (8, 33.75), 4.0


def _hash(x, y):
    h = (x.astype(np.uint64) * np.uint64(374761393) + y.astype(np.uint64) * np.uint64(668265263)) & np.uint64(0xFFFFFFFF)
    h = ((h ^ (h >> np.uint64(13))) * np.uint64(1274126177)) & np.uint64(0xFFFFFFFF)
    return ((h ^ (h >> np.uint64(16))) & np.uint64(0xFFFFFFFF)).astype(np.float64) / 4294967296.0


def _noise(x, y):
    xi, yi = np.floor(x), np.floor(y)
    fx, fy = x - xi, y - yi
    sx, sy = fx * fx * (3 - 2 * fx), fy * fy * (3 - 2 * fy)
    xi, yi = xi.astype(np.int64) + 1000, yi.astype(np.int64) + 1000
    a, b, c, d = _hash(xi, yi), _hash(xi + 1, yi), _hash(xi, yi + 1), _hash(xi + 1, yi + 1)
    return a + (b - a) * sx + (c - a) * sy + (a - b - c + d) * sx * sy


def ground(x, y):
    """Color of the lot at points (x, y) in meters: asphalt, stall lines, crosswalks, an arrow, center dashes."""
    g = 84 + 22 * _noise(x * 0.7, y * 0.7) + 10 * _noise(x * 3.1, y * 3.1)
    rgb = np.stack([g, g, g + 3], axis=-1)
    white, w = np.array([236, 236, 230.0]), 0.07
    stall_x = np.abs(((x - 14) % 2.7) - 0.06) < w
    rows = (((y >= 2) & (y <= 7.5) & (np.abs(y - 2) < w)) | ((y >= 38.5) & (y <= 44) & (np.abs(y - 44) < w)) |
            (((y >= 2) & (y <= 7.5)) | ((y >= 38.5) & (y <= 44))) & stall_x) & (x >= 14) & (x <= 58.1)
    island = (x >= 16) & (x <= 54.1) & (y >= 17) & (y <= 29)
    isl_line = island & ((np.abs(((x - 16) % 2.7) - 0.06) < w) | (np.abs(y - 23) < w))
    zebra = (((x >= 26) & (x <= 30) & (y >= 9.5) & (y <= 15)) & (np.floor((x - 26) / 0.5) % 2 == 0)) | \
            (((x >= 34) & (x <= 38) & (y >= 31.2) & (y <= 36.3)) & (np.floor((x - 34) / 0.5) % 2 == 0))
    ax, ay = x - 46, np.abs(y - 12.25)
    arrow = ((ax > -2) & (ax <= 0.2) & (ay < 0.18)) | ((ax > 0.2) & (ax < 1.2) & (ay < 0.6 * (1.2 - ax)))
    for m in (rows, isl_line, zebra, arrow):
        rgb[m] = white
    rgb[island & ~isl_line & (y > 22.6) & (y < 23.4)] = [150, 160, 110]
    blue = (x >= 41.5) & (x <= 44.2) & (y >= 2.1) & (y <= 7.4) & ~stall_x
    rgb[blue] = [40, 92, 170]
    dash = ((np.abs(y - 12.25) < w) | (np.abs(y - 33.75) < w)) & (x > 13) & (x < 57) & (((x - 13) % 4) < 2)
    dash |= ((np.abs(x - 8) < w) | (np.abs(x - 62) < w)) & (y > 17) & (y < 29) & (((y - 17) % 4) < 2)
    rgb[dash] = [222, 190, 70]
    edge = (x < 0.3) | (x > LOT_W - 0.3) | (y < 0.3) | (y > LOT_H - 0.3)
    rgb[edge] = [200, 170, 60]
    return rgb


def _segments():
    L1, L2, Q = (P1[0] - P0[0]) - 2 * RC, (P2[1] - P1[1]) - 2 * RC, math.pi / 2 * RC
    return [("line", (P0[0] + RC, P0[1], 0.0), L1), ("arc", (P1[0] - RC, P1[1] + RC, -math.pi / 2), Q),
            ("line", (P1[0], P1[1] + RC, math.pi / 2), L2), ("arc", (P2[0] - RC, P2[1] - RC, 0.0), Q),
            ("line", (P2[0] - RC, P2[1], math.pi), L1), ("arc", (P3[0] + RC, P3[1] - RC, math.pi / 2), Q),
            ("line", (P3[0], P3[1] - RC, -math.pi / 2), L2), ("arc", (P0[0] + RC, P0[1] + RC, math.pi), Q)]


LOOP_M = sum(s[2] for s in _segments())


def pose(s):
    """Pose of base_link after driving s meters: x, y (m) and heading (rad, counterclockwise from east)."""
    s = s % LOOP_M
    for kind, a, length in _segments():
        if s <= length:
            if kind == "line":
                return a[0] + s * math.cos(a[2]), a[1] + s * math.sin(a[2]), a[2]
            ang = a[2] + s / RC
            return a[0] + RC * math.cos(ang), a[1] + RC * math.sin(ang), ang + math.pi / 2
        s -= length
    return P0[0], P0[1], 0.0


def footprint(step=0.15):
    """Ground points ahead of the camera, in its own frame: f forward (4 to 11 m), l to the left, inside the 72 degree view."""
    pts = []
    for f in np.arange(4.0, 11.0 + 1e-9, step):
        half = 0.68 * f
        for l in np.arange(-half, half + 1e-9, step):
            pts.append((f, l))
    return np.array(pts)


def mosaic(speed=4.0, heading_err_deg=0.0, offset_ms=0.0, pitch_err_deg=0.0, ignore_lever=False, res=0.2, step=0.15):
    """Blend every image into one map. Returns the RGB map (rows north up) and misplacement statistics."""
    cols, rows_ = int(LOT_W / res), int(LOT_H / res)
    total = np.zeros((rows_ * cols, 3)); count = np.zeros(rows_ * cols)
    fp = footprint(step)
    f, l = fp[:, 0], fp[:, 1]
    dp, dpsi, dt = math.radians(pitch_err_deg), math.radians(heading_err_deg), offset_ms / 1000.0
    fe = CAM_H / np.tan(np.arctan(CAM_H / f) + dp)      # where a wrong pitch puts each ground point
    ratio = fe / f
    frames = int(LOOP_M / (speed / FPS))
    xs, ys, mx, my = [], [], [], []
    err5, err10 = [], []
    for k in range(frames):
        s = k * speed / FPS
        bx, by, ps = pose(s)
        ax, ay, pa = pose(s + speed * dt)                  # the pose the late time stamp picks
        antx, anty = ax + ANT_FWD * math.cos(pa), ay + ANT_FWD * math.sin(pa)
        pe = pa + dpsi
        lever = 0.0 if ignore_lever else ANT_FWD
        ex, ey = antx - lever * math.cos(pe), anty - lever * math.sin(pe)
        c, sn, ce, se = math.cos(ps), math.sin(ps), math.cos(pe), math.sin(pe)
        xs.append(bx + c * (CAM_FWD + f) - sn * l); ys.append(by + sn * (CAM_FWD + f) + c * l)
        mx.append(ex + ce * (CAM_FWD + fe) - se * l * ratio); my.append(ey + se * (CAM_FWD + fe) + ce * l * ratio)
        for d, bucket in ((5.0, err5), (10.0, err10)):
            de = CAM_H / math.tan(math.atan(CAM_H / d) + dp)
            gx, gy = bx + c * (CAM_FWD + d), by + sn * (CAM_FWD + d)
            bucket.append(math.hypot(ex + ce * (CAM_FWD + de) - gx, ey + se * (CAM_FWD + de) - gy))
    X, Y, MX, MY = (np.concatenate(v) for v in (xs, ys, mx, my))
    color = ground(X, Y)                                   # what the camera saw
    px, py = np.floor(MX / res).astype(int), rows_ - 1 - np.floor(MY / res).astype(int)   # where it was placed
    ok = (px >= 0) & (px < cols) & (py >= 0) & (py < rows_)
    idx = py[ok] * cols + px[ok]
    for ch in range(3):
        total[:, ch] = np.bincount(idx, weights=color[ok, ch], minlength=rows_ * cols)
    count = np.bincount(idx, minlength=rows_ * cols)
    img = np.zeros((rows_ * cols, 3))
    seen = count > 0
    img[seen] = total[seen] / count[seen, None]
    gy_, gx_ = np.divmod(np.arange(rows_ * cols), cols)
    faint = ground(gx_[~seen] * res, (rows_ - 1 - gy_[~seen]) * res)
    img[~seen] = 237 + 0.06 * (faint - 120)                # the parts no image covered, very faint
    return img.reshape(rows_, cols, 3).clip(0, 255).astype(np.uint8), {
        "frames": frames, "err_5m": float(np.mean(err5)), "err_10m": float(np.mean(err10)),
        "coverage": float(seen.mean())}


def demo_api(params_json):
    """Web page entry point. Input: {"speed", "heading", "offset_ms", "pitch", "ignore_lever"}."""
    p = json.loads(params_json or "{}")
    img, st = mosaic(float(p.get("speed", 4)), float(p.get("heading", 0)), float(p.get("offset_ms", 0)),
                     float(p.get("pitch", 0)), bool(p.get("ignore_lever", False)))
    st.update({"W": img.shape[1], "H": img.shape[0], "rgb": base64.b64encode(img.tobytes()).decode(), "res": 0.2})
    return json.dumps(st)


def main(argv):
    cases = [("no errors", {}), ("heading 0.5 deg", {"heading_err_deg": 0.5}), ("time offset 50 ms", {"offset_ms": 50}),
             ("pitch 0.2 deg", {"pitch_err_deg": 0.2}), ("lever arm ignored", {"ignore_lever": True})]
    for name, kw in cases:
        img, st = mosaic(**kw)
        print("{:18s} misplaced {:5.0f} cm at 5 m, {:5.0f} cm at 10 m ahead".format(name, st["err_5m"] * 100, st["err_10m"] * 100))
        if "--save" in argv:
            import matplotlib.pyplot as plt
            plt.imsave("mosaic_" + name.replace(" ", "_").replace(".", "") + ".png", img)


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