#!/usr/bin/env python3
"""MVP verification of a positive finite-rank Markov spectral layer on Z_M.

The cyclic grid is a finite compact-group analogue of SO(2).  DFT convolution
is exactly equivariant to cyclic shifts.  The Fejer kernel is constructed as a
squared Dirichlet kernel, so positivity and unit mass are not assumptions.
"""
import json
import math
from pathlib import Path
import numpy as np

SEED = 1203
rng = np.random.default_rng(SEED)
M = 256


def fejer_kernel(M, L):
    """Nonnegative unit-mass kernel with 2L-1 retained Fourier modes."""
    j = np.arange(M)
    # squared finite geometric sum; works stably at j=0
    z = np.exp(2j * np.pi * j / M)
    d = np.zeros(M, dtype=np.complex128)
    for r in range(L):
        d += z ** r
    k = (np.abs(d) ** 2) / (M * L)
    return k.real


def centered_mask(M, L):
    freqs = np.fft.fftfreq(M) * M
    return (np.abs(freqs) < L).astype(float)


def signed_kernel(M, L):
    return np.fft.ifft(centered_mask(M, L)).real


def conv(x, k):
    return np.fft.ifft(np.fft.fft(x) * np.fft.fft(k)).real


def shift(x, s):
    return np.roll(x, s)


def induced_inf_norm(k):
    # For a circulant convolution, this is exactly its l_infinity operator norm.
    return float(np.sum(np.abs(k)))


def spectral_coefficients(k):
    return np.fft.fft(k).real


def run():
    Ls = [2, 4, 8, 16, 32, 64]
    kernel_rows = []
    attenuation_rows = []
    random_rows = []
    equivariance_errors = []
    order_errors = []

    # Prediction 1: positivity + mass imply Markov properties for every L.
    # Prediction 2: the induced sup norm is exactly one for Fejer, while signed
    # truncation has a growing norm (and can amplify a worst-case input).
    for L in Ls:
        kf = fejer_kernel(M, L)
        ks = signed_kernel(M, L)
        af = spectral_coefficients(kf)
        raw_norm = induced_inf_norm(ks)
        fejer_norm = induced_inf_norm(kf)
        worst = np.sign(np.roll(ks, np.argmax(np.abs(ks))))
        # A direct worst-case witness for the first convolution row.
        witness_ratio = np.max(np.abs(conv(worst, ks))) / np.max(np.abs(worst))
        f = rng.uniform(-1, 1, M)
        q = rng.uniform(-1, 1, M)
        markov_ratio = np.max(np.abs(conv(f, kf) - conv(q, kf))) / np.max(np.abs(f-q))
        # Deterministic order test q=f+r, r>=0: positivity predicts F(q)-F(f)>=0.
        f0 = rng.normal(size=M)
        r0 = rng.random(M)
        order_errors.append(float(max(0.0, -np.min(conv(f0 + r0, kf) - conv(f0, kf)))))
        kernel_rows.append({
            "L": L, "min_fejer_kernel": float(kf.min()),
            "mass_error": float(abs(kf.sum() - 1)),
            "fejer_inf_norm": fejer_norm, "signed_inf_norm": raw_norm,
            "signed_worst_case_ratio": float(witness_ratio),
            "fejer_random_ratio": float(markov_ratio),
        })

        # Prediction 3: Fejer eigenvalue at frequency k is max(1-|k|/L,0).
        checks = []
        for k in [0, max(1, L // 2), L - 1, L, min(M // 2, L + 3)]:
            expected = max(1.0 - abs(k) / L, 0.0)
            actual = af[k % M]
            checks.append({"k": k, "observed": float(actual), "predicted": expected,
                           "abs_error": float(abs(actual - expected))})
        attenuation_rows.append({"L": L, "checks": checks})

        # Exact equivariance test: F(T_s x) = T_s F(x).
        x = rng.normal(size=M)
        s = int(rng.integers(0, M))
        equivariance_errors.append(float(np.max(np.abs(conv(shift(x, s), kf) - shift(conv(x, kf), s)))))

        # Average bounded random-input behavior, plus 20 repeated layers.
        vals = []
        vals_s = []
        for _ in range(100):
            x = rng.uniform(-1, 1, M)
            vals.append(np.max(np.abs(conv(x, kf))) / np.max(np.abs(x)))
            vals_s.append(np.max(np.abs(conv(x, ks))) / np.max(np.abs(x)))
        xf = rng.uniform(-1, 1, M)
        xs = xf.copy()
        for _ in range(20):
            xf, xs = conv(xf, kf), conv(xs, ks)
        random_rows.append({"L": L, "mean_fejer_ratio": float(np.mean(vals)),
                            "mean_signed_ratio": float(np.mean(vals_s)),
                            "20_layer_signed_norm_ratio": float(np.max(np.abs(xs)) / np.max(np.abs(xf)) if np.max(np.abs(xf)) else 0.0)})

    # Correct, reproducible stacked-layer comparison with a common initial signal.
    stack_rows = []
    for L in Ls:
        kf, ks = fejer_kernel(M, L), signed_kernel(M, L)
        x0 = rng.uniform(-1, 1, M)
        xf, xs = x0.copy(), x0.copy()
        for _ in range(20):
            xf, xs = conv(xf, kf), conv(xs, ks)
        stack_rows.append({"L": L, "fejer_20_layer_sup_ratio": float(np.max(abs(xf))/np.max(abs(x0))),
                           "signed_20_layer_sup_ratio": float(np.max(abs(xs))/np.max(abs(x0)))})

    result = {
        "seed": SEED, "M": M,
        "predictions": [
            "For every L, min(K)>=0 and sum(K)=1, hence induced sup norm and random error ratios are <=1.",
            "Matched-rank signed truncation has induced sup norm sum(abs(K)); its worst-case ratio equals that value and generally grows with rank.",
            "Fejer Fourier eigenvalue at integer frequency k is max(1-|k|/L,0), giving linear attenuation to the cutoff and zero beyond it."
        ],
        "kernel_sweep": kernel_rows,
        "attenuation_sweep": attenuation_rows,
        "equivariance_max_errors": equivariance_errors,
        "order_violation_max": max(order_errors),
        "random_and_stack": random_rows,
        "stacked_20_layers": stack_rows,
        "summary": {
            "max_fejer_mass_error": max(r["mass_error"] for r in kernel_rows),
            "min_fejer_kernel": min(r["min_fejer_kernel"] for r in kernel_rows),
            "max_equivariance_error": max(equivariance_errors),
            "max_attenuation_error": max(c["abs_error"] for row in attenuation_rows for c in row["checks"]),
            "signed_norms": [r["signed_inf_norm"] for r in kernel_rows],
            "fejer_norms": [r["fejer_inf_norm"] for r in kernel_rows]
        }
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result["summary"], indent=2))
    print("kernel_sweep")
    for r in kernel_rows:
        print(r)
    print("stacked_20_layers")
    for r in stack_rows:
        print(r)

if __name__ == "__main__":
    run()