import json
import numpy as np
from experiment import integrate


def run(lam, T, ns):
    f = lambda t, y: lam * (1.0 - y)
    box = (np.array([0.0]), np.array([1.0]))
    exact = 1.0 - np.exp(-lam * T)
    out = []
    for n in ns:
        row = {"n": n}
        for method in ["unconstrained", "clip", "constrained"]:
            try:
                y, violation, residual, iters = integrate(
                    method, [0.0], f, T, n, box)
                row[method] = {
                    "final": float(y[0]),
                    "abs_error": float(abs(y[0] - exact)),
                    "max_constraint_violation": float(violation),
                    "mean_residual": float(residual),
                    "inner_iterations": int(iters),
                }
            except Exception as exc:
                row[method] = {"error": str(exc)}
        out.append(row)
    return out


if __name__ == "__main__":
    result = {
        "lambda_10": run(10.0, 1.0, [2, 4, 8, 16, 32]),
        "lambda_40": run(40.0, 1.0, [2, 4, 8, 16, 32]),
    }
    with open("sweep_results.json", "w") as fp:
        json.dump(result, fp, indent=2)
    print(json.dumps(result, indent=2))