"""
block_matching.py: find the disparity of every pixel in a stereo pair by block matching.

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

After rectification (Lecture 13), the match of a left pixel lies on the same
row of the right image, shifted left by its disparity d. Block matching tries
every shift d = 0, 1, ..., max_disp - 1, compares a small window around the
pixel by the sum of absolute differences (SAD), and keeps the cheapest shift.
Depth then follows from Z = f B / d.

The scene is a small synthetic pair with a known true disparity everywhere:
    a textured wall far away (d = 5), a plain painted wall (d = 11),
    a near box (d = 20), and a fence with stripes every 8 pixels (d = 8).

Two checks that real matchers use mark doubtful pixels as invalid instead of guessing:
    uniqueness    the best cost must beat the second best (away from it) by a margin
    left-right    matching right to left must give the same d back

Run it from a terminal:
    python3 block_matching.py            accuracy for several window sizes
"""
import base64
import json
import sys

import numpy as np

W, H, MAX_DISP = 200, 120, 32
# region in the left image (x0, x1, y0, y1), true disparity, and texture
LAYERS = [
    {"name": "far wall", "box": (0, W, 0, H), "d": 5, "tex": "noise"},
    {"name": "plain wall", "box": (14, 70, 18, 100), "d": 11, "tex": "plain"},
    {"name": "near box", "box": (92, 152, 30, 92), "d": 20, "tex": "noise2"},
    {"name": "fence", "box": (160, 196, 10, 110), "d": 8, "tex": "stripes"},
]


def _hash(x, y, salt):
    """Repeatable pseudo-random values in [0, 1) for integer grids."""
    h = (x.astype(np.uint64) * np.uint64(374761393) + y.astype(np.uint64) * np.uint64(668265263) + np.uint64(salt)) & 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 _value_noise(x, y, salt):
    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 = _hash(xi, yi, salt), _hash(xi + 1, yi, salt)
    c, d = _hash(xi, yi + 1, salt), _hash(xi + 1, yi + 1, salt)
    return a + (b - a) * sx + (c - a) * sy + (a - b - c + d) * sx * sy


def texture(layer, X, Y):
    """Brightness of a layer at scene column X (left-image coordinates) and row Y."""
    t = layer["tex"]
    if t == "plain":
        return np.full(X.shape, 150.0)
    if t == "stripes":
        return np.where((np.floor(X) % 8) < 4, 205.0, 70.0)
    s = 0.35 if t == "noise" else 0.5
    return 60 + 105 * _value_noise(X * s, Y * s, layer["d"] * 13) + 45 * _hash(np.floor(X).astype(np.int64) + 500, Y.astype(np.int64), layer["d"])


