"""Interval forward-mode AD (value + 6 partials, all Arb balls) for the degree-4 gradient formula."""
from flint import arb
import itertools
NV=6
PI=arb.pi()
INF=arb(0,1e30)
class D:
    __slots__=("v","d")
    def __init__(s,v,d=None): s.v=v if isinstance(v,arb) else arb(v); s.d=d if d is not None else [arb(0)]*NV
    def _c(o): return o if isinstance(o,D) else D(o)
    def __add__(a,b): b=D._c(b); return D(a.v+b.v,[x+y for x,y in zip(a.d,b.d)])
    __radd__=__add__
    def __sub__(a,b): b=D._c(b); return D(a.v-b.v,[x-y for x,y in zip(a.d,b.d)])
    def __rsub__(a,b): return D._c(b)-a
    def __neg__(a): return D(-a.v,[-x for x in a.d])
    def __mul__(a,b):
        b=D._c(b); return D(a.v*b.v,[x*b.v+a.v*y for x,y in zip(a.d,b.d)])
    __rmul__=__mul__
    def __truediv__(a,b):
        b=D._c(b); q=a.v/b.v; return D(q,[(x-q*y)/b.v for x,y in zip(a.d,b.d)])
    def __rtruediv__(a,b): return D._c(b)/a
    def sqrt(a):
        r=a.v.sqrt(); return D(r,[x/(2*r) for x in a.d])
def asinD(x):
    if not x.v.is_finite():            # nan or unbounded
        return D(arb(0).union(PI/2).union(-PI/2),[INF]*NV)
    lo=x.v.lower(); hi=x.v.upper()     # exact endpoints, rounded outward
    l=lo if lo>-1 else arb(-1); h=hi if hi<1 else arb(1)   # the true value lies in [-1,1]
    if l>h: return D(arb(0).union(PI/2).union(-PI/2),[INF]*NV)
    val=l.asin().union(h.asin())
    if lo>-1 and hi<1:
        dv=1/(1-x.v*x.v).sqrt(); return D(val,[dv*y for y in x.d])
    return D(val,[INF]*NV)
HALF=D(arb(0.25,0.25),[INF]*NV)
def cond_orth(R,i,j):
    n=len(R); K=[k for k in range(n) if k not in (i,j)]; rij=R[i][j]; det=1-rij*rij
    if not det.v>0: return HALF
    def cov(a,b): return R[a][b]-(R[a][i]*R[b][i]+R[a][j]*R[b][j]-rij*(R[a][i]*R[b][j]+R[a][j]*R[b][i]))/det
    S=[[cov(a,b) for b in K] for a in K]
    if not all(S[t][t].v>0 for t in range(len(K))): return HALF
    pc=lambda x,y: S[x][y]/(S[x][x]*S[y][y]).sqrt()
    if len(K)==2: return 0.25+asinD(pc(0,1))*(1/(2*PI))
    return 0.125+(asinD(pc(0,1))+asinD(pc(0,2))+asinD(pc(1,2)))*(1/(4*PI))
PAIRS=list(itertools.combinations(range(4),2))
def grad_formula(rsa,eq,cD):
    """cD: list of 6 D inputs. returns list of 6 D (gradient components), or None entries for sqrt-type coords"""
    s=[(1-r*r).sqrt() for r in rsa]
    R=[[D(1)]*5 for _ in range(5)]
    for i in range(4): R[0][i+1]=R[i+1][0]=D(rsa[i])
    for (p,q),v in zip(PAIRS,cD): R[p+1][q+1]=R[q+1][p+1]=v*(s[p]*s[q])+rsa[p]*rsa[q]
    Rp=[row[:] for row in R]
    for i in range(1,5): Rp[0][i]=-R[0][i]; Rp[i][0]=-R[i][0]
    out=[]
    for m,(p,q) in enumerate(PAIRS):
        P,Q=p+1,q+1
        d=-cond_orth(R,P,Q)-cond_orth(Rp,P,Q)*3
        for T in itertools.combinations(range(1,5),3):
            if P in T and Q in T:
                idx=[0]+list(T); sub=[[Rp[x][y] for y in idx] for x in idx]; d=d+cond_orth(sub,idx.index(P),idx.index(Q))
        S=R[P][Q]
        out.append(d*(s[p]*s[q])/(((1-S)*(1+S)).sqrt()*(2*PI)))
    return out
def grad_box_centered(rsa,eq,lo,hi):
    """centered-form enclosure of grad G over the box; also returns plain enclosure; intersection used."""
    cen=[(lo[m]+hi[m])/2 for m in range(NV)]; rad=[(hi[m]-lo[m])/2 for m in range(NV)]
    # point evaluation (value only matters)
    gp=grad_formula(rsa,eq,[D(arb(cen[m])) for m in range(NV)])
    # box evaluation with derivatives (Hessian enclosure)
    unit=lambda m:[arb(1) if k==m else arb(0) for k in range(NV)]
    gb=grad_formula(rsa,eq,[D(arb(cen[m],rad[m]),unit(m)) for m in range(NV)])
    res=[]
    for m in range(NV):
        plain=gb[m].v
        if all(h.is_finite() for h in gb[m].d) and gp[m].v.is_finite():
            cf=gp[m].v+sum((h*arb(0,rad[k]) for k,h in enumerate(gb[m].d)),arb(0))
            # intersect
            lo_=max(float(cf.lower()),float(plain.lower())) if plain.is_finite() else float(cf.lower())
            hi_=min(float(cf.upper()),float(plain.upper())) if plain.is_finite() else float(cf.upper())
            res.append(arb((lo_+hi_)/2,(hi_-lo_)/2*(1+1e-12)+1e-300))
        else:
            res.append(plain)
    return res