import json
import math
import numpy as np


def fejer_coeffs(K):
    """p(t)=(1/(K+1))*sum_{j=0}^K t^j."""
    c = np.zeros(K + 1)
    c[:] = 1.0 / (K + 1)
    return c


def jackson_coeffs(K):
    """A Jackson-type polynomial: normalized square of a Fejer amplitude.

    For N=floor(K/2)+1, q(t)=sum_{j=0}^{N-1}t^j/N and p=q^2.
    It has p(1)=1, degree <=K, and on the unit circle its magnitude is
    the squared Fejer-kernel amplitude, giving quadratic suppression away
    from t=1.
    """
    N = K // 2 + 1
    q = np.ones(N) / N
    c = np.convolve(q, q)
    assert len(c) - 1 <= K
    return c


def polyval(c, z):
    out = np.zeros_like(z, dtype=np.complex128)
    for a in c[::-1]:
        out = out * z + a
    return out


def rotation_spectrum(s, n=4096, phase_max=math.pi):
    """J eigenvalues=(1+exp(i phi))/2, with dist(1,sigma(J))=s.
    Dense phases expose the worst-case envelope rather than lucky zeros.
    """
    phi0 = 2.0 * math.asin(min(1.0, s))
    phis = np.linspace(phi0, phase_max, n)
    # Include conjugate pairs implicitly since magnitudes are symmetric.
    return (1.0 + np.exp(1j * phis)) / 2.0, phis


def residual_ratio(lam, coeffs):
    t = 2.0 * lam - 1.0
    p = polyval(coeffs, t)
    return float(np.max(np.abs((1.0 - lam) * p)))


def weighted_residual(lam, coeffs, seed=0):
    rng = np.random.default_rng(seed)
    y = rng.normal(size=len(lam)) + 1j * rng.normal(size=len(lam))
    r0 = np.linalg.norm((1-lam)*y)
    p = polyval(coeffs, 2*lam-1)
    return float(np.linalg.norm((1-lam)*p*y) / r0)


def apply_filter(lam, y, coeffs):
    """Horner application of p(2J-I) to a vector, using J diagonalized here."""
    t = 2*lam - 1
    z = np.zeros_like(y, dtype=np.complex128)
    for a in coeffs[::-1]:
        z = t*z + a*y
    return z


def math_checks():
    checks = {}
    for K in [3, 7, 15, 31]:
        f, j = fejer_coeffs(K), jackson_coeffs(K)
        checks[str(K)] = {
            "fejer_p1_error": abs(polyval(f, np.array([1+0j]))[0]-1),
            "jackson_p1_error": abs(polyval(j, np.array([1+0j]))[0]-1),
            "jackson_degree": len(j)-1,
            "jackson_nonnegative_coeffs": bool(np.min(j) >= 0),
        }
    # Direct polynomial-vs-iterates identity on a random diagonal operator.
    rng = np.random.default_rng(3)
    lam = (1 + np.exp(1j*rng.uniform(-math.pi, math.pi, 20))) / 2
    y = rng.normal(size=20)+1j*rng.normal(size=20)
    K=15; c=fejer_coeffs(K)
    direct = apply_filter(lam, y, c)
    z=y.copy(); avg=np.zeros_like(y)
    for _ in range(K+1):
        avg += z/(K+1)
        z=(2*lam-1)*z
    checks["iterate_identity_error"] = float(np.linalg.norm(direct-avg))
    return checks


def run_sweep():
    rows=[]
    # Include both critical sK=O(1) and increasingly gapped regimes.
    for K in [7, 15, 31, 63]:
        for s in [0.25/K, 0.5/K, 1.0/K, 1.5/K, 2.0/K, 4.0/K, 8.0/K, 16.0/K]:
            s=min(s, .95)
            lam,_=rotation_spectrum(s)
            f=fejer_coeffs(K); j=jackson_coeffs(K)
            rf=residual_ratio(lam,f); rj=residual_ratio(lam,j)
            # The rule in the proposal, with candidate K fixed for this row.
            chosen = "fejer" if s*K < 2 else "jackson"
            ra = rf if chosen=="fejer" else rj
            rows.append({"K":K,"s":s,"sK":s*K,"fejer":rf,"jackson":rj,
                         "adaptive":ra,"choice":chosen,
                         "jackson_scaled":rj*K*K*s,
                         "fejer_scaled":rf*(K+1)})
    return rows


def summary(rows):
    # Ratios are robust envelope quantities on the same dense spectrum.
    critical=[r for r in rows if r["sK"] <= 1.5]
    gapped=[r for r in rows if r["sK"] >= 4]
    # Fit log slope of Jackson residual vs s at each K in gapped range.
    slopes=[]
    for K in sorted(set(r["K"] for r in gapped)):
        q=[r for r in gapped if r["K"]==K]
        slopes.append(float(np.polyfit(np.log([r["s"] for r in q]), np.log([r["jackson"] for r in q]),1)[0]))
    # Candidate adaptive degree selection, measuring best of {3,7,15} with rule.
    adaptive_gain=[]
    for s in [0.01,0.03,0.06,0.12,0.25,0.5]:
        vals=[]
        for K in [3,7,15]:
            lam,_=rotation_spectrum(s)
            f=fejer_coeffs(K); j=jackson_coeffs(K)
            vals.append((residual_ratio(lam,f) if s*K<2 else residual_ratio(lam,j),K))
        adaptive_gain.append({"s":s,"selected":min(vals)[1],"residual":min(vals)[0]})
    return {
      "critical_mean_fejer_times_Kplus1": float(np.mean([r["fejer_scaled"] for r in critical])),
      "critical_max_fejer_times_Kplus1": float(np.max([r["fejer_scaled"] for r in critical])),
      "gapped_jackson_K2s_mean": float(np.mean([r["jackson_scaled"] for r in gapped])),
      "gapped_jackson_log_s_slopes": slopes,
      "adaptive_candidates": adaptive_gain,
      "rows":rows
    }

if __name__ == "__main__":
    out={"math_checks":math_checks()}
    rows=run_sweep(); out["summary"]=summary(rows)
    with open("results.json","w") as f: json.dump(out,f,indent=2)
    print(json.dumps(out["math_checks"],indent=2))
    print(json.dumps({k:v for k,v in out["summary"].items() if k!="rows"},indent=2))


def adaptive_coeffs(K, gap_hat, baseline_residual=None, candidate_residual=None,
                    reject_ratio=1.10):
    """Select Fejer below s*K=2, otherwise Jackson, with residual safeguard."""
    f = fejer_coeffs(K)
    if gap_hat * K < 2.0:
        return f, "fejer", False
    j = jackson_coeffs(K)
    rejected = (baseline_residual is not None and candidate_residual is not None
                and candidate_residual > reject_ratio * baseline_residual)
    return (f, "fejer", True) if rejected else (j, "jackson", False)


def adaptive_apply(lam, y, K, gap_hat, reject_ratio=1.10):
    """Apply the adaptive rule on a diagonal toy operator."""
    f, j = fejer_coeffs(K), jackson_coeffs(K)
    r0 = np.linalg.norm((1-lam)*y)
    yf = apply_filter(lam, y, f)
    rf = np.linalg.norm((1-lam)*yf) / r0
    if gap_hat*K < 2:
        return yf, "fejer", False, rf
    yj = apply_filter(lam, y, j)
    rj = np.linalg.norm((1-lam)*yj) / r0
    if rj > reject_ratio*rf:
        return yf, "fejer", True, rf
    return yj, "jackson", False, rj