"""Independent checks for convex-position probabilities -- numberdb.org/T468.

Run: sage -python checks.py
Run against the draft/table: NUMBERDB_API_KEY=... sage -python checks.py --stored
The campaign wrapper accepts NUMBERDB_KEY_FROM_STDIN=1 for the second command.
Keep generate.py and check.py beside this file. The simulation is a separate,
finite-sample diagnostic, not the justification for the 100 stored digits.

Disk integration uses Theorem 1(i') and L0,...,L5 from Marckert's preprint,
sharing no polynomial-in-pi coefficients with generate.py. It independently
checks n=3,...,7 at high precision. At n=8 the high-precision cross-check uses
the same published formula rearranged; only the simulation independently
checks its geometric meaning. The disk four-point row is also checked by a
Blaschke-Petkantschin reduction of the mean area of a random triangle, with the
normalization of the line-coordinate change checked first.
"""
import itertools
import json
import math
import os
import sys
from pathlib import Path

import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
import mpmath as mp
import numpy as np

from generate import ConvexPosition, REGIONS, N_MIN, N_MAX
from check import exactness, measure, agrees_with, names_its_rings


def plain(value):
    return value['number'] if isinstance(value, dict) else value


def dyck_count(k):
    """Enumerate balanced walks that never go negative, without binomials."""
    count = 0
    for ups in itertools.combinations(range(2*k), k):
        chosen = set(ups)
        height = 0
        for j in range(2*k):
            height += 1 if j in chosen else -1
            if height < 0:
                break
        else:
            count += 1
    return count


def factorial_probability(key):
    region, n = key
    if region == 'equilateral-triangle':
        return QQ(2**n*math.factorial(3*n-3))/QQ(math.factorial(n-1)**3*math.factorial(2*n))
    return (QQ(math.comb(2*n-2,n-1))/QQ(math.factorial(n)))**2


def disk_formula(n):
    """Independent mpmath evaluation in the source's complement form."""
    t = mp.pi**2
    formulas = {3: mp.mpf(1), 4: 1-35/(12*t), 5: 1-305/(48*t),
                6: 1-(146400*t-473473)/(11520*t**2),
                7: 1-(512400*t-2900611)/(23040*t**2),
                8: 1-(62664108221+1721664000*t**2-18670881600*t)/(48384000*t**3)}
    return formulas[n]


def segment_L(k, t):
    """Marckert equation (5) and the displayed functions on preprint page 4."""
    s = mp.sin(t/2)
    z = mp.sin(t)
    if k == 0:
        return s
    if k == 1:
        return s*(t-z)
    if k == 2:
        return s*(3*t*t+z*z-16*s*s)/6
    if k == 3:
        return s*(2*s**4*z+9*t**3+27*s*s*t+7*s*s*z+105*(z-t))/54
    if k == 4:
        return s*(160*s**6+48*s**8+60*s**4+540*t**4-13725*t*t
                  -7200*z*t+83700*s*s)/12960
    if k == 5:
        return s*(40500*z*t*t-584*s**4*z+12000*s**4*t-272*s**6*z
                  -549000*t**3-2745000*s*s*t+44270*s*s*z
                  +7102095*(t-z)-64*s**8*z+10800*t**5)/1296000
    raise ValueError(k)


def disk_integral(n):
    f = lambda t: sum(segment_L(k,t)*segment_L(n-2-k,2*mp.pi-t)
                      for k in range(n-1))
    return mp.factorial(n-2)*mp.quad(f,[0,mp.pi,2*mp.pi])/(2**(n-2)*mp.pi**(n-1))


