import json
import numpy as np

GAMMA = 0.25
A = np.array([[0, 0, 0, 0, 0], [0.5, 0, 0, 0, 0],
              [17/50, -1/25, 0, 0, 0],
              [371/1360, -137/2720, 15/544, 0, 0],
              [25/24, -49/48, 125/16, -85/12, 0]], float)


def residual(z, b, h, rates):
    # r=b-Az for A=I+h*gamma*diag(rates), i.e. F(y)=-diag(rates)y.
    return b - (z + h * GAMMA * rates * z)


def chebyshev_solve(b, h, rates, K, inflate=1.10, supplied_beta=False):
    """Chebyshev semi-iteration for the diagonal implicit stage system.

    supplied_beta=True reproduces the idea_context literally.  The default uses
    the dimensionally consistent momentum beta=d^2*alpha_k*alpha_{k-1}/4.
    """
    lo = 1.0
    hi = (1.0 + h * GAMMA * float(np.max(rates))) * inflate
    c, d = (lo + hi) / 2.0, (hi - lo) / 2.0
    z = b.copy(); zprev = z.copy(); alpha_prev = None
    rs = []
    for k in range(K):
        r = residual(z, b, h, rates); rs.append(float(np.linalg.norm(r)))
        alpha = 1.0 / c if k == 0 else 1.0 / (c - d*d*alpha_prev/4.0)
        beta = 0.0 if k == 0 else d*d*alpha/4.0
        if not supplied_beta and k > 0:
            beta *= alpha_prev
        znew = z + alpha*r + beta*(z-zprev)
        zprev, z, alpha_prev = z, znew, alpha
    rs.append(float(np.linalg.norm(residual(z, b, h, rates))))
    return z, np.asarray(rs)


def fixed_point_solve(b, h, rates, K):
    z = b.copy(); rs = []
    for _ in range(K):
        r = residual(z, b, h, rates); rs.append(float(np.linalg.norm(r))); z = z + r
    rs.append(float(np.linalg.norm(residual(z, b, h, rates))))
    return z, np.asarray(rs)


def sdirk_step(y, h, rates, method="cheb", K=24):
    fvals = []
    for i in range(5):
        b = y + h * sum(A[i, j] * fvals[j] for j in range(i)) if i else y.copy()
        z, _ = (chebyshev_solve(b, h, rates, K) if method == "cheb"
                else fixed_point_solve(b, h, rates, K))
        fvals.append(-rates * z)
    return z


def rk4_step(y, h, rates):
    f = lambda x: -rates*x
    k1 = f(y); k2 = f(y+h*k1/2); k3 = f(y+h*k2/2); k4 = f(y+h*k3)
    return y + h*(k1+2*k2+2*k3+k4)/6


def main():
    np.random.seed(7)
    rates = np.geomspace(1.0, 1000.0, 16)
    y0 = np.random.randn(rates.size)
    h = 0.20; b = y0.copy()
    _, literal = chebyshev_solve(b, h, rates, 20, supplied_beta=True)
    _, corrected = chebyshev_solve(b, h, rates, 20)
    _, fixed = fixed_point_solve(b, h, rates, 20)
    report = {"rates": rates.tolist(), "checks": {
        "h": h, "literal_beta_final_over_initial": float(literal[-1]/literal[0]),
        "corrected_beta_final_over_initial": float(corrected[-1]/corrected[0]),
        "fixed_final_over_initial": float(fixed[-1]/fixed[0]),
        "corrected_residuals": corrected.tolist(),
        "literal_residuals": literal.tolist(),
        "fixed_residuals": fixed.tolist(),
        "corrected_monotone_after_transients": bool(np.all(np.diff(corrected[2:]) <= 1e-12))}}
    report["steps"] = []
    for h in [0.001, 0.01, 0.05, 0.2]:
        exact = y0 * np.exp(-rates*h)
        yc = sdirk_step(y0, h, rates, "cheb")
        yf = sdirk_step(y0, h, rates, "fixed")
        yr = rk4_step(y0, h, rates)
        report["steps"].append({"h": h, "cheb_error": float(np.linalg.norm(yc-exact)),
            "fixed_error": float(np.linalg.norm(yf-exact)), "rk4_error": float(np.linalg.norm(yr-exact)),
            "cheb_norm": float(np.linalg.norm(yc)), "fixed_norm": float(np.linalg.norm(yf)),
            "rk4_norm": float(np.linalg.norm(yr)), "rk4_finite": bool(np.all(np.isfinite(yr)))})
    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()