"""Rigorous (Arb) enclosure of grad_theta G over a box, via Plackett's identity (closed form)."""
from flint import arb, ctx
ctx.prec=64
PI=arb.pi()
IDX=((0,1),(0,2),(1,2))          # residual pairs (Y indices 1..3 -> 0..2)
def ball(lo,hi): return arb(lo).union(arb(hi))   # contains [lo,hi]
def asin_ball(x):
    if not x.is_finite(): return arb('nan')             # undetermined: callers fall back to the Lipschitz bound
    lo=x.lower(); hi=x.upper()                       # exact endpoints, rounded outward
    lo=lo if lo>-1 else arb(-1); hi=hi if hi<1 else arb(1)   # the true value lies in [-1,1]
    return lo.asin().union(hi.asin())
def cond_orth(R,i,j,k,l):
    # P(X_k>0,X_l>0 | X_i=X_j=0) = 1/4 + asin(pc)/(2pi), pc = partial corr of k,l given i,j
    rij=R[i][j]; det=1-rij*rij
    def cov(a,b): # conditional covariance
        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
    skk=cov(k,k); sll=cov(l,l); skl=cov(k,l)
    if not (skk>0 and sll>0): return arb(0.5,0.5)   # unknown: conditional probability in [0,1]
    pc=skl/(skk*sll).sqrt()
    return arb(1)/4+asin_ball(pc)/(2*PI)
def grad_box(rs,lo,hi):
    """rs: list of 3 arb; lo,hi: float lists. returns list of 3 arb balls for dG/dtheta_m"""
    s=[(1-r*r).sqrt() for r in rs]
    th=[ball(lo[m],hi[m]) for m in range(3)]
    C={}
    for m,(i,j) in enumerate(IDX): C[(i,j)]=C[(j,i)]=th[m].cos()
    R=[[arb(1)]*4 for _ in range(4)]
    for a in range(3): R[0][a+1]=R[a+1][0]=rs[a]
    for a in range(3):
        for b in range(3):
            if a!=b: R[a+1][b+1]=rs[a]*rs[b]+s[a]*s[b]*C[(a,b)]
    Rp=[row[:] for row in R]
    for a in range(1,4): Rp[0][a]=-R[0][a]; Rp[a][0]=-R[a][0]
    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)]
        Rij=R[i][j]
        if abs(float((rs[a]-rs[b]).mid()))<1e-30:   # equal r: s*sqrt(1+cos)/(2pi sqrt(1+R))
            fac=s[a]*(1+th[m].cos()).sqrt()/(2*PI*(1+Rij).sqrt())
        else:
            fac=s[a]*s[b]*th[m].sin()/(2*PI*(1-Rij*Rij).sqrt())
        c=cond_orth(R,i,j,k,l); cp=cond_orth(Rp,i,j,k,l)
        # dG/dtheta = -s_a s_b sin th * phi2 * (cp - c)   [dG/dR_ij = phi2*(c - ... ) see notes]
        g.append(-fac*(cp-c))
    return g