"""
tame_q.py -- the coefficient triangle and the tame readout over an
unramified base of degree f > 1: Teichmuller digits, and a kernel that
splits each class.

QUESTION. Both statements were read at f = 1. Over W = W(F_q), q =
p^f, K totally ramified of degree e over W, F = x^e + sum p b_i x^i -
p d: (a) is the string of u = p/pi^e, in Teichmuller digits, still a
triangular bijection of the coefficient digits with diagonal +/-
ubar^(k+2); (b) where does a unit of the class (c, m) land at an odd p
with s = e/(p - 1) an integer?

DERIVATION (by hand, frozen before code). (a) The f = 1 fixed point
uses no f: a Teichmuller digit c at level r = i + ke enters D = d - sum
b_i pi^i as -/+ omega(c) u^k pi^r, so u_r = s_r c_r + g_r over F_q, s_r
= +/- ubar^(k+2), - at the d levels. (b) As at f = 1, y = z^(p^m) =
1 + pi^s tau, tau = t^(p^m)(1 + eta), v(eta) > p^m, and y^p - 1 =
pi^(ps) tau B, B = tau^(p-1) + u + u sum_(2<=k<=p-1) C(p,k)/p (pi^s
tau)^(k-1). In F_q, a^(p-1) = 1 fails (TRANSPLANT): with b = a^(p^m),
a the residue of t, B's residue is b^(p-1) + ubar, the graded p-th
power at the bend (Fesenko-Vostokov's theta_0 = ubar). (T1) b^(p-1)
!= -ubar: landing p s. (T2) the kernel b^(p-1) = -ubar, nonempty iff
N(-ubar) = 1 in F_p, then p - 1 residues b: B = (u - omega(ubar)) +
omega(ubar)(t_1/a)^(p^m) pi^(p^m) + terms free of t_1 at level p^m
(the k = 2 term when c = 1) or deeper, so with the DEPARTURE l =
v(u - omega(ubar)), the least level of a nonzero Teichmuller digit of
(F - x^e + p omega(dbar))/p, kernel units land at p s + min(l, p^m),
all of them when l < p^m, all but one t_1 in F_q per residue a when
l >= p^m. At f =
1 and ubar = -1, omega(-1) = -1 and l is v(u + 1); at f = 1 and
ubar != -1, v(u + 1) is 0 and the kernel is empty.

DESIGN. O_K = Z_q[pi]/(F) mod p^N, Z_q = Z_p[y]/(h), h a monic lift of
an irreducible of degree f mod p; omega(x) = x^(q^N). (a) p = 3, f = 2,
e = 2, every cell to level 3 (5832); p = 2, f = 2, e = 3, every cell to
level 4 (768); the diagonal at 60 random cells per grid point, each at
one random level up to 3e and a random new digit. (b) (p, f, e) in (3, 2, 2), (3, 2, 4),
(3, 2, 6), (3, 2, 12), (5, 2, 4), (3, 3, 2): every ubar in F_q^x at random digits
to level 2e + s, and every ubar with N(-ubar) = 1 at designed
departures l = 1 .. s + 2 and none; units z = 1 + omega(a) pi^c +
omega(t_1) pi^(c+1), every a != 0 and t_1, every class.

PREDICTIONS. (Pa) no two cells print one string; every diagonal move
equals s_r (c' - c). (Pb1) every non-kernel unit lands at p s. (Pb2)
kernel units at p s + min(l, p^m) when l < p^m; at least p s + p^m,
with equality at q - 1 of the q values of t_1 for each a, when l >=
p^m. (Pb3) kernel nonempty iff N(-ubar) = 1, then p - 1 values of a.

KILL. Printed counts: (b) readings off (Pb1)/(Pb2) plus (poly, class,
a) with l >= p^m and an equality count other than q - 1; (a) colliding
strings plus diagonal moves off s_r. Either nonzero at any grid point
fires.

CONTROLS (read first). C1 f = 1 against tame.landings unit for unit on
tame.py's designed polynomials at ubar = -1 (integer digits there). C2
plant: the departure read as v(u + 1) at f = 2, ubar != -1 with N(-ubar)
= 1: the (Pb2) count must be nonzero. C3 omega(x)^q = omega(x) mod
p^N with residue x, every x in F_q.

CHANGED AFTER THE FIRST RUN. The first run's (b) read the branch l <
p^m only at (3, 2, 6)'s class (1, 1); (3, 2, 12) was added (class (2,
1), p^m = 3) and every grid point now prints the readings per branch.
Both runs missed the kill.

FINDINGS (copied from the printed run). Controls: C3 0 bad; C1 176
units, 0 differ; C2 the plant at ubar = 1, (p, f, e) = (3, 2, 4), reads
l = 0 and puts 18 kernel units off (Pb2). (a) 5832 cells, 5832 strings
(p 3, f 2, e 2, level 3); 768 and 768 (p 2, f 2, e 3, level 4); 300
diagonal moves, 0 off s_r. (b) 0 readings off (Pb1), (Pb2) or the
equality rule at every grid point; kernel size p - 1 exactly where
N(-ubar) = 1 and 0 elsewhere; every designed departure read back as
v(u - omega(ubar)). Read per branch (off the kernel units, kernel units
at l < p^m, (poly, class, a) at l >= p^m): (3, 2, 2) 1368, 0, 40; (3,
2, 4) 1584, 0, 48; (3, 2, 6) 3600, 198, 90; (3, 2, 12) 4896, 216, 136;
(5, 2, 4) 25800, 0, 120; (3, 3, 2) 51246, 0, 130. KILL MISSES. Run: 17.1
s, peak commit 11.2 MB.

VERDICT. Both statements pass to f > 1 with the residue field read as
F_q: the triangle verbatim in Teichmuller digits, and the tame readout
with two changes, both forced by a^(p-1) != 1 in F_q. The bend's graded
map b -> b^p + ubar b has the kernel b^(p-1) = -ubar, nonempty iff
N(-ubar) = 1 (iff zeta_p lies in K), so a class splits: units off the
kernel land at p s, the kernel's p - 1 residues at least at p s +
min(l, p^m), every unit there when l < p^m and all but one t_1 per
residue when l >= p^m, l = v(u - omega(ubar)) the first nonzero Teichmuller coefficient
digit of (F - x^e + p omega(dbar))/p. At f = 1 the kernel is the whole
class at ubar = -1 and empty elsewhere: the f = 1 readout. At f > 1
the kernel's condition ubar = -1 widens to N(-ubar) = 1, (q - 1)/(p -
1) residues; a head, every level-1 unit past the bend, still needs the
kernel to be all of F_q^x, f = 1. No class with a nonempty kernel is
rigid at f > 1.
"""

