import numpy as np
from scipy.stats import linregress

SEED = 3164


def correlated_field(n, beta=1.5, k0=1.0, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    w = rng.normal(size=n)
    fk = np.fft.rfft(w)
    k = np.arange(len(fk), dtype=float)
    filt = (k + k0) ** (-beta / 2.0)
    g = np.fft.irfft(fk * filt, n=n)
    return (g - g.mean()) / (g.std() + 1e-12)


def make_mixer(n, alpha=1.3, beta=1.5, correlated=True, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    g = correlated_field(n, beta=beta, rng=rng) if correlated else rng.normal(size=n)
    sigma = np.exp(g - 0.5 * np.var(g))
    signs = rng.choice([-1.0, 1.0], size=(n, n))
    signs = np.triu(signs, 1)
    signs = signs + signs.T
    d = np.abs(np.arange(n)[:, None] - np.arange(n)[None, :])
    M = signs * sigma[:, None] * sigma[None, :] * (1.0 + d) ** (-alpha)
    np.fill_diagonal(M, 0.0)
    raw_norm = np.linalg.svd(M, compute_uv=False)[0]
    return M / (raw_norm + 1e-12), raw_norm, sigma


def local_mixer(n, width=2, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    d = np.abs(np.arange(n)[:, None] - np.arange(n)[None, :])
    M = rng.normal(size=(n, n)) * (d <= width) * (d > 0)
    return M / (np.linalg.svd(M, compute_uv=False)[0] + 1e-12)


def singular_gap(M):
    s = np.linalg.svd(M, compute_uv=False)
    return s[-1], 1.0 - s[1] / (s[0] + 1e-12)


def participation_entropy(v):
    p = np.asarray(v) ** 2
    p = p / (p.sum() + 1e-12)
    return float(-np.sum(p * np.log(p + 1e-15)))


def propagation_entropy(M, steps=8, trials=80, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    vals = []
    # Absolute influence profile after repeated normalized residual propagation.
    A = np.eye(M.shape[0]) + 0.5 * M
    for _ in range(trials):
        x = rng.normal(size=M.shape[0])
        x /= np.linalg.norm(x)
        for _ in range(steps):
            x = A @ x
            x /= np.linalg.norm(x) + 1e-12
        vals.append(participation_entropy(x))
    return float(np.mean(vals))


def fit_scaling(ns, ys):
    fit = linregress(np.log(ns), np.log(np.maximum(ys, 1e-14)))
    return -fit.slope, fit.rvalue ** 2


def main():
    rng = np.random.default_rng(SEED)
    print('CORE MATH CHECK')
    n = 96
    M, raw_norm, _ = make_mixer(n, correlated=True, rng=rng)
    norm = np.linalg.svd(M, compute_uv=False)[0]
    print(f'raw_norm={raw_norm:.6f} normalized_norm={norm:.6f}')
    for c in (0.25, 0.5, 0.9):
        A = np.eye(n) + (c / (norm + 1e-12)) * M
        # The residual update has bounded one-step amplification <= 1+c.
        amp = np.linalg.svd(A, compute_uv=False)[0]
        print(f'c={c:.2f} observed_amp={amp:.6f} bound={1+c:.6f}')
    # Correlation check: edges sharing a node have correlated magnitudes only for the field model.
    def shared_edge_corr(correlated):
        a, _, s = make_mixer(128, correlated=correlated, rng=rng)
        # Remove deterministic distance effect by examining adjacent edges from a common node.
        x = np.abs(a[64, 1:64])
        y = np.abs(a[64, 65:128])
        return np.corrcoef(x, y)[0, 1]
    print(f'shared_endpoint_abs_corr iid={shared_edge_corr(False):.4f} correlated={shared_edge_corr(True):.4f}')

    print('FINITE SIZE PROPAGATION ENTROPY')
    ns = [32, 48, 64, 96, 128, 192]
    result = {}
    for name in ('local', 'iid', 'correlated'):
        ents = []
        gaps = []
        for n in ns:
            vals_e, vals_g = [], []
            for rep in range(12):
                rr = np.random.default_rng(SEED + 1000 * rep + n)
                if name == 'local':
                    mm = local_mixer(n, rng=rr)
                else:
                    mm, _, _ = make_mixer(n, beta=1.5, correlated=(name == 'correlated'), rng=rr)
                vals_e.append(propagation_entropy(mm, steps=8, trials=12, rng=rr))
                vals_g.append(singular_gap(mm)[0])
            ents.append(np.mean(vals_e)); gaps.append(np.mean(vals_g))
        # Compare the two proposed entropy forms using residual sum of squares.
        x = np.log(np.asarray(ns, float))
        X1 = np.column_stack([x, np.ones_like(x)])
        X2 = np.column_stack([x*x, x, np.ones_like(x)])
        rss_log = np.sum((np.asarray(ents) - X1 @ np.linalg.lstsq(X1, ents, rcond=None)[0]) ** 2)
        rss_log2 = np.sum((np.asarray(ents) - X2 @ np.linalg.lstsq(X2, ents, rcond=None)[0]) ** 2)
        z, r2 = fit_scaling(ns, gaps)
        result[name] = (ents, gaps, z, r2, rss_log, rss_log2)
        print(f'{name:10s} entropy=' + ','.join(f'{v:.4f}' for v in ents))
        print(f'{name:10s} gap=' + ','.join(f'{v:.6f}' for v in gaps) + f' z={z:.3f} R2={r2:.3f} RSS_log={rss_log:.6g} RSS_log2={rss_log2:.6g}')

    print('STABILITY WITHOUT NORMALIZATION')
    mm, raw_norm, _ = make_mixer(96, correlated=True, rng=np.random.default_rng(SEED))
    raw = mm * raw_norm
    x = np.random.default_rng(SEED + 9).normal(size=96); x /= np.linalg.norm(x)
    for c in (0.5, 1.0):
        y = x.copy()
        for _ in range(30):
            y = y + c * raw @ y
        print(f'raw_update_c={c:.1f} final_norm={np.linalg.norm(y):.3e}')


if __name__ == '__main__':
    main()