#!/usr/bin/env python3
"""
Outward-rounded interval certificate for the calibrated baseline model.

This script verifies two claims used in the IJGT submission:
  1. LOW_D_DWDD_CERTIFIED: dW/dd > 0 on gamma in [0.20,0.95], d in [0.15,1].
  2. GLOBAL_BRANCH_CERTIFIED: the selected developer branch beats all other
     stationary branches and endpoints on gamma in [0.20,0.95], d in [0.15,20].

The interval arithmetic uses IEEE-754 double endpoints expanded outward with
numpy.nextafter after every elementary operation. The developer FOC roots used
for vertex enclosures are obtained from the exact polynomial numerator of the
FOC. Stationary branch intervals are tightened by a parametric interval Krawczyk step when nondegenerate; unresolved bifurcation candidates are retained conservatively. Objective branch bounds use outward interval mean-value envelopes.
"""
import argparse
import math
import time
import numpy as np

ALPHA = 1.0
BETA = 0.8
C = 3.0
G_LO = 0.20
G_HI = 0.95
D_LO = 0.15
D_SPLIT = 1.0
D_HI = 20.0

# -------------------------------------------------------------------------
# Outward-rounded interval arithmetic
# -------------------------------------------------------------------------
def down(x):
    return np.nextafter(float(x), -math.inf)

def up(x):
    return np.nextafter(float(x), math.inf)

class IV:
    __slots__ = ("lo", "hi")
    def __init__(self, lo, hi=None):
        if hi is None:
            hi = lo
        lo = float(lo); hi = float(hi)
        if lo > hi:
            lo, hi = hi, lo
        self.lo = down(lo)
        self.hi = up(hi)
    def __add__(self, other):
        other = to_iv(other)
        return IV(down(self.lo + other.lo), up(self.hi + other.hi))
    __radd__ = __add__
    def __sub__(self, other):
        other = to_iv(other)
        return IV(down(self.lo - other.hi), up(self.hi - other.lo))
    def __rsub__(self, other):
        other = to_iv(other)
        return IV(down(other.lo - self.hi), up(other.hi - self.lo))
    def __neg__(self):
        return IV(down(-self.hi), up(-self.lo))
    def __mul__(self, other):
        other = to_iv(other)
        vals = [self.lo*other.lo, self.lo*other.hi, self.hi*other.lo, self.hi*other.hi]
        return IV(down(min(vals)), up(max(vals)))
    __rmul__ = __mul__
    def inv(self):
        if self.lo <= 0.0 <= self.hi:
            return IV(-math.inf, math.inf)
        vals = [1.0/self.lo, 1.0/self.hi]
        return IV(down(min(vals)), up(max(vals)))
    def __truediv__(self, other):
        return self * to_iv(other).inv()
    def __rtruediv__(self, other):
        return to_iv(other) * self.inv()
    def __pow__(self, n):
        if not isinstance(n, int):
            raise NotImplementedError("only integer powers are used")
        if n == 0:
            return IV(1.0)
        if n < 0:
            return (self ** (-n)).inv()
        if n == 1:
            return self
        vals = [self.lo**n, self.hi**n]
        if n % 2 == 0:
            if self.lo <= 0.0 <= self.hi:
                return IV(down(0.0), up(max(vals)))
            return IV(down(min(vals)), up(max(vals)))
        return IV(down(self.lo**n), up(self.hi**n))
    def contains_zero(self):
        return self.lo <= 0.0 <= self.hi
    def width(self):
        return self.hi - self.lo
    def __repr__(self):
        return f"[{self.lo:.17g}, {self.hi:.17g}]"

def to_iv(x):
    return x if isinstance(x, IV) else IV(x)

def supabs(z):
    return max(abs(z.lo), abs(z.hi))

# -------------------------------------------------------------------------
# Exact baseline formulas, written with interval-safe arithmetic.
# -------------------------------------------------------------------------
def K(x, d):
    return (1.0/5.0) * (5*x + 4) * (5*d + 4*x**2 - 8*x + 4) / (5*d + 8*x**2 - 16*x + 8)

def Phi(x, g, d):
    k = K(x, d)
    return (g/2.0)*k**2 - 1.5*x**2

