import numpy as np
from numba import njit


N = 3_691_560
M = 10
BCF10_DIM = 120
MAX_ARANK = 7
RM2_DIM = 56
RM6_DIM = 848
RM6_DEGREE = 6
RM6_WORDS = 14
ENUM_BINS = 65
TRUTH_WORDS = 16
WORD_BITS = 64
ENUM_STEP = 8
MAX_RM6_WEIGHT = 32
FULL_MASK = (1 << M) - 1
RM3_10_LENGTH = 1024
RM3_11_LENGTH = 2048
RM3_10_BINS = RM3_10_LENGTH // ENUM_STEP + 1
RM3_11_BINS = RM3_11_LENGTH // ENUM_STEP + 1
RM3_11_HALF_BINS = RM3_11_LENGTH // (2 * ENUM_STEP) + 1
MOD_PRIME_COUNT = 8


@njit
def dims_ok(forms, dims):
    for row in range(N):
        support = 0
        bit = 0

        for i in range(M):
            for j in range(i + 1, M):
                for k in range(j + 1, M):
                    if bit < WORD_BITS:
                        coeff = (forms[row, 0] >> bit) & np.uint64(1)
                    else:
                        coeff = (forms[row, 1] >> (bit - WORD_BITS)) & np.uint64(1)

                    if coeff != 0:
                        support |= (1 << i) | (1 << j) | (1 << k)
                    bit += 1

        count = 0
        for i in range(M):
            count += (support >> i) & 1

        if count > dims[row]:
            return False

    return True


def stabilizer_identity_ok(dims, stabilizers):
    gl = []
    for n in range(M + 1):
        order = 1
        for i in range(n):
            order *= (1 << n) - (1 << i)
        gl.append(order)

    orbit_sum = 0
    for m, s in zip(dims, stabilizers):
        m = int(m)
        full_stabilizer = int(s) * (1 << (m * (M - m))) * gl[M - m]
        orbit_sum += gl[M] // full_stabilizer

    return orbit_sum == (1 << BCF10_DIM) - 1


def catalog_ok(forms, stabilizers):
    monomial_to_bit = {}
    bit = 0
    for i in range(M):
        for j in range(i + 1, M):
            for k in range(j + 1, M):
                monomial_to_bit[f"{i}{j}{k}"] = bit
                bit += 1

    row = 0
    with open("catalog.txt", encoding="ascii") as file:
        for line in file:
            fields = line.split()
            if not fields or not fields[0].isdigit():
                continue

            if len(fields) != 3 or int(fields[0]) != row + 1:
                return False

            lo = 0
            hi = 0
            for term in fields[1].split("+"):
                bit = monomial_to_bit.get(term, -1)
                if bit < 0:
                    return False
                if bit < WORD_BITS:
                    lo |= 1 << bit
                else:
                    hi |= 1 << (bit - WORD_BITS)

            if lo != int(forms[row, 0]) or hi != int(forms[row, 1]) or int(fields[2]) != int(stabilizers[row]):
                return False

            row += 1

    return row == N


def is_prime(n):
    if n < 2:
        return False
    if n % 2 == 0:
        return n == 2
    d = 3
    while d * d <= n:
        if n % d == 0:
            return False
        d += 2
    return True


def verification_primes():
    primes = []
    n = (1 << 31) - 1
    while len(primes) < MOD_PRIME_COUNT:
        if is_prime(n):
            primes.append(n)
        n -= 2
    return primes


def gl_orders():
    orders = []
    for n in range(M + 1):
        order = 1
        for i in range(n):
            order *= (1 << n) - (1 << i)
        orders.append(order)
    return orders


_orbit_unique = None
_orbit_inverse = None
_orbit_mod_cache = {}


