"""Reproduce the manuscript's finite scalar tables with outward integer arithmetic.

This independently transcribes Appendix B of OpenAI's cubic-torus manuscript
at commit adc7f1241b42e322a6451854ab7e4b4c146bf78a. It checks the arithmetic
tables, not the geometric reduction or the analytic interpolation bounds.
Run with Python assertions enabled. JSON is written to standard output.
"""

from fractions import Fraction as F
from math import comb, isqrt
import hashlib
import json
from pathlib import Path

N = 10**40


def ceil_div(a, b):
    return -((-a) // b)


class I:
    def __init__(self, value=0, *, bounds=None):
        if bounds is not None:
            self.lo, self.hi = bounds
        else:
            q = F(value)
            self.lo = (N * q.numerator) // q.denominator
            self.hi = ceil_div(N * q.numerator, q.denominator)
        assert self.lo <= self.hi

    def __add__(self, other):
        other = interval(other)
        return I(bounds=(self.lo + other.lo, self.hi + other.hi))

    __radd__ = __add__

    def __neg__(self):
        return I(bounds=(-self.hi, -self.lo))

    def __sub__(self, other):
        return self + (-interval(other))

    def __rsub__(self, other):
        return interval(other) - self

    def __mul__(self, other):
        other = interval(other)
        products = [a * b for a in (self.lo, self.hi)
                    for b in (other.lo, other.hi)]
        return I(bounds=(min(products) // N, ceil_div(max(products), N)))

    __rmul__ = __mul__

    def reciprocal(self):
        if self.hi < 0:
            return -(-self).reciprocal()
        assert self.lo > 0
        return I(bounds=(N * N // self.hi, ceil_div(N * N, self.lo)))

    def __truediv__(self, other):
        return self * interval(other).reciprocal()

    def __rtruediv__(self, other):
        return interval(other) / self

    def __pow__(self, power):
        assert isinstance(power, int) and power >= 0
        if power == 0:
            return I(1)
        values = [self.lo**power, self.hi**power]
        if power % 2 == 0 and self.lo <= 0 <= self.hi:
            values.append(0)
        scale = N ** (power - 1)
        return I(bounds=(min(values) // scale, ceil_div(max(values), scale)))

    def sqrt(self):
        assert self.lo >= 0
        low = isqrt(self.lo * N)
        high = isqrt(self.hi * N)
        high += high * high < self.hi * N
        return I(bounds=(low, high))


def interval(x):
    return x if isinstance(x, I) else I(x)


def minimum(a, b):
    a, b = interval(a), interval(b)
    return I(bounds=(min(a.lo, b.lo), min(a.hi, b.hi)))


def strict_between(value, low, high):
    assert F(value.lo, N) > F(low)
    assert F(value.hi, N) < F(high)


def decimal(integer):
    sign = '-' if integer < 0 else ''
    integer = abs(integer)
    return f'{sign}{integer // N}.{integer % N:040d}'


def enclosure(value):
    return [decimal(value.lo), decimal(value.hi)]


PI_LO, PI_HI = F('3.14159265358979'), F('3.14159265358980')
pi = I(bounds=(I(PI_LO).lo, I(PI_HI).hi))
lam, ell, dm = I('.98'), I('.28'), I('.21')
D = I(48).sqrt() / 7
a0, b0 = 1 / pi, 1 - 1 / pi


def atan_bounds(z, degree):
    assert degree % 2 == 1
    partial = sum((-1)**j * z**(2*j+1) / (2*j+1)
                  for j in range(degree + 1))
    return partial, partial + z**(2*degree+3) / (2*degree+3)


def asin(z):
    def half(x):
        return x / (2 + 2 * (1 - x*x).sqrt()).sqrt()
    j = half(half(z))
    assert F(j.hi, N) < F('.36')
    series = sum((I(F(comb(2*k, k), 4**k * (2*k+1))) * j**(2*k+1)
                  for k in range(11)), I(0))
    low = 4 * series
    high = low + 4 * j**23 / (1 - j*j)
    return I(bounds=(low.lo, high.hi))


def moments(z):
    beta = (1 - z*z).sqrt()
    L = [asin(z), 1 - beta]
    for n in range(2, 5):
        L.append(((n-1)*L[n-2] - z**(n-1)*beta) / n)
    J = []
    for j in range(4):
        boundary = z**j / beta - (1 if j == 0 else 0)
        if j:
            boundary -= j * L[j-1]
        J.append(lam*L[j+1] - (1-lam)*boundary)
    return J


def values(h, use_gain):
    t = 2 / (1 + (1 + 4*h*h/pi).sqrt())
    x, hb = h*a0, h*b0
    assert hb.hi < D.lo or hb.lo > D.hi
    y = minimum(D, hb)
    Jx, Jy = moments(x), moments(y)
    base = lam*h/pi + ell*t - 1 + pi/(h*h)*Jx[2] + (Jy[1]-Jx[1])/h
    if hb.hi < D.lo:
        k = (b0 + D/h) / 2
        sigma = pi/4
        v = 1 + sigma*k*k
        root = v.sqrt()
        r0 = pi.sqrt()*k*root - pi*k*k/2
        r1 = pi.sqrt()*(1+2*sigma*k*k)/root - pi*k
        r2 = (pi.sqrt()*sigma*k*(3+2*sigma*k*k)/(v*root)-pi)/2
        r3 = pi.sqrt()*sigma/(2*v*v*root)
        coefficients = [r0-k*r1+k*k*r2-k**3*r3,
                        r1-2*k*r2+3*k*k*r3, r2-3*k*r3, r3]
        JD = moments(D)
        base += sum((coefficients[j]/h**j*(JD[j]-Jy[j])
                     for j in range(4)), I(0))
    gain = None
    if use_gain:
        u = minimum(b0, I('.96')/h)
        P = u - 1 + 2*h*t/pi
        Q = 2/pi * (pi + 2*h*h*t).sqrt()
        C = minimum(I('.134')*h**3, I('.36'))
        q = 2*P/(1 + (1+6*C*Q*Q*P/ell).sqrt())
        gain = minimum(C*dm**3, ell*((P-q)/Q)**2+C*q**3)
    return base, gain


KNOTS = [
    ('.64',1744,None), ('.74',1200,None), ('.84',735,None),
    ('.92',432,None), ('.98',247,None), ('1.02',144,76),
    ('1.06',57,98), ('1.10',-15,123), ('1.14',-74,151),
    ('1.18',-119,182), ('1.22',-153,215), ('1.26',-178,247),
    ('1.30',-196,271), ('1.34',-211,297), ('1.38',-226,325),
    ('1.42',-247,332), ('1.46',-275,332), ('1.50',-260,318),
    ('1.54',-193,293), ('1.58',-76,271), ('1.62',86,251),
    ('1.70',536,None), ('1.90',2275,None), ('2.25',6853,None),
]
GAPS = [4416,2298,1859,1313,967,1037,581,280,145,137,137,
        209,269,379,369,90,90,100,520,1470,250,1899,14266]


def main():
    assert __debug__, 'Run with assertions enabled'
    a, b = atan_bounds(F(1, 5), 31), atan_bounds(F(1, 239), 9)
    assert PI_LO < 16*a[0]-4*b[1] < 16*a[1]-4*b[0] < PI_HI
    knots = []
    for h, b_lower, s_lower in KNOTS:
        base, gain = values(I(h), s_lower is not None)
        strict_between(10**5*base-b_lower, '1.10', '1.95')
        row = {'h': h, 'B_tilde': enclosure(base), 'b_lower': b_lower}
        if gain is not None:
            strict_between(10**5*gain-s_lower, '1.00', '1.66')
            row.update(S=enclosure(gain), s_lower=s_lower)
        if h == '1.46':
            strict_between(10**5*base, '-273.170', '-273.169')
            assert gain.lo == gain.hi == I('.00333396').lo
        knots.append(row)
    D_plus = F('.989743318611')
    assert 49*D_plus**2 > 48
    b_minus = 1-1/PI_LO

    def error(h):
        return F('.008')*max(D_plus/h-b_minus, 0)**4

    gaps = []
    for left, right, stated in zip(KNOTS, KNOTS[1:], GAPS):
        a, b = F(left[0]), F(right[0])
        interior = a >= F('1.02') and b <= F('1.62')
        K = F('2.4') if interior else 2/a**2
        margin = min(F(row[1]+(row[2] if interior else 0), 10**5)-error(F(row[0]))
                     for row in (left, right))-K*(b-a)**2/8
        assert margin > 0 and 10**6*margin >= stated
        gaps.append({'interval':[left[0],right[0]], 'exact_margin':str(margin),
                     'certified_lower_bound':str(F(stated, 10**6))})
    assert min(F(g['certified_lower_bound']) for g in gaps) == F(9, 100000)
    report = {
        'family':'354', 'source_commit':'adc7f1241b42e322a6451854ab7e4b4c146bf78a',
        'implementation':'Independent transcription of the manuscript arithmetic prescription',
        'script_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        'scale':str(N), 'assertions_enabled':__debug__, 'all_checks_pass':True,
        'scope':'Finite arithmetic enclosures and interpolation margins only; geometric reduction and analytic derivative bounds are not verified by this program.',
        'pi_enclosure':[str(PI_LO),str(PI_HI)], 'knots':knots, 'intervals':gaps,
    }
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()