import random
import sys
import time

import tame

CHECKS = []
INF = 10 ** 9
SEED = 1699


def check(name, ok, detail=""):
    CHECKS.append(bool(ok))
    tail = f"  ({detail})" if detail else ""
    print(f"  {'PASS' if ok else 'FAIL'}  {name}{tail}")


def irreducible(p, f):
    """A monic irreducible of degree f <= 3 over F_p (no root suffices),
    coefficients low to high."""
    if f == 1:
        return [0, 1]
    for n in range(p ** f):
        h = [(n // p ** i) % p for i in range(f)] + [1]
        if all(sum(c * x ** i for i, c in enumerate(h)) % p for x in range(p)):
            return h
    raise ValueError


class Ring:
    """O_K = Z_q[pi]/(F) mod p^N; a Z_q element is a tuple of f ints, an
    O_K element a list of e of them."""

    def __init__(self, p, f, e, b, d, N):
        self.p, self.f, self.e, self.N = p, f, e, N
        self.P = p ** N
        self.q = p ** f
        self.h = irreducible(p, f)
        self.b, self.d = b, d
        if d is None:
            return
        # pi^e = p d - sum_i p b_i pi^i
        self.top = [self.zmul((p,) + (0,) * (f - 1), d)] + \
            [self.zneg(self.zmul((p,) + (0,) * (f - 1), b[i]))
             for i in range(1, e)]

    # Z_q mod p^N
    def zmul(self, x, y):
        f, P, h = self.f, self.P, self.h
        pr = [0] * (2 * f - 1)
        for i, a in enumerate(x):
            if a:
                for j, c in enumerate(y):
                    pr[i + j] += a * c
        for k in range(2 * f - 2, f - 1, -1):
            c = pr[k]
            if c:
                pr[k] = 0
                for j in range(f):
                    pr[k - f + j] -= c * h[j]
        return tuple(c % P for c in pr[:f])

    def zadd(self, x, y):
        return tuple((a + c) % self.P for a, c in zip(x, y))

    def zneg(self, x):
        return tuple((-a) % self.P for a in x)

    def zpow(self, x, n):
        r = (1,) + (0,) * (self.f - 1)
        while n:
            if n & 1:
                r = self.zmul(r, x)
            x = self.zmul(x, x)
            n >>= 1
        return r

    def zval(self, x):
        best = INF
        for a in x:
            if a:
                v = 0
                while a % self.p == 0:
                    a //= self.p
                    v += 1
                best = min(best, v)
        return best

    def omega(self, x):
        """Teichmuller lift of the residue tuple x."""
        return self.zpow(tuple(a % self.P for a in x), self.q ** self.N)

    # O_K mod p^N
    def zero(self):
        return [(0,) * self.f for _ in range(self.e)]

    def one(self):
        z = self.zero()
        z[0] = (1,) + (0,) * (self.f - 1)
        return z

    def mul(self, x, y):
        e = self.e
        pr = [(0,) * self.f for _ in range(2 * e - 1)]
        for i, a in enumerate(x):
            if any(a):
                for j, c in enumerate(y):
                    if any(c):
                        pr[i + j] = self.zadd(pr[i + j], self.zmul(a, c))
        for k in range(2 * e - 2, e - 1, -1):
            c = pr[k]
            if any(c):
                pr[k] = (0,) * self.f
                for j in range(e):
                    pr[k - e + j] = self.zadd(pr[k - e + j],
                                              self.zmul(c, self.top[j]))
        return pr[:e]

    def add(self, x, y):
        return [self.zadd(a, c) for a, c in zip(x, y)]

    def scal(self, z, x):
        return [self.zmul(z, a) for a in x]

    def power(self, x, n):
        r = self.one()
        while n:
            if n & 1:
                r = self.mul(r, x)
            x = self.mul(x, x)
            n >>= 1
        return r

    def pi_pow(self, r):
        x = self.zero()
        x[1 % self.e] = (1,) + (0,) * (self.f - 1)
        if self.e == 1:
            raise ValueError
        return self.power(x, r) if r else self.one()

    def val(self, x):
        best = INF
        for i, a in enumerate(x):
            v = self.zval(a)
            if v < INF:
                best = min(best, self.e * v + i)
        return min(best, self.e * self.N)

    def inverse(self, z):
        """Newton from the residue's inverse in F_q."""
        r0 = self.fq_inv(tuple(a % self.p for a in z[0]))
        y = self.zero()
        y[0] = r0
        two = self.zero()
        two[0] = (2,) + (0,) * (self.f - 1)
        for _ in range(2 * self.N.bit_length() + 6):
            y = self.mul(y, self.add(two, [self.zneg(a) for a in self.mul(z, y)]))
        return y

    # F_q as residue tuples
    def fq_mul(self, x, y):
        return tuple(a % self.p for a in self.zmul(x, y))

    def fq_pow(self, x, n):
        return tuple(a % self.p for a in self.zpow(x, n))

    def fq_inv(self, x):
        return self.fq_pow(x, self.q - 2)


def fq_all(p, f):
    return [tuple((n // p ** i) % p for i in range(f)) for n in range(p ** f)]


def fq_sub(p, x, y):
    return tuple((a - c) % p for a, c in zip(x, y))


def norm_is_one(R, x):
    """N_(F_q/F_p)(x) = x^((q-1)/(p-1)) == 1."""
    return R.fq_pow(x, (R.q - 1) // (R.p - 1)) == (1,) + (0,) * (R.f - 1)


def build(p, f, e, dbar, digs, N):
    """F from Teichmuller coefficient digits: digs[r - 1] at level r =
    i + k e (b_i's p^k digit, d's at i = 0), d's p^0 digit dbar."""
    R0 = Ring(p, f, e, [None] * e, None, N)
    b = [(0,) * f for _ in range(e)]
    d = R0.omega(dbar)
    for r, c in enumerate(digs, start=1):
        if not any(c):
            continue
        i, k = r % e, r // e
        t = R0.zmul(R0.omega(c), (p ** k,) + (0,) * (f - 1))
        if i == 0:
            d = R0.zadd(d, t)
        else:
            b[i] = R0.zadd(b[i], t)
    return Ring(p, f, e, b, d, N)


def u_of(R):
    D = R.zero()
    D[0] = R.d
    for i in range(1, R.e):
        D[i] = R.zneg(R.b[i])
    return R.inverse(D)


def digits(R, u, L):
    """Teichmuller digits u_0..u_L of u."""
    p, e = R.p, R.e
    ubar = tuple(a % p for a in u[0])
    z = [a for a in u]
    out = []
    for r in range(L + 1):
        i, k = r % e, r // e
        w = tuple((a // p ** k) % p for a in z[i])
        dig = R.fq_mul(w, R.fq_pow(ubar, k))
        out.append(dig)
        if any(dig):
            z = R.add(z, [R.zneg(a) for a in R.scal(R.omega(dig), R.pi_pow(r))])
    return out


# ------------------------------------------------------------ the triangle

def section_triangle(rng):
    print("\n(a) THE TRIANGLE in Teichmuller digits")
    ctot = 0
    for (p, f, e, r) in ((3, 2, 2, 3), (2, 2, 3, 4)):
        N = r // e + 3
        F = fq_all(p, f)
        seen, coll, n = {}, 0, 0
        for d0 in F[1:]:
            for idx in range(len(F) ** r):
                digs = [F[(idx // len(F) ** j) % len(F)] for j in range(r)]
                R = build(p, f, e, d0, digs, N)
                s = tuple(digits(R, u_of(R), r))
                coll += s in seen
                seen[s] = 1
                n += 1
        print(f"  p {p} f {f} e {e} to level {r}: {n} cells, "
              f"{len(seen)} strings, {coll} collisions")
        check(f"Pa bijection p {p} f {f} e {e}", coll == 0 and n == len(seen))
        ctot += coll
    off, tot = 0, 0
    for (p, f, e) in ((3, 2, 2), (3, 2, 4), (2, 2, 3), (5, 2, 4), (3, 3, 2)):
        F = fq_all(p, f)
        L = 3 * e
        N = L // e + 3
        for _ in range(60):
            d0 = rng.choice(F[1:])
            digs = [rng.choice(F) for _ in range(L)]
            R = build(p, f, e, d0, digs, N)
            u0 = digits(R, u_of(R), L)
            ubar = u0[0]
            lv = rng.randrange(1, L + 1)
            new = rng.choice([c for c in F if c != digs[lv - 1]])
            d2 = list(digs)
            d2[lv - 1] = new
            R2 = build(p, f, e, d0, d2, N)
            u1 = digits(R2, u_of(R2), L)
            i, k = lv % e, lv // e
            s = R.fq_pow(ubar, k + 2)
            if i == 0:
                s = tuple((-a) % p for a in s)
            want = R.fq_mul(s, fq_sub(p, new, digs[lv - 1]))
            ok = u1[:lv] == u0[:lv] and fq_sub(p, u1[lv], u0[lv]) == want
            off += not ok
            tot += 1
    print(f"  diagonal: {tot} (cell, level, change) reads, {off} off s_r")
    check("Pa diagonal s_r = +/- ubar^(k+2)", off == 0)
    return ctot + off


# ------------------------------------------------------------ the readout

def departure(R, u):
    ubar = tuple(a % R.p for a in u[0])
    w = R.zero()
    w[0] = R.omega(ubar)
    return R.val(R.add(u, [R.zneg(a) for a in w]))


def classes(p, e):
    return tame.classes(p, e)


def landings_q(R, c, m, units):
    """{(a, t1): v(z^(p^(m+1)) - 1)} with the levels below checked."""
    out, lev_ok = {}, True
    pc, pc1 = R.pi_pow(c), R.pi_pow(c + 1)
    for (a, t1) in units:
        z = R.add(R.one(), R.add(R.scal(a, pc), R.scal(t1, pc1)))
        y = z
        for j in range(m + 1):
            ym = [a_ for a_ in y]
            ym[0] = R.zadd(ym[0], R.zneg(R.one()[0]))
            lev_ok &= R.val(ym) == c * R.p ** j
            y = R.power(y, R.p)
        y[0] = R.zadd(y[0], R.zneg(R.one()[0]))
        out[(a, t1)] = R.val(y)
    return out, lev_ok


def read_poly(R, cls, plant=False):
    """Counts off (Pb1), (Pb2) and the equality rule, and the kernel
    sizes, over every class and unit."""
    p, f, e, q = R.p, R.f, R.e, R.q
    s = e // (p - 1)
    u = u_of(R)
    ubar = tuple(a % p for a in u[0])
    if plant:
        w = R.zero()
        w[0] = R.zneg(R.one()[0])
        l = R.val(R.add(u, [R.zneg(a) for a in w]))
    else:
        l = departure(R, u)
    F = fq_all(p, f)
    mubar = tuple((-a) % p for a in ubar)
    om = {x: R.omega(x) for x in F}
    off1 = off2 = offeq = 0
    nb = [0, 0, 0]
    kers = set()
    lev_all = True
    for (c, m) in cls:
        units = [(om[a], om[t]) for a in F[1:] for t in F]
        land, lev_ok = landings_q(R, c, m, units)
        lev_all &= lev_ok
        ker = [a for a in F[1:]
               if R.fq_pow(R.fq_pow(a, p ** m), p - 1) == mubar]
        kers.add(len(ker))
        for a in F[1:]:
            vals = [land[(om[a], om[t])] for t in F]
            if a not in ker:
                nb[0] += len(vals)
                off1 += sum(v != p * s for v in vals)
            elif l < p ** m:
                nb[1] += len(vals)
                off2 += sum(v != p * s + l for v in vals)
            else:
                nb[2] += 1
                off2 += sum(v < p * s + p ** m for v in vals)
                offeq += sum(v == p * s + p ** m for v in vals) != q - 1
    return dict(l=l, ubar=ubar, n1=norm_is_one(R, mubar), off1=off1,
                off2=off2, offeq=offeq, kers=kers, lev=lev_all, nb=nb)


def designed_digits(rng, F, l, top):
    digs = []
    for r in range(1, top + 1):
        if r < l:
            digs.append(F[0])
        elif r == l:
            digs.append(rng.choice(F[1:]))
        else:
            digs.append(rng.choice(F))
    return digs


GRID = [(3, 2, 2), (3, 2, 4), (3, 2, 6), (3, 2, 12), (5, 2, 4),
        (3, 3, 2)]


def section_readout(rng):
    print("\n(b) THE TAME READOUT at f > 1: per (p, f, e) the polynomials "
          "read, kernel sizes, counts off (Pb1), (Pb2), equality")
    tot = 0
    for (p, f, e) in GRID:
        s = e // (p - 1)
        cls = classes(p, e)
        top = 2 * e + s
        N = (p * s + s + e + 2) // e + 3
        F = fq_all(p, f)
        rows = []
        for ub in F[1:]:
            R0 = Ring(p, f, e, [None] * e, None, N)
            dbar = R0.fq_inv(ub)
            rows.append(("rand", build(p, f, e, dbar,
                                       [rng.choice(F) for _ in range(top)], N)))
            if norm_is_one(R0, tuple((-a) % p for a in ub)):
                for l in list(range(1, s + 3)) + [INF]:
                    dg = designed_digits(rng, F, l, top) if l < INF \
                        else [F[0]] * top
                    rows.append((l, build(p, f, e, dbar, dg, N)))
        o1 = o2 = oe = 0
        nb = [0, 0, 0]
        lmis = 0
        kerok = True
        levok = True
        nn = 0
        for tag, R in rows:
            r = read_poly(R, cls)
            o1 += r["off1"]
            o2 += r["off2"]
            oe += r["offeq"]
            nb = [x + y for x, y in zip(nb, r["nb"])]
            levok &= r["lev"]
            nn += r["n1"]
            want = {p - 1} if r["n1"] else {0}
            kerok &= r["kers"] == want
            if tag != "rand" and r["l"] != (tag if tag < INF else
                                            e * R.N):
                lmis += 1
        tot += o1 + o2 + oe
        print(f"  p {p} f {f} e {e} (s {s}, classes {cls}): {len(rows)} "
              f"polys, {nn} with N(-ubar) = 1; off Pb1 {o1}, Pb2 {o2}, "
              f"equality {oe}; designed l missed {lmis}; kernel sizes "
              f"{'p - 1 or 0 as N(-ubar)' if kerok else 'OFF'}; levels "
              f"{'ok' if levok else 'OFF'}")
        print(f"      read: {nb[0]} units off the kernel, {nb[1]} kernel "
              f"units at l < p^m, {nb[2]} (poly, class, a) at l >= p^m")
        check(f"Pb p {p} f {f} e {e}", o1 + o2 + oe == 0 and kerok and levok
              and lmis == 0)
    return tot


# ------------------------------------------------------------ controls

def section_controls(rng):
    print("CONTROLS")
    # C3 Teichmuller
    bad = 0
    for (p, f) in ((3, 2), (5, 2), (3, 3), (2, 2)):
        R = Ring(p, f, 2, [None] * 2, None, 5)
        for x in fq_all(p, f):
            w = R.omega(x)
            bad += R.zpow(w, R.q) != w or tuple(a % p for a in w) != x
    check("C3 omega(x)^q = omega(x), residue x", bad == 0, f"{bad} bad")
    # C1 f = 1 against tame.landings
    mism, n = 0, 0
    for (p, e) in ((3, 2), (3, 4), (5, 4), (3, 6)):
        s = e // (p - 1)
        for l in (1, s, s + 1, INF):
            b, d = tame.designed(rng, p, e, l, 2 * e + s)
            loc = tame.local(p, e, b, d)
            Nt = (p * s + s + e + 2) // e + 2
            R = Ring(p, 1, e, [(x % p ** Nt,) for x in b], (d % p ** Nt,), Nt)
            for (c, m) in tame.classes(p, e):
                ref, _, _ = tame.landings(loc, c, m, 2)
                units = []
                for idx in range((p - 1) * p):
                    units.append(((1 + idx % (p - 1),), ((idx // (p - 1)),)))
                mine, _ = landings_q(R, c, m, units)
                for idx, (a, t) in enumerate(units):
                    n += 1
                    mism += mine[(a, t)] != ref[idx]
    check("C1 f = 1 landings equal tame.landings", mism == 0,
          f"{n} units, {mism} differ")
    # C2 plant: departure read as v(u + 1) at f = 2, ubar != -1, N(-ubar) = 1
    p, f, e = 3, 2, 4
    s = e // (p - 1)
    N = (p * s + s + e + 2) // e + 3
    R0 = Ring(p, f, e, [None] * e, None, N)
    F = fq_all(p, f)
    ub = next(x for x in F[1:] if x != (p - 1,) + (0,) * (f - 1)
              and norm_is_one(R0, tuple((-a) % p for a in x)))
    R = build(p, f, e, R0.fq_inv(ub), designed_digits(rng, F, 1, 2 * e + s), N)
    r = read_poly(R, tame.classes(p, e), plant=True)
    check("C2 plant: departure v(u + 1) breaks Pb2", r["off2"] > 0,
          f"ubar {ub}, plant l {r['l']}, off Pb2 {r['off2']}")


def main():
    t0 = time.time()
    rng = random.Random(SEED)
    section_controls(rng)
    ka = section_triangle(rng)
    kb = section_readout(rng)
    print(f"\nKILL: (b) off-counts {kb}, (a) off-counts {ka}: "
          f"{'FIRES' if ka or kb else 'MISSES'}")
    print(f"\n{sum(CHECKS)}/{len(CHECKS)} checks pass  "
          f"[{time.time() - t0:.1f} s]")
    return 0 if all(CHECKS) else 1


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