"""Affine-arithmetic enclosure of grad_theta G over a theta-box (degree 3), same formula as grad.py.
Uses the rigorous affine forms of aa.py (Arb coefficients, bounded nonlinear remainders), plus
cos/sin of an affine form x = c + r e via Taylor: f(c+d) = f(c) + f'(c) d + R, |R| <= d^2/2 for f = cos, sin."""
from flint import arb
from aa import AA, recip, sqrt, asin_interval, UNIT
PI=arb.pi()
IDX=((0,1),(0,2),(1,2))
def _trig(cen,rad,m,n,fn):
    # fn='cos' or 'sin' of the single-noise form cen + rad e_m
    c=arb(cen); r=arb(rad)
    if fn=="cos": f0=c.cos(); f1=-c.sin()
    else: f0=c.sin(); f1=c.cos()
    return AA(f0,[f1*r if k==m else arb(0) for k in range(n)],(r*r/2).upper())
def grad_box_aa(rs,lo,hi):
    n=3
    ball=[arb(lo[m]).union(arb(hi[m])) for m in range(n)]
    cen=[ball[m].mid() for m in range(n)]
    rad=[abs(ball[m]-cen[m]).upper() for m in range(n)]
    s=[(1-r*r).sqrt() for r in rs]
    C={}
    for m,(i,j) in enumerate(IDX): C[(i,j)]=C[(j,i)]=_trig(cen[m],rad[m],m,n,"cos")
    def mk(sign):
        R=[[None]*4 for _ in range(4)]
        for a in range(4): R[a][a]=AA.const(1,n)
        for a in range(3): R[0][a+1]=R[a+1][0]=AA.const(sign*rs[a],n)
        for a in range(3):
            for b in range(3):
                if a!=b: R[a+1][b+1]=C[(a,b)]*(s[a]*s[b])+rs[a]*rs[b]
        return R
    R=mk(1); Rp=mk(-1)
    def cond(Rm,i,j,k,l):
        rij=Rm[i][j]; idet=recip(1-rij*rij)
        if idet is None: return None
        def cov(a,b): return Rm[a][b]-(Rm[a][i]*Rm[b][i]+Rm[a][j]*Rm[b][j]-rij*(Rm[a][i]*Rm[b][j]+Rm[a][j]*Rm[b][i]))*idet
        skk=cov(k,k); sll=cov(l,l); skl=cov(k,l)
        q=sqrt(skk*sll)
        if q is None: return None
        iq=recip(q)
        if iq is None: return None
        return arb(1)/4+asin_interval((skl*iq).interval())/(2*PI)
    g=[]
    for m,(a,b) in enumerate(IDX):
        i,j=a+1,b+1; k,l=[x for x in range(4) if x not in (i,j)]
        th=ball[m]; Rij=R[i][j].interval()
        if abs(float((rs[a]-rs[b]).mid()))<1e-30:
            fac=s[a]*(1+th.cos()).sqrt()/(2*PI*(1+Rij).sqrt())
        else:
            fac=s[a]*s[b]*th.sin()/(2*PI*(1-Rij*Rij).sqrt())
        c=cond(R,i,j,k,l); cp=cond(Rp,i,j,k,l)
        if c is None or cp is None or not fac.is_finite(): g.append(arb("nan")); continue
        g.append(-fac*(cp-c))
    return g