def area_integral_check(values):
    # For an unoriented line at signed distance t from the disk center, its
    # chord length is l(t)=2 sqrt(1-t^2). The pair-coordinate Jacobian is |s1-s2|.
    # Normalization: pi integral l(t)^3/3 dt = pi^2, the product of two disk areas.
    with mp.workdps(140):
        normalization = mp.quad(lambda t: (2*mp.sqrt(1-t*t))**3/3,[-1,0,1])/mp.pi
        assert abs(normalization-1) < mp.mpf('1e-130')
    # E area / area = [integral l(t)^4 J(t) dt]/(12*pi^3),
    # J(t)=integral |u-t| 2 sqrt(1-u^2) du. Swap u,t and integrate the polynomial
    # 16(1-t^2)^2 |u-t| in t exactly. Semicircle moments are pi*C_k/4^k.
    P = PolynomialRing(QQ,'u'); u = P.gen()
    W = 16*(u-QQ(2)/3*u**3+QQ(1)/5*u**5)
    V = 16*(QQ(1)/2*u**2-QQ(1)/2*u**4+QQ(1)/6*u**6)
    inner = u*(2*W-W(-1)-W(1))-2*V+V(-1)+V(1)
    normalized_integral = sum(c*QQ(dyck_count(k//2))/4**(k//2)
                              for k,c in enumerate(inner) if k%2 == 0)
    assert all(inner[k] == 0 for k in range(1,inner.degree()+1,2))
    assert normalized_integral == QQ(35)/4
    with mp.workdps(140):
        area_ratio = mp.mpf(str(normalized_integral/12))/mp.pi**2
        expected = 1-4*area_ratio
        assert abs(expected-disk_formula(4)) < mp.mpf('1e-130')
    return {'line_coordinate_normalization': '1 to 130 decimal places',
            'integrated_polynomial': str(inner),
            'semicircle_integral_divided_by_pi': str(normalized_integral),
            'normalized_mean_triangle_area': '35/(48*pi^2)'}


def convex_flags(points):
    """A generic planar n-point sample is convex iff it has n supporting edges."""
    m,n,_ = points.shape
    edges = np.zeros(m,dtype=np.int16)
    for i in range(n):
        for j in range(i+1,n):
            dx = points[:,j,0]-points[:,i,0]
            dy = points[:,j,1]-points[:,i,1]
            positive = np.ones(m,dtype=bool)
            negative = np.ones(m,dtype=bool)
            for k in range(n):
                if k in (i,j):
                    continue
                cross = dx*(points[:,k,1]-points[:,i,1])-dy*(points[:,k,0]-points[:,i,0])
                positive &= cross > 0
                negative &= cross < 0
            edges += positive | negative
    return edges == n


def simulations(values, trials=500000):
    controls = np.array([[[0.,0.],[1.,0.],[1.,1.],[0.,1.]],
                         [[0.,0.],[2.,0.],[0.,2.],[.4,.3]]])
    assert convex_flags(controls).tolist() == [True,False]
    rng = np.random.default_rng(20261024211700)
    out = {}
    for region in REGIONS:
        counts = np.zeros(N_MAX+1,dtype=np.int64)
        area_sum = 0.; area_sq = 0.
        moment_sums = np.zeros(4); moment_sq = np.zeros(4)
        for start in range(0,trials,10000):
            m = min(10000,trials-start)
            points = rng.random((m,N_MAX,2))
            if region == 'equilateral-triangle':
                fold = points.sum(axis=2)>1
                points[fold] = 1-points[fold]
                targets = np.array([1/3,1/3,1/6,1/6])
                area = .5
            elif region == 'unit-square':
                targets = np.array([.5,.5,1/3,1/3]); area = 1.
            else:
                radius=np.sqrt(points[:,:,0]); theta=2*np.pi*points[:,:,1]
                points=np.stack((radius*np.cos(theta),radius*np.sin(theta)),axis=2)
                targets = np.array([0.,0.,.25,.25]); area=np.pi
            first = points[:,0,:]
            moments = np.column_stack((first,first**2))
            moment_sums += moments.sum(axis=0); moment_sq += (moments**2).sum(axis=0)
            p0,p1,p2=points[:,0],points[:,1],points[:,2]
            areas = np.abs((p1[:,0]-p0[:,0])*(p2[:,1]-p0[:,1])
                           -(p1[:,1]-p0[:,1])*(p2[:,0]-p0[:,0]))/(2*area)
            area_sum+=areas.sum();area_sq+=(areas**2).sum()
            for n in range(N_MIN,N_MAX+1):
                flags=convex_flags(points[:,:n,:]);counts[n]+=flags.sum()
                if start==0:
                    # An invertible shear, unequal scaling and translation.
                    aff=points[:1000,:n,:]@np.array([[2.,1.],[1.,3.]])+[5.,-2.]
                    assert np.array_equal(flags[:1000],convex_flags(aff))
        means=moment_sums/trials
        se=np.sqrt(np.maximum(0,moment_sq/trials-means**2)/trials)
        assert np.all(np.abs(means-targets)<6*se+1e-12)
        target_area={'equilateral-triangle':1/12,'unit-square':11/144,
                     'unit-disk':35/(48*np.pi**2)}[region]
        area_mean=area_sum/trials
        area_se=math.sqrt((area_sq/trials-area_mean**2)/trials)
        assert abs(area_mean-target_area)<6*area_se
        rows={}
        for n in range(N_MIN,N_MAX+1):
            v=values[(region,n)]; truth=float(v.center() if hasattr(v,'center') else v)
            observed=counts[n]/trials
            sigma=math.sqrt(truth*(1-truth)/trials)
            assert abs(observed-truth) <= 6*sigma+1/trials,(region,n,observed,truth)
            rows[str(n)]={'successes':int(counts[n]),'estimate':observed,
                          'reference':truth,'standard_errors':(observed-truth)/sigma if sigma else 0.}
        out[region]={'trials':trials,'moments':means.tolist(),'expected_moments':targets.tolist(),
                     'normalized_triangle_area':area_mean,'expected_triangle_area':target_area,
                     'probabilities':rows,'affine_samples_per_n':1000}
        print('Simulation passed:',region,flush=True)
    return out


def stored_checks(g, values):
    report=g.verify(sample=None)
    assert report.ok,report
    document=numberdb.table('T468')
    from numberdb._generate import _stored_entries
    entries=_stored_entries(document)
    assert len(entries)==18
    out={}
    from decimal import Decimal
    for identity, written in entries.items():
        region,n=identity.split(',');key=(region,int(n))
        out[key]=written
        v=values[key]
        if key[0]!='unit-disk' or key[1]==3:
            assert QQ(written)==v
        else:
            # NumberDB plain decimals represent +- one unit in their last place.
            center=Decimal(written);ulp=Decimal(1).scaleb(center.as_tuple().exponent)
            with mp.workdps(160):
                x=mp.mpf(written);rad=mp.mpf(str(ulp));ref=disk_formula(key[1])
                assert abs(ref-x)<rad
                if key[1]<=7:
                    assert abs(disk_integral(key[1])-x)<rad
    # Check recurrences using stored exact numbers, not the generator's values.
    for region in REGIONS[:2]:
        for n in range(3,8):
            if region=='unit-square':ratio=(QQ(2*(2*n-1))/QQ(n*(n+1)))**2
            else:ratio=QQ(2*(3*n)*(3*n-1)*(3*n-2))/QQ(n**3*(2*n+2)*(2*n+1))
            assert QQ(out[(region,n+1)])==QQ(out[(region,n)])*ratio
    return {'verify':str(report),'readback_entries':len(out),'document':document}


def main():
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    g=ConvexPosition();values={(p['region'],int(p['n'])):plain(g.value(p,100)) for p in g.enumerate()}
    assert not exactness(values)
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    polys={k:v for k,v in values.items() if k[0] != 'unit-disk'}
    assert not agrees_with(polys,factorial_probability)
    catalan_document = numberdb.table('T264')
    QR = PolynomialRing(QQ,'q')
    for n in range(4,9):
        polynomial = QR(catalan_document['Numbers'][str(n-1)])
        assert (QQ(polynomial(1))/math.factorial(n-1))**2 == polys[('unit-square',n)]
    for (region,n),v in polys.items():
        if region=='unit-square':
            assert v==(QQ(dyck_count(n-1))/math.factorial(n-1))**2
    for region in REGIONS:
        assert values[(region,3)]==1
        for n in range(3,8):
            assert values[(region,n)]>values[(region,n+1)]>0
    ball_report={}
    with mp.workdps(160):
        for n in range(4,9):
            ball=values[('unit-disk',n)]
            center=mp.mpf(str(ball.center()));rad=mp.mpf(str(ball.rad()))
            assert mp.isfinite(center) and 0 < rad < mp.mpf('1e-120')
            assert abs(center-disk_formula(n))<rad
            ball_report[str(n)]={'radius':str(ball.rad()),'relative_decimal_accuracy':float(-mp.log10(rad/abs(center)))}
    # The integral with n=3 is a normalization control with known answer 1.
    integrals={}
    for precision in (130,170):
        with mp.workdps(precision):
            assert abs(disk_integral(3)-1)<mp.mpf(10)**(-precision+10)
            for n in range(3,8):
                value=disk_integral(n)
                assert abs(value-disk_formula(n))<mp.mpf(10)**(-precision+10),(precision,n)
                integrals[f'{n}@{precision}']=mp.nstr(value,precision-5)
    evidence={'table':'T468','measure':measure(values),'exact_rational_entries':13,
              'ball_entries':5,'balls':ball_report,'disk_convolution_integrals':integrals,
              'disk_area_integral':area_integral_check(values)}
    print('Exact formulas, Dyck counts, balls and disk integrals passed.',flush=True)
    evidence['simulation']=simulations(values)
    if '--stored' in sys.argv or os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        evidence['stored']=stored_checks(g,values)
    print('EVIDENCE_JSON_BEGIN')
    print(json.dumps(evidence,indent=2))
    print('EVIDENCE_JSON_END')


if __name__=='__main__':
    main()