"""Original bounded RC timing uncertainty and common-reference examples."""
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def crossing(high, threshold, baseline=0.0, tau=1e-6):
    assert np.all(np.asarray(baseline) < threshold)
    assert np.all(np.asarray(threshold) < high)
    return tau * np.log((high - baseline) / (high - threshold))


def exact_uncertainty(rho):
    z = np.array([-1, -1, 1, 1])
    w = np.array([-1, 1, -1, 1])
    dh = 0.02 * z
    dt = 0.01 * (rho * z + np.sqrt(max(0, 1 - rho**2)) * w)
    assert np.isclose(np.mean(dh**2), 0.02**2)
    assert np.isclose(np.mean(dt**2), 0.01**2)
    assert np.isclose(np.mean(dh * dt), 0.02 * 0.01 * rho)
    times = crossing(1 + dh, 0.5 + dt)
    return np.std(times) * 1e9


def main():
    rho = np.linspace(-1, 1, 601)
    approximate = np.sqrt(8e-16 * (1 - rho)) * 1e9
    cases = np.array([-1.0, 0.0, 0.9, 1.0])
    exact = np.array([exact_uncertainty(r) for r in cases])
    prediction = np.sqrt(8e-16 * (1 - cases)) * 1e9
    assert np.allclose(exact[:-1], prediction[:-1], rtol=0.01)
    assert exact[-1] < 1e-6
    epsilon = np.linspace(-0.1, 0.1, 501)
    common_gain = crossing(1 + epsilon, 0.5 * (1 + epsilon))
    two_level_offset = crossing(1 + epsilon, 0.5 + epsilon)
    three_level_offset = crossing(1 + epsilon, 0.5 + epsilon, baseline=epsilon)
    nominal = crossing(1.0, 0.5)
    assert np.allclose(common_gain, nominal, rtol=0, atol=1e-20)
    assert np.allclose(three_level_offset, nominal, rtol=0, atol=1e-20)
    assert np.isclose((crossing(1.000001, 0.500001) - crossing(0.999999, 0.499999))
                      / 0.000002, 1e-6, rtol=1e-7)

    fig, (uncertainty, reference) = plt.subplots(1, 2, figsize=(10, 4.4), dpi=100,
                                               layout="constrained")
    fig.set_facecolor("white")
    uncertainty.plot(rho, approximate, color="#286f99", label="First-order uncertainty")
    uncertainty.scatter(cases, exact, color="#9f4b27", s=35, zorder=5,
                        label="Exact bounded four-state model")
    uncertainty.set(xlim=(-1.05, 1.05), ylim=(-1, 44), xlabel="Correlation coefficient ρ",
                    ylabel="Timing standard uncertainty (ns)",
                    title="Covariance changes the crossing-time uncertainty")
    uncertainty.legend(loc="upper right", fontsize=8)
    uncertainty.annotate("Matched common gain: zero", (1, exact[-1]), xytext=(-0.05, 5),
                          fontsize=9, arrowprops={"arrowstyle": "->", "color": "#444444"})
    reference.plot(epsilon * 100, common_gain * 1e6, color="#286f99", linewidth=2.7,
                   label="Common gain; baseline remains zero")
    reference.plot(epsilon * 100, three_level_offset * 1e6, "--", color="#52762a",
                   label="Common offset of all three levels")
    reference.plot(epsilon * 100, two_level_offset * 1e6, color="#9f4b27", linewidth=2,
                   label="Offset of high/threshold; baseline fixed")
    reference.set(xlim=(-10, 10), ylim=(0.57, 0.82), xlabel="Common perturbation ε (%)",
                  ylabel="Crossing time (µs)", title="Which levels share the error matters")
    reference.legend(loc="lower right", fontsize=7.8)
    for ax in (uncertainty, reference):
        ax.grid(alpha=0.15)
    fig.savefig(Path(__file__).with_suffix(".png").name, facecolor="white")
    plt.close(fig)


if __name__ == "__main__":
    main()