#!/usr/bin/env python3
"""Checks for Section 16 (Hecke symmetry and transverse test functions).

 (1) Lemma "The two cotangent projections", step 3: the pairing
       (v_+, v_-) -> Res tr( [a(t), v_+] v_- ) dt
     between g_in/i and g_out/i is nondegenerate when ad a(0) is invertible
     on g/m_j, and degenerate when it is not.   (gl_n, lambda = m*varpi_j)
 (2) Lemma "A transverse homogeneous test function", step 1: for a graded
     symplectic space with form of weight two and two graded Lagrangians,
     the decomposition exists and, after the enlargement, Lambda is a
     graded Lagrangian complement of both enlarged subspaces.
 (3) Lemma "The microstalk is the Whittaker coefficient", step 4: a
     Laurent-polynomial truncation has an injective matrix with
     one-dimensional cokernel for each operator used in the proof.
 (4) Lemma "Quadratic stabilization": the universal conic rs + t w^2 = 0
     is smooth and is the blowup of P^1 x A^1 at one point (chart formulas).
"""
import random, itertools, sys
from sympy import Matrix, Rational, zeros, eye, symbols, expand, diff

random.seed(19)

# ---------------------------------------------------------------- (1)
def hessian_rank(n, j, m, a_coeffs):
    """a_coeffs[k] = block-diagonal n x n matrix, coefficient of t^k (k<m)."""
    up = [(p, q) for p in range(j) for q in range(j, n)]      # weight +m
    lo = [(q, p) for (p, q) in up]                            # weight -m
    plus = [(pq, e) for pq in up for e in range(m)]           # E_pq t^e, 0<=e<m
    minus = [(qp, e) for qp in lo for e in range(-m, 0)]      # E_qp t^e, -m<=e<0
    def E(p, q):
        M = zeros(n, n); M[p, q] = 1; return M
    G = zeros(len(plus), len(minus))
    for i, ((p, q), e) in enumerate(plus):
        for k, ((q2, p2), e2) in enumerate(minus):
            # residue: coefficient of t^{-1} in a(t) t^e t^e2  ->  a_c with c+e+e2=-1
            c = -1 - e - e2
            if 0 <= c < m:
                A = a_coeffs[c]
                br = A * E(p, q) - E(p, q) * A
                G[i, k] = (br * E(q2, p2)).trace()
    return G.rank(), len(plus)

def rand_block(n, j, spec1, spec2):
    """block diagonal with prescribed eigenvalues (upper triangular + conj)."""
    def blk(spec):
        k = len(spec)
        T = zeros(k, k)
        for i in range(k):
            T[i, i] = spec[i]
            for l in range(i + 1, k):
                T[i, l] = random.randint(-2, 2)
        while True:
            P = Matrix(k, k, lambda *_: random.randint(-2, 2))
            if P.det() != 0: break
        return P * T * P.inv()
    A = zeros(n, n)
    A[:j, :j] = blk(spec1); A[j:, j:] = blk(spec2)
    return A

def rand_m(n, j):
    A = zeros(n, n)
    A[:j, :j] = Matrix(j, j, lambda *_: random.randint(-3, 3))
    A[j:, j:] = Matrix(n - j, n - j, lambda *_: random.randint(-3, 3))
    return A

ok = True
for (n, j, m) in [(2, 1, 2), (3, 1, 3), (3, 2, 2), (4, 2, 2), (4, 1, 3)]:
    for trial in range(3):
        a0 = rand_block(n, j, [1, 3, 5][:j], [2, 4, 6, 8][:n - j])   # disjoint spectra
        coeffs = [a0] + [rand_m(n, j) for _ in range(m - 1)]
        r, d = hessian_rank(n, j, m, coeffs)
        good = (r == d == 2 * 0 + j * (n - j) * m)
        ok &= good
    # negative control: a common eigenvalue
    a0 = rand_block(n, j, [1, 3, 5][:j], [1, 4, 6, 8][:n - j])
    coeffs = [a0] + [rand_m(n, j) for _ in range(m - 1)]
    r, d = hessian_rank(n, j, m, coeffs)
    ok &= (r < d)
    print("(1) gl_%d, j=%d, m=%d: dim a=%d, pairing nondegenerate; degenerate control rank %d<%d"
          % (n, j, m, d, r, d))
assert ok
print("(1) OK")

# ---------------------------------------------------------------- (2)
def span_basis(vecs, dim):
    """independent subset of column vectors (list of Matrix dim x 1)."""
    out = []; M = zeros(dim, 0)
    for v in vecs:
        M2 = M.row_join(v)
        if M2.rank() > M.cols:
            M = M2; out.append(v)
    return out

