"""Pair model of GL_n/(GL_{n-p} x GL_p) as rank-(n-p) idempotents P (P = projection onto U along V).

Exact arithmetic in F = Q (prime ell, uniformizer ell, e = val(2) = 1 if ell == 2 else 0)
or F = Q(sqrt2) at the prime above 2 (uniformizer sqrt2, e = val(2) = 2).

Checks, for random data:
 (1) label(t^lam x0) = phi(lam) = (lam_i - lam_{n+1-i})_{i<=p}   for dominant lam;
 (2) label(k t^lam h x0) <=_dom phi(lam)   (partial sums, k,h in K random);
 (3) Cartan type of eps_x = 2P - 1 equals mu((label - e)^+)   for random x = g x0, g in GL_n(F);
 (4) label of the involution base point (1+J)/2 is (e,...,e);
 (5) n = 2: label of t^(a,0) x0 is a and Cartan label of tau(t^(a,0) x0) is max(a - e, 0).
label(P) = polar part of the Smith exponents of P, computed from determinantal divisors (minors).
Usage: python3 -I pair_tau.py
"""
from fractions import Fraction as Fr
import itertools, random

INF = 10 ** 9


class Q2:
    """a + b*sqrt2 with rational a, b."""
    __slots__ = ('a', 'b')

    def __init__(self, a, b=0):
        self.a = Fr(a); self.b = Fr(b)

    def __add__(s, o):
        o = o if isinstance(o, Q2) else Q2(o); return Q2(s.a + o.a, s.b + o.b)

    __radd__ = __add__

    def __sub__(s, o):
        o = o if isinstance(o, Q2) else Q2(o); return Q2(s.a - o.a, s.b - o.b)

    def __rsub__(s, o):
        return Q2(o) - s

    def __neg__(s):
        return Q2(-s.a, -s.b)

    def __mul__(s, o):
        o = o if isinstance(o, Q2) else Q2(o)
        return Q2(s.a * o.a + 2 * s.b * o.b, s.a * o.b + s.b * o.a)

    __rmul__ = __mul__

    def inv(s):
        n = s.a * s.a - 2 * s.b * s.b
        return Q2(s.a / n, -s.b / n)

    def __truediv__(s, o):
        o = o if isinstance(o, Q2) else Q2(o); return s * o.inv()

    def iszero(s):
        return s.a == 0 and s.b == 0


def vfr(x, ell):
    if x == 0:
        return INF
    n, d, v = x.numerator, x.denominator, 0
    while n % ell == 0:
        n //= ell; v += 1
    while d % ell == 0:
        d //= ell; v -= 1
    return v


class Field:
    def __init__(self, kind, ell=2):
        self.kind, self.ell = kind, ell
        if kind == 'Q':
            self.unif = Q2(ell); self.e = 1 if ell == 2 else 0
        else:
            self.unif = Q2(0, 1); self.e = 2

    def val(self, x):
        if x.iszero():
            return INF
        if self.kind == 'Q':
            return vfr(x.a, self.ell)
        return min(2 * vfr(x.a, 2) if x.a != 0 else INF, 2 * vfr(x.b, 2) + 1 if x.b != 0 else INF)

    def upow(self, k):
        r = Q2(1)
        base = self.unif if k >= 0 else self.unif.inv()
        for _ in range(abs(k)):
            r = r * base
        return r

    def rand_int(self, R=3):
        if self.kind == 'Q':
            return Q2(random.randint(-R, R))
        return Q2(random.randint(-R, R), random.randint(-R, R))


def mat_mul(A, B):
    n, m, l = len(A), len(B), len(B[0])
    return [[sum((A[i][k] * B[k][j] for k in range(m)), Q2(0)) for j in range(l)] for i in range(n)]


def det(M):
    M = [row[:] for row in M]
    n = len(M); d = Q2(1)
    for c in range(n):
        piv = next((r for r in range(c, n) if not M[r][c].iszero()), None)
        if piv is None:
            return Q2(0)
        if piv != c:
            M[c], M[piv] = M[piv], M[c]; d = -d
        d = d * M[c][c]
        inv = M[c][c].inv()
        for r in range(c + 1, n):
            if not M[r][c].iszero():
                f = M[r][c] * inv
                M[r] = [M[r][j] - f * M[c][j] for j in range(n)]
    return d


def inverse(M):
    n = len(M)
    A = [M[i][:] + [Q2(1 if i == j else 0) for j in range(n)] for i in range(n)]
    for c in range(n):
        piv = next(r for r in range(c, n) if not A[r][c].iszero())
        A[c], A[piv] = A[piv], A[c]
        inv = A[c][c].inv()
        A[c] = [x * inv for x in A[c]]
        for r in range(n):
            if r != c and not A[r][c].iszero():
                f = A[r][c]
                A[r] = [A[r][j] - f * A[c][j] for j in range(2 * n)]
    return [row[n:] for row in A]


def smith_exponents(M, F):
    """ascending Smith exponents via determinantal divisors; INF for the kernel part."""
    n = len(M)
    deltas = [0]
    for j in range(1, n + 1):
        best = INF
        for I in itertools.combinations(range(n), j):
            for J in itertools.combinations(range(n), j):
                v = F.val(det([[M[i][k] for k in J] for i in I]))
                best = min(best, v)
        deltas.append(best)
    ex = []
    for j in range(1, n + 1):
        ex.append(INF if deltas[j] >= INF else deltas[j] - deltas[j - 1])
    return ex


