#!/usr/bin/env python3
"""Exact rational checks using only the Python standard library."""
from __future__ import annotations
import csv
import json
from collections import Counter
from fractions import Fraction as F
from functools import lru_cache
from itertools import product
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
MAX_N = 12


def zero(r, c=None):
    return [[F(0) for _ in range(r if c is None else c)] for _ in range(r)]


def eye(n):
    return [[F(i == j) for j in range(n)] for i in range(n)]


def tr(a):
    return [list(row) for row in zip(*a)]


def mul(a, b):
    return [[sum((x*y for x, y in zip(row, col)), F(0))
             for col in zip(*b)] for row in a]


def scale(c, a):
    return [[F(c)*x for x in row] for row in a]


def add(a, b):
    return [[x+y for x, y in zip(ra, rb)] for ra, rb in zip(a, b)]


def power(a, k):
    result = eye(len(a))
    while k:
        if k & 1:
            result = mul(result, a)
        a = mul(a, a)
        k //= 2
    return result


def diag(*blocks):
    size = sum(len(b) for b in blocks)
    result = zero(size)
    start = 0
    for block in blocks:
        for i, row in enumerate(block):
            for j, value in enumerate(row):
                result[start+i][start+j] = value
        start += len(block)
    return result


def reversal(m):
    return [[F(i+j == m-1) for j in range(m)] for i in range(m)]


def jordan(m):
    return [[F(j == i+1) for j in range(m)] for i in range(m)]


def rank(a):
    b = [row[:] for row in a]
    rows, cols = len(b), len(b[0])
    pivot_row = 0
    for col in range(cols):
        pivot = next((i for i in range(pivot_row, rows) if b[i][col]), None)
        if pivot is None:
            continue
        b[pivot_row], b[pivot] = b[pivot], b[pivot_row]
        value = b[pivot_row][col]
        b[pivot_row] = [x/value for x in b[pivot_row]]
        for i in range(pivot_row+1, rows):
            value = b[i][col]
            b[i] = [x-value*y for x, y in zip(b[i], b[pivot_row])]
        pivot_row += 1
        if pivot_row == rows:
            break
    return pivot_row


def inertia(a):
    """Exact rational congruence elimination, including isotropic pivots."""
    assert a == tr(a)
    b = [row[:] for row in a]
    p = q = z = 0
    while b:
        n = len(b)
        pivot = next((i for i in range(n) if b[i][i]), None)
        if pivot is None:
            pair = next(((i, j) for i in range(n) for j in range(i+1, n)
                         if b[i][j]), None)
            if pair is None:
                z += n
                break
            i, j = pair
            change = eye(n)
            change[j][i] = 1
            b = mul(tr(change), mul(b, change))
            pivot = i
        order = [pivot] + [i for i in range(n) if i != pivot]
        b = [[b[i][j] for j in order] for i in order]
        value = b[0][0]
        assert value
        p += int(value > 0)
        q += int(value < 0)
        b = [[b[i][j]-b[i][0]*b[0][j]/value
              for j in range(1, n)] for i in range(1, n)]
    assert p+q+z == len(a)
    return p, q, z


@lru_cache(None)
def block_inertia(m, sign, j):
    n = jordan(m)
    nj = power(n, j)
    result = inertia(mul(scale(sign, reversal(m)), nj))
    assert result[0]+result[1] == rank(nj)
    return result


def ordinary_partitions(total, minimum=1):
    if total == 0:
        yield ()
        return
    for first in range(minimum, total+1):
        for tail in ordinary_partitions(total-first, first):
            yield (first,)+tail


def signed_partitions(n):
    for partition in ordinary_partitions(n):
        counts = Counter(partition)
        lengths = sorted(counts)
        for positives in product(*(range(counts[m]+1) for m in lengths)):
            yield {m: (a, counts[m]-a) for m, a in zip(lengths, positives)}


def matrix_profile(n, blocks):
    result = []
    for j in range(n+2):
        p = q = 0
        for m, (cp, cm) in blocks.items():
            for sign, count in ((1, cp), (-1, cm)):
                if count:
                    bp, bq, _ = block_inertia(m, sign, j)
                    p += count*bp
                    q += count*bq
        result.append((p, q))
    return result


def product_coefficients(max_n):
    """Independent unbounded-coin extraction of the bivariate product."""
    coeff = {(p, q): 0 for p in range(max_n+1)
             for q in range(max_n+1-p)}
    coeff[0, 0] = 1
    for m in range(1, max_n+1):
        a, b = (m+1)//2, m//2
        for positive, negative in ((a, b), (b, a)):
            for p in range(positive, max_n+1):
                for q in range(negative, max_n+1-p):
                    coeff[p, q] += coeff[p-positive, q-negative]
    return coeff


def cmul(x, y):
    a, b = x
    c, d = y
    return a*c-b*d, a*d+b*c


def hermitian_pair(h, x, y):
    total = (F(0), F(0))
    for i in range(len(h)):
        for j in range(len(h)):
            value = cmul((x[i][0], -x[i][1]), cmul((h[i][j], F(0)), y[j]))
            total = total[0]+value[0], total[1]+value[1]
    return total


