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

SEED = 1452

def sphere(rng, n, k):
    x = rng.normal(size=(k, n))
    return x / np.linalg.norm(x, axis=1, keepdims=True)


def estimate_jacobian(theta, beta_fn, n, m, delta, bsp, bmin, sigma, rng):
    u = sphere(rng, n, bsp)
    bp, bm = [], []
    for ui in u:
        ep = rng.normal(0.0, sigma / math.sqrt(bmin), size=m)
        em = rng.normal(0.0, sigma / math.sqrt(bmin), size=m)
        bp.append(beta_fn(theta + delta * ui) + ep)
        bm.append(beta_fn(theta - delta * ui) + em)
    bp, bm = np.asarray(bp), np.asarray(bm)
    return (n / bsp) * ((bp - bm) / (2 * delta)).T @ u


def mean_mse(theta, beta, J, n, m, delta, bsp, bmin, sigma, reps, seed):
    rng = np.random.default_rng(seed)
    vals = []
    for _ in range(reps):
        h = estimate_jacobian(theta, beta, n, m, delta, bsp, bmin, sigma, rng)
        vals.append(float(np.sum((h - J) ** 2)))
    return float(np.mean(vals)), float(np.std(vals) / math.sqrt(reps))


def mechanism_checks():
    rng = np.random.default_rng(SEED)
    n, m = 8, 3
    A = rng.normal(size=(m, n)) / math.sqrt(n)
    theta = rng.normal(size=n)
    q = rng.normal(size=n)
    c = 0.7
    beta_lin = lambda x: A @ x
    beta_non = lambda x: A @ x + c * (q @ x) ** 3 * np.ones(m)
    Jlin = A
    Jnon = A + 3 * c * (q @ theta) ** 2 * np.ones((m, 1)) @ q[None, :]
    out = {}

    # Prediction 1: with a linear response and negligible rollout noise, error falls ~1/b_sp.
    vals = []
    for k in [2, 4, 8, 16, 32, 64]:
        e, se = mean_mse(theta, beta_lin, Jlin, n, m, 0.08, k, 100000, 0.0, 250, SEED + k)
        vals.append((k, e, se))
    slope = np.polyfit(np.log([x[0] for x in vals]), np.log([x[1] for x in vals]), 1)[0]
    out['sphere_scaling'] = {'prediction': 'MSE proportional to 1/b_sp (slope -1)', 'observed': vals, 'loglog_slope': float(slope)}

    # Prediction 2: rollout noise contribution falls as 1/(delta^2*b_sp*b_min).
    vals = []
    for bmin in [4, 16, 64, 256]:
        e, se = mean_mse(theta, beta_lin, Jlin, n, m, 0.10, 16, bmin, 0.8, 300, SEED + bmin)
        vals.append((bmin, e, se))
    slope = np.polyfit(np.log([x[0] for x in vals]), np.log([x[1] for x in vals]), 1)[0]
    out['noise_scaling'] = {'prediction': 'noise MSE proportional to 1/b_min', 'observed': vals, 'loglog_slope': float(slope)}

    # Prediction 3: central-difference nonlinear bias has MSE proportional to delta^4.
    vals = []
    for d in [0.01, 0.02, 0.04, 0.08, 0.16]:
        e, se = mean_mse(theta, beta_non, Jnon, n, m, d, 2000, 100000, 0.0, 30, SEED + int(d * 10000))
        vals.append((d, e, se))
    slope = np.polyfit(np.log([x[0] for x in vals]), np.log([x[1] for x in vals]), 1)[0]
    out['delta_bias'] = {'prediction': 'finite-difference MSE proportional to delta^4', 'observed': vals, 'loglog_slope': float(slope)}
    return out


def optimizer_comparison():
    # Performative quadratic: ordinary loss has gradient theta-y; deployed mean y=beta(theta)=a theta.
    # True objective is 0.5*(theta-a theta)^2 + 0.5*sigma_y^2, whose gradient is (1-a)^2 theta.
    rng = np.random.default_rng(SEED + 99)
    n, m, a = 6, 2, 0.70
    B = np.zeros((m, n)); B[:, :m] = a * np.eye(m)
    target = np.ones(m)
    theta0 = np.full(n, 1.5)
    alpha = 0.18
    steps = 80
    def beta(x): return B @ x
    def run(method):
        theta = theta0.copy(); losses=[]
        for t in range(steps):
            # Current-environment sample mean; loss = .5 ||theta[:m]-y||^2.
            ymean = beta(theta)
            gtheta = np.zeros(n); gtheta[:m] = theta[:m] - ymean - target
            # g_beta for .5||theta-y-beta target||^2 with y=beta+target convention.
            # Use a fixed differentiable surrogate: loss(beta,theta)=.5||theta[:m]-beta-target||^2.
            gbeta = -(theta[:m] - ymean - target)
            if method == 'sgd':
                g = gtheta
            elif method == 'oracle':
                g = gtheta + B.T @ gbeta
            else:
                h = estimate_jacobian(theta, beta, n, m, 0.04, 16, 128, 0.20, rng)
                g = gtheta + h.T @ gbeta
            theta -= alpha * g
            residual = theta[:m] - beta(theta) - target
            losses.append(0.5 * float(residual @ residual))
        return losses
    results = {k: run(k) for k in ['sgd', 'sphere', 'oracle']}
    return {'final_loss': {k: float(v[-1]) for k,v in results.items()},
            'loss_at_20': {k: float(v[19]) for k,v in results.items()},
            'a': a, 'steps': steps, 'alpha': alpha}


def main():
    checks = mechanism_checks()
    comparison = optimizer_comparison()
    report = {'seed': SEED, 'checks': checks, 'optimizer_comparison': comparison}
    Path('results.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()