def label(P, F, p):
    ex = smith_exponents(P, F)
    lab = sorted((-c for c in ex if c < 0), reverse=True)
    assert len(lab) <= p
    return tuple(lab + [0] * (p - len(lab)))


def cartan(g, F):
    ex = smith_exponents(g, F)
    assert all(c < INF for c in ex)
    return tuple(sorted(ex, reverse=True))


def mu(b, n):
    p = len(b)
    return tuple(list(b) + [0] * (n - 2 * p) + [-x for x in reversed(b)])


def dom_le(b, a):
    sa = sb = 0
    for x, y in zip(a, b):
        sa += x; sb += y
        if sb > sa:
            return False
    return True


def P0(n, p):
    P = [[Q2(0)] * n for _ in range(n)]
    for i in range(p):
        P[i][i] = Q2(1); P[n - 1 - i][i] = Q2(1)
    for j in range(p, n - p):
        P[j][j] = Q2(1)
    return P


def PJ(n, p):
    P = [[Q2(0)] * n for _ in range(n)]
    for i in range(p):
        for a in (i, n - 1 - i):
            for b in (i, n - 1 - i):
                P[a][b] = Q2(Fr(1, 2))
    for j in range(p, n - p):
        P[j][j] = Q2(1)
    return P


def rand_K(n, F):
    while True:
        k = [[F.rand_int() for _ in range(n)] for _ in range(n)]
        if F.val(det(k)) == 0:
            return k


def rand_G(n, F):
    while True:
        g = [[F.rand_int() * F.upow(random.randint(-2, 2)) for _ in range(n)] for _ in range(n)]
        if not det(g).iszero():
            return g


def tl(lam, F):
    n = len(lam)
    return [[F.upow(lam[i]) if i == j else Q2(0) for j in range(n)] for i in range(n)]


def conj(g, P):
    return mat_mul(mat_mul(g, P), inverse(g))


def rand_dom(n, lo=-2, hi=3):
    return tuple(sorted((random.randint(lo, hi) for _ in range(n)), reverse=True))


def main():
    random.seed(20261017)
    fields = [Field('Q', 2), Field('Q', 3), Field('Q2')]
    names = ['Q_2 (e=1)', 'Q_3 (e=0)', 'Q_2(sqrt2) (e=2)']
    ok = True
    for F, name in zip(fields, names):
        cnt = [0, 0, 0, 0]
        stats = {'support: label = phi': 0, 'support: label < phi': 0, 'tau: some 0<l_i<=e': 0}
        for (n, p) in ((2, 1), (3, 1), (4, 1), (4, 2), (5, 1), (5, 2), (6, 3)):
            Pb = P0(n, p)
            # (4) involution base point
            lj = label(PJ(n, p), F, p)
            if lj != tuple([F.e] * p):
                ok = False; print('J-label FAIL', name, n, p, lj)
            trials = {2: 60, 3: 40, 4: 15, 5: 6, 6: 3}[n]
            for _ in range(trials):
                lam = rand_dom(n)
                phi = tuple(lam[i] - lam[n - 1 - i] for i in range(p))
                # (1)
                l1 = label(conj(tl(lam, F), Pb), F, p)
                if l1 != phi:
                    ok = False; print('top FAIL', name, n, p, lam, l1, phi)
                cnt[0] += 1
                # (2)
                g = mat_mul(mat_mul(rand_K(n, F), tl(lam, F)), rand_K(n, F))
                P = conj(g, Pb)
                l2 = label(P, F, p)
                if not dom_le(l2, phi):
                    ok = False; print('support FAIL', name, n, p, lam, l2, phi)
                cnt[1] += 1
                stats['support: label = phi' if l2 == phi else 'support: label < phi'] += 1
                # (3)
                P = conj(rand_G(n, F), Pb)
                l3 = label(P, F, p)
                eps = [[2 * P[i][j] - (1 if i == j else 0) for j in range(n)] for i in range(n)]
                c = cartan(eps, F)
                expect = mu(tuple(max(x - F.e, 0) for x in l3), n)
                if c != expect:
                    ok = False; print('tau FAIL', name, n, p, l3, c, expect)
                cnt[2] += 1
                if any(0 < x <= F.e for x in l3):
                    stats['tau: some 0<l_i<=e'] += 1
        # (5) n = 2 explicit
        for a in range(0, 5):
            P = conj(tl((a, 0), F), P0(2, 1))
            la = label(P, F, 1)
            eps0 = [[2 * P0(2, 1)[i][j] - (1 if i == j else 0) for j in range(2)] for i in range(2)]
            eps = [[2 * P[i][j] - (1 if i == j else 0) for j in range(2)] for i in range(2)]
            tau = mat_mul(eps, inverse(eps0))
            ct = cartan(tau, F)
            if la != (a,) or ct != (max(a - F.e, 0), -max(a - F.e, 0)):
                ok = False; print('n=2 FAIL', name, a, la, ct)
            cnt[3] += 1
        print(name, ': checks run (top, support, tau-formula, n=2):', cnt, stats)
    print('ALL CHECKS PASSED' if ok else 'SOME CHECK FAILED')


if __name__ == '__main__':
    main()