"""Best available rigorous gradient enclosure of Psi4 over a box: dual-number centered form,
affine arithmetic, and affine arithmetic on a 2^6 sub-box grid where needed; intersected."""
import itertools
import dualK, aa
def _meet(x,y):
    if x is None: return y
    if y is None: return x
    try: return x.intersection(y)
    except Exception: return x if float(x.rad())<=float(y.rad()) else y
def grad_best(rsa,lo,hi,need=None,refine_tol=None,refine=True):
    """need: coordinates that must be finite (others may stay None); refine with sub-boxes when a
    needed coordinate is None or wider than refine_tol."""
    g1=dualK.grad_box(rsa,list(lo),list(hi)); g2=aa.grad_aa(rsa,list(lo),list(hi))
    g=[_meet(a,b) for a,b in zip(g1,g2)]
    if not refine: return g
    need=range(6) if need is None else need
    bad=[m for m in need if g[m] is None or (refine_tol is not None and float(g[m].rad())>refine_tol)]
    if not bad: return g
    w=[hi[m]-lo[m] for m in range(6)]; order=sorted(range(6),key=lambda m:-w[m])
    split=[m for m in order[:4] if w[m]>=0.5*w[order[0]]]
    mids=[(lo[m]+hi[m])/2 for m in range(6)]
    acc=[None]*6; failed=[False]*6
    for idx in itertools.product((0,1),repeat=len(split)):
        sl=list(lo); sh=list(hi)
        for b,m in zip(idx,split):
            if b==0: sh[m]=mids[m]
            else: sl[m]=mids[m]
        gs=aa.grad_aa(rsa,sl,sh,only=bad)
        for m in bad:
            if gs[m] is None: failed[m]=True
            elif not failed[m]: acc[m]=gs[m] if acc[m] is None else acc[m].union(gs[m])
    for m in bad:
        if not failed[m]: g[m]=_meet(g[m],acc[m])
    return g