"""Binder cumulants of finite Ising rectangles -- numberdb.org/T594.

Run with SageMath and the numberdb package:
    sage -pip install numberdb
    sage -python generate.py                  # verify every stored entry
    sage -python generate.py --fill           # fill the existing draft

The --fill action requires NUMBERDB_API_KEY. It sends entries and attachments;
it does not offer, review, or make the table public. In the campaign wrapper,
NUMBERDB_KEY_FROM_STDIN=1 reads the key and NUMBERDB_PUBLISH=1 selects --fill.

All 2 <= m <= n <= 6 are enumerated lexicographically. See argument-selection.md.
The spin cut counts are exact integers. Packed polynomial coefficients use
base 2^(N+1), exceeding the total 2^N spin assignments, so addition never carries
between coefficients. The exponent index is k + (N+1)*d: k counts down spins
and d counts disagreement bonds. Row transfer adds each bond exactly once.

U is computed exactly in Q(sqrt(2)). A rational result is returned exactly;
otherwise its two rational coefficients are enclosed using real ball arithmetic
at bits(digits, losing=96). The guard is measured in checks.py; finite values
must retain at least bits(digits) bits before serialization.
"""
import os
import sys
from functools import lru_cache
from pathlib import Path

import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.number_field.number_field import QuadraticField
from sage.rings.real_arb import RealBallField

WORKING_GUARD = 96
FIELD = QuadraticField(2, 'a')
A = FIELD.gen()
BC = A - 1


@lru_cache(None)
def cut_counts(m, n):
    """Return {(down_spin_count, cut_bonds): multiplicity} for P_m square P_n."""
    m, n = int(m), int(n)
    if m < 1 or n < 1:
        raise ValueError('Side lengths must be positive integers')
    sites = m * n
    stride = sites + 1
    bits = sites + 1
    mask = (1 << bits) - 1
    states = range(1 << m)
    downs = [s.bit_count() for s in states]
    internal = [sum(((s >> i) ^ (s >> (i+1))) & 1
                    for i in range(m-1)) for s in states]
    previous = [1 << (bits * (downs[s] + stride * internal[s])) for s in states]
    for _ in range(1, n):
        current = []
        for s in states:
            packed = 0
            for t in states:
                shift = bits * (downs[s] + stride * (internal[s] + (s ^ t).bit_count()))
                packed += previous[t] << shift
            current.append(packed)
        previous = current
    packed = sum(previous)
    counts = {}
    index = 0
    while packed:
        coefficient = packed & mask
        if coefficient:
            counts[index % stride, index // stride] = coefficient
        packed >>= bits
        index += 1
    return counts


def moment_sums(m, n, bond_weight=BC):
    """Unnormalized sums of M^0, M^2, M^4 with weight b^d."""
    sites = int(m) * int(n)
    edges = 2*sites-int(m)-int(n)
    powers = [bond_weight**d for d in range(edges+1)]
    return tuple(sum(count * (sites-2*k)**r * powers[d]
                     for (k, d), count in cut_counts(int(m), int(n)).items())
                 for r in (0, 2, 4))


@lru_cache(None)
def exact_cumulant(m, n):
    s0, s2, s4 = moment_sums(m, n)
    return FIELD(1) - s0*s4/(3*s2*s2)


def enclose_quadratic(value, digits):
    p, q = QQ(value[0]), QQ(value[1])
    if q == 0:
        return p
    balls = RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))
    result = balls(p) + balls(q)*balls(2).sqrt()
    assert result.is_finite()
    assert result.accuracy() >= numberdb.bits(digits)
    return result


class BinderCumulants(numberdb.Generator):
    table = 'T594'
    parameters = ('m', 'n')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'argument-selection.md',
             'numerical-evidence.json', 'cut-counts.json', 'skill-evidence.json')

    def enumerate(self, max_side=6):
        for m in range(2, int(max_side)+1):
            for n in range(m, int(max_side)+1):
                yield {'m': m, 'n': n}

    def value(self, params, digits):
        m, n = int(params['m']), int(params['n'])
        if not 2 <= m <= n:
            raise ValueError('The table requires 2 <= m <= n')
        return enclose_quadratic(exact_cumulant(m, n), digits)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        key = sys.stdin.read().strip()
        if '=' in key and key.split('=', 1)[0].replace('_', '').isalnum():
            key = key.split('=', 1)[1].strip().strip('\"\'')
        os.environ['NUMBERDB_API_KEY'] = key
    generator = BinderCumulants()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Compute critical finite-rectangle Binder cumulants from exact spin counts'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)