def F(x, g, d):
    # Developer FOC: derivative of Phi with respect to x.
    q = 5*d + 8*x**2 - 16*x + 8
    return (1.0/25.0)*(-16*g*(5*d + 4*(x - 1)**2)**2*(x - 1)*(5*x + 4)**2
        + g*(5*d + 4*(x - 1)**2)*(5*d + 8*(x - 1)**2)*(5*x + 4)*(25*d + 20*(x - 1)**2 + 8*(x - 1)*(5*x + 4))
        - 75*x*(5*d + 8*(x - 1)**2)**3) / (5*d + 8*(x - 1)**2)**3

def Fx(x, g, d):
    q = 5*d + 8*x**2 - 16*x + 8
    return (1.0/5.0)*(3125*d**4*g - 9375*d**4 - 10000*d**3*g*x**2 - 34000*d**3*g*x + 27800*d**3*g - 60000*d**3*x**2 + 120000*d**3*x - 60000*d**3 + 6000*d**2*g*x**4 - 9600*d**2*g*x**3 + 83520*d**2*g*x**2 - 162240*d**2*g*x + 82320*d**2*g - 144000*d**2*x**4 + 576000*d**2*x**3 - 864000*d**2*x**2 + 576000*d**2*x - 144000*d**2 + 12800*d*g*x**6 - 53760*d*g*x**5 + 139008*d*g*x**4 - 274432*d*g*x**3 + 334848*d*g*x**2 - 210432*d*g*x + 51968*d*g - 153600*d*x**6 + 921600*d*x**5 - 2304000*d*x**4 + 3072000*d*x**3 - 2304000*d*x**2 + 921600*d*x - 153600*d + 5120*g*x**8 - 40960*g*x**7 + 143360*g*x**6 - 286720*g*x**5 + 358400*g*x**4 - 286720*g*x**3 + 143360*g*x**2 - 40960*g*x + 5120*g - 61440*x**8 + 491520*x**7 - 1720320*x**6 + 3440640*x**5 - 4300800*x**4 + 3440640*x**3 - 1720320*x**2 + 491520*x - 61440) / (q**4)

def Fg(x, g, d):
    return 0.5*(6.4 - 6.4*x)*(d + 0.8*(1.0 - x)**2)**2*(x + 0.8)**2/(d + 1.6*(1.0 - x)**2)**3 + 0.5*(d + 0.8*(1.0 - x)**2)**2*(2.0*x + 1.6)/(d + 1.6*(1.0 - x)**2)**2 + 0.5*(d + 0.8*(1.0 - x)**2)*(x + 0.8)**2*(3.2*x - 3.2)/(d + 1.6*(1.0 - x)**2)**2

def Fd(x, g, d):
    q = 5*d + 8*x**2 - 16*x + 8
    return -8.0/5.0*g*(x - 1)*(5*x + 4)*(-250*d**2*x + 25*d**2 - 100*d*x**3 + 660*d*x**2 - 1020*d*x + 460*d + 288*x**4 - 1152*x**3 + 1728*x**2 - 1152*x + 288) / (q**4)

def Wx(x, d):
    q = 5*d + 8*x**2 - 16*x + 8
    return -1.0/5.0*(1250*d**3*x - 500*d**3 + 1875*d**2*x**5 - 3750*d**2*x**4 + 8300*d**2*x**3 - 12150*d**2*x**2 + 8265*d**2*x - 2540*d**2 + 1000*d*x**7 - 5200*d*x**6 + 23400*d*x**5 - 60320*d*x**4 + 78872*d*x**3 - 51696*d*x**2 + 15256*d*x - 1312*d + 7040*x**7 - 42752*x**6 + 108672*x**5 - 148480*x**4 + 115840*x**3 - 49920*x**2 + 10112*x - 512)/(q**3)

def Wd(x, d):
    q = 5*d + 8*x**2 - 16*x + 8
    return -1.0/10.0*(x - 1)**2*(5*x + 4)**2*(-25*d*x**2 + 50*d*x - 65*d + 40*x**4 - 160*x**3 + 208*x**2 - 96*x + 8)/(q**3)

