"""Synthetic phase-stepped Gaussian interferometry; save the PNG to cwd."""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def frames(zeta):
    zeta = np.asarray(zeta)
    phase = -np.arctan(zeta)
    signal = 1 / (1 + zeta**2)
    background = 1 + signal
    amplitude = 2 * np.sqrt(signal)
    steps = np.arange(4) * np.pi / 2
    return background[..., None] + amplitude[..., None] * np.cos(phase[..., None] + steps)


def recover(intensities):
    return np.arctan2(intensities[..., 3] - intensities[..., 1],
                      intensities[..., 0] - intensities[..., 2])


def main():
    zeta = np.linspace(-3, 3, 601)
    intensities = frames(zeta)
    np.testing.assert_allclose(recover(intensities), -np.arctan(zeta), atol=1e-14)
    np.testing.assert_allclose(frames(1), [2.5, 2.5, .5, .5], atol=1e-14)
    np.testing.assert_allclose(frames(-1), [2.5, .5, .5, 2.5], atol=1e-14)
    assert np.all(intensities >= -1e-14)
    np.testing.assert_allclose(intensities[:, 0], intensities[::-1, 0], atol=1e-14)

    plt.rcParams.update({"font.size": 11, "figure.facecolor": "white", "axes.facecolor": "white"})
    fig, axes = plt.subplots(1, 2, figsize=(9.6, 4.2), dpi=100)
    colors = ["#276FBF", "#BD5C18", "#518548", "#8E4C96"]
    labels = ["δ = 0", "δ = π/2", "δ = π", "δ = 3π/2"]
    for i, (color, label) in enumerate(zip(colors, labels)):
        axes[0].plot(zeta, intensities[:, i], color=color, linewidth=1.8, label=label)
    axes[0].set_title("Four phase-stepped intensities", fontsize=12)
    axes[0].set_ylabel("Normalized detector intensity")
    axes[0].set_ylim(-.12, 4.25)
    axes[0].legend(fontsize=9, loc="upper right", framealpha=.95)

    axes[1].plot(zeta, -np.arctan(zeta), color="#343434", linewidth=2,
                 label="Model: −atan(z/zR)")
    sample = np.linspace(-3, 3, 25)
    axes[1].scatter(sample, recover(frames(sample)), marker="o", s=24,
                    facecolor="white", edgecolor="#276FBF", zorder=3,
                    label="Retrieved from four frames")
    for location in [-1, 1]:
        axes[1].axvline(location, color="#777777", linestyle=":", linewidth=.8)
    axes[1].set_yticks([-np.pi / 2, -np.pi / 4, 0, np.pi / 4, np.pi / 2],
                      ["−π/2", "−π/4", "0", "π/4", "π/2"])
    axes[1].set_ylim(-1.65, 1.65)
    axes[1].set_title("Recovered relative phase", fontsize=12)
    axes[1].set_ylabel("Phase (rad)")
    axes[1].legend(fontsize=9, loc="upper right")
    for ax in axes:
        ax.set_xlim(-3, 3)
        ax.set_xlabel("Axial coordinate z/zR")
        ax.grid(alpha=.18)
        ax.spines[["top", "right"]].set_visible(False)
    fig.suptitle("Synthetic Gaussian/reference interferometry", fontsize=13, y=.98)
    fig.text(.5, .025, "The δ = 0 frame is identical at z = ±zR; the quadratures determine the phase sign.",
             ha="center", fontsize=10)
    fig.subplots_adjust(left=.075, right=.98, top=.84, bottom=.18, wspace=.30)
    output = Path(__file__).with_suffix(".png").name
    fig.savefig(output, dpi=100, facecolor="white")
    plt.close(fig)
    print(f"{output}: 960×420; all phase-retrieval and frame-symmetry checks passed")


if __name__ == "__main__":
    main()