#!/usr/bin/env python3
"""Numerical verification of Proposition 3.1 of

    B. R. Davis, "Плоские модели систем со связностью: характеризация и
    препятствия" (Flat models of systems with connection: characterization
    and obstructions).

Proposition 3.1 states that for orthonormal X, Y in T_x M and gamma_eps the
boundary of the eps-square spanned by them,

    tau_{gamma_eps} = exp(-eps^2 F_x(X,Y) + O(eps^3)),
    d_G(tau_{gamma_eps}, e) = eps^2 ||F_x(X,Y)|| + O(eps^3),

and in particular  lim_{eps->0} eps^-2 d_G(tau_{gamma_eps}, e) = ||F_x(X,Y)||.

This script computes the left-hand side by numerically transporting around the
square and compares it to the curvature computed independently from the gauge
field.  G = SU(2); the base is the (u,v) plane with X, Y the coordinate
directions.

Three cases are run, plus a self-check on the third:

  curved    a connection with non-vanishing curvature, including a non-zero
            commutator term, so the test exercises the full
            F = dA + (1/2)[A,A] and not merely dA.

  flat      a gauge transform of the zero connection, A = g^-1 dg.  Non-zero,
            non-abelian, non-constant, and flat.  The measured ratio must go to
            zero.  This is the control: it fails loudly if the script is
            picking up a generic "bigger loop, bigger number" signal rather
            than curvature.

  abelian-only
            the curved connection measured against a DELIBERATELY WRONG
            prediction that drops the [A_u, A_v] term.  This must NOT converge.
            It is here so that a passing run cannot be explained by the test
            being insensitive to the non-abelian part of the curvature.

  3b        the Case 3 gate run against the CORRECT curvature, which must make
            it fail.  A control that passes whatever it is shown establishes
            nothing, so this is reported rather than left to be taken on trust.

Conventions follow the paper.  The inner product on su(2) is <A,B> = -tr(AB),
Ad-invariant and positive definite on anti-Hermitian traceless matrices; d_G is
the distance of the bi-invariant metric it induces, so d_G(exp z, e) = ||z||
within the injectivity radius.  The SAME inner product is used for ||F|| and
for d_G -- the identity is normalisation-independent only if that is true.

Run:  python flat_models_prop31_check.py
Exits 0 on success, 1 if any case fails.  numpy only.

Copyright (c) 2026 Bee Rosa Davis.  Licensed CC BY-NC 4.0 --
https://creativecommons.org/licenses/by-nc/4.0/ -- free to use, share and adapt
for education and research with attribution; commercial use requires permission.
Deposited alongside the paper; the Zenodo record carries the canonical terms.
"""

from __future__ import annotations

import sys

import numpy as np

# ----------------------------------------------------------------------------
# su(2): basis, inner product, exponential, logarithm, distance
# ----------------------------------------------------------------------------

SIGMA = (
    np.array([[0, 1], [1, 0]], dtype=complex),
    np.array([[0, -1j], [1j, 0]], dtype=complex),
    np.array([[1, 0], [0, -1]], dtype=complex),
)

# T_a = -(i/2) sigma_a.  Anti-Hermitian, traceless, and [T_a, T_b] = eps_abc T_c.
T = tuple(-0.5j * s for s in SIGMA)

I2 = np.eye(2, dtype=complex)


def inner(a: np.ndarray, b: np.ndarray) -> float:
    """Ad-invariant inner product on su(2):  <A,B> = -tr(AB), real."""
    return float(np.real(-np.trace(a @ b)))


def norm(a: np.ndarray) -> float:
    q = inner(a, a)
    return float(np.sqrt(q)) if q > 0.0 else 0.0


def su2_exp(z: np.ndarray) -> np.ndarray:
    """exp of an su(2) element, in closed form.

    Writing z = x . T with r = |x|, we have z^2 = -(r^2/4) I, hence
    exp(z) = cos(r/2) I + (2 sin(r/2) / r) z, continuous at r = 0.
    """
    # ||z||^2 = r^2 / 2  for this basis, so r = sqrt(2) * ||z||.
    r = np.sqrt(2.0) * norm(z)
    if r < 1e-14:
        # second order is already far below double precision here
        return I2 + z + 0.5 * (z @ z)
    return np.cos(r / 2.0) * I2 + (2.0 * np.sin(r / 2.0) / r) * z


