#!/usr/bin/env python3
"""Exact checks for the manuscript in submission.md.

Requires Python 3, SymPy, and mpmath. No network access. Exits 0 on success.
"""

from __future__ import annotations

import platform
import sys

import mpmath as mp
import sympy as sp


def fail(message: str) -> None:
    raise SystemExit(message)


def main() -> None:
    chain_store()
    wright_n2()
    linkage_and_replicator()
    revision_attraction()
    buyer_seller()
    owner_intruder()
    rock_scissors_paper()
    sunspot()
    print(
        "PASS"
        f" python={platform.python_version()}"
        f" sympy={sp.__version__}"
        f" mpmath={mp.__version__}"
    )


def chain_store() -> None:
    x, y, lam = sp.symbols("x y lam")
    xdot = x * (1 - x) * (2 * y - 1)
    ydot = y * (1 - y) * (2 * x)
    jac = sp.Matrix(
        [
            [sp.diff(xdot, x), sp.diff(xdot, y)],
            [sp.diff(ydot, x), sp.diff(ydot, y)],
        ]
    ).subs(x, 0)
    jac = sp.simplify(jac)
    expected = sp.Matrix([[2 * y - 1, 0], [2 * y * (1 - y), 0]])
    if sp.simplify(jac - expected) != sp.zeros(2):
        fail("chain-store Jacobian")
    poly = sp.simplify(jac.charpoly(lam).as_expr())
    if sp.expand(poly - lam * (lam - 2 * y + 1)) != 0:
        fail("chain-store characteristic polynomial")


def wright_n2() -> None:
    p00, p01, p10 = sp.symbols("p00 p01 p10", positive=True)
    p11 = p01 * p10 / p00
    mass = p00 + p01 + p10 + p11
    if sp.simplify(p00 * mass - (p00 + p01) * (p00 + p10)) != 0:
        fail("binary independence identity")
    u0, u1, v0, v1 = sp.symbols("u0 u1 v0 v1", positive=True)
    joint = sp.Matrix([[u0 * v0, u0 * v1], [u1 * v0, u1 * v1]])
    if sp.simplify(joint[0, 0] * joint[1, 1] - joint[0, 1] * joint[1, 0]) != 0:
        fail("product distributions leave the Wright manifold")


def linkage_and_replicator() -> None:
    def mat(name: str) -> sp.Matrix:
        return sp.Matrix(
            [[sp.symbols(f"{name}{i}{j}") for j in range(2)] for i in range(2)]
        )

    A, B, C, D = mat("A"), mat("B"), mat("C"), mat("D")
    p = sp.Matrix(sp.symbols("p00 p01 p10 p11"))
    pure = [(0, 0), (0, 1), (1, 0), (1, 1)]

    def payoff(a: int, b: int, c: int, d: int) -> sp.Expr:
        return A[a, c] + B[a, d] + C[b, c] + D[b, d]

    pi = [
        sp.expand(sum(p[k] * payoff(a, b, c, d) for k, (c, d) in enumerate(pure)))
        for a, b in pure
    ]
    if sp.simplify(pi[0] + pi[3] - pi[1] - pi[2]) != 0:
        fail("role-additive payoffs")
    mean = sp.expand(sum(p[i] * pi[i] for i in range(4)))
    field = [sp.expand(p[i] * (pi[i] - mean)) for i in range(4)]
    K = sp.symbols("K")
    defect = p[0] * p[3] - K * p[1] * p[2]
    derivative = sp.expand(sum(sp.diff(defect, p[i]) * field[i] for i in range(4)))
    expected = sp.expand(defect * (pi[0] + pi[3] - 2 * mean))
    if sp.simplify(derivative - expected) != 0:
        fail("replicator linkage derivative")

    P = sp.Matrix(sp.symbols("q1:4(1:4)", positive=True)).reshape(3, 3)

    def cross(i: int, j: int, k: int, l: int) -> sp.Expr:
        return (P[i, j] * P[k, l]) / (P[i, l] * P[k, j])

    if sp.simplify(cross(0, 1, 2, 0) - cross(2, 0, 0, 1)) != 0:
        fail("cross ratio symmetry")
    if sp.simplify(cross(0, 1, 2, 0) * cross(0, 0, 2, 1) - 1) != 0:
        fail("cross ratio inversion")


def revision_attraction() -> None:
    p00, p01, p10, b0, c0 = sp.symbols("p00 p01 p10 b0 c0")
    b1, c1 = 1 - b0, 1 - c0
    p11 = 1 - p00 - p01 - p10
    p1 = {0: p00 + p01, 1: p10 + p11}
    p2 = {0: p00 + p10, 1: p01 + p11}
    b, c = {0: b0, 1: b1}, {0: c0, 1: c1}
    pure = [(0, 0), (0, 1), (1, 0), (1, 1)]
    coords = [p00, p01, p10, p11]
    dots = []
    for (a, j), coord in zip(pure, coords):
        dots.append(
            sp.expand(
                sp.Rational(1, 2) * (p2[j] * b[a] - coord)
                + sp.Rational(1, 2) * (p1[a] * c[j] - coord)
            )
        )
    defect = sp.expand(p1[0] * p2[0] - p00)
    derivative = sp.expand(
        sum(sp.diff(defect, var) * dots[i] for i, var in enumerate((p00, p01, p10)))
    )
    if sp.simplify(derivative + defect) != 0:
        fail("revision defect decay")


