"""MVP for a locally-private spectral whitening front-end.

The Gaussian noise mechanism here is the implementation-plan provisional
mechanism, not the paper's optimal Gaussian-specific LDP estimator.
"""
import json
from pathlib import Path
import numpy as np

SEED = 145
rng = np.random.default_rng(SEED)


def ar1(n, t, phi=0.85, rng=None):
    rng = np.random.default_rng() if rng is None else rng
    x = np.zeros((n, t))
    x[:, 0] = rng.normal(size=n)
    e = rng.normal(size=(n, t))
    for k in range(1, t):
        x[:, k] = phi * x[:, k - 1] + np.sqrt(1 - phi * phi) * e[:, k]
    return x


def autocov_features(x, max_lag, B):
    """Clipped, centered per-sequence autocovariances, lags 0..max_lag."""
    z = np.clip(x, -B, B)
    z = z - z.mean(axis=1, keepdims=True)
    n, t = z.shape
    out = np.empty((n, max_lag + 1))
    for h in range(max_lag + 1):
        out[:, h] = np.mean(z[:, :t-h] * z[:, h:], axis=1)
    return out


def estimate_spectrum(x, alpha=None, B=2.5, smooth=3, max_lag=None, rng=None):
    """Estimate gamma then FFT it; alpha=None is the nonprivate control."""
    rng = np.random.default_rng() if rng is None else rng
    n, t = x.shape
    max_lag = t // 2 if max_lag is None else max_lag
    g = autocov_features(x, max_lag, B)
    if alpha is not None:
        # Clipping makes each feature bounded by B^2. This is a deliberately
        # simple Gaussian-noise local mechanism, as proposed for the MVP.
        sigma = 2.0 * B * B / alpha
        g = g + rng.normal(scale=sigma, size=g.shape)
    gamma = np.zeros(t)
    gamma[:max_lag + 1] = g.mean(axis=0)
    # A real even covariance sequence is required for a real nonnegative PSD.
    for h in range(1, max_lag + 1):
        if h < t:
            gamma[-h] = gamma[h]
    f = np.real(np.fft.fft(gamma))
    if smooth > 1:
        pad = np.r_[f[-smooth//2:], f, f[:smooth//2]]
        f = np.convolve(pad, np.ones(smooth) / smooth, mode='valid')[:t]
    f = 0.5 * (f + f[::-1])  # numerical symmetrization (FFT-bin reversal)
    return np.maximum(f, 1e-3)


def whiten(x, f, eps=0.05):
    z = x - x.mean(axis=1, keepdims=True)
    w = 1.0 / np.sqrt(np.maximum(f, eps))
    y = np.real(np.fft.ifft(np.fft.fft(z, axis=1) * w[None, :], axis=1))
    return y


def periodogram(x):
    z = x - x.mean(axis=1, keepdims=True)
    return np.mean(np.abs(np.fft.fft(z, axis=1)) ** 2, axis=0) / x.shape[1]


def autocorr_energy(x, max_lag=10):
    z = x - x.mean(axis=1, keepdims=True)
    vals = []
    for h in range(1, max_lag + 1):
        a = np.mean(z[:, :z.shape[1]-h] * z[:, h:])
        v = np.mean(z[:, :z.shape[1]-h] ** 2)
        vals.append(float(a / max(v, 1e-12)))
    return float(np.mean(np.square(vals))), vals


def verify_math():
    x = ar1(1, 128, phi=.7, rng=np.random.default_rng(7))[0]
    z = x - x.mean()
    # Wiener-Khinchin: circular autocorrelation FFT equals |FFT|^2.
    lhs = np.real(np.fft.ifft(np.abs(np.fft.fft(z)) ** 2)) / len(z)
    rhs = np.array([np.mean(z * np.roll(z, -h)) for h in range(len(z))])
    identity_err = float(np.max(np.abs(lhs - rhs)))
    f = np.maximum(np.real(np.fft.fft(rhs)), 1e-5)
    y = np.real(np.fft.ifft(np.fft.fft(z) / np.sqrt(f)))
    fy = np.abs(np.fft.fft(y)) ** 2 / len(y)
    # With the exact spectrum, each occupied Fourier bin is flattened.
    flatten_ratio = float(np.std(fy) / np.mean(fy))
    return identity_err, flatten_ratio


def alpha_probe():
    # Repeated independent datasets estimate the same AR spectrum. The
    # private variance should worsen rapidly as alpha decreases; fit slope.
    t, n, reps = 48, 256, 28
    true = np.array([0.85 ** h for h in range(t // 2 + 1)])
    trueg = np.zeros(t); trueg[:len(true)] = true
    for h in range(1, len(true)): trueg[-h] = true[h]
    truef = np.maximum(np.real(np.fft.fft(trueg)), 1e-3)
    out = {}
    for a in [0.5, 1.0, 2.0, 4.0]:
        errs = []
        for r in range(reps):
            xx = ar1(n, t, rng=np.random.default_rng(1000 + r))
            est = estimate_spectrum(xx, alpha=a, B=2.5, smooth=3,
                                    rng=np.random.default_rng(2000 + r))
            errs.append(np.mean((est - truef) ** 2))
        out[str(a)] = float(np.mean(errs))
    # compare alpha^-4 endpoints, a descriptive check rather than a theorem test
    ratio = out['0.5'] / out['4.0']
    return out, float(ratio)


def main():
    identity_err, exact_flatten = verify_math()
    t, ntrain, ntest = 64, 512, 512
    train = ar1(ntrain, t, rng=np.random.default_rng(11))
    test = ar1(ntest, t, rng=np.random.default_rng(12))
    rows = {}
    for name, alpha in [('raw', None), ('nonprivate', None), ('private_alpha1', 1.0), ('private_alpha4', 4.0)]:
        if name == 'raw':
            y = test
        else:
            f = estimate_spectrum(train, alpha=alpha, B=2.5, smooth=5,
                                  rng=np.random.default_rng(300))
            y = whiten(test, f, eps=.05)
        energy, ac = autocorr_energy(y)
        p = periodogram(y)
        rows[name] = {'lag_autocorr_energy': energy,
                      'mean_abs_lag1_10': float(np.mean(np.abs(ac))),
                      'periodogram_cv': float(np.std(p) / max(np.mean(p), 1e-12))}
    alpha_mse, endpoint_ratio = alpha_probe()
    result = {'math': {'fft_identity_max_error': identity_err,
                       'exact_filter_spectrum_cv': exact_flatten},
              'heldout_ar1': rows,
              'alpha_probe_mse': alpha_mse,
              'alpha_0.5_to_alpha_4_mse_ratio': endpoint_ratio,
              'config': {'T': t, 'train_sequences': ntrain, 'test_sequences': ntest,
                         'phi': .85, 'B': 2.5, 'epsilon': .05, 'seed': SEED,
                         'noise': 'Gaussian sigma=2*B^2/alpha'}}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()