"""Rigorous affine arithmetic over Arb, for gradient enclosures of Psi4 over a box.

A form x = c + sum_i a_i e_i + [-err, err] with noise symbols e_i in [-1,1] (one per box coordinate)
represents, for every e in [-1,1]^n, a set containing the true value at the point with those noise
values. Coefficients are Arb balls; err is a nonnegative Arb upper bound. Nonlinear operations use
linear approximations whose error is bounded rigorously (convexity / AM-GM), so the represented
set always contains the true value.
"""
from flint import arb
import itertools
PI=arb.pi()
UNIT=arb(0,1)
class AA:
    __slots__=("c","a","e")
    def __init__(self,c,a,e=None):
        self.c=c if isinstance(c,arb) else arb(c); self.a=a; self.e=e if e is not None else arb(0)
    @staticmethod
    def const(v,n): return AA(v,[arb(0)]*n)
    def rad(self):  # upper bound of |x - c| (as arb, exact upper)
        return (sum((abs(t) for t in self.a),arb(0))+self.e).upper()
    def interval(self):
        return self.c+sum((t*UNIT for t in self.a),arb(0))+self.e*UNIT
    def __add__(self,o):
        if not isinstance(o,AA): return AA(self.c+o,list(self.a),self.e)
        return AA(self.c+o.c,[x+y for x,y in zip(self.a,o.a)],self.e+o.e)
    __radd__=__add__
    def __neg__(self): return AA(-self.c,[-x for x in self.a],self.e)
    def __sub__(self,o): return self+(-o)
    def __rsub__(self,o): return (-self)+o
    def __mul__(self,o):
        if not isinstance(o,AA): return AA(self.c*o,[x*o for x in self.a],(self.e*abs(o)).upper())
        lin=[self.c*y+o.c*x for x,y in zip(self.a,o.a)]
        ra=sum((abs(t) for t in self.a),arb(0)); rb=sum((abs(t) for t in o.a),arb(0))
        err=abs(self.c)*o.e+abs(o.c)*self.e+(ra+self.e)*(rb+o.e)
        return AA(self.c*o.c,lin,err.upper())
    __rmul__=__mul__
    def _affine(self,alpha,zeta,delta):
        # result = alpha*x + zeta + [-delta,delta], alpha exact (point arb), zeta ball, delta >=0
        r=self*alpha
        return AA(r.c+zeta,r.a,(r.e+delta).upper())
def recip(x):
    """1/x for x>0 on its range; None if the range is not certified positive."""
    I=x.interval(); l=I.lower(); u=I.upper()
    if not l>0: return None
    alpha=arb((-1/(u*u)).mid())            # exact point, < 0
    g=lambda t: 1/t-alpha*t                 # convex on t>0
    hi=g(l).union(g(u)).upper()            # max of a convex function at the endpoints
    lo=(2*(-alpha).sqrt()).lower()         # AM-GM: 1/t - alpha t >= 2 sqrt(-alpha)
    zeta=(arb(lo)+arb(hi))/2; delta=((arb(hi)-arb(lo))/2).upper()
    return x._affine(alpha,zeta,delta+zeta.rad())
def sqrt(x):
    """sqrt(x) for x>=0 on its range; None if the range is not certified nonnegative."""
    I=x.interval(); l=I.lower(); u=I.upper()
    if not l>=0: return None
    if not u>l: return AA(I.sqrt(),[arb(0)]*len(x.a))
    alpha=arb(((u.sqrt()-l.sqrt())/(u-l)).mid())
    if not alpha>0: return None
    g=lambda t: t.sqrt()-alpha*t            # concave on t>=0
    lo=g(l).union(g(u)).lower()            # min of a concave function at the endpoints
    hi=(1/(4*alpha)).upper()               # sqrt(t) - alpha t <= 1/(4 alpha)
    zeta=(arb(lo)+arb(hi))/2; delta=((arb(hi)-arb(lo))/2).upper()
    return x._affine(alpha,zeta,delta+zeta.rad())
def asin_interval(I):
    if not I.is_finite(): return arb(0).union(PI/2).union(-PI/2)
    l=I.lower(); u=I.upper()
    l=l if l>-1 else arb(-1); u=u if u<1 else arb(1)
    if l>u: return arb(0).union(PI/2).union(-PI/2)
    return l.asin().union(u.asin())

PAIRS=list(itertools.combinations(range(4),2)); COMP={p:tuple(k for k in range(4) if k not in p) for p in PAIRS}
def grad_aa(rsa,lo,hi,only=None):
    """Enclosures (arb) of dPsi4/dc_pq over the box [lo,hi] intersected with the interior of the
    elliptope, or None per coordinate if the affine forms cannot certify the formula."""
    n=6
    ball=[arb(lo[m]).union(arb(hi[m])) for m in range(n)]
    cen=[arb(ball[m].mid()) for m in range(n)]
    rad=[(ball[m]-cen[m]).upper()*1 for m in range(n)]
    # c_m = cen_m + rad'_m e_m with rad' an exact upper bound of the half-width
    rd=[abs(ball[m]-cen[m]).upper() for m in range(n)]
    c=[AA(cen[m],[rd[m] if k==m else arb(0) for k in range(n)]) for m in range(n)]
    s=[(1-r*r).sqrt() for r in rsa]
    S=[[None]*4 for _ in range(4)]
    for i in range(4): S[i][i]=AA.const(1,n)
    for (p,q),cm in zip(PAIRS,c): S[p][q]=S[q][p]=cm*(s[p]*s[q])+rsa[p]*rsa[q]
    out=[]
    for mm,(p,q) in enumerate(PAIRS):
        if only is not None and mm not in only: out.append(None); continue
        k,l=COMP[(p,q)]; rpq=S[p][q]
        det=1-rpq*rpq
        idet=recip(det)
        if idet is None: out.append(None); continue
        def cov(a,b):
            return S[a][b]-(S[a][p]*S[b][p]+S[a][q]*S[b][q]-rpq*(S[a][p]*S[b][q]+S[a][q]*S[b][p]))*idet
        skk=cov(k,k); sll=cov(l,l); skl=cov(k,l)
        prod=skk*sll; sq=sqrt(prod)
        if sq is None: out.append(None); continue
        isq=recip(sq)
        if isq is None: out.append(None); continue
        pc=(skl*isq).interval()
        A=asin_interval(pc)
        sd=sqrt(det)
        if sd is None: out.append(None); continue
        isd=recip(sd)
        if isd is None: out.append(None); continue
        g=A*isd.interval()*(-(s[p]*s[q])/(2*PI*PI))
        out.append(g if g.is_finite() else None)
    return out