def make_pair(camera_noise=3.0, seed=1):
    """Left image, right image, and the true disparity of every left pixel."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[0:H, 0:W].astype(float)
    left, right = np.zeros((H, W)), np.zeros((H, W))
    truth = np.zeros((H, W))
    best_d = np.full((H, W), -1.0)
    for L in LAYERS:                                   # later layers are in front in the left image
        x0, x1, y0, y1 = L["box"]
        m = (xx >= x0) & (xx < x1) & (yy >= y0) & (yy < y1)
        left[m] = texture(L, xx, yy)[m]
        truth[m] = L["d"]
        # right image: the point at left column x + d shows at right column x; the nearest layer wins
        xl = xx + L["d"]
        mr = (xl >= x0) & (xl < x1) & (yy >= y0) & (yy < y1) & (L["d"] > best_d)
        right[mr] = texture(L, xl, yy)[mr]
        best_d[mr] = L["d"]
    right[best_d < 0] = 120.0
    left += rng.normal(0, camera_noise, left.shape)
    right += rng.normal(0, camera_noise, right.shape)
    return np.clip(left, 0, 255), np.clip(right, 0, 255), truth


def box_sum(img, r):
    """Sum over a (2r+1) x (2r+1) window at every pixel, with an integral image. Borders are inf."""
    S = np.zeros((img.shape[0] + 1, img.shape[1] + 1))
    S[1:, 1:] = img.cumsum(0).cumsum(1)
    out = np.full(img.shape, np.inf)
    k = 2 * r + 1
    out[r:img.shape[0] - r, r:img.shape[1] - r] = S[k:, k:] - S[:-k, k:] - S[k:, :-k] + S[:-k, :-k]
    return out


def cost_volume(a, b, window, max_disp, direction=1):
    """SAD cost for every pixel and every shift. direction 1: left to right, -1: right to left."""
    r = (window - 1) // 2
    cost = np.full((max_disp,) + a.shape, np.inf)
    for d in range(max_disp):
        diff = np.full(a.shape, 60.0)                  # outside the other image: a fixed, poor cost
        if direction == 1:
            diff[:, d:] = np.abs(a[:, d:] - b[:, :a.shape[1] - d])
        else:
            diff[:, :a.shape[1] - d] = np.abs(a[:, :a.shape[1] - d] - b[:, d:])
        cost[d] = box_sum(diff, r)
    return cost


def block_matching(left, right, window=5, max_disp=MAX_DISP, uniqueness=0.0, lr_check=False):
    """Disparity for every left pixel; -1 where the match is invalid."""
    cost = cost_volume(left, right, window, max_disp, 1)
    disp = np.argmin(cost, axis=0)
    best = np.take_along_axis(cost, disp[None], 0)[0]
    disp = disp.astype(float)
    disp[~np.isfinite(best)] = -1
    if uniqueness > 0:
        # the best cost must beat every cost more than 1 shift away by the margin
        dd = np.arange(max_disp)[:, None, None]
        far = np.abs(dd - disp[None]) > 1
        second = np.where(far, cost, np.inf).min(axis=0)
        disp[best * (1 + uniqueness) >= second] = -1
    if lr_check:
        cr = cost_volume(right, left, window, max_disp, -1)
        dr = np.argmin(cr, axis=0)
        cols = np.arange(left.shape[1])[None, :] - np.maximum(disp, 0).astype(int)
        rows = np.arange(left.shape[0])[:, None].repeat(left.shape[1], 1)
        ok = cols >= 0
        back = np.full(disp.shape, -99.0)
        back[ok] = dr[rows[ok], cols[ok]]
        disp[(np.abs(back - disp) > 1) & (disp >= 0)] = -1
    return disp


def evaluate(disp, truth):
    """Share of valid pixels within 1 px of the truth, overall and in each kind of region."""
    yy, xx = np.mgrid[0:H, 0:W]
    def region(x0, x1, y0, y1):
        return (xx >= x0) & (xx < x1) & (yy >= y0) & (yy < y1)
    box = region(92, 152, 30, 92)
    regions = {"all": np.ones_like(truth, dtype=bool), "plain wall": region(14, 70, 18, 100),
               "fence": region(160, 196, 10, 110), "box edges": box & ~region(96, 148, 34, 88)}
    out = {}
    valid = disp >= 0
    for k, m in regions.items():
        v = m & valid
        out[k] = {"correct": float((np.abs(disp - truth) <= 1)[v].mean()) if v.any() else None,
                  "valid": float(v.sum() / m.sum())}
    return out


def to_b64(img):
    return base64.b64encode(np.clip(img, 0, 255).astype(np.uint8).tobytes()).decode()


def demo_api(params_json):
    """Web page entry point. Input: {"window": odd int, "uniqueness": 0..0.5, "lr": bool}."""
    p = json.loads(params_json or "{}")
    window = int(p.get("window", 5)) | 1
    left, right, truth = make_pair()
    disp = block_matching(left, right, window, MAX_DISP, float(p.get("uniqueness", 0.0)), bool(p.get("lr", False)))
    code = np.where(disp < 0, 255, disp * 8)           # 0..248 for d = 0..31, 255 for invalid
    return json.dumps({"W": W, "H": H, "left": to_b64(left), "right": to_b64(right),
                       "disp": to_b64(code), "truth": to_b64(truth * 8), "stats": evaluate(disp, truth), "window": window})


def main(argv):
    left, right, truth = make_pair()
    for w in (3, 5, 9, 15, 21):
        s = evaluate(block_matching(left, right, w), truth)
        print("window {:2d}: right within 1 px {:4.0%}   plain wall {:4.0%}   fence {:4.0%}   box edges {:4.0%}".format(
            w, s["all"]["correct"], s["plain wall"]["correct"], s["fence"]["correct"], s["box edges"]["correct"]))
    s = evaluate(block_matching(left, right, 9, uniqueness=0.15, lr_check=True), truth)
    print("window  9 with uniqueness and left-right checks: right {:4.0%} of the {:4.0%} still valid; plain wall valid {:4.0%}".format(
        s["all"]["correct"], s["all"]["valid"], s["plain wall"]["valid"]))


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