import json
import numpy as np


def make_adapter(C, b, M, tol=1e-12):
    n = C.shape[1]
    Minv = np.linalg.inv(M)
    G = C @ Minv @ C.T
    up = Minv @ C.T @ np.linalg.solve(G, b)
    # Euclidean nullspace, then weighted orthonormalize via Gram eigendecomposition.
    U, s, Vh = np.linalg.svd(C, full_matrices=True)
    rank = int(np.sum(s > tol * (s[0] if len(s) else 1.0)))
    N = Vh[rank:].T
    gram = N.T @ M @ N
    ew, Q = np.linalg.eigh(gram)
    keep = ew > tol * max(1.0, ew.max())
    V = N @ Q[:, keep] @ np.diag(1.0 / np.sqrt(ew[keep]))
    return up, V, rank, s


def mass_norm(x, M):
    return float(np.sqrt(max(0.0, x @ M @ x)))


def run():
    rng = np.random.default_rng(7)
    n, m = 12, 3
    # A nontrivial SPD mass matrix and full-row-rank sampled trace operator.
    A = rng.normal(size=(n, n))
    M = A.T @ A + 0.5 * np.eye(n)
    C = np.zeros((m, n))
    C[0, 0] = 1; C[0, 1] = .4
    C[1, 4] = 1; C[1, 7] = -.25
    C[2, 9] = 1; C[2, 10] = .3; C[2, 11] = -.2
    b = np.array([0.7, -0.2, 0.35])
    up, V, rank, s = make_adapter(C, b, M)
    # Core algebra verification.
    algebra = {
        "constraint_up": float(np.linalg.norm(C @ up - b)),
        "constraint_V": float(np.linalg.norm(C @ V)),
        "weighted_orthogonality": float(np.linalg.norm(V.T @ M @ V - np.eye(V.shape[1]))),
        "rank": rank, "latent_dim": int(V.shape[1]),
    }

    # Frozen ambient mechanism: deliberately has normal and tangent components.
    # The tangent part is a skew operator, hence reduced dynamics conserve M-energy.
    R = rng.normal(size=(n, n)); K = R - R.T
    # F(u)=M^{-1}K(u-up)+ normal forcing; tangent projection removes normal forcing.
    Minv = np.linalg.inv(M)
    normal = rng.normal(size=n)
    def F(u):
        return Minv @ (K @ (u - up)) + 2.0 * normal
    z0 = rng.normal(size=V.shape[1])
    u0 = up + V @ z0
    dt, steps = 0.035, 600

    def reduced_rollout():
        z = z0.copy(); us=[]
        for _ in range(steps):
            u = up + V @ z
            z = z + dt * (V.T @ M @ F(u))
            us.append(up + V @ z)
        return np.asarray(us)

    # Same ambient Euler update, then exact affine projection in the M metric.
    # Projection is the standard control and uses the same F and timestep.
    def project(u):
        return u + Minv @ C.T @ np.linalg.solve(C @ Minv @ C.T, b - C @ u)
    def projected_rollout():
        u=u0.copy(); us=[]
        for _ in range(steps):
            u = project(u + dt * F(u)); us.append(u.copy())
        return np.asarray(us)

    # Soft penalty control: ambient update with a finite penalty and no projection.
    # The penalty is applied in M^{-1} C^T coordinates.
    def penalty_rollout(lam=30.0):
        u=u0.copy(); us=[]
        for _ in range(steps):
            u = u + dt * (F(u) + lam * (Minv @ C.T @ (b - C @ u)))
            us.append(u.copy())
        return np.asarray(us)

    red, proj, pen = reduced_rollout(), projected_rollout(), penalty_rollout()
    def stats(arr):
        residual=np.linalg.norm((arr @ C.T) - b[None,:], axis=1)
        e0=mass_norm(u0,M)
        drift=np.abs(np.array([mass_norm(x,M) for x in arr])-e0)/e0
        return {"max_constraint_residual":float(residual.max()),
                "final_constraint_residual":float(residual[-1]),
                "max_relative_mass_norm_drift":float(drift.max()),
                "final_relative_mass_norm_drift":float(drift[-1])}
    # With a complete weighted nullspace basis, reduced Euler equals the
    # corresponding M-metric projection of one ambient Euler step.
    u_test = u0 + dt * F(u0)
    z_test = z0 + dt * (V.T @ M @ F(u0))
    one_step_difference = np.linalg.norm((up + V @ z_test) - project(u_test))
    penalty_sweep = {str(lam): stats(penalty_rollout(lam))
                     for lam in (1.0, 10.0, 30.0, 100.0)}
    result={"seed":7,"dt":dt,"steps":steps,"algebra":algebra,
            "one_step_reduced_vs_projection_difference":float(one_step_difference),
            "reduced_coordinate":stats(red),"projection_control":stats(proj),
            "soft_penalty_lambda_30":stats(pen),"soft_penalty_sweep":penalty_sweep}
    # A deliberately unconstrained version demonstrates why the adapter matters.
    u=u0.copy(); raw=[]
    for _ in range(steps):
        u=u+dt*F(u); raw.append(u.copy())
    result["unconstrained_ambient"] = stats(np.asarray(raw))
    print(json.dumps(result, indent=2))

if __name__ == '__main__': run()