def nullspace_of(rows, dim):
    if not rows: return [eye(dim)[:, i] for i in range(dim)]
    return Matrix.vstack(*rows).nullspace()

def intersect(A, B, dim):
    """basis of span(A) cap span(B)."""
    if not A or not B: return []
    MA = Matrix.hstack(*A); MB = Matrix.hstack(*B)
    ns = MA.row_join(-MB).nullspace()
    return span_basis([MA * v[:len(A), 0] for v in ns], dim)

def graded_test(wts):
    N = len(wts); D = 2 * N
    # basis: e_i (weight wts[i]), f_i (weight 2-wts[i]); omega(e_i,f_j)=delta
    J = zeros(D, D)
    for i in range(N):
        J[i, N + i] = 1; J[N + i, i] = -1
    om = lambda u, v: (u.T * J * v)[0, 0]
    wt = [wts[i] for i in range(N)] + [2 - wts[i] for i in range(N)]
    allw = sorted(set(wt))
    E_w = {a: [eye(D)[:, k] for k in range(D) if wt[k] == a] for a in allw}
    # random graded Lagrangian L
    L = {}
    for a in allw:
        if a < 1:
            basis = E_w[a]; r = random.randint(0, len(basis))
            vs = []
            for _ in range(r):
                v = zeros(D, 1)
                for b in basis: v += random.randint(-2, 2) * b
                vs.append(v)
            L[a] = span_basis(vs, D)
    for a in allw:
        if a > 1:
            La = L.get(2 - a, [])
            # annihilator of L_{2-a} inside E_a
            basis = E_w[a]
            if not basis: L[a] = []; continue
            B = Matrix.hstack(*basis)
            rows = [(u.T * J * B) for u in La]
            L[a] = [B * c for c in nullspace_of(rows, len(basis))]
    if 1 in allw:
        idx = [i for i in range(N) if wts[i] == 1]
        A = [i for i in idx if random.random() < 0.5]
        rest = [i for i in idx if i not in A]
        S = {(i, k): random.randint(-2, 2) for i in rest for k in rest if i <= k}
        vs = [eye(D)[:, N + i] for i in A]
        for i in rest:
            v = eye(D)[:, i]
            for k in rest:
                v += S[(min(i, k), max(i, k))] * eye(D)[:, N + k]
            vs.append(v)
        L[1] = vs
    Lall = [v for a in allw for v in L.get(a, [])]
    assert len(Lall) == N and all(om(u, v) == 0 for u in Lall for v in Lall), "L not Lagrangian"
    V = {a: [eye(D)[:, N + i] for i in range(N) if 2 - wts[i] == a] for a in allw}
    # I, U+, U-  (weight by weight)
    I = {}; Up = {}; Um = {}
    for a in allw:
        I[a] = intersect(L.get(a, []), V.get(a, []), D)
        Up[a] = span_basis(I[a] + L.get(a, []), D)[len(I[a]):]
        Um[a] = span_basis(I[a] + V.get(a, []), D)[len(I[a]):]
    Upl = [(a, v) for a in allw for v in Up[a]]
    Uml = [(a, v) for a in allw for v in Um[a]]
    Il = [(a, v) for a in allw for v in I[a]]
    assert len(Upl) == len(Uml)
    if Upl:
        G = Matrix(len(Upl), len(Uml), lambda i, k: om(Upl[i][1], Uml[k][1]))
        assert G.det() != 0, "U+ x U- not perfect"
    # orthogonal of U+ + U-, graded complement of I in it, made isotropic
    Jc = []
    for a in allw:
        basis = E_w[a]
        if not basis: continue
        B = Matrix.hstack(*basis)
        rows = [(u.T * J * B) for (_, u) in Upl + Uml]
        perp = [B * c for c in nullspace_of(rows, len(basis))]
        comp = span_basis(I[a] + perp, D)[len(I[a]):]
        Jc += [(a, v) for v in comp]
    assert len(Jc) == len(Il)
    if Jc:
        P = Matrix(len(Il), len(Jc), lambda i, k: om(Il[i][1], Jc[k][1]))
        assert P.det() != 0
        # phi(x) in I with om(phi(x), y) = beta(x,y)/2 for y in J
        newJ = []
        for (a, x) in Jc:
            rhs = Matrix([Rational(1, 2) * om(x, y) for (_, y) in Jc])
            coef = P.T.solve(rhs)
            phi = zeros(D, 1)
            for c, (_, i) in zip(coef, Il): phi += c * i
            # phi must be homogeneous of the same weight as x
            assert all(phi[k] == 0 for k in range(D) if wt[k] != a)
            newJ.append((a, x - phi))
        Jc = newJ
        assert all(om(u, v) == 0 for (_, u) in Jc for (_, v) in Jc)
    # enlargement: new symplectic space U-new + U+new with form = -old
    m = len(Upl); D2 = D + 2 * m
    J2 = zeros(D2, D2); J2[:D, :D] = J
    for i in range(m):
        for k in range(m):
            val = om(Upl[i][1], Uml[k][1])
            J2[D + i, D + m + k] = -val       # (u+new_i , u-new_k)
            J2[D + m + k, D + i] = val
    om2 = lambda u, v: (u.T * J2 * v)[0, 0]
    ext = lambda v: v.col_join(zeros(2 * m, 1))
    unit = lambda k: eye(D2)[:, k]
    Lam = [ext(v) for (_, v) in Jc] \
        + [ext(Upl[i][1]) + unit(D + i) for i in range(m)] \
        + [ext(Uml[k][1]) + unit(D + m + k) for k in range(m)]
    assert len(Lam) == D2 // 2
    assert Matrix.hstack(*Lam).rank() == D2 // 2
    assert all(om2(u, v) == 0 for u in Lam for v in Lam), "Lambda not isotropic"
    nilp = [ext(v) for v in Lall] + [unit(D + m + k) for k in range(m)]
    vert = [ext(eye(D)[:, N + i]) for i in range(N)] + [unit(D + i) for i in range(m)]
    assert Matrix.hstack(*(Lam + nilp)).rank() == D2, "Lambda meets nilp"
    assert Matrix.hstack(*(Lam + vert)).rank() == D2, "Lambda meets vert"
    return len(Il), m

