import json
import numpy as np
from scipy.optimize import root, minimize

# Two-stage Gauss-Legendre Runge-Kutta method, formal order four.
A = np.array([[0.25, 0.25 - np.sqrt(3.0) / 6.0],
              [0.25 + np.sqrt(3.0) / 6.0, 0.25]])
b = np.array([0.5, 0.5])
c = np.array([0.5 - np.sqrt(3.0) / 6.0,
              0.5 + np.sqrt(3.0) / 6.0])


def gauss_unconstrained(y, t, h, f):
    y = np.asarray(y, dtype=float)
    d = y.size

    def residual(x):
        Y = x.reshape(2, d)
        vals = np.array([f(t + c[i] * h, Y[i]) for i in range(2)])
        return (Y - y[None, :] - h * A.dot(vals)).ravel()

    sol = root(residual, np.tile(y, 2), method="hybr")
    if not sol.success:
        raise RuntimeError("implicit solve failed: " + sol.message)
    Y = sol.x.reshape(2, d)
    vals = np.array([f(t + c[i] * h, Y[i]) for i in range(2)])
    z = y + h * b.dot(vals)
    return z, Y, float(np.linalg.norm(residual(sol.x))), 2


def gauss_constrained(y, t, h, f, lower, upper, rho=1.0):
    """Paper objective with explicit box feasibility on stages and final state."""
    y = np.asarray(y, dtype=float)
    lower = np.broadcast_to(np.asarray(lower, dtype=float), y.shape)
    upper = np.broadcast_to(np.asarray(upper, dtype=float), y.shape)
    d = y.size

    def unpack(x):
        return x[:2*d].reshape(2, d), x[2*d:]

    def residuals(x):
        Y, z = unpack(x)
        vals = np.array([f(t + c[i] * h, Y[i]) for i in range(2)])
        rs = Y - y[None, :] - h * A.dot(vals)
        rz = z - y - h * b.dot(vals)
        return rs, rz

    def objective(x):
        rs, rz = residuals(x)
        return 0.5 * (np.sum(rs * rs) + rho * np.sum(rz * rz))

    # Explicit bounds are equivalent to projected optimization for this box.
    x0 = np.concatenate([np.tile(np.clip(y, lower, upper), 2),
                         np.clip(y, lower, upper)])
    bounds = [(float(lower[j]), float(upper[j])) for _ in range(3) for j in range(d)]
    sol = minimize(objective, x0, method="L-BFGS-B", bounds=bounds,
                   options={"maxiter": 250, "ftol": 1e-14, "gtol": 1e-10, "maxls": 30})
    Y, z = unpack(sol.x)
    rs, rz = residuals(sol.x)
    return z, Y, float(np.sqrt(np.sum(rs * rs) + rho * np.sum(rz * rz))), 2, int(sol.nit), bool(sol.success)


def integrate(method, y0, f, T, n, box=None):
    y = np.asarray(y0, dtype=float).copy()
    t = 0.0
    max_violation = 0.0
    residual_sum = 0.0
    iterations = 0
    for _ in range(n):
        h = T / n
        if method == "unconstrained":
            z, Y, r, _ = gauss_unconstrained(y, t, h, f)
        elif method == "clip":
            z, Y, r, _ = gauss_unconstrained(y, t, h, f)
            z = np.clip(z, box[0], box[1])
            Y = np.clip(Y, box[0], box[1])
        elif method == "constrained":
            z, Y, r, _, nit, ok = gauss_constrained(y, t, h, f, box[0], box[1])
            iterations += nit
            if not ok:
                raise RuntimeError("constrained optimizer failed")
        else:
            raise ValueError(method)
        if box is not None:
            max_violation = max(max_violation,
                float(max(np.max(box[0] - Y), np.max(Y - box[1]),
                          np.max(box[0] - z), np.max(z - box[1]), 0.0)))
        residual_sum += r
        y, t = z, t + T / n
    return y, max_violation, residual_sum / n, iterations


def main():
    np.random.seed(7)
    results = {"math_check": [], "stiff_test": []}

    # Claimed high-order behavior away from the boundary: y' = y, exact exp(T).
    f_exp = lambda t, y: y
    exact = np.exp(1.0)
    errors = []
    for n in [2, 4, 8, 16, 32]:
        y, _, _, _ = integrate("unconstrained", [1.0], f_exp, 1.0, n)
        errors.append([n, float(abs(y[0] - exact))])
    orders = [float(np.log(errors[i-1][1] / errors[i][1]) / np.log(2.0))
              for i in range(1, len(errors))]
    results["math_check"] = {"errors": errors, "observed_orders": orders,
                              "expected_order": 4}

    # Stiff positive scalar benchmark. Exact solution remains in [0, 1].
    # The large positive rate makes unconstrained RK stages leave the box.
    lam = 40.0
    f_stiff = lambda t, y: lam * (1.0 - y)
    T, n, y0 = 1.0, 10, np.array([0.0])
    box = (np.array([0.0]), np.array([1.0]))
    exact_final = 1.0 - np.exp(-lam * T)
    for method in ["unconstrained", "clip", "constrained"]:
        y, violation, residual, iters = integrate(method, y0, f_stiff, T, n, box)
        results["stiff_test"].append({
            "method": method, "final": float(y[0]),
            "abs_error": float(abs(y[0] - exact_final)),
            "max_constraint_violation": violation,
            "mean_residual": residual, "inner_iterations": iters})

    with open("results.json", "w") as fp:
        json.dump(results, fp, indent=2)
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()