"""Pair model X = {(U,V): F^n = U (+) V, dim U = q, dim V = p} over O = F_ell[[t]].

Base point x_0 = (U_0, V_0) split for L_0 = O^n (U_0 = first q coordinates, V_0 = last p;
any O-split base point is in the same K-orbit).  For g in K t^lam K (lam dominant, >= 0)
the point g x_0 has label  b = type( pi_V(L')/(L' cap V_0) ),  L' = g^{-1} L_0, and L' runs
over all lattices containing L_0 with L'/L_0 of type lam.  So supp(T_lam * phi_0) is the set
of labels b of submodules M of E = t^{-k}L_0/L_0 = (O/t^k)^n of type lam, with
b = type(pi_V(M) / (M cap E_V)).

Checks:
  (A) b <=_dom phi(lam) = (lam_i - lam_{n+1-i})_{i<=p} for every dominant lam >= 0 tested
      (function-level H5(a) for the pair model, here in residue characteristic ell);
  (B) for lam = (a_1..a_p,0..0) also b <= a componentwise (node A5-4 (iii)), and the
      label a occurs (A5-4 (ii));
  (C) the label equals the negative part of the Smith exponents of the projector pi_x
      in a basis of L' (checked by an independent Laurent-polynomial computation).

Usage: python3 -I pairmodel.py ell p q maxpart
"""
import sys, os, itertools
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fplin import rref


def setup(n, k):
    idx = [(i, j) for i in range(n) for j in range(-k, 0)]
    pos = {x: a for a, x in enumerate(idx)}
    tmap = [pos.get((i, j + 1)) for (i, j) in idx]
    return idx, pos, tmap


def tvec(v, tmap, ell):
    w = [0] * len(v)
    for a, x in enumerate(v):
        if x and tmap[a] is not None:
            w[tmap[a]] = (w[tmap[a]] + x) % ell
    return tuple(w)


def tstable_span(gens, tmap, ell):
    out = []
    for v in gens:
        while any(v):
            out.append(v)
            v = tvec(v, tmap, ell)
    return rref(out, ell)


def mtype(W, tmap, ell):
    dims = []
    cur = list(W)
    while cur:
        R = rref(cur, ell)
        dims.append(len(R))
        cur = [w for w in (tvec(v, tmap, ell) for v in R) if any(w)]
    dims.append(0)
    ge = [dims[s] - dims[s + 1] for s in range(len(dims) - 1)]
    lam = []
    for size in range(len(ge), 0, -1):
        lam += [size] * (ge[size - 1] - (ge[size] if size < len(ge) else 0))
    return tuple(lam)


def submodules_of_type(n, k, lam, ell):
    """All t-stable subspaces of (O/t^k)^n of module type lam (BFS with pruning)."""
    idx, pos, tmap = setup(n, k)
    N = len(idx)
    lamp = list(lam) + [0] * n
    target_dim = sum(lam)
    basis = [tuple(1 if b == a else 0 for b in range(N)) for a in range(N)]
    levels = {0: {tuple()}}
    found = set()
    for d in range(0, target_dim):
        for W in levels.get(d, ()):
            for coeffs in itertools.product(range(ell), repeat=N):
                if not any(coeffs) or next(c for c in coeffs if c) != 1:
                    continue
                v = coeffs
                W2 = tstable_span(list(W) + [v], tmap, ell)
                if len(W2) == len(W) or len(W2) > target_dim:
                    continue
                ty = mtype(W2, tmap, ell)
                typ = list(ty) + [0] * n
                if any(typ[i] > lamp[i] for i in range(n)):
                    continue
                if len(W2) == target_dim:
                    if ty == tuple(x for x in lam if x):
                        found.add(W2)
                else:
                    levels.setdefault(len(W2), set()).add(W2)
    return found, idx, tmap


def label(M, idx, tmap, ell, q):
    """type of pi_V(M)/(M cap E_V); E_V = coordinates i >= q."""
    N = len(idx)
    Vcoords = [a for a, (i, j) in enumerate(idx) if i >= q]
    Ucoords = [a for a, (i, j) in enumerate(idx) if i < q]
    MV = rref([tuple(v[a] if a in set(Vcoords) else 0 for a in range(N)) for v in M], ell)
    # M cap E_V: vectors of M with zero U-part: solve via span of M plus kernel
    # build matrix of M's basis restricted to U coords and find combinations giving 0
    rows = [[v[a] for a in Ucoords] for v in M]
    # left kernel of rows: combinations c with sum c_r rows_r = 0
    m = len(M)
    if m == 0:
        return tuple()
    # transpose solve
    from fplin import nullspace
    T = [[rows[r][cidx] for r in range(m)] for cidx in range(len(Ucoords))]
    ker = nullspace(T, m, ell)
    NV = rref([tuple(sum(c[r] * M[r][a] for r in range(m)) % ell for a in range(N)) for c in ker], ell) if ker else tuple()
    dNV = len(NV)
    dims = []
    cur = list(MV)
    s = 0
    while True:
        Wd = len(rref(list(cur) + list(NV), ell)) - dNV
        dims.append(Wd)
        if Wd == 0:
            break
        cur = [w for w in (tvec(v, tmap, ell) for v in rref(cur, ell)) if any(w)]
    ge = [dims[s] - dims[s + 1] for s in range(len(dims) - 1)]
    lab = []
    for size in range(len(ge), 0, -1):
        lab += [size] * (ge[size - 1] - (ge[size] if size < len(ge) else 0))
    return tuple(lab)


def dom_le(b, a):
    """b <=_dom a as labels of Lambda_S^+ (p-tuples, nonneg): partial sums."""
    n = max(len(a), len(b))
    a = list(a) + [0] * (n - len(a)); b = list(b) + [0] * (n - len(b))
    sa = sb = 0
    for x, y in zip(a, b):
        sa += x; sb += y
        if sb > sa:
            return False
    return True


def dominant_lams(n, maxpart):
    for lam in itertools.product(range(maxpart, -1, -1), repeat=n):
        if all(lam[i] >= lam[i + 1] for i in range(n - 1)) and lam[-1] == 0 and sum(lam) > 0:
            yield lam


if __name__ == '__main__':
    ell = int(sys.argv[1]); p = int(sys.argv[2]); q = int(sys.argv[3]); maxpart = int(sys.argv[4])
    n = p + q
    ok = True
    for lam in dominant_lams(n, maxpart):
        k = lam[0]
        found, idx, tmap = submodules_of_type(n, k, lam, ell)
        phi = tuple(lam[i] - lam[n - 1 - i] for i in range(p))
        labs = {}
        for M in found:
            b = label(M, idx, tmap, ell, q)
            b = tuple(list(b) + [0] * (p - len(b)))
            labs[b] = labs.get(b, 0) + 1
            if not dom_le(b, phi):
                ok = False
                print('H5(a) VIOLATION', lam, b, phi)
        is_lift = all(x == 0 for x in lam[p:])
        if is_lift:
            a = lam[:p]
            for b in labs:
                if any(bb > aa for bb, aa in zip(b, a)):
                    ok = False
                    print('componentwise VIOLATION', lam, b)
            if a not in labs:
                ok = False
                print('top label missing', lam)
        print('lam', lam, 'phi', phi, 'lift' if is_lift else '', 'labels:', dict(sorted(labs.items(), reverse=True)))
    print('ALL PAIR-MODEL CHECKS PASSED' if ok else 'SOME CHECK FAILED')