import json
import numpy as np

OMEGA = 1.0 / 16.0


def heis(p, q):
    x, y, z = p
    X, Y, Z = q
    return np.array([x + X, y + Y, z + Z + 0.5 * (x * Y - y * X)])


def four(v):
    return -OMEGA*v[0] + (0.5 + OMEGA)*(v[1] + v[2]) - OMEGA*v[3]


def refine(seq, geometric=True):
    """Odd/even refinement. Boundary inserted points use linear interpolation."""
    n, d = seq.shape
    out = np.empty((2*n-1, d), dtype=float)
    out[::2] = seq
    for j in range(n-1):
        if 1 <= j <= n-3:
            w = seq[j-1:j+3]
            a, b = four(w[:, 0]), four(w[:, 1])
            z = four(w[:, 2])
            if geometric:
                z += 0.5 * (seq[j, 0] * b - seq[j, 1] * a)
            out[2*j+1] = [a, b, z]
        else:
            out[2*j+1] = 0.5 * (seq[j] + seq[j+1])
    return out


def path_z(x, y):
    """Discrete signed area accumulated along a fine path, starting at zero."""
    z = np.zeros(len(x))
    z[1:] = np.cumsum(0.5 * (x[:-1]*y[1:] - y[:-1]*x[1:]))
    return z


def make_paths(seed=7, count=100, n=65):
    rng = np.random.default_rng(seed)
    t = np.linspace(0, 1, n)
    paths = []
    for _ in range(count):
        # Smooth closed-ish controls with nontrivial winding and signed area.
        phase = rng.uniform(0, 2*np.pi, 4)
        ax, ay = rng.uniform(.7, 1.3, 2)
        x = ax*np.cos(2*np.pi*t + phase[0]) + .18*np.sin(6*np.pi*t + phase[1])
        y = ay*np.sin(2*np.pi*t + phase[2]) + .18*np.cos(4*np.pi*t + phase[3])
        z = path_z(x, y)
        paths.append(np.stack([x, y, z], axis=1))
    return paths


def main():
    # Core algebra: associativity should hold to floating point precision.
    rng = np.random.default_rng(123)
    assoc = []
    for _ in range(1000):
        p, q, r = rng.normal(size=(3, 3))
        assoc.append(np.max(np.abs(heis(heis(p,q),r)-heis(p,heis(q,r)))))
    assoc_err = max(assoc)

    # Reversal: reversing the traversal changes the signed polygonal area sign.
    x = np.array([0., 1., 1., 0., 0.]); y = np.array([0., 0., 1., 1., 0.])
    area = path_z(x, y)[-1]
    reverse_area = path_z(x[::-1], y[::-1])[-1]

    paths = make_paths()
    errors = {"linear": [], "four_point": [], "heisenberg": []}
    z_ranges = {k: [] for k in errors}
    for fine in paths:
        coarse = fine[::2]
        # Linear control is standard interpolation; ordinary four-point is the direct baseline.
        lin = np.empty_like(fine); lin[::2] = coarse
        for j in range(len(coarse)-1): lin[2*j+1] = .5*(coarse[j]+coarse[j+1])
        fp = refine(coarse, geometric=False)
        hp = refine(coarse, geometric=True)
        for name, pred in [("linear",lin),("four_point",fp),("heisenberg",hp)]:
            errors[name].append(np.sqrt(np.mean((pred[:,2]-fine[:,2])**2)))
            z_ranges[name].append(float(np.max(np.abs(pred[:,2]))))
    summary = {
        "associativity_max_abs_error": assoc_err,
        "unit_square_area": area,
        "reversed_unit_square_area": reverse_area,
        "reversal_sum": area + reverse_area,
        "rmse_mean": {k: float(np.mean(v)) for k,v in errors.items()},
        "rmse_median": {k: float(np.median(v)) for k,v in errors.items()},
        "max_abs_z_mean": {k: float(np.mean(v)) for k,v in z_ranges.items()},
        "n_paths": len(paths), "fine_length": len(paths[0]), "coarse_length": len(paths[0][::2]),
    }
    print(json.dumps(summary, indent=2))

if __name__ == '__main__':
    main()