"""Hilbert–Schmidt volumes of quantum states of fixed rank -- numberdb.org/T571.

Run with SageMath (install numberdb once with sage -pip install numberdb):
    sage -python generate.py                  # compare with the stored draft
    sage -python generate.py --fill           # write entries, never publish
    sage -python checks.py                    # independent private checks

The metric is Tr(d rho^2). The stratum includes its support subspace.
Source: Zyczkowski and Sommers, quant-ph/0302197v3, equations (4.3),
(5.3), (7.8), (3.10), and (A27)-(A29). Nullity in that paper is N-r.
Gamma factors are reduced to exact rationals and powers of sqrt(pi).
Only the final rational * sqrt(integer) * pi^integer uses Arb arithmetic.
A 64-bit guard is measured by checks.py, not treated as an error bound.
The 270 entries cover both fields, 2 <= N <= 16, and every 1 <= r <= N.
See argument-selection.md for the discrete families and finite cutoff.
"""
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from functools import lru_cache
from math import factorial
import os
import sys

WORKING_GUARD = 64


def gamma_half(twice_argument):
    """Gamma(t/2) = q*pi^(h/2), returning exact (q,h)."""
    t = int(twice_argument)
    assert t > 0
    if t % 2 == 0:
        return QQ(factorial(t // 2 - 1)), 0
    m = (t - 1) // 2
    return QQ(factorial(2*m)) / QQ(4**m * factorial(m)), 1


@lru_cache(None)
def exact_factors(field, N, r):
    """Return (q,s,p) such that V = q*sqrt(s)*pi^p."""
    assert field in ('real', 'complex') and N >= 2 and 1 <= r <= N
    beta = 1 if field == 'real' else 2
    d = N - r
    alpha = 1 + beta*d
    A = alpha*r + beta*r*(r-1)//2
    coefficient = QQ(1) / QQ(factorial(A-1)*factorial(r))
    half_pi = 0
    den_q, den_h = gamma_half(2+beta)
    for j in range(1, r+1):
        for t in (2+beta*j, 2*alpha+beta*(j-1)):
            q,h = gamma_half(t)
            coefficient *= q
            half_pi += h
        coefficient /= den_q
        half_pi -= den_h
    assert half_pi == 0   # the eigenvalue integral is rational for beta=1,2
    m = r*(2*N-r-1)//2
    if beta == 2:
        coefficient *= 2**m
        for k in range(d+1,N+1):
            coefficient /= QQ(factorial(k-1))
        radicand, pi_power = r, m
    else:
        coefficient *= 2**(m//2)
        radicand = r * 2**(m%2)
        for k in range(d+1,N+1):
            q,h = gamma_half(k)
            coefficient /= q
            half_pi += k-h
        assert half_pi % 2 == 0
        pi_power = half_pi//2
    assert pi_power > 0  # nonzero algebraic multiple of a positive power of pi
    return coefficient, radicand, pi_power


class HilbertSchmidtVolumes(numberdb.Generator):
    table = 'T571'
    parameters = ('field','N','rank')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py','checks.py','argument-selection.md','check.py')

    def enumerate(self, max_N=16):
        for field in ('real','complex'):
            for N in range(2,max_N+1):
                for r in range(1,N+1):
                    yield {'field':field,'N':N,'rank':r}

    def value(self, params, digits):
        q,s,p = exact_factors(str(params['field']),int(params['N']),int(params['rank']))
        R = RealBallField(numberdb.bits(digits,losing=WORKING_GUARD))
        value = R(q)*R(s).sqrt()*R.pi()**p
        assert value.is_finite() and value > 0
        return value


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    g = HilbertSchmidtVolumes()
    if '--fill' in sys.argv:
        print(g.publish(message='Fill 270 independently checked intrinsic rank-stratum volumes.'))
    else:
        result = g.verify(sample=None)
        print(result)
        sys.exit(0 if result.ok else 1)