"""Rigorous Arb enclosure of K4(S) = E[prod sgn] via the Plackett path, integrand using one acb_mat inverse."""
from flint import arb, acb, acb_mat, ctx
import itertools
ctx.prec=64
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}
PI=arb.pi(); C4=4/(PI*PI)
def asin_a(z,a):
    i=acb(0,1); return -i*(i*z+(1-z*z).sqrt(analytic=a)).log(analytic=a)
def integrand(S):
    Sa=[[acb(S[p][q]) for q in range(4)] for p in range(4)]
    def f(t,a):
        M=acb_mat(4,4,[ (acb(1) if p==q else acb(0))*(1-t)+t*Sa[p][q] for p in range(4) for q in range(4)])
        try: P=M.inv()
        except ZeroDivisionError: return acb(float("nan"))   # not certified analytic here
        tot=acb(0)
        for (p,q) in PAIRS:
            k,l=COMP[(p,q)]; pc=-P[k,l]/(P[k,k]*P[l,l]).sqrt(analytic=a); rho=t*Sa[p][q]
            tot+=Sa[p][q]*C4*asin_a(pc,a)/(1-rho*rho).sqrt(analytic=a)
        return tot
    return f
def integrand_poly(S):
    """Same integrand as integrand(S), without a matrix inverse. Along R_t=(1-t)I+tS, for the pair
    (p,q) with complement (k,l) and D=1-t^2 S_pq^2, the conditional covariances given (p,q) satisfy
      D*cov(a,b) = A_ab(t) := t S_ab D' - t^2 (S_ap S_bp + S_aq S_bq) + t^3 S_pq (S_ap S_bq + S_aq S_bp),
    with t S_ab D' replaced by D when a=b (unit diagonal). Hence pc = A_kl / sqrt(A_kk A_ll)."""
    Sa=[[acb(S[p][q]) for q in range(4)] for p in range(4)]
    coef=[]
    for (p,q) in PAIRS:
        k,l=COMP[(p,q)]; spq=Sa[p][q]
        def poly(a,b):
            c2=-(Sa[a][p]*Sa[b][p]+Sa[a][q]*Sa[b][q]); c3=spq*(Sa[a][p]*Sa[b][q]+Sa[a][q]*Sa[b][p])
            if a==b: return (acb(1),acb(0),c2-spq*spq,c3)          # 1 + c2' t^2 + c3 t^3
            return (acb(0),Sa[a][b],c2,c3-Sa[a][b]*spq*spq)         # t S_ab (1 - t^2 spq^2) + ...
        coef.append((spq,poly(k,k),poly(l,l),poly(k,l)))
    def f(t,a):
        tot=acb(0)
        for spq,Pk,Pl,Pkl in coef:
            ev=lambda P: ((P[3]*t+P[2])*t+P[1])*t+P[0]
            Akk=ev(Pk); All=ev(Pl); Akl=ev(Pkl)
            pc=Akl/(Akk*All).sqrt(analytic=a); rho=t*spq
            tot+=spq*C4*asin_a(pc,a)/(1-rho*rho).sqrt(analytic=a)
        return tot
    return f
def tailK(S,eta):
    # integrand term S_pq*(4/pi^2)*asin(.)/sqrt(1-t^2 S_pq^2), |asin|<=pi/2
    tot=arb(0)
    for (p,q) in PAIRS:
        rho=S[p][q]; ar=abs(rho)
        if not ar>0: continue
        a1=ar if ar<1 else arb(1)
        I=(a1.asin()-((1-eta)*a1).asin())/a1
        bound=ar*C4*(PI/2)*I
        tot+=bound*arb(0,1)
    return tot
def K4(S,eta=arb(10)**-8,tol=arb(2)**-34):
    """K = int_0^1 f(t) dt.  On [0,1-eta] substitute t = 1-exp(-s), dt = exp(-s) ds, s in [0, ln(1/eta)]:
    a pole or branch point of f at t = 1+lam (lam>0 small, from a nearly singular endpoint) moves to
    Im s = pi, so the integrand in s stays analytic in a wide strip. The piece [1-eta,1] is bounded by tailK."""
    f=integrand_poly(S)
    def g(s,a):
        e=(-s).exp(); return f(1-e,a)*e
    v=acb.integral(g,0,(1/eta).log(),rel_tol=tol,abs_tol=tol,eval_limit=10**6,depth_limit=10**4)
    return v.real+tailK(S,eta)
def SigmaY(rsa,c):
    s=[(1-r*r).sqrt() for r in rsa]; S=[[arb(1)]*4 for _ in range(4)]
    for (p,q),v in zip(PAIRS,c): S[p][q]=S[q][p]=rsa[p]*rsa[q]+s[p]*s[q]*arb(v)
    return S
def Psi(rsa,c,tol=None):
    const=sum((r.acos() for r in rsa),arb(0))/(4*PI)+arb(1)/8
    if tol is None: return const-K4(SigmaY(rsa,c))/8
    return const-K4(SigmaY(rsa,c),tol=arb(tol))/8