"""Symmetry-resolved Fourier bifurcation monitor and toy experiment."""
import json
import numpy as np


def cyclic_shift(n):
    S = np.zeros((n, n), dtype=complex)
    for i in range(n):
        S[(i + 1) % n, i] = 1.0
    return S


def spatial_fourier(n):
    j = np.arange(n)
    # columns are eigenvectors of the cyclic shift
    return np.exp(2j * np.pi * np.outer(j, j) / n) / np.sqrt(n)


def circulant_from_modes(modes):
    n = len(modes)
    U = spatial_fourier(n)
    return U @ np.diag(modes) @ U.conj().T


def sector_projectors(n):
    U = spatial_fourier(n)
    return [np.outer(U[:, k], U[:, k].conj()) for k in range(n)]


def restricted_spectrum(M, n=None):
    """Return each cyclic sector's scalar block and margin |1-lambda|."""
    if n is None:
        n = M.shape[0]
    U = spatial_fourier(n)
    blocks = np.diag(U.conj().T @ M @ U)
    margins = np.abs(1.0 - blocks)
    return blocks, margins


def temporal_fft_symbols(period_map):
    """Temporal Fourier coefficients of a sequence of one-step matrices."""
    period_map = np.asarray(period_map)
    return np.fft.fft(period_map, axis=0) / len(period_map)


def winding_number(characteristic_values):
    """Discrete phase winding around a closed contour (integer after rounding)."""
    z = np.asarray(characteristic_values)
    phase = np.unwrap(np.angle(z))
    return int(np.rint((phase[-1] - phase[0]) / (2 * np.pi)))


def circle_winding(eigenvalue, radius=0.15, points=512):
    """Winding of chi(mu)=mu-lambda as mu traverses a circle centered at 1."""
    theta = np.linspace(0, 2*np.pi, points + 1)
    mu = 1 + radius * np.exp(1j * theta)
    return winding_number(mu - eigenvalue)


def run():
    rng = np.random.default_rng(7)
    n = 8
    # A nontrivial equivariant ring: each spatial Fourier sector has a known gain.
    # Sector 3 is deliberately the first one to hit the discrete-time unit circle.
    base = np.array([0.62, 0.71, 0.80, 0.93, 0.77, 0.80, 0.71, 0.62])
    boundary_rows = []
    for gamma in np.linspace(0.70, 1.30, 25):
        M = circulant_from_modes(gamma * base)
        lam, margins = restricted_spectrum(M, n)
        boundary_rows.append((gamma, float(np.max(np.abs(lam))), int(np.argmin(np.abs(1-lam)))))
    # Predicted first unit-circle crossing is gamma=1/max(base)=1/base[3].
    predicted_boundary = 1.0 / np.max(base)
    observed_boundary = min(boundary_rows, key=lambda x: abs(x[1]-1))[0]

    # Margin law: for the critical real sector, m=|1-gamma*a|, slope a.
    gammas = np.linspace(0.55, 1.02, 48)
    margins = []
    for gamma in gammas:
        lam, mm = restricted_spectrum(circulant_from_modes(gamma * base), n)
        margins.append(mm[3])
    margins = np.asarray(margins)
    fit = np.polyfit(gammas[gammas < predicted_boundary], margins[gammas < predicted_boundary], 1)
    predicted_slope = -base[3]
    observed_slope = fit[0]

    # Critical-sector prediction: add a small perturbation and measure the dominant FFT mode.
    gamma = 1.03
    M = circulant_from_modes(gamma * base)
    x = rng.normal(size=n)
    amplitudes = []
    for _ in range(20):
        x = M.real @ x
        amplitudes.append(np.abs(np.fft.fft(x)))
    dominant_sector = int(np.argmax(np.mean(amplitudes[-8:], axis=0)))

    # Winding transition: circle centered at 1 encloses lambda iff its distance is < radius.
    w_inside = circle_winding(1.0 + 0.05j, radius=0.15)
    w_outside = circle_winding(1.0 + 0.30j, radius=0.15)

    # Mini comparison: optimize a scalar gain per sector toward target gains.
    # Baseline uses one global LR multiplier from the worst full-Jacobian margin;
    # blockwise monitor damps only the near-critical sector.
    target = np.array([0.45, 0.52, 0.58, 0.985, 0.50, 0.58, 0.52, 0.45])
    init = np.array([0.55, 0.60, 0.66, 1.08, 0.63, 0.66, 0.60, 0.55])
    def optimize(blockwise):
        p = init.copy()
        losses = []
        for _ in range(80):
            loss = float(np.mean((p-target)**2))
            losses.append(loss)
            grad = 2*(p-target)/n
            # monitor return-map margins; near one gets damping, but no clipping of params
            mm = np.abs(1-p)
            if blockwise:
                scale = np.minimum(1.0, mm/0.20)
                # The critical sector is slowed, other sectors retain the base step.
                step = 0.22 * scale
            else:
                scale = min(1.0, float(np.min(mm))/0.20)
                step = np.full(n, 0.22 * scale)
            p -= step * grad
        return losses, p
    loss_global, final_global = optimize(False)
    loss_block, final_block = optimize(True)

    result = {
        "predictions": {
            "unit_circle_boundary": {"predicted": predicted_boundary, "observed_grid": observed_boundary, "relative_error": abs(observed_boundary-predicted_boundary)/predicted_boundary},
            "margin_slope": {"predicted": predicted_slope, "observed_fit": float(observed_slope), "relative_error": abs(observed_slope-predicted_slope)/abs(predicted_slope)},
            "critical_spatial_sector": {"predicted": 3, "observed_dominant_fft": dominant_sector},
            "winding": {"inside_circle": w_inside, "outside_circle": w_outside}
        },
        "comparison": {
            "global_monitor_final_loss": loss_global[-1],
            "blockwise_monitor_final_loss": loss_block[-1],
            "global_monitor_steps_to_1e-4": next((i for i,x in enumerate(loss_global) if x < 1e-4), None),
            "blockwise_monitor_steps_to_1e-4": next((i for i,x in enumerate(loss_block) if x < 1e-4), None),
            "global_final": final_global.tolist(), "blockwise_final": final_block.tolist()
        },
        "checks": {"boundary_ok_5pct": bool(abs(observed_boundary-predicted_boundary)/predicted_boundary < .05),
                   "slope_ok_10pct": bool(abs(observed_slope-predicted_slope)/abs(predicted_slope) < .10),
                   "sector_ok": bool(dominant_sector == 3),
                   "winding_ok": bool(w_inside == 1 and w_outside == 0)}
    }
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    run()