def role_payoff(eps: sp.Expr) -> sp.Matrix:
    hawk = sp.Matrix([[-1, 2], [0, 1]])
    buyer = sp.Matrix([[5, 1], [4, 2]])
    seller = sp.Matrix([[4, 0], [6, -2]])
    within = eps * hawk
    pure = [(0, 0), (0, 1), (1, 0), (1, 1)]
    matrix = sp.zeros(4)
    for i, (a, b) in enumerate(pure):
        for j, (c, d) in enumerate(pure):
            matrix[i, j] = within[a, c] + buyer[a, d] + seller[b, c] + within[b, d]
    return sp.expand(matrix)


def reduced_replicator(matrix: sp.Matrix, weights: sp.Matrix):
    z = sp.Matrix(sp.symbols("z0 z1 z2"))
    state = sp.Matrix([z[0], z[1], z[2], 1 - z[0] - z[1] - z[2]])
    payoff = sp.expand(matrix * state)
    mean = sp.expand(sum(state[i] * payoff[i] for i in range(4)))
    field = [sp.expand(state[i] * (payoff[i] - mean)) for i in range(3)]
    jac = sp.Matrix(
        [[sp.diff(field[i], z[j]) for j in range(3)] for i in range(3)]
    )
    point = dict(zip(z, weights))
    return (
        state.subs(point),
        [sp.simplify(component.subs(point)) for component in field],
        sp.simplify(jac.subs(point)),
    )


def buyer_seller() -> None:
    eps, lam, t = sp.symbols("eps lam t")
    matrix = role_payoff(eps)
    weights = sp.Matrix(
        [
            sp.Rational(1, 4) + lam,
            sp.Rational(1, 4) - lam,
            sp.Rational(1, 4) - lam,
        ]
    )
    _state, field, jac = reduced_replicator(matrix, weights)
    for component in field:
        if sp.simplify(component) != 0:
            fail("buyer-seller fiber is not at rest")
    poly = sp.factor(sp.simplify(jac).charpoly(t).as_expr())
    quad = (
        t**2
        + (eps + 2 * lam) * t
        + (eps**2 * (1 - 16 * lam**2) - 32 * lam**2 + 2) / 4
    )
    if sp.expand(poly - t * quad) != 0:
        fail("buyer-seller characteristic polynomial")
    disc = sp.factor((eps + 2 * lam) ** 2 - 4 * (quad.subs(t, 0)))
    expected = 16 * eps**2 * lam**2 + 36 * lam**2 + 4 * eps * lam - 2
    if sp.expand(disc - expected) != 0:
        fail("buyer-seller discriminant")
    point = {eps: sp.Rational(1, 10), lam: sp.Rational(-1, 10)}
    if sp.simplify(expected.subs(point)) != -sp.Rational(1049, 625):
        fail("buyer-seller sample radicand")
    if sp.simplify((-eps / 2 - lam).subs(point)) != sp.Rational(1, 20):
        fail("buyer-seller sample real part")


def owner_intruder() -> None:
    eps = sp.symbols("eps", real=True)
    hawk = sp.Matrix([[-1, 2], [0, 1]])
    pure = [(0, 0), (0, 1), (1, 0), (1, 1)]

    def entry(fo: int, fi: int, oo: int, oi: int) -> sp.Expr:
        return (
            eps * hawk[fo, oo]
            + (sp.Rational(1, 2) - eps) * hawk[fo, oi]
            + (sp.Rational(1, 2) - eps) * hawk[fi, oo]
            + eps * hawk[fi, oi]
        )

    matrix = sp.zeros(4)
    for i, focal in enumerate(pure):
        for j, other in enumerate(pure):
            matrix[i, j] = sp.expand(entry(*focal, *other))
    t = sp.symbols("t")
    weights = sp.Matrix(
        [sp.Rational(1, 4) - t, sp.Rational(1, 4) + t, sp.Rational(1, 4) + t]
    )
    _state, field, jac = reduced_replicator(matrix, weights)
    for component in field:
        if sp.simplify(component) != 0:
            fail("owner-intruder fiber is not at rest")
    jac0 = sp.simplify(jac.subs(t, 0))
    expected = sp.Matrix(
        [
            [sp.Rational(-1, 4), sp.Rational(-1, 8), sp.Rational(-1, 8)],
            [0, sp.Rational(1, 8) - eps / 2, eps / 2 - sp.Rational(1, 8)],
            [0, eps / 2 - sp.Rational(1, 8), sp.Rational(1, 8) - eps / 2],
        ]
    )
    if sp.simplify(jac0 - expected) != sp.zeros(3):
        fail("owner-intruder Jacobian")
    s = sp.symbols("s")
    poly = sp.expand(jac0.charpoly(s).as_expr())
    expected_poly = sp.expand(
        (s + sp.Rational(1, 4)) * s * (s - (sp.Rational(1, 4) - eps))
    )
    if sp.expand(poly - expected_poly) != 0:
        fail("owner-intruder characteristic polynomial")
    column = sp.Matrix([sp.simplify(matrix[i, 1]) for i in range(4)])
    gaps = [sp.simplify(column[1] - column[k]) for k in (0, 2, 3)]
    if gaps != [sp.Rational(1, 2) - 2 * eps, 1 - 4 * eps, sp.Rational(1, 2) - 2 * eps]:
        fail("owner-intruder strict-equilibrium gaps")


