"""Analytic layer-normalization response; writes a PNG in the working directory.

The three input coordinates are (-t, 0, t), with unit gains and zero biases.
No data or illustrations from another publication are used.
"""

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


def response(t, epsilon):
    return t / np.sqrt(2 * t**2 / 3 + epsilon)


def slope(t, epsilon):
    return epsilon / (2 * t**2 / 3 + epsilon) ** 1.5


def verify():
    for epsilon in (0.04, 0.25, 1.0):
        np.testing.assert_allclose(slope(0.0, epsilon), 1 / np.sqrt(epsilon))
        for t in (-0.8, -0.1, 0.0, 0.2, 0.7):
            step = 1e-6
            finite_difference = (
                response(t + step, epsilon) - response(t - step, epsilon)
            ) / (2 * step)
            np.testing.assert_allclose(
                finite_difference, slope(t, epsilon), rtol=1e-8, atol=1e-9
            )
        z = np.array([-0.3, 0.0, 0.3])
        centered = z - z.mean()
        q = np.mean(centered**2)
        u = centered / np.sqrt(q + epsilon)
        np.testing.assert_allclose(np.mean(u**2), q / (q + epsilon))
        projection = np.eye(3) - np.ones((3, 3)) / 3
        jacobian = projection / np.sqrt(q + epsilon) - np.outer(
            centered, centered
        ) / (3 * (q + epsilon) ** 1.5)
        np.testing.assert_allclose(jacobian @ np.ones(3), 0, atol=1e-14)
        np.testing.assert_allclose(
            jacobian @ centered,
            epsilon * centered / (q + epsilon) ** 1.5,
            atol=1e-14,
        )
        assert np.linalg.norm(jacobian, 2) <= 1 / np.sqrt(epsilon) + 1e-12


def main():
    verify()
    plt.rcParams.update({"font.size": 10, "axes.titlesize": 11})
    fig, (forward, derivative) = plt.subplots(
        1, 2, figsize=(8.6, 4.1), dpi=100, facecolor="white"
    )
    fig.subplots_adjust(left=0.075, right=0.98, bottom=0.18, top=0.79, wspace=0.3)
    fig.suptitle("What epsilon changes in layer normalization", y=0.97, fontsize=14)
    fig.text(
        0.5, 0.88,
        r"Input $(-t,0,t)$; plot the third normalized coordinate; gain = 1, bias = 0",
        ha="center", fontsize=10,
    )
    t = np.linspace(-3, 3, 1201)
    near_zero = np.linspace(-1, 1, 1201)
    for epsilon, color in zip((0.04, 0.25, 1.0), ("#0072b2", "#d55e00", "#009e73")):
        label = rf"$\varepsilon={epsilon:g}$"
        forward.plot(t, response(t, epsilon), color=color, lw=2, label=label)
        derivative.plot(near_zero, slope(near_zero, epsilon), color=color, lw=2)
    asymptote = np.sqrt(3 / 2)
    for sign in (-1, 1):
        forward.axhline(sign * asymptote, color="#888888", ls="--", lw=0.8)
    forward.axhline(0, color="#999999", lw=0.7)
    forward.axvline(0, color="#999999", lw=0.7)
    forward.set(xlim=(-3, 3), ylim=(-1.4, 1.4), xlabel="Input amplitude t",
                ylabel="Normalized coordinate", title="Forward response")
    forward.legend(loc="lower right", framealpha=0.95, fontsize=9)
    derivative.set(xlim=(-1, 1), ylim=(0, 6.15), xlabel="Input amplitude t",
                   ylabel="Derivative of normalized coordinate",
                   title="Sensitivity near a constant input")
    derivative.text(0.5, 0.96, r"Slope at $t=0$: $1/\sqrt{\varepsilon}$",
                    transform=derivative.transAxes, ha="center", va="top", fontsize=10)
    for ax in (forward, derivative):
        ax.set_facecolor("white")
        ax.grid(alpha=0.22)
        ax.spines[["top", "right"]].set_visible(False)
    fig.savefig(Path.cwd() / "layer-normalization-epsilon-response.png",
                dpi=100, facecolor="white", transparent=False)
    plt.close(fig)


if __name__ == "__main__":
    main()