"""
stereo_baseline.py: pick the stereo baseline from the range and accuracy you need.

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

A rectified stereo pair measures depth from disparity (Lecture 13):

    Z = f B / d            f: focal length in pixels, B: baseline in meters, d: disparity in pixels

Matching is never perfect. If d is off by delta_d pixels, the depth is off by

    delta_Z = Z^2 delta_d / (f B)

so twice as far means four times the error, and twice the baseline (or focal
length) means half. A wider baseline also has a cost: the matcher only searches
d = 0 to (search - 1) pixels, so anything closer than Z_min = f B / search is
out of range.

Camera numbers are the course rig: Blackfly S, 4.5 um pixels, 1616 pixels wide.
A 5 mm lens gives f = 5 / 0.0045 = 1111 px.

Run it from a terminal:
    python3 stereo_baseline.py                 the baseline table for the 5 mm lens
    python3 stereo_baseline.py 25 0.10         baseline needed for 10 cm at 25 m
"""
import json
import math
import sys

PIXEL_MM = 0.0045            # 4.5 um pixels
WIDTH_PX = 1616
GARAGE_TESTS_M = [3, 5, 10, 15]


def focal_px(f_mm):
    return f_mm / PIXEL_MM


def hfov_deg(f_px):
    return math.degrees(2 * math.atan(WIDTH_PX / 2 / f_px))


def depth_error(Z, f, B, match_px=0.25):
    """1-sigma depth error (m) at distance Z (m)."""
    return Z * Z * match_px / (f * B)


def disparity(Z, f, B):
    return f * B / Z


def closest_depth(f, B, search=256):
    """Nearest distance whose disparity still fits in the search range."""
    return f * B / search


def max_range(accuracy_m, f, B, match_px=0.25):
    """Farthest distance where the depth error is still within accuracy_m."""
    return math.sqrt(accuracy_m * f * B / match_px)


def baseline_for(Z, accuracy_m, f, match_px=0.25):
    """Smallest baseline that gives accuracy_m at distance Z."""
    return Z * Z * match_px / (f * accuracy_m)


def table(f, baselines=(0.12, 0.30, 0.50, 0.80, 1.20), distances=(3, 5, 10, 15, 20, 40), match_px=0.25, search=256):
    return [{"B": B, "errors_m": [depth_error(Z, f, B, match_px) for Z in distances],
             "closest_m": closest_depth(f, B, search)} for B in baselines]


def demo_api(params_json):
    """Web page entry point. Input: {"f_mm", "B", "match_px", "search", "range_m", "acc_m"}."""
    p = json.loads(params_json or "{}")
    f = focal_px(float(p.get("f_mm", 5.0)))
    B, e = float(p.get("B", 0.5)), float(p.get("match_px", 0.25))
    search, R, acc = int(p.get("search", 256)), float(p.get("range_m", 15.0)), float(p.get("acc_m", 0.12))
    zs = [0.25 * k for k in range(1, 161)]
    zmin = closest_depth(f, B, search)
    out = {
        "f_px": f, "hfov_deg": hfov_deg(f), "B": B, "match_px": e, "search": search,
        "curve": {"z": zs, "err": [depth_error(z, f, B, e) for z in zs]},
        "closest_m": zmin, "err_at_range": depth_error(R, f, B, e), "disp_at_range": disparity(R, f, B),
        "meets": depth_error(R, f, B, e) <= acc and R >= zmin, "max_range_m": max_range(acc, f, B, e),
        "baseline_needed": baseline_for(R, acc, f, e),
        "closest_with_needed": closest_depth(f, baseline_for(R, acc, f, e), search),
        "garage": [{"Z": Z, "d": disparity(Z, f, B), "in_range": Z >= zmin, "err_m": depth_error(Z, f, B, e)} for Z in GARAGE_TESTS_M],
        "table": table(f, match_px=e, search=search), "distances": [3, 5, 10, 15, 20, 40],
    }
    return json.dumps(out)


def main(argv):
    f = focal_px(5.0)
    if len(argv) > 2:
        Z, acc = float(argv[1]), float(argv[2])
        B = baseline_for(Z, acc, f)
        print("for {:.0f} cm at {:.0f} m: baseline {:.2f} m, closest usable depth {:.1f} m".format(acc * 100, Z, B, closest_depth(f, B)))
        return
    print("f = {:.0f} px, {:.0f} deg wide; matching error 0.25 px; search 256 px".format(f, hfov_deg(f)))
    print("  B (m)   " + "".join("{:>9}".format(str(z) + " m") for z in (3, 5, 10, 15, 20, 40)) + "   closest")
    for row in table(f):
        print("  {:5.2f}   ".format(row["B"]) + "".join("{:>9}".format("{:.1f} cm".format(v * 100) if v < 1 else "{:.2f} m".format(v)) for v in row["errors_m"])
              + "   {:5.1f} m".format(row["closest_m"]))


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