"""Centered-form enclosure of grad Psi4 (sign-moment form) over a box in C-space, Arb dual numbers."""
from flint import arb
import itertools
from dual import D, asinD, NV, PI
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_formula(rsa,cD):
    s=[(1-r*r).sqrt() for r in rsa]
    S=[[D(1)]*4 for _ in range(4)]
    for (p,q),v in zip(PAIRS,cD): S[p][q]=S[q][p]=v*(s[p]*s[q])+rsa[p]*rsa[q]
    out=[]
    for (p,q) in PAIRS:
        k,l=COMP[(p,q)]; rpq=S[p][q]; det=1-rpq*rpq
        if not det.v>0: 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]))/det
        skk=cov(k,k); sll=cov(l,l); skl=cov(k,l)
        if not (skk.v>0 and sll.v>0): out.append(None); continue
        pc=skl/(skk*sll).sqrt()
        a=asinD(pc)
        # dPsi/dc_pq = -(1/8) s_p s_q (4/pi^2) asin(pc)/sqrt(1-S_pq^2)
        out.append(a*(-(s[p]*s[q])/(2*PI*PI))/(det.sqrt()))
    return out
def grad_box(rsa,lo,hi):
    """rigorous enclosure of grad Psi over the box: centered form intersected with the plain enclosure"""
    ball=[arb(lo[m]).union(arb(hi[m])) for m in range(NV)]          # contains [lo,hi]
    cen=[float(ball[m].mid()) for m in range(NV)]                    # exact float centre point
    gp=grad_formula(rsa,[D(arb(cen[m])) for m in range(NV)])
    unit=lambda m:[arb(1) if k==m else arb(0) for k in range(NV)]
    gb=grad_formula(rsa,[D(ball[m],unit(m)) for m in range(NV)])
    res=[]
    for m in range(NV):
        if gb[m] is None: res.append(None); continue
        plain=gb[m].v
        if gp[m] is not None and all(h.is_finite() for h in gb[m].d) and gp[m].v.is_finite():
            cf=gp[m].v+sum((h*(ball[k]-arb(cen[k])) for k,h in enumerate(gb[m].d)),arb(0))
            if plain.is_finite():
                I=cf.intersection(plain)
                res.append(I)
            else: res.append(cf)
        else: res.append(plain if plain.is_finite() else None)
    return res
def sqrt_K(rsa,p,q,lo_c):
    """|dPsi/dc_pq| <= K/sqrt(1-c) on the box (equal r): K = s/(4 pi sqrt(1+S_min))"""
    s=(1-rsa[p]*rsa[p]).sqrt(); Smin=rsa[p]*rsa[q]+s*s*arb(lo_c)
    return s/(4*PI*(1+Smin).sqrt())