"""Draw original ideal stereo ray geometry; write the PNG to the current directory.

Run with ``make _media/rectified-stereo-triangulation.png``. The drawing
uses virtual image planes in front of pinholes, so coordinates are upright.
The diagram has arbitrary common length units, rather than pixel coordinates.
"""
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def main():
    fig, ax = plt.subplots(figsize=(9.2, 5.5), layout="constrained")
    fig.patch.set_facecolor("white")
    baseline, focal, depth, point_x = 2.0, 1.3, 5.0, 1.5
    colors = ("#176b91", "#bb5425")
    for center, color, name in zip((0.0, baseline), colors, ("Left", "Right")):
        image_x = center + focal * (point_x - center) / depth
        ax.plot([center, point_x], [0, depth], color=color, lw=2)
        ax.plot([center, center], [0, depth + .3], "--", color=".65", lw=1)
        ax.plot([center - .62, center + .62], [focal, focal], color=color, lw=3)
        ax.scatter([center], [0], color=color, s=60, zorder=3)
        ax.scatter([image_x], [focal], color=color, s=45, zorder=4)
        ax.text(center, -.2, f"{name} pinhole", ha="center", va="top", fontsize=11)
        ax.annotate("", (image_x, focal + .18), (center, focal + .18),
                    arrowprops={"arrowstyle": "<->", "color": color})
        ax.text((center + image_x) / 2, focal + .28,
                r"$u_L$" if center == 0 else r"$u_R<0$", ha="center", color=color)
    ax.scatter([point_x], [depth], s=65, color="#3e5e32", zorder=4)
    ax.text(point_x + .15, depth, r"Scene point $P=(X,Z)$", va="center", fontsize=12)
    ax.annotate("", (baseline, -.65), (0, -.65), arrowprops={"arrowstyle": "<->"})
    ax.text(baseline / 2, -.79, r"Camera baseline $B$", ha="center", va="top", fontsize=12)
    for x, top, label in ((-.95, focal, r"$f$"), (-1.35, depth, r"$Z$")):
        ax.annotate("", (x, top), (x, 0), arrowprops={"arrowstyle": "<->"})
        ax.text(x - .10, top / 2, label, ha="right", va="center", fontsize=13)
    ax.text(2.75, focal, "Virtual image planes\n(front projection)", va="center", fontsize=11)
    ax.text(2.75, 3.1, r"$d=u_L-u_R=fB/Z$" + "\n" + r"$Z=fB/d$", fontsize=15)
    ax.set(xlim=(-1.8, 5.0), ylim=(-1.1, 5.8), aspect="equal")
    ax.set_title("Stereo depth from two viewing directions", fontsize=15)
    ax.axis("off")
    fig.savefig(Path("rectified-stereo-triangulation.png"), dpi=100, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()