cnt = 0
for wts in [[0], [2], [1], [0, 2], [1, 1], [0, 1, 2], [-1, 3, 1], [0, 0, 2, 2],
            [2, 2, 1, 0], [-2, 4, 0, 2, 1], [3, -1, 1, 1, 0]]:
    for trial in range(6):
        dI, m = graded_test(wts); cnt += 1
print("(2) %d random graded Lagrangian pairs: decomposition exists, Lambda is a"
      " graded Lagrangian complement of both" % cnt)
# the obstruction of the Remark: weights 0 and 2 in a plane
print("(2) OK")

# ---------------------------------------------------------------- (3)
def laurent_check(shift, Nmax=12):
    """operator D(x^n) = (n - shift) x^{n-1} - x^n on span{x^n, -N<=n<=N}."""
    # functional c with c_n = (n - shift) c_{n-1}
    c = {}
    if shift == Rational(1, 2):
        c[0] = Rational(1)
        for n in range(1, Nmax + 1): c[n] = (n - shift) * c[n - 1]
        for n in range(0, -Nmax - 1, -1): c[n - 1] = c[n] / (n - shift)
    else:   # shift = 0: c_n = n c_{n-1}; c_n = 0 for n>=0, c_{-1} = 1
        for n in range(0, Nmax + 1): c[n] = Rational(0)
        c[-1] = Rational(1)
        for n in range(-1, -Nmax - 1, -1): c[n - 1] = c[n] / n
    # c kills the image
    for n in range(-Nmax + 1, Nmax + 1):
        assert (n - shift) * c[n - 1] - c[n] == 0
    # truncated matrix: columns x^n (-N+1..N), rows x^k (-N..N): rank = number of columns
    cols = list(range(-Nmax + 1, Nmax + 1)); rows = list(range(-Nmax, Nmax + 1))
    M = zeros(len(rows), len(cols))
    for jn, n in enumerate(cols):
        M[rows.index(n), jn] = -1
        M[rows.index(n - 1), jn] = n - shift
    assert M.rank() == len(cols)          # injective
    assert len(rows) - M.rank() == 1      # cokernel of the truncation is a line
    return True
laurent_check(Rational(1, 2)); laurent_check(Rational(0))
print("(3) Laurent-polynomial truncations: injective matrices,"
      " cokernels of dimension one   OK")

# ---------------------------------------------------------------- (4)
r, s, w, t = symbols('r s w t')
Gc = r * s + t * w**2
# in each of the three standard charts one partial derivative is identically 1,
# so the total space is smooth over Z
for chart, sub, var in [("w=1", {w: 1}, t), ("r=1", {r: 1}, s), ("s=1", {s: 1}, r)]:
    assert diff(Gc.subs(sub), var) == 1, chart
# chart w=1: t = -r s, and the ideal (r/w, t) = (r) is invertible
assert expand((Gc.subs({w: 1})).subs(t, -r * s)) == 0
# chart s=1: r = -t w^2, r/w = -t w, the ideal (r/w, t) = (t) is invertible
assert expand((Gc.subs({s: 1})).subs(r, -t * w**2)) == 0
print("(4) universal conic: total space smooth over Z in all three charts;"
      " blowup chart formulas   OK")
print("ALL OK")