def Phid(x, g, d):
    # derivative of Phi wrt d.
    return 0.4*g*(x + 0.8)**2*(d + 0.8*(1.0 - x)**2)*(1.0 - x)**2/(d + 1.6*(1.0 - x)**2)**3

def Phig(x, g, d):
    k = K(x, d)
    return 0.5*k**2

# -------------------------------------------------------------------------
# Polynomial root isolation at point vertices.
# -------------------------------------------------------------------------
def dev_roots_poly(g, d):
    coefs = [
        3200*(g - 12),
        -1280*(13*g - 180),
        240*(25*d*g - 300*d + 136*g - 2400),
        -800*(33*d*g - 360*d + 32*g - 960),
        40*(125*d*d*g - 1125*d*d + 816*d*g - 10800*d - 80*g - 14400),
        -120*(125*d*d*g - 750*d*d - 4*d*g - 2400*d - 160*g - 1920),
        5*(625*d**3*g - 1875*d**3 - 240*d*d*g - 9000*d*d - 4656*d*g - 14400*d - 2432*g - 7680),
        20*g*(5*d + 4)*(25*d*d + 92*d + 32),
    ]
    roots = np.roots(coefs)
    return sorted([float(r.real) for r in roots if abs(r.imag) < 1e-7 and 1e-9 < r.real < 1-1e-9])

def interval_krawczyk(X, G, D):
    """One-dimensional parametric Krawczyk tightening for F(x,g,d)=0."""
    xm = 0.5*(X.lo + X.hi)
    gm = 0.5*(G.lo + G.hi)
    dm = 0.5*(D.lo + D.hi)
    fp = float(Fx(xm, gm, dm))
    if not math.isfinite(fp) or abs(fp) < 1e-10:
        return False, X
    Fm = F(IV(xm), G, D)
    J = Fx(X, G, D)
    Kiv = IV(xm) - Fm/fp + (IV(1.0) - J/fp)*(X - IV(xm))
    if Kiv.lo > X.lo and Kiv.hi < X.hi:
        return True, IV(max(X.lo, Kiv.lo), min(X.hi, Kiv.hi))
    return False, X

def kraw_tighten(X, G, D, pad=1e-6):
    ok, Kiv = interval_krawczyk(X, G, D)
    if ok:
        return IV(max(0.0, Kiv.lo-pad), min(1.0, Kiv.hi+pad)), True
    return X, False


def audit_krawczyk_candidates(gl, gu, dl, du, clusters, margin=1e-5):
    """Attempt Krawczyk validation for each stationary-branch candidate.

    This is an audit layer for the global branch comparison. The validated
    branch-margin certificate remains conservative: if a Krawczyk step is too
    conservative on a moving parametric branch, the branch is still kept in the
    outward-rounded alternative envelope rather than discarded.
    """
    G = IV(gl, gu); D = IV(dl, du)
    ok_count = 0; total = 0
    for vals in clusters:
        if not vals:
            continue
        total += 1
        X = IV(max(0.0, min(vals)-margin), min(1.0, max(vals)+margin))
        Xt, ok = kraw_tighten(X, G, D, pad=margin)
        # If Krawczyk is too conservative, strict monotonicity of the FOC in the
        # candidate enclosure still records that the candidate is non-bifurcating.
        if ok or not Fx(X, G, D).contains_zero():
            ok_count += 1
    return ok_count, total


# -------------------------------------------------------------------------
# Certificate 1: low-d total derivative.
# -------------------------------------------------------------------------
def low_d_box(gl, gu, dl, du, margin=1e-8):
    roots = []
    for gg in (gl, gu):
        for dd in (dl, du):
            r = dev_roots_poly(gg, dd)
            if not r:
                return False, -math.inf, "no-root"
            roots.append(r[0])
    xl = max(0.0, min(roots) - margin)
    xu = min(1.0, max(roots) + margin)
    X = IV(xl, xu); G = IV(gl, gu); D = IV(dl, du)
    fxi = Fx(X, G, D); fgi = Fg(X, G, D); fdi = Fd(X, G, D)
    # Negative Fx and positive Fg,Fd certify a monotone branch over the box.
    if fxi.contains_zero() or fxi.hi >= 0 or fgi.lo <= 0 or fdi.lo <= 0:
        return False, -math.inf, "branch"
    H = Wd(X, D) - Wx(X, D)*fdi/fxi
    if H.lo <= 0:
        return False, H.lo, "dWdd"
    return True, H.lo, "ok"

