import json
import numpy as np


def steering(points, omega, kappa):
    points = np.asarray(points)
    return np.exp(1j * kappa * (omega @ points.T)) / np.sqrt(len(omega))


def orthonormal_columns(A):
    Q, _ = np.linalg.qr(A)
    return Q[:, :A.shape[1]]


def adversarial_rotation(U, eps, rng):
    """Rotate every signal direction toward an orthogonal complement by angle asin(eps)."""
    M, s = U.shape
    R = rng.standard_normal((M, s)) + 1j * rng.standard_normal((M, s))
    R = R - U @ (U.conj().T @ R)
    W = orthonormal_columns(R)
    theta = np.arcsin(eps)
    V = np.cos(theta) * U + np.sin(theta) * W
    V = orthonormal_columns(V)
    return V


def scores(U, phi):
    corr = U.conj().T @ phi
    power = np.sum(np.abs(corr) ** 2, axis=0)
    residual = 1.0 - power
    # A simple cosine-style control: retain only the strongest learned direction.
    cosine = 1.0 - np.max(np.abs(corr) ** 2, axis=0)
    return residual.real, cosine.real


def nearest_local_minimum(values, grid, target, radius=0.18):
    d = np.linalg.norm(grid - target[None, :], axis=1)
    eligible = d <= radius
    idx = np.where(eligible)[0][np.argmin(values[eligible])]
    return grid[idx], float(values[idx])


def main():
    rng = np.random.default_rng(3137)
    M, s, kappa = 64, 3, 18.0
    omega = rng.normal(size=(M, 2))
    omega /= np.linalg.norm(omega, axis=1, keepdims=True)
    targets = np.array([[-0.42, -0.22], [0.08, 0.31], [0.43, -0.34]])
    U = orthonormal_columns(steering(targets, omega, kappa))
    P = U @ U.conj().T

    # Cheap direct verification of |q_tilde-q| <= ||P_tilde-P||.
    check_rows = []
    for eps in [0.0, 0.02, 0.05, 0.10, 0.20, 0.35]:
        V = adversarial_rotation(U, eps, rng)
        Pt = V @ V.conj().T
        projector_distance = np.linalg.norm(Pt - P, 2)
        test_points = rng.uniform(-0.8, 0.8, size=(300, 2))
        ph = steering(test_points, omega, kappa)
        q, _ = scores(U, ph)
        qt, _ = scores(V, ph)
        max_delta = float(np.max(np.abs(qt - q)))
        check_rows.append({"requested_eps": eps, "projector_distance": projector_distance,
                           "max_score_change": max_delta,
                           "bound_slack": projector_distance - max_delta})

    # Localization on a common dense candidate grid, using all three known wells.
    axis = np.linspace(-0.8, 0.8, 81)
    xx, yy = np.meshgrid(axis, axis)
    grid = np.stack([xx.ravel(), yy.ravel()], axis=1)
    grid_phi = steering(grid, omega, kappa)
    eps = 0.20
    V = adversarial_rotation(U, eps, rng)
    q_idea, c_idea = scores(V, grid_phi)
    # Baseline uses the same estimated subspace but a max-cosine similarity rather than
    # accumulated projection residual, which is a common single-component control.
    idea_err, base_err = [], []
    idea_vals, base_vals = [], []
    for t in targets:
        pi, vi = nearest_local_minimum(q_idea, grid, t)
        pb, vb = nearest_local_minimum(c_idea, grid, t)
        idea_err.append(float(np.linalg.norm(pi - t)))
        base_err.append(float(np.linalg.norm(pb - t)))
        idea_vals.append(vi)
        base_vals.append(vb)

    # Repeat rotations to assess graceful degradation, keeping the same Fourier setup.
    degradation = []
    for eps2 in [0.0, 0.05, 0.10, 0.20, 0.35, 0.50]:
        V2 = adversarial_rotation(U, eps2, rng)
        q2, c2 = scores(V2, grid_phi)
        ei, eb = [], []
        for t in targets:
            pi, _ = nearest_local_minimum(q2, grid, t)
            pb, _ = nearest_local_minimum(c2, grid, t)
            ei.append(np.linalg.norm(pi - t)); eb.append(np.linalg.norm(pb - t))
        degradation.append({"eps": eps2, "projector_distance": float(np.linalg.norm(V2@V2.conj().T-P,2)),
                            "idea_mean_localization_error": float(np.mean(ei)),
                            "cosine_mean_localization_error": float(np.mean(eb))})

    out = {"config": {"M": M, "rank": s, "kappa": kappa, "targets": targets.tolist()},
           "math_check": check_rows,
           "localization_eps_0p2": {"idea_mean_error": float(np.mean(idea_err)),
                                     "cosine_mean_error": float(np.mean(base_err)),
                                     "idea_errors": idea_err, "cosine_errors": base_err,
                                     "idea_scores": idea_vals, "cosine_scores": base_vals},
           "degradation": degradation}
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()