def rock_scissors_paper() -> None:
    A = sp.Matrix([[0, 6, -4], [-4, 0, 4], [2, -2, 0]])
    p0 = sp.Matrix([sp.Rational(10, 29), sp.Rational(8, 29), sp.Rational(11, 29)])
    payoff = sp.simplify(A * p0)
    if payoff[0] != payoff[1] or payoff[1] != payoff[2] or payoff[0] != sp.Rational(4, 29):
        fail("RSP equilibrium")
    system = sp.Matrix([[2, 3, -4], [1, -4, 2], [1, 1, 1]])
    if system.det() != -29:
        fail("RSP equilibrium is not a unique linear solution")
    p1, p2 = sp.symbols("p1 p2")
    p = sp.Matrix([p1, p2, 1 - p1 - p2])
    Ap = A * p
    mean = sp.simplify(p.dot(Ap))
    field = [sp.simplify(p[i] * (Ap[i] - mean)) for i in range(2)]
    jac = sp.Matrix(
        [
            [sp.diff(field[0], p1), sp.diff(field[0], p2)],
            [sp.diff(field[1], p1), sp.diff(field[1], p2)],
        ]
    ).subs({p1: p0[0], p2: p0[1]})
    values = list(sp.simplify(jac.eigenvals()).keys())
    root = 2 * sp.sqrt(879) * sp.I / 29
    expected = {-sp.Rational(2, 29) - root, -sp.Rational(2, 29) + root}
    if set(values) != expected:
        fail("RSP eigenvalues")


def sunspot() -> None:
    A = sp.Matrix([[0, 6, -4], [-4, 0, 4], [2, -2, 0]])
    p0 = [sp.Rational(10, 29), sp.Rational(8, 29), sp.Rational(11, 29)]
    xvec, yvec = [-4, 1, 3], [-2, -1, 3]
    index = [(a, b) for a in range(3) for b in range(3)]
    payoff = [
        [
            sp.Rational(1, 2) * (A[a, c] + A[b, d])
            for c, d in index
        ]
        for a, b in index
    ]
    r = sp.symbols("r")
    state = [
        sp.simplify(p0[a] * p0[b] + r * xvec[a] * yvec[b]) for a, b in index
    ]
    if sp.simplify(sum(state) - 1) != 0:
        fail("sunspot mass")
    marginal = [sp.simplify(sum(state[i] for i, (a, _b) in enumerate(index) if a == k)) for k in range(3)]
    if marginal != p0:
        fail("sunspot marginals")

    mp.mp.dps = 40
    pay = [[mp.mpf(entry) for entry in row] for row in payoff]

    def oscillator(value: sp.Rational):
        point = [
            mp.mpf(p0[a] * p0[b]) + mp.mpf(value) * xvec[a] * yvec[b]
            for a, b in index
        ]
        if min(point) <= 0:
            fail("sunspot state leaves the simplex")

        def reduced(coords):
            full = coords + [1 - sum(coords)]
            gained = [
                sum(pay[i][j] * full[j] for j in range(9)) for i in range(9)
            ]
            mean = sum(full[i] * gained[i] for i in range(9))
            return [full[i] * (gained[i] - mean) for i in range(8)]

        base = point[:8]
        step = mp.mpf("1e-18")
        jac = mp.matrix(8)
        for col in range(8):
            plus, minus = base.copy(), base.copy()
            plus[col] += step
            minus[col] -= step
            fp, fm = reduced(plus), reduced(minus)
            for row in range(8):
                jac[row, col] = (fp[row] - fm[row]) / (2 * step)
        values = mp.eig(jac, left=False, right=False)
        return max(values, key=lambda item: abs(mp.im(item)))

    at_zero = oscillator(sp.Integer(0))
    exact_real = mp.mpf(-1) / 29
    exact_imag = mp.sqrt(879) / 29
    if abs(mp.re(at_zero) - exact_real) > mp.mpf("1e-9"):
        fail("sunspot real part at linkage equilibrium")
    if abs(abs(mp.im(at_zero)) - exact_imag) > mp.mpf("1e-9"):
        fail("sunspot imaginary part at linkage equilibrium")
    low = oscillator(sp.Rational(9, 1000))
    high = oscillator(sp.Rational(1, 100))
    if mp.re(low) >= 0 or mp.re(high) <= 0:
        fail("sunspot real-part sign change")


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