def su2_log(g: np.ndarray) -> np.ndarray:
    """Principal logarithm of g in SU(2), valid for r < pi.

    Taking the traceless part P = g - (tr g / 2) I kills the cos(r/2) term and
    leaves P = (2 sin(r/2) / r) z, so ||P||_F = sqrt(2) sin(r/2).  Recovering r
    through arcsin rather than arccos(tr g / 2) matters: near the identity
    tr g / 2 = 1 - r^2/8, and for the small loops this script measures that is
    1 - O(1e-15), where arccos has no precision left.  arcsin near zero is
    well conditioned.
    """
    p = g - 0.5 * np.trace(g) * I2
    s = float(np.linalg.norm(p, "fro")) / np.sqrt(2.0)  # = sin(r/2)
    s = min(s, 1.0)
    r = 2.0 * np.arcsin(s)
    if np.real(np.trace(g)) < 0.0:  # r > pi, outside the injectivity radius
        raise ValueError("holonomy outside the injectivity radius; shrink eps")
    if r < 1e-14:
        return p
    return (r / (2.0 * np.sin(r / 2.0))) * p


def dist_to_identity(g: np.ndarray) -> float:
    """d_G(g, e) for the bi-invariant metric induced by <.,.>."""
    return norm(su2_log(g))


# ----------------------------------------------------------------------------
# Connections on the (u,v) plane, valued in su(2)
# ----------------------------------------------------------------------------


class Connection:
    """A gauge field A = A_u du + A_v dv with a known curvature at the origin."""

    name = "connection"

    def A(self, u: float, v: float) -> tuple[np.ndarray, np.ndarray]:
        raise NotImplementedError

    def curvature_at_origin(self) -> np.ndarray:
        """F(X,Y) = d_u A_v - d_v A_u + [A_u, A_v], evaluated at (0,0)."""
        raise NotImplementedError


class CurvedConnection(Connection):
    """A_u = alpha v T1 + gamma T3,   A_v = beta u T2 + delta T1.

    Affine in (u, v), so the partial derivatives are exact constants and no
    finite differencing enters the predicted curvature.  At the origin the
    commutator [A_u, A_v] = gamma delta [T3, T1] = gamma delta T2 is non-zero,
    which is the point: the prediction being tested is not merely dA.
    """

    name = "curved"

    def __init__(self, alpha=0.7, beta=1.3, gamma=0.9, delta=1.1):
        self.alpha, self.beta, self.gamma, self.delta = alpha, beta, gamma, delta

    def A(self, u, v):
        a_u = self.alpha * v * T[0] + self.gamma * T[2]
        a_v = self.beta * u * T[1] + self.delta * T[0]
        return a_u, a_v

    def curvature_at_origin(self):
        d_u_A_v = self.beta * T[1]          # d/du of (beta u T2 + delta T1)
        d_v_A_u = self.alpha * T[0]         # d/dv of (alpha v T1 + gamma T3)
        a_u, a_v = self.A(0.0, 0.0)
        return d_u_A_v - d_v_A_u + (a_u @ a_v - a_v @ a_u)

    def curvature_without_commutator(self):
        """The wrong answer: dA alone, dropping (1/2)[A,A]."""
        return self.beta * T[1] - self.alpha * T[0]


class FlatConnection(Connection):
    """A = h^-1 dh for h(u,v) = exp(u T1) exp(v T2).

    Pure gauge, hence identically flat, but non-zero, non-constant and
    non-abelian -- so a script that merely responds to the size of A, or that
    mishandles path ordering, will not pass this one.
    """

    name = "flat"

    @staticmethod
    def _h(u, v):
        return su2_exp(u * T[0]) @ su2_exp(v * T[1])

    def A(self, u, v):
        h = self._h(u, v)
        h_inv = np.conjugate(h).T  # SU(2): inverse is the conjugate transpose
        # d/du h = T1 h   (left factor), so h^-1 d_u h = h^-1 T1 h
        d_u_h = T[0] @ h
        # d/dv h = exp(u T1) T2 exp(v T2) = h T2
        d_v_h = h @ T[1]
        return h_inv @ d_u_h, h_inv @ d_v_h

    def curvature_at_origin(self):
        return np.zeros((2, 2), dtype=complex)


# ----------------------------------------------------------------------------
# Parallel transport
# ----------------------------------------------------------------------------


