import json
import numpy as np


def make_basis(n=32):
    v = np.linspace(-1.0, 1.0, n)
    V1, V2 = np.meshgrid(v, v, indexing="ij")
    return np.stack([np.ones((n, n)), V1, V2, 0.5 * (V1**2 + V2**2)])


def moments(x, phi):
    return np.einsum("kij,ij->k", phi, x)


def gram(phi):
    return np.einsum("kij,lij->kl", phi, phi)


def conserve(target, compressed, phi, ridge=1e-12):
    """Smallest Frobenius-norm correction of compressed to target moments."""
    G = gram(phi)
    delta = moments(target, phi) - moments(compressed, phi)
    coeff = np.linalg.solve(G + ridge * np.eye(len(phi)), delta)
    return compressed + np.einsum("k,kij->ij", coeff, phi)


def rank_svd(x, rank):
    u, s, vt = np.linalg.svd(x, full_matrices=False)
    return (u[:, :rank] * s[:rank]) @ vt[:rank]


def zero_moment_part(x, phi):
    """Orthogonal projection of x onto the nullspace of all moments."""
    return conserve(np.zeros_like(x), x, phi)


def main():
    rng = np.random.default_rng(329)
    n, rank, steps = 32, 2, 200
    phi = make_basis(n)

    # Core verification: exact moment preservation and minimum-norm property.
    x, y = rng.normal(size=(2, n, n))
    corrected = conserve(x, y, phi)
    residual = np.max(np.abs(moments(corrected, phi) - moments(x, phi)))
    optimal_distance = np.linalg.norm(corrected - y)
    random_distances = []
    for _ in range(100):
        q0 = zero_moment_part(rng.normal(size=(n, n)), phi)
        random_distances.append(np.linalg.norm(corrected + q0 - y))

    v = np.linspace(-1, 1, n)
    V1, V2 = np.meshgrid(v, v, indexing="ij")
    # Deliberately not low rank: several separated smooth components.
    exact = (np.exp(-18 * ((V1 + .55)**2 + (V2 - .35)**2))
             + .8 * np.exp(-14 * ((V1 - .35)**2 + (V2 + .45)**2))
             + .15 * np.sin(7 * V1 + 2 * V2) * np.cos(5 * V2))
    target0 = moments(exact, phi)
    baseline, idea = exact.copy(), exact.copy()
    base_drift, idea_drift, base_err, idea_err = [], [], [], []
    for _ in range(steps):
        # Exact update conserves the four selected moments.
        increment = zero_moment_part(rng.normal(size=(n, n)), phi)
        increment *= 0.004 * np.linalg.norm(exact) / np.linalg.norm(increment)
        exact = exact + increment
        baseline = rank_svd(baseline + increment, rank)
        candidate = idea + increment
        idea = conserve(candidate, rank_svd(candidate, rank), phi)
        base_drift.append(np.linalg.norm(moments(baseline, phi) - target0))
        idea_drift.append(np.linalg.norm(moments(idea, phi) - target0))
        base_err.append(np.linalg.norm(baseline - exact) / np.linalg.norm(exact))
        idea_err.append(np.linalg.norm(idea - exact) / np.linalg.norm(exact))

    dense = n * n
    factor_units = rank * (2 * n + 1)
    corrected_units = (rank + len(phi)) * (2 * n + 1)
    result = {
        "math_check": {
            "max_moment_residual": float(residual),
            "min_random_feasible_distance_minus_optimal": float(min(random_distances) - optimal_distance),
            "passed": bool(residual < 1e-9 and min(random_distances) > optimal_distance),
        },
        "experiment": {
            "n": n, "rank": rank, "steps": steps,
            "baseline_final_moment_drift_l2": float(base_drift[-1]),
            "idea_final_moment_drift_l2": float(idea_drift[-1]),
            "baseline_max_moment_drift_l2": float(max(base_drift)),
            "idea_max_moment_drift_l2": float(max(idea_drift)),
            "baseline_mean_relative_state_error": float(np.mean(base_err)),
            "idea_mean_relative_state_error": float(np.mean(idea_err)),
            "dense_storage_units": dense,
            "baseline_factor_storage_units": factor_units,
            "idea_factor_storage_upper_bound": corrected_units,
            "baseline_storage_ratio": factor_units / dense,
            "idea_storage_ratio_upper_bound": corrected_units / dense,
        },
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()