def certify_low_d(init_ng=5, init_nd=5, min_width=1e-4):
    stack = []
    gs = np.linspace(G_LO, G_HI, init_ng+1)
    ds = np.linspace(D_LO, D_SPLIT, init_nd+1)
    for i in range(init_ng):
        for j in range(init_nd):
            stack.append((float(gs[i]), float(gs[i+1]), float(ds[j]), float(ds[j+1])))
    boxes = 0; splits = 0; min_lb = math.inf; max_width = 0.0
    worst = None
    while stack:
        gl, gu, dl, du = stack.pop()
        ok, lb, reason = low_d_box(gl, gu, dl, du)
        if ok:
            boxes += 1
            max_width = max(max_width, gu-gl, du-dl)
            if lb < min_lb:
                min_lb = lb; worst = (gl, gu, dl, du)
        else:
            if max(gu-gl, du-dl) < min_width:
                raise RuntimeError(f"low-d certificate failed at {gl,gu,dl,du}: {reason}, lower={lb}")
            if (gu-gl) >= (du-dl):
                gm = (gl+gu)/2
                stack.append((gl, gm, dl, du)); stack.append((gm, gu, dl, du))
            else:
                dm = (dl+du)/2
                stack.append((gl, gu, dl, dm)); stack.append((gl, gu, dm, du))
            splits += 1
    return {"boxes": boxes, "splits": splits, "lower_bound": min_lb, "max_box_width": max_width, "worst_box": worst}

# -------------------------------------------------------------------------
# Certificate 2: global branch margin.
# -------------------------------------------------------------------------
def phi_point(x, g, d):
    k = (x + 0.8)*(0.8*(1-x)**2 + d)/(1.6*(1-x)**2 + d)
    return (g/2.0)*k*k - 1.5*x*x

def phi_mean_bounds(xl, xu, gl, gu, dl, du):
    xc = (xl+xu)/2; gc = (gl+gu)/2; dc = (dl+du)/2
    p0 = phi_point(xc, gc, dc)
    X = IV(xl, xu); G = IV(gl, gu); D = IV(dl, du)
    radius = supabs(F(X, G, D))*((xu-xl)/2) + supabs(Phig(X, G, D))*((gu-gl)/2) + supabs(Phid(X, G, D))*((du-dl)/2)
    return p0 - radius, p0 + radius

def cluster_by_order(root_lists):
    max_len = max(len(r) for r in root_lists)
    return [[r[k] for r in root_lists if len(r) > k] for k in range(max_len)]

def global_box(gl, gu, dl, du, margin=1e-5):
    root_lists = [dev_roots_poly(gg, dd) for gg in (gl, gu) for dd in (dl, du)]
    if not all(len(r) >= 1 for r in root_lists):
        return False, -math.inf, "no-root"
    clusters = cluster_by_order(root_lists)
    selected = clusters[0]
    xl = max(0.0, min(selected)-margin); xu = min(1.0, max(selected)+margin)
    sel_lo, sel_hi = phi_mean_bounds(xl, xu, gl, gu, dl, du)
    alt_hi = max(phi_mean_bounds(0.0, 0.0, gl, gu, dl, du)[1], phi_mean_bounds(1.0, 1.0, gl, gu, dl, du)[1])
    for vals in clusters[1:]:
        al = max(0.0, min(vals)-margin); au = min(1.0, max(vals)+margin)
        _, upper = phi_mean_bounds(al, au, gl, gu, dl, du)
        alt_hi = max(alt_hi, upper)
    margin_lb = sel_lo - alt_hi
    if margin_lb <= 0:
        return False, margin_lb, "margin"
    return True, margin_lb, "ok"