def transport_segment(conn, start, delta, steps):
    """Transport along the straight segment c(t) = start + t * delta, t in [0,1].

    Solves  g'(t) = -A_{c(t)}(c'(t)) g(t),  g(0) = e,  by the midpoint product
    of exponentials.  Every factor is an exponential of an su(2) element, so
    the result stays exactly in SU(2) -- no drift off the group to re-project.
    """
    g = I2.copy()
    h = 1.0 / steps
    for k in range(steps):
        t_mid = (k + 0.5) * h
        u = start[0] + t_mid * delta[0]
        v = start[1] + t_mid * delta[1]
        a_u, a_v = conn.A(u, v)
        m = delta[0] * a_u + delta[1] * a_v  # A_c(cdot)
        g = su2_exp(-m * h) @ g
    return g


def holonomy_square(conn, eps, steps=512):
    """Transport around the boundary of the eps-square, as the paper orders it.

    c1: 0 -> eps X,  c2: eps X -> eps X + eps Y,
    c3: eps X + eps Y -> eps Y,  c4: eps Y -> 0.

    Transport operators compose on the left, so tau = T4 T3 T2 T1.
    """
    t1 = transport_segment(conn, (0.0, 0.0), (eps, 0.0), steps)
    t2 = transport_segment(conn, (eps, 0.0), (0.0, eps), steps)
    t3 = transport_segment(conn, (eps, eps), (-eps, 0.0), steps)
    t4 = transport_segment(conn, (0.0, eps), (0.0, -eps), steps)
    return t4 @ t3 @ t2 @ t1


# ----------------------------------------------------------------------------
# The measurement
# ----------------------------------------------------------------------------


def run_case(conn, predicted_F, label, exponents, steps=512):
    """Measure eps^-2 d_G(tau, e) against ||predicted_F|| over a range of eps."""
    target = norm(predicted_F)
    print(f"\n  {label}")
    print(f"  predicted ||F|| = {target:.12f}")
    print()
    print("      eps         eps^-2 d_G(tau,e)     |ratio - ||F|| |    "
          "rel. err of log(tau)/(-eps^2) vs F")
    print("    " + "-" * 94)

    rows = []
    for k in exponents:
        eps = 2.0 ** (-k)
        tau = holonomy_square(conn, eps, steps)
        ratio = dist_to_identity(tau) / eps**2
        gap = abs(ratio - target)

        # Sharper than the norm: the full matrix statement of Prop 3.1,
        # log(tau) = -eps^2 F + O(eps^3), tested in su(2) and not just in R.
        f_measured = su2_log(tau) / (-(eps**2))
        denom = target if target > 1e-12 else 1.0
        matrix_rel = norm(f_measured - predicted_F) / denom

        rows.append((eps, ratio, gap, matrix_rel))
        print(f"    {eps:>10.3e}   {ratio:>18.12f}   {gap:>12.3e}    {matrix_rel:>12.3e}")

    return target, rows


def gate(conn, prediction, eps, steps=512):
    """Evaluate Case 3's pass condition for one prediction at one eps.

    Returns (scalar gap, relative matrix error, whether the gate passes).
    """
    target = norm(prediction)
    tau = holonomy_square(conn, eps, steps)
    gap = abs(dist_to_identity(tau) / eps**2 - target)
    denom = target if target > 1e-12 else 1.0
    mat = norm(su2_log(tau) / (-(eps**2)) - prediction) / denom
    return gap, mat, (gap > 1e-2 and mat > 1e-2)


def report_order(rows, target):
    """O(eps) residual means the gap roughly halves as eps halves."""
    print()
    if target > 1e-12:
        print("    successive gap ratios (expect ~0.5 for an O(eps) residual):")
        line = []
        for (_, _, g_prev, _), (_, _, g_now, _) in zip(rows, rows[1:]):
            line.append("  -  " if g_prev == 0 else f"{g_now / g_prev:.3f}")
        print("      " + "  ".join(line))


