import json
import math
import numpy as np


def fold(p, j):
    """Apply the stated local rational fold to complex polygon coordinate j."""
    q = np.array(p, copy=True)
    n = len(q)
    a, b, c = p[(j - 1) % n], p[j], p[(j + 1) % n]
    den = (b - c) - (a - b)  # 2*b-a-c
    q[j] = ((b - c) * a - (a - b) * c) / den
    return q, abs(den)


def apply_schedule(p, schedule):
    q = np.array(p, copy=True)
    denoms = []
    for j in schedule:
        q, d = fold(q, int(j))
        denoms.append(d)
    return q, np.asarray(denoms)


def relerr(a, b):
    return float(np.linalg.norm(a - b) / max(np.linalg.norm(b), 1e-30))


def polygon(seed, n=8, dtype=np.complex128):
    rng = np.random.default_rng(seed)
    # A separated, mildly perturbed convex polygon avoids accidental singularities.
    theta = np.arange(n) * (2 * np.pi / n)
    radius = 1.0 + .04 * rng.normal(size=n)
    p = radius * np.exp(1j * theta) + .02 * (rng.normal(size=n) + 1j*rng.normal(size=n))
    return p.astype(dtype)


def fold_math_check():
    rng = np.random.default_rng(10)
    errs = []
    for _ in range(2000):
        a, c = rng.normal(size=2) + 1j*rng.normal(size=2)
        b = rng.normal() + 1j*rng.normal()
        if abs(2*b-a-c) < .1:
            continue
        q, _ = fold(np.array([a, b, c], dtype=np.complex128), 1)
        r, _ = fold(q, 1)
        errs.append(abs(r[1]-b) / max(abs(b), 1e-12))
    return {"samples": len(errs), "max_relative_error": max(errs), "median_relative_error": float(np.median(errs))}


def schedule_check(dtype, schedule, repeats=1):
    p = polygon(3, dtype=dtype)
    sched = list(schedule) * repeats
    y, den = apply_schedule(p, sched)
    z, den2 = apply_schedule(y, list(reversed(sched)))
    return {"dtype": str(dtype), "steps": len(sched), "reconstruction_relative_error": relerr(z, p),
            "minimum_denominator": float(min(np.min(den), np.min(den2))),
            "max_abs_state": float(max(np.max(np.abs(p)), np.max(np.abs(y)), np.max(np.abs(z))))}


def denominator_sweep():
    # Directly approach the singular locus 2b-a-c=0 while retaining nonzero a,b,c.
    a, c = 0.2 + .4j, -0.7 + .1j
    out = []
    for eps in [1e-1, 1e-3, 1e-5, 1e-7, 1e-9]:
        b = (a+c)/2 + eps*(1+0.3j)
        p = np.array([a,b,c], dtype=np.complex128)
        q, d = fold(p, 1)
        r, _ = fold(q, 1)
        out.append({"epsilon": eps, "denominator": float(d),
                    "folded_abs": float(abs(q[1])), "involution_error": float(abs(r[1]-b))})
    return out


def main():
    schedule = [0, 2, 4, 6, 1, 3, 5, 7]
    results = {
        "math_check": fold_math_check(),
        "fixed_schedule": [schedule_check(np.complex128, schedule, k) for k in [1, 5, 20]],
        "float32_schedule": [schedule_check(np.complex64, schedule, k) for k in [1, 5, 20]],
        "singularity_sweep": denominator_sweep(),
        "memory_note": {"reversible_saved_fold_states": 0,
                        "ordinary_unroll_saved_fold_states_20_steps": 20,
                        "state_size_points": 8},
    }
    with open("results.json", "w") as f:
        json.dump(results, f, indent=2)
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()