def orbit_residues(dims, stabilizers, modulus):
    global _orbit_unique, _orbit_inverse

    if _orbit_unique is None:
        pairs = np.empty((N, 2), dtype=np.uint64)
        pairs[:, 0] = dims
        pairs[:, 1] = stabilizers
        _orbit_unique, _orbit_inverse = np.unique(pairs, axis=0, return_inverse=True)

    if modulus not in _orbit_mod_cache:
        gl = gl_orders()
        gl10 = gl[M]
        residues = np.empty(len(_orbit_unique), dtype=np.uint64)

        for i, (m, s) in enumerate(_orbit_unique):
            m = int(m)
            s = int(s)
            full_stabilizer = s * (1 << (m * (M - m))) * gl[M - m]
            residues[i] = (gl10 // full_stabilizer) % modulus

        _orbit_mod_cache[modulus] = residues[_orbit_inverse]

    return _orbit_mod_cache[modulus]


def exact_div(a, b):
    if a % b != 0:
        raise ValueError("non-exact division")
    return a // b


def rank_quadratic_count(m, j):
    if j == 0:
        return 1

    numerator = 1
    for i in range(2 * j):
        numerator *= (1 << m) - (1 << i)

    denominator = 1 << (j * j)
    for i in range(1, j + 1):
        denominator *= (1 << (2 * i)) - 1

    return exact_div(numerator, denominator)


def rm2_10_weight_enumerator():
    weights = [0] * RM3_10_BINS

    for j in range(M // 2 + 1):
        nforms = rank_quadratic_count(M, j)
        rank_factor = 1 << (2 * j)
        low = (1 << (M - 1)) - (1 << (M - j - 1))
        mid = 1 << (M - 1)
        high = (1 << (M - 1)) + (1 << (M - j - 1))

        weights[low // ENUM_STEP] += nforms * rank_factor
        weights[mid // ENUM_STEP] += nforms * 2 * ((1 << M) - rank_factor)
        weights[high // ENUM_STEP] += nforms * rank_factor

    if sum(weights) != 1 << RM2_DIM:
        raise ValueError("RM(2,10) weight enumerator sum mismatch")

    return weights


def parse_weight_distribution(path, length):
    bins = length // ENUM_STEP + 1
    weights = [0] * bins

    with open(path, encoding="ascii") as file:
        for line in file:
            line = line.split("#", 1)[0].strip()
            if not line:
                continue

            if line.endswith("+"):
                line = line[:-1].strip()

            if line == "1":
                weights[0] = 1
            elif line.startswith("x^"):
                weights[int(line[2:]) // ENUM_STEP] = 1
            else:
                coeff, monomial = line.split()
                weights[int(monomial[2:]) // ENUM_STEP] = int(coeff)

    return weights


@njit
def assemble_rm3_10_mod(enumerators, orbit_mod, zero_mod, modulus):
    out = np.zeros(RM3_10_BINS, dtype=np.uint64)

    for t in range(RM3_10_BINS):
        out[t] = zero_mod[t]

    for row in range(N):
        orbit = orbit_mod[row]
        for t in range(ENUM_BINS):
            c = enumerators[row, t] % modulus
            if c != 0:
                add = (orbit * c) % modulus
                out[t] = (out[t] + add) % modulus
                if t != ENUM_BINS - 1:
                    mirror = RM3_10_BINS - 1 - t
                    out[mirror] = (out[mirror] + add) % modulus

    return out


@njit
def assemble_rm3_11_mod(enumerators, orbit_mod, zero_mod, modulus):
    out = np.zeros(RM3_11_HALF_BINS, dtype=np.uint64)

    for i in range(RM3_10_BINS):
        ci = zero_mod[i]
        if ci == 0:
            continue
        for j in range(RM3_10_BINS):
            t = i + j
            if t >= RM3_11_HALF_BINS:
                continue
            cj = zero_mod[j]
            if cj != 0:
                out[t] = (out[t] + ((ci * cj) % modulus)) % modulus

    idx = np.empty(RM3_10_BINS, dtype=np.uint16)
    val = np.empty(RM3_10_BINS, dtype=np.uint64)

    for row in range(N):
        nitems = 0

        for t in range(ENUM_BINS):
            c = enumerators[row, t] % modulus
            if c != 0:
                idx[nitems] = t
                val[nitems] = c
                nitems += 1

                if t != ENUM_BINS - 1:
                    idx[nitems] = RM3_10_BINS - 1 - t
                    val[nitems] = c
                    nitems += 1

        orbit = orbit_mod[row]

        for a in range(nitems):
            ia = idx[a]
            va = val[a]
            for b in range(nitems):
                t = ia + idx[b]
                if t < RM3_11_HALF_BINS:
                    add = (((va * val[b]) % modulus) * orbit) % modulus
                    out[t] = (out[t] + add) % modulus

    return out


def enumerator_rm3_10_ok(enumerators, dims, stabilizers):
    target = parse_weight_distribution("enumerator_rm3_10.txt", RM3_10_LENGTH)
    zero = rm2_10_weight_enumerator()

    for modulus in verification_primes():
        orbit_mod = orbit_residues(dims, stabilizers, modulus)
        zero_mod = np.array([x % modulus for x in zero], dtype=np.uint64)
        target_mod = np.array([x % modulus for x in target], dtype=np.uint64)

        if not np.array_equal(assemble_rm3_10_mod(enumerators, orbit_mod, zero_mod, np.uint64(modulus)), target_mod):
            return False

    return True


def enumerator_rm3_11_ok(enumerators, dims, stabilizers):
    target = parse_weight_distribution("enumerator_rm3_11.txt", RM3_11_LENGTH)
    zero = rm2_10_weight_enumerator()

    for modulus in verification_primes():
        orbit_mod = orbit_residues(dims, stabilizers, modulus)
        zero_mod = np.array([x % modulus for x in zero], dtype=np.uint64)
        target_mod = np.array([x % modulus for x in target[:RM3_11_HALF_BINS]], dtype=np.uint64)

        if not np.array_equal(assemble_rm3_11_mod(enumerators, orbit_mod, zero_mod, np.uint64(modulus)), target_mod):
            return False

    return True


@njit
def enumerators_sum_ok(enumerators):
    coset_size = np.uint64(1 << RM2_DIM)

    for row in range(N):
        total = enumerators[row, WORD_BITS]
        for col in range(WORD_BITS):
            total += np.uint64(2) * enumerators[row, col]
        if total != coset_size:
            return False

    return True


@njit
def enumerators_divisibility_ok(enumerators, dims):
    for row in range(N):
        m = dims[row]
        mask = np.uint64((1 << m) - 1)
        for col in range(ENUM_BINS):
            if (enumerators[row, col] & mask) != 0:
                return False

    return True


@njit
def decompositions_ok(forms, decompositions):
    triples = np.empty((BCF10_DIM, 3), dtype=np.uint8)
    bit = 0
    for i in range(M):
        for j in range(i + 1, M):
            for k in range(j + 1, M):
                triples[bit, 0] = i
                triples[bit, 1] = j
                triples[bit, 2] = k
                bit += 1

    one = np.uint64(1)

    for row in range(N):
        rec_lo = np.uint64(0)
        rec_hi = np.uint64(0)

        for col in range(MAX_ARANK):
            term = decompositions[row, col]
            if term != 0:
                r0 = term & 0x3FF
                r1 = (term >> 10) & 0x3FF
                r2 = (term >> 20) & 0x3FF

                for bit in range(BCF10_DIM):
                    i = triples[bit, 0]
                    j = triples[bit, 1]
                    k = triples[bit, 2]
                    a0 = (r0 >> i) & 1
                    a1 = (r1 >> i) & 1
                    a2 = (r2 >> i) & 1
                    b0 = (r0 >> j) & 1
                    b1 = (r1 >> j) & 1
                    b2 = (r2 >> j) & 1
                    c0 = (r0 >> k) & 1
                    c1 = (r1 >> k) & 1
                    c2 = (r2 >> k) & 1
                    coeff = (
                        (a0 & ((b1 & c2) ^ (b2 & c1)))
                        ^ (b0 & ((a1 & c2) ^ (a2 & c1)))
                        ^ (c0 & ((a1 & b2) ^ (a2 & b1)))
                    )
                    if coeff:
                        if bit < WORD_BITS:
                            rec_lo ^= one << bit
                        else:
                            rec_hi ^= one << (bit - WORD_BITS)

        if rec_lo != forms[row, 0] or rec_hi != forms[row, 1]:
            return False

    return True


@njit
def corrections_rm2_ok(forms, enumerators, corrections):
    one = np.uint64(1)
    monomials = np.empty(BCF10_DIM + RM2_DIM, dtype=np.uint16)
    bit = 0

    for i in range(M):
        for j in range(i + 1, M):
            for k in range(j + 1, M):
                monomials[bit] = (1 << i) | (1 << j) | (1 << k)
                bit += 1

    for i in range(M):
        for j in range(i + 1, M):
            monomials[bit] = (1 << i) | (1 << j)
            bit += 1

    for i in range(M):
        monomials[bit] = 1 << i
        bit += 1

    monomials[bit] = 0

    truth_tables = np.zeros((BCF10_DIM + RM2_DIM, TRUTH_WORDS), dtype=np.uint64)
    for m in range(BCF10_DIM + RM2_DIM):
        monomial = monomials[m]
        for x in range(1 << M):
            if (x & monomial) == monomial:
                truth_tables[m, x >> 6] |= one << (x & 63)

    words = np.empty(TRUTH_WORDS, dtype=np.uint64)

    for row in range(N):
        words.fill(0)
        f0 = forms[row, 0]
        f1 = forms[row, 1]

        for bit in range(WORD_BITS):
            if ((f0 >> bit) & one) != 0:
                for word in range(TRUTH_WORDS):
                    words[word] ^= truth_tables[bit, word]

        for bit in range(WORD_BITS, BCF10_DIM):
            if ((f1 >> (bit - WORD_BITS)) & one) != 0:
                for word in range(TRUTH_WORDS):
                    words[word] ^= truth_tables[bit, word]

        q = corrections[row]
        for bit in range(RM2_DIM):
            if ((q >> bit) & one) != 0:
                for word in range(TRUTH_WORDS):
                    words[word] ^= truth_tables[BCF10_DIM + bit, word]

        weight = 0
        for word in range(TRUTH_WORDS):
            x = words[word]
            while x != 0:
                weight += 1
                x &= x - one

        distance = 0
        for col in range(ENUM_BINS):
            if enumerators[row, col] != 0:
                distance = ENUM_STEP * col
                break

        if weight != distance:
            return False

    return True


@njit
def corrections_rm6_ok(forms, corrections):
    one = np.uint64(1)
    monomials = np.empty(BCF10_DIM + RM6_DIM, dtype=np.uint16)
    bit = 0

    for i in range(M):
        for j in range(i + 1, M):
            for k in range(j + 1, M):
                monomials[bit] = FULL_MASK ^ ((1 << i) | (1 << j) | (1 << k))
                bit += 1

    monomials[bit] = 0
    bit += 1

    comb = np.empty(RM6_DEGREE, dtype=np.uint8)
    for degree in range(1, RM6_DEGREE + 1):
        for i in range(degree):
            comb[i] = i

        while True:
            monomial = 0
            for i in range(degree):
                monomial |= 1 << comb[i]
            monomials[bit] = monomial
            bit += 1

            pos = degree - 1
            while pos >= 0 and comb[pos] == M - degree + pos:
                pos -= 1
            if pos < 0:
                break

            comb[pos] += 1
            for i in range(pos + 1, degree):
                comb[i] = comb[i - 1] + 1

    truth_tables = np.zeros((BCF10_DIM + RM6_DIM, TRUTH_WORDS), dtype=np.uint64)
    for m in range(BCF10_DIM + RM6_DIM):
        monomial = monomials[m]
        for x in range(1 << M):
            if (x & monomial) == monomial:
                truth_tables[m, x >> 6] |= one << (x & 63)

    words = np.empty(TRUTH_WORDS, dtype=np.uint64)

    for row in range(N):
        words.fill(0)
        f0 = forms[row, 0]
        f1 = forms[row, 1]

        for bit in range(WORD_BITS):
            if ((f0 >> bit) & one) != 0:
                for word in range(TRUTH_WORDS):
                    words[word] ^= truth_tables[bit, word]

        for bit in range(WORD_BITS, BCF10_DIM):
            if ((f1 >> (bit - WORD_BITS)) & one) != 0:
                for word in range(TRUTH_WORDS):
                    words[word] ^= truth_tables[bit, word]

        for w in range(RM6_WORDS):
            coeffs = corrections[row, w]
            base = w * WORD_BITS
            stop = WORD_BITS
            if base + stop > RM6_DIM:
                stop = RM6_DIM - base

            for shift in range(stop):
                if ((coeffs >> shift) & one) != 0:
                    bit = base + shift
                    for word in range(TRUTH_WORDS):
                        words[word] ^= truth_tables[BCF10_DIM + bit, word]

        weight = 0
        for word in range(TRUTH_WORDS):
            x = words[word]
            while x != 0:
                weight += 1
                if weight > MAX_RM6_WEIGHT:
                    return False
                x &= x - one

    return True


def status(ok):
    return "OK" if ok else "FAILED"


def format_histogram(hist, step=1):
    return " ".join(f"{i * step}:{int(count)}" for i, count in enumerate(hist) if count)


def print_checks(rows):
    check_width = max(len(description) for description, _, _ in rows)

    print(f"{'check':<{check_width}}  status", flush=True)
    print(f"{'-' * check_width}  {'-' * 6}", flush=True)
    for description, fn, args in rows:
        ok = fn(*args)
        print(f"{description:<{check_width}}  {status(ok):>6}", flush=True)


def print_histograms(rows):
    name_width = max(len(name) for name, _ in rows)

    print(f"{'histogram':<{name_width}}  counts")
    print(f"{'-' * name_width}  {'-' * 6}")
    for name, histogram in rows:
        print(f"{name:<{name_width}}  {histogram}")


forms = np.load("forms.npy", mmap_mode="r")
dims = np.load("dims.npy", mmap_mode="r")
stabilizers = np.load("stabilizers.npy", mmap_mode="r")
enumerators = np.load("enumerators.npy", mmap_mode="r")
decompositions = np.load("decompositions.npy", mmap_mode="r")
corrections_rm2 = np.load("corrections_rm2.npy", mmap_mode="r")
corrections_rm6 = np.load("corrections_rm6.npy", mmap_mode="r")

checks = [
    ("forms use at most their natural dimension", dims_ok, (forms, dims)),
    ("orbit-stabilizer sum is 2^120 - 1", stabilizer_identity_ok, (dims, stabilizers)),
    ("catalog.txt matches forms.npy and stabilizers.npy", catalog_ok, (forms, stabilizers)),
    ("each enumerator has total coset size 2^56", enumerators_sum_ok, (enumerators,)),
    ("enumerators are divisible by 2^dim", enumerators_divisibility_ok, (enumerators, dims)),
    ("enumerators assemble enumerator_rm3_10.txt", enumerator_rm3_10_ok, (enumerators, dims, stabilizers)),
    ("enumerators assemble enumerator_rm3_11.txt", enumerator_rm3_11_ok, (enumerators, dims, stabilizers)),
    ("alternating decompositions reconstruct forms", decompositions_ok, (forms, decompositions)),
    ("RM(2,10) corrections attain d2", corrections_rm2_ok, (forms, enumerators, corrections_rm2)),
    ("RM(6,10) corrections give weight <= 32", corrections_rm6_ok, (forms, corrections_rm6)),
]

print_checks(checks)
print()
print_histograms(
    [
        ("dims", format_histogram(np.bincount(dims, minlength=M + 1))),
        (
            "d2 degenerate",
            format_histogram(
                np.bincount(np.argmax(enumerators[dims < M] != 0, axis=1), minlength=ENUM_BINS),
                ENUM_STEP,
            ),
        ),
        (
            "d2 nondegenerate",
            format_histogram(
                np.bincount(np.argmax(enumerators[dims == M] != 0, axis=1), minlength=ENUM_BINS),
                ENUM_STEP,
            ),
        ),
        ("alternating rank", format_histogram(np.bincount(np.count_nonzero(decompositions != 0, axis=1), minlength=MAX_ARANK + 1))),
    ]
)