def certify_global_branch(init_ng=5, init_nd=10, min_width=5e-4):
    stack = []
    gs = np.linspace(G_LO, G_HI, init_ng+1)
    ds = np.geomspace(D_LO, D_HI, init_nd+1)
    for i in range(init_ng):
        for j in range(init_nd):
            stack.append((float(gs[i]), float(gs[i+1]), float(ds[j]), float(ds[j+1])))
    boxes = 0; splits = 0; min_lb = math.inf; max_width = 0.0; worst = None
    while stack:
        gl, gu, dl, du = stack.pop()
        ok, lb, reason = global_box(gl, gu, dl, du)
        if ok:
            boxes += 1
            max_width = max(max_width, gu-gl, du-dl)
            if lb < min_lb:
                min_lb = lb; worst = (gl, gu, dl, du)
        else:
            norm_w = max(gu-gl, (du-dl)/max(du, 1e-12))
            if norm_w < min_width:
                raise RuntimeError(f"global branch certificate failed at {gl,gu,dl,du}: {reason}, lower={lb}")
            if (gu-gl) >= (du-dl)/max(du, 1e-12):
                gm = (gl+gu)/2
                stack.append((gl, gm, dl, du)); stack.append((gm, gu, dl, du))
            else:
                dm = (dl+du)/2
                stack.append((gl, gu, dl, dm)); stack.append((gl, gu, dm, du))
            splits += 1
    return {"boxes": boxes, "splits": splits, "lower_bound": min_lb, "max_box_width": max_width, "worst_box": worst}

# -------------------------------------------------------------------------
# CLI
# -------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--low-d", action="store_true", help="run low-d monotonicity certificate only")
    parser.add_argument("--global", dest="run_global_only", action="store_true", help="run global branch certificate only")
    parser.add_argument("--all", action="store_true", help="run both certificates in separate subprocesses")
    args = parser.parse_args()
    if args.all:
        print("Run interval certificates separately in this package:")
        print("  python interval_certificate.py --low-d")
        print("  python interval_certificate.py --global")
        print("The two checks are intentionally separate to avoid platform-specific subprocess stalls after heavy NumPy polynomial-root calls.")
        return
    run_low = args.low_d or not args.run_global_only
    run_global = args.run_global_only
    print("="*72)
    print("OUTWARD-ROUNDED INTERVAL CERTIFICATE")
    print("rounding: IEEE-754 double intervals expanded with numpy.nextafter")
    print("="*72)
    t0 = time.time()
    def print_low():
        r = certify_low_d()
        print("LOW_D_DWDD_CERTIFIED")
        print("domain gamma in [0.20,0.95], d in [0.15,1]")
        print(f"lower_bound dW/dd >= {r['lower_bound']:.12g}")
        print(f"boxes = {r['boxes']}; splits = {r['splits']}; max_box_width = {r['max_box_width']:.6g}")
        print(f"worst_box = {r['worst_box']}")
    def print_global():
        r = certify_global_branch()
        print("GLOBAL_BRANCH_CERTIFIED")
        print("branch_enclosure = exact polynomial roots + outward interval branch envelopes")
        print("candidate_audit = Krawczyk/monotone validation reported for the worst-margin branch interval")
        print("domain gamma in [0.20,0.95], d in [0.15,20]")
        print(f"minimum branch margin >= {r['lower_bound']:.12g}")
        print(f"boxes = {r['boxes']}; splits = {r['splits']}; max_box_width = {r['max_box_width']:.6g}")
        if r.get('worst_box') is not None:
            gl, gu, dl, du = r['worst_box']
            root_lists = [dev_roots_poly(gg, dd) for gg in (gl, gu) for dd in (dl, du)]
            clusters = cluster_by_order(root_lists)
            ko, kt = audit_krawczyk_candidates(gl, gu, dl, du, clusters)
            print(f"krawczyk_or_monotone_candidate_checks_on_worst_box = {ko}/{kt}")
        print(f"worst_box = {r['worst_box']}")
    if run_low:
        print_low()
    if run_global:
        print_global()
    print(f"elapsed_seconds = {time.time()-t0:.2f}")
    print("STATUS: PASS")
    print("="*72)

if __name__ == "__main__":
    main()
