import json
import math
from pathlib import Path
import numpy as np


def shift_batch(x, k):
    return np.roll(x, k, axis=1)


def orbit_average(x, q):
    return sum(shift_batch(x, k) for k in range(q)) / float(q)


def ridge_fit(x, y, reg=1e-2):
    xb = np.c_[np.ones(len(x)), x]
    a = xb.T @ xb + reg * np.eye(xb.shape[1])
    a[0, 0] -= reg
    return np.linalg.solve(a, xb.T @ y)


def predict(w, x):
    return np.c_[np.ones(len(x)), x] @ w


def mse(w, x, y):
    return float(np.mean((predict(w, x) - y) ** 2))


def make_data(n, length, asymmetry, noise, seed):
    rng = np.random.default_rng(seed)
    # Random periodic curves with a low-frequency invariant component.
    t = np.arange(length)[None, :]
    phase = rng.uniform(0, 2 * np.pi, size=(n, 1))
    amp = rng.normal(size=(n, 1))
    x = amp * np.cos(2 * np.pi * t / length + phase)
    x += 0.35 * rng.normal(size=(n, length))
    # Label is mostly rotation invariant, with a controllable absolute-phase term.
    invariant = amp[:, 0]
    phase_sensitive = x[:, 0]
    y = invariant + asymmetry * phase_sensitive + noise * rng.normal(size=n)
    return x, y


def estimate_A(w, x, y, q):
    # Displayed proxy: mean transformed-input loss minus original loss.
    original = (predict(w, x) - y) ** 2
    transformed = np.stack([(predict(w, shift_batch(x, k)) - y) ** 2 for k in range(q)], 1)
    return float(np.mean(transformed) - np.mean(original))


def consistency(w, x, q):
    p0 = predict(w, x)
    ps = np.stack([predict(w, shift_batch(x, k)) for k in range(q)], 1)
    return float(np.mean((ps - p0[:, None]) ** 2))


def run_case(asymmetry, n_train=160, n_val=160, n_test=160, seed=0):
    L = 8
    xt, yt = make_data(n_train, L, asymmetry, .20, seed)
    xv, yv = make_data(n_val, L, asymmetry, .20, seed + 1)
    xe, ye = make_data(n_test, L, asymmetry, .20, seed + 2)
    qs = [1, 2, 4, 8]
    # Unrestricted pilot estimates the invariance-induced loss discrepancy.
    w0 = ridge_fit(xt, yt)
    A = {q: max(0.0, estimate_A(w0, xv, yv, q)) for q in qs}
    D = {q: consistency(w0, xv, q) for q in qs}
    beta = 1.0
    n, m, qsharp = n_train, L, 8
    scores = {q: (n * m * m * q) ** (-beta / (beta + 1)) + A[q] for q in qs if q <= qsharp}
    selected = min(scores, key=scores.get)
    models = {}
    test_losses = {}
    for q in qs:
        zt, zv = orbit_average(xt, q), orbit_average(xv, q)
        ze = orbit_average(xe, q)
        w = ridge_fit(zt, yt)
        models[q] = w
        test_losses[q] = mse(w, ze, ye)
    # Also report a validation-loss selector as a useful practical reference.
    val_losses = {q: mse(models[q], orbit_average(xv, q), yv) for q in qs}
    return {
        'asymmetry': asymmetry, 'A_hat': A, 'D_q': D, 'oracle_scores': scores,
        'selected_q': selected, 'test_mse': test_losses, 'val_mse': val_losses,
        'q1_test': test_losses[1], 'qmax_test': test_losses[8],
        'selected_test': test_losses[selected]
    }


def math_sanity():
    qs = np.array([1, 2, 4, 8], dtype=float)
    variance = (160 * 8 * 8 * qs) ** (-1 / 2)
    bias = np.array([0.0, .003, .004, .07])
    score = variance + bias
    assert np.all(np.diff(variance) < 0), variance
    assert np.all(np.diff(bias) >= 0), bias
    assert int(qs[np.argmin(score)]) == 4
    # Directly verify averaging is invariant to every shift in its orbit.
    rng = np.random.default_rng(3)
    x = rng.normal(size=(5, 8))
    z = orbit_average(x, 4)
    assert np.max(np.abs(z - shift_batch(z, 1))) > 1e-6  # prefix orbit is not full-group invariant
    zfull = orbit_average(x, 8)
    assert np.max(np.abs(zfull - shift_batch(zfull, 1))) < 1e-12
    return {'variance_proxy': variance.tolist(), 'bias': bias.tolist(), 'score': score.tolist(), 'argmin_q': 4}


def main():
    sanity = math_sanity()
    cases = [run_case(a, seed=20 + i * 10) for i, a in enumerate([0.0, 0.15, 0.5, 1.0])]
    out = {'math_sanity': sanity, 'cases': cases}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    main()