def main() -> int:
    exponents = list(range(3, 13))  # eps from 1/8 down to 1/4096
    failures = []

    print("=" * 100)
    print("Proposition 3.1:  lim_{eps->0} eps^-2 d_G(tau_{gamma_eps}, e) = ||F_x(X,Y)||")
    print("G = SU(2);  <A,B> = -tr(AB);  d_G from the induced bi-invariant metric")
    print("=" * 100)

    # ---- 1. curved -------------------------------------------------------
    curved = CurvedConnection()
    f_true = curved.curvature_at_origin()
    target, rows = run_case(curved, f_true, "CASE 1 -- curved connection", exponents)
    report_order(rows, target)

    gap_fine = rows[-1][2]
    mat_fine = rows[-1][3]
    tol = 1e-3
    ok = gap_fine < tol and mat_fine < tol
    print(f"\n    -> finest gap {gap_fine:.3e}, matrix rel. err {mat_fine:.3e}, "
          f"tolerance {tol:.0e}:  {'PASS' if ok else 'FAIL'}")
    if not ok:
        failures.append("curved")

    # ---- 2. flat control -------------------------------------------------
    flat = FlatConnection()
    target_f, rows_f = run_case(
        flat, flat.curvature_at_origin(), "CASE 2 -- flat control (A = h^-1 dh)", exponents
    )
    ratio_fine = rows_f[-1][1]
    ok_flat = ratio_fine < 1e-6
    print(f"\n    -> ratio at the finest eps is {ratio_fine:.3e}, must vanish:  "
          f"{'PASS' if ok_flat else 'FAIL'}")
    if not ok_flat:
        failures.append("flat")

    # ---- 3. the deliberately wrong prediction ----------------------------
    f_wrong = curved.curvature_without_commutator()
    target_w, rows_w = run_case(
        curved, f_wrong, "CASE 3 -- same connection, prediction missing [A_u, A_v]", exponents
    )
    report_order(rows_w, target_w)
    gap_wrong = rows_w[-1][2]
    mat_wrong = rows_w[-1][3]
    # Both must stay large.  Asserting only the scalar gap would leave the
    # control half-blind: it is the Lie-algebra column that actually sees the
    # dropped commutator as a direction and not merely as a change in length.
    ok_wrong = gap_wrong > 1e-2 and mat_wrong > 1e-2
    scalar_plateau = abs(norm(f_true) - norm(f_wrong))
    commutator = norm(f_true - f_wrong)
    print(f"\n    -> finest scalar gap {gap_wrong:.3e}, matrix error {mat_wrong:.3e}; "
          f"neither may converge, so both above 1e-2 is the pass:  "
          f"{'PASS' if ok_wrong else 'FAIL'}")
    print("       These plateau at two different quantities, which is worth keeping")
    print("       distinct.  The scalar column measures lengths only, so its gap")
    print(f"       tends to  | ||F|| - ||F_wrong|| | = {scalar_plateau:.12f}.")
    print("       The matrix column compares elements of su(2), so its error tends")
    print(f"       to the dropped commutator itself, ||F - F_wrong|| = "
          f"||[A_u, A_v]|| = {commutator:.12f}")
    print(f"       (printed above relative to ||F_wrong||, hence {mat_wrong:.4f}).")
    if not ok_wrong:
        failures.append("abelian-only")

    # ---- 3b. does the Case 3 gate actually discriminate? -----------------
    print("\n  CASE 3b -- the Case 3 gate, turned on itself")
    print("  A control that passes whatever it is shown is not a control.  Case 3")
    print("  passes only because its prediction omits the commutator; hand the same")
    print("  gate the CORRECT curvature and it has to go red.  Both lines below come")
    print("  from the same code path that decides Case 3.\n")
    eps_fine = 2.0 ** (-exponents[-1])
    g_w, m_w, gate_w = gate(curved, f_wrong, eps_fine)
    g_r, m_r, gate_r = gate(curved, f_true, eps_fine)
    print(f"      wrong prediction  (Case 3 as shipped):   gap={g_w:.3e}   "
          f"matrix={m_w:.3e}   gate={'PASS' if gate_w else 'FAIL'}")
    print(f"      correct prediction (mechanism removed):  gap={g_r:.3e}   "
          f"matrix={m_r:.3e}   gate={'PASS' if gate_r else 'FAIL'}")
    ok_discriminates = gate_w and not gate_r
    print(f"\n    -> the gate must pass the first and fail the second:  "
          f"{'PASS' if ok_discriminates else 'FAIL'}")
    if not ok_discriminates:
        failures.append("gate-discrimination")

    print("\n" + "=" * 100)
    if failures:
        print("FAILED: " + ", ".join(failures))
        return 1
    print("All cases pass.")
    print("Proposition 3.1 is reproduced numerically, the flat control vanishes,")
    print("the measurement is sensitive to the non-abelian term of the curvature,")
    print("and the control that establishes that last point is itself falsifiable.")
    print("=" * 100)
    return 0


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