import json
import math
from pathlib import Path
import numpy as np


def rotation(phi):
    return np.array([[math.cos(phi), -math.sin(phi)],
                     [math.sin(phi), math.cos(phi)]], dtype=float)


def certificate_scalar_rotation(a, phi, theta):
    """Exact min_gamma ||gamma exp(-j theta) a R(phi)-I||_2.

    A planar rotation is a complex scalar. For phase error d in (-pi,pi],
    the positive-gamma optimum is gamma=cos(d)/a when cos(d)>0 and the
    optimum boundary is gamma -> 0 otherwise.
    """
    d = (phi - theta + math.pi) % (2 * math.pi) - math.pi
    c = math.cos(d)
    if c > 0:
        gamma = c / a
        z = abs(math.sin(d))
    else:
        gamma = 0.0
        z = 1.0
    # The certificate is only meaningful for z < 1.
    Gamma = math.asin(min(1.0, z))
    return gamma, z, Gamma, d


def run():
    rng = np.random.default_rng(2193)
    rows = []

    # Prediction 1: Gamma_theta(a R(phi)) = |wrapped(phi-theta)|
    phase_errors = np.linspace(-0.9, 0.9, 19)
    err = []
    for d in phase_errors:
        _, z, g, wrapped = certificate_scalar_rotation(1.7, d, 0.0)
        err.append(abs(g - abs(d)))
    pred1 = float(max(err))

    # Prediction 2: phase composition is additive for products.
    n = 8
    a = 0.91
    phis = rng.uniform(-0.28, 0.28, size=n)
    P = np.eye(2)
    for phi in phis:
        P = (a * rotation(phi)) @ P
    direct_angle = math.atan2(P[1, 0], P[0, 0])
    expected_angle = ((float(np.sum(phis)) + math.pi) % (2 * math.pi)) - math.pi
    pred2 = abs(((direct_angle - expected_angle + math.pi) % (2 * math.pi)) - math.pi)

    # Prediction 3: for phase sum pi, I+P loses invertibility precisely at a^N=1.
    # The predicted distance of the exact product SRG point from -1 equals |1-a^N|.
    gains = np.linspace(0.82, 1.18, 73)
    sigma = []
    margin = []
    for gain in gains:
        P = (gain * rotation(math.pi / n))
        Pn = np.linalg.matrix_power(P, n)
        sigma.append(float(np.linalg.svd(np.eye(2) + Pn, compute_uv=False)[-1]))
        margin.append(abs(1.0 - gain ** n))
    sigma = np.asarray(sigma)
    margin = np.asarray(margin)
    pred3 = float(np.max(np.abs(sigma - margin)))
    observed_cross = float(gains[np.argmin(sigma)])
    predicted_cross = 1.0

    # Prediction 4 / sharp boundary: a product with phase pi has bounded
    # powers below gain 1 and amplification above gain 1.
    horizons = np.arange(1, 13)
    below, above = 0.97, 1.03
    below_norm = np.array([below ** k for k in horizons])
    above_norm = np.array([above ** k for k in horizons])
    # Fit log growth rates; predicted slopes are log(gain).
    slope_below = float(np.polyfit(horizons, np.log(below_norm), 1)[0])
    slope_above = float(np.polyfit(horizons, np.log(above_norm), 1)[0])

    # Small optimization demo: phase-margin penalty moves phase away from the
    # dangerous product angle pi, while a vanilla objective leaves it there.
    # Parameters are block phases; target is deliberately unstable at pi.
    N = 10
    target = np.full(N, math.pi / N)
    def optimize(use_penalty, steps=250, lr=0.08, lam=3.0):
        x = target.copy()
        for _ in range(steps):
            # synthetic fitting pressure toward the dangerous target
            grad = 2.0 * (x - target)
            total = float(np.sum(x))
            # keep total phase at least 0.25 rad from pi (a differentiable
            # hinge proxy for the estimated -1 exclusion margin)
            if use_penalty:
                dist = abs(((total - math.pi + math.pi) % (2 * math.pi)) - math.pi)
                if dist < 0.25:
                    sign = 1.0 if total >= math.pi else -1.0
                    grad += lam * 2.0 * (0.25 - dist) * sign
            x -= lr * grad
        P = np.linalg.matrix_power(rotation(float(np.mean(x))), N)
        return x, float(np.linalg.svd(np.eye(2) + P, compute_uv=False)[-1]), float(abs(1 + np.exp(1j * np.sum(x))))
    x0, s0, m0 = optimize(False)
    x1, s1, m1 = optimize(True)

    result = {
        "certificate_phase_max_abs_error": pred1,
        "composition_angle_abs_error": pred2,
        "invertibility_margin_max_abs_error": pred3,
        "invertibility_crossing_observed": observed_cross,
        "invertibility_crossing_predicted": predicted_cross,
        "power_log_slope_below_observed_predicted": [slope_below, math.log(below)],
        "power_log_slope_above_observed_predicted": [slope_above, math.log(above)],
        "optimization_baseline_sigma_min": s0,
        "optimization_phase_margin_sigma_min": s1,
        "optimization_baseline_margin": m0,
        "optimization_phase_margin_margin": m1,
        "n_blocks": n,
        "gain_sweep": {"gains": gains.tolist(), "sigma_min": sigma.tolist(), "exact_margin": margin.tolist()},
        "phase_sweep": {"errors": phase_errors.tolist()}
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    run()