def verify_examples():
    h2, j2 = reversal(2), jordan(2)
    assert inertia(mul(h2, j2)) == (1, 0, 1)
    assert inertia(mul(h2, scale(-1, j2))) == (0, 1, 1)
    h6 = diag(reversal(4), reversal(2))
    n6a = diag(jordan(4), scale(-1, jordan(2)))
    n6b = diag(scale(-1, jordan(4)), jordan(2))
    assert [inertia(mul(h6, power(n6a, j))) for j in range(4)] == [
        (3, 3, 0), (2, 2, 2), (1, 1, 4), (1, 0, 5)]
    assert [inertia(mul(h6, power(n6b, j))) for j in range(4)] == [
        (3, 3, 0), (2, 2, 2), (1, 1, 4), (0, 1, 5)]

    h = [[F(1), F(0)], [F(0), F(-1)]]
    t = [[F(0), F(1)], [F(-1), F(0)]]
    assert mul(tr(t), h) == mul(h, t)
    assert power(t, 2) == scale(-1, eye(2))
    vp = [(F(1), F(0)), (F(0), F(1))]
    vm = [(F(1), F(0)), (F(0), F(-1))]
    assert hermitian_pair(h, vp, vp) == (0, 0)
    assert hermitian_pair(h, vm, vm) == (0, 0)
    assert hermitian_pair(h, vp, vm) == (2, 0)

    h4, n1 = zero(4), zero(4)
    for i in range(2):
        h4[i][i+2] = h4[i+2][i] = F(1)
        n1[i][i+2] = F(1)
    pair_checks = 0
    for value in (F(2), F(3), F(5)):
        n2 = zero(4)
        n2[0][2], n2[1][3] = F(1), value
        assert rank(n1) == rank(n2) == 2
        assert power(n1, 2) == power(n2, 2) == mul(n1, n2) == mul(n2, n1) == zero(4)
        assert mul(h4, n1) == mul(tr(n1), h4)
        assert mul(h4, n2) == mul(tr(n2), h4)
        assert inertia(mul(h4, n1)) == inertia(mul(h4, n2)) == (2, 0, 2)
        for parameter in (F(1, 2), F(1), F(3, 2), value, value+1):
            expected = (sum(x > parameter for x in (1, value)),
                        sum(x < parameter for x in (1, value)),
                        2+sum(x == parameter for x in (1, value)))
            assert inertia(mul(h4, add(n2, scale(-parameter, n1)))) == expected
            pair_checks += 1

    embedding = [[F(1), F(0)], [F(0), F(3, 5)], [F(0), F(4, 5)]]
    effects = [[[F(1), F(0)], [F(0), F(9, 25)]],
               [[F(0), F(0)], [F(0), F(16, 25)]]]
    p1, p2 = zero(3), zero(3)
    p1[0][0] = p1[1][1] = p2[2][2] = F(1)
    assert mul(tr(embedding), embedding) == eye(2)
    for effect, projection in zip(effects, (p1, p2)):
        assert mul(tr(embedding), mul(projection, embedding)) == effect
    assert sum(rank(e) for e in effects) == 3
    return {"power_inertia_examples": 2, "nonreal_spectrum_example": True,
            "positive_pair_pencil_checks": pair_checks,
            "minimal_naimark_dimension": 3}


def main():
    data = ROOT / "data"
    data.mkdir(exist_ok=True)
    enumerated = Counter({(0, 0): 1})
    profile_counts = {}
    tested = 0
    for n in range(1, MAX_N+1):
        seen = set()
        for blocks in signed_partitions(n):
            profile = matrix_profile(n, blocks)
            p, q = profile[0]
            assert p+q == n
            r = [a+b for a, b in profile]
            d = [a-b for a, b in profile]
            for m in range(1, n+1):
                total = r[m-1]-2*r[m]+r[m+1]
                difference = d[m-1]-d[m+1]
                assert total >= 0 and abs(difference) <= total
                assert (total-difference) % 2 == 0
                recovered = ((total+difference)//2, (total-difference)//2)
                assert recovered == blocks.get(m, (0, 0))
            key = tuple(profile)
            assert key not in seen
            seen.add(key)
            enumerated[p, q] += 1
            tested += 1
        profile_counts[n] = len(seen)
    assert dict(enumerated) == product_coefficients(MAX_N)
    assert enumerated[1, 1] == 3 and enumerated[2, 1] == 4 and enumerated[2, 2] == 10
    for (p, q), value in enumerated.items():
        assert value == enumerated[q, p]
        if min(p, q) == 0:
            assert value == 1
    with (data / "branch_counts.csv").open("w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(["dimension", "positive_index", "negative_index",
                         "complete_theories_with_one_nilpotent"])
        for n in range(MAX_N+1):
            for p in range(n+1):
                writer.writerow([n, p, n-p, enumerated[p, n-p]])
    lines = [r"\begin{table}[htbp]", r"\centering", r"\begin{tabular}{rll}",
             r"\toprule", r"\(n\) & \(b_{p,n-p}\) for \(0\leq p\leq\lfloor n/2\rfloor\) & \(\sum_{p=0}^n b_{p,n-p}\) \\",
             r"\midrule"]
    for n in range(1, 9):
        values = ", ".join(str(enumerated[p, n-p]) for p in range(n//2+1))
        total = sum(enumerated[p, n-p] for p in range(n+1))
        lines.append(f"{n} & {values} & {total} " + r"\\")
    lines.extend([r"\bottomrule", r"\end{tabular}",
                  r"\caption{Numbers of complete theories with one named self-adjoint nilpotent operator. The omitted signatures follow from \(b_{p,q}=b_{q,p}\).}",
                  r"\label{tab:branch-counts}", r"\end{table}", ""])
    (data / "branch_table.tex").write_text("\n".join(lines))
    report = {"arithmetic": "exact rational congruence elimination",
              "maximum_enumerated_dimension": MAX_N,
              "signed_profiles_verified": tested,
              "profiles_by_dimension": profile_counts,
              "independent_generating_function_check": True,
              "block_power_inertias": "computed from rational matrices by congruence elimination",
              "examples": verify_examples(), "result": "all checks passed"}
    (data / "verification.json").write_text(json.dumps(report, indent=2)+"\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()