import json, math, time, random
import numpy as np
from scipy.linalg import hadamard

SEED = 155
np.random.seed(SEED); random.seed(SEED)
N = 16
M = 16

# A has columns a_k; a_k are binary sign vectors. Hadamard rows provide
# A A^T / N = I, hence exactly N evaluations suffice here.
H = hadamard(N).astype(np.float64)
A_exact = H.T  # [N coefficients, N selected evaluations]
lambda_exact = np.full(M, 1.0 / N, dtype=np.float64)

# Cheap algebraic and numerical verification first.
residual = A_exact @ np.diag(lambda_exact) @ A_exact.T - np.eye(N)
rng = np.random.default_rng(SEED)
C = rng.normal(size=(2000, N))
exact_E = ((C @ A_exact) ** 2 * lambda_exact).sum(axis=1)
true_E = (C*C).sum(axis=1)
relative_error = np.max(np.abs(exact_E - true_E) / np.maximum(true_E, 1e-12))
# Compare random uniform sign subsampling, unbiased with the same number of evals.
A_random = rng.choice([-1.0, 1.0], size=(N, M))
random_trials = np.empty((500, len(C)))
for t in range(500):
    Ar = rng.choice([-1.0, 1.0], size=(N, M))
    random_trials[t] = ((C @ Ar) ** 2).mean(axis=1)
rand_rel_rmse = float(np.sqrt(np.mean((random_trials - true_E[None,:])**2)) / np.sqrt(np.mean(true_E**2)))
rand_rel_std_mean = float(np.mean(np.std(random_trials, axis=0) / np.maximum(true_E, 1e-12)))
# Float16 accumulation drift, versus float32 accumulation.
Cf = C.astype(np.float16); Af = A_exact.astype(np.float16)
E16 = ((Cf @ Af).astype(np.float32)**2 * (1.0/N)).sum(axis=1)
mixed_rel = float(np.max(np.abs(E16-true_E)/np.maximum(true_E,1e-12)))

# Torch mini experiment. The only stochasticity in the baseline estimator is
# the uniformly random sign design drawn independently at each forward pass.
import torch
from torch import nn

def get_device():
    if torch.cuda.is_available():
        try:
            torch.zeros(1, device='cuda')
            return torch.device('cuda')
        except Exception:
            pass
    return torch.device('cpu')

device = get_device()
torch.manual_seed(SEED)
if device.type == 'cuda': torch.cuda.manual_seed_all(SEED)

g = torch.Generator(device='cpu').manual_seed(SEED)
X = torch.randn(96, 8, generator=g)
true_w = torch.randn(8, 1, generator=g)
y = torch.tanh(X @ true_w) + 0.05*torch.randn(96,1,generator=g)
X, y = X.to(device), y.to(device)
Aex_t = torch.tensor(A_exact, dtype=torch.float32, device=device)

def make_model():
    torch.manual_seed(SEED + 7)
    return nn.Sequential(nn.Linear(8, 32), nn.Tanh(), nn.Linear(32, N), nn.Tanh(), nn.Linear(N, 1)).to(device)

def run(exact, steps=220):
    model = make_model(); opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    losses=[]; start=time.perf_counter()
    for step in range(steps):
        opt.zero_grad(set_to_none=True)
        c = model[0:4](X)  # coefficient-producing trunk, excluding output layer
        if exact:
            z = c @ Aex_t
            E = (z*z).mean(dim=1, keepdim=True)
        else:
            # Same 16 activation evaluations and unbiased normalization statistic.
            Ar = torch.randint(0, 2, (N,M), device=device, dtype=torch.float32)*2-1
            E = ((c @ Ar)**2).mean(dim=1, keepdim=True)
        cn = c / torch.sqrt(E + 1e-5)
        pred = model[4](cn)
        loss = ((pred-y)**2).mean()
        loss.backward(); opt.step(); losses.append(float(loss.detach().cpu()))
    if device.type == 'cuda': torch.cuda.synchronize()
    return {'final_loss': float(np.mean(losses[-20:])), 'best_loss': float(np.min(losses)),
            'first_loss': losses[0], 'time_sec': time.perf_counter()-start}

# Gradient variance at one fixed c, directly measuring the statistic-induced
# noise that the identity claims to remove.
c0 = torch.randn(32, N, device=device, generator=torch.Generator(device=device).manual_seed(SEED+11), requires_grad=True)
target = torch.randn_like(c0)
def grad_samples(exact, count=80):
    vals=[]; grads=[]
    for i in range(count):
        c = c0.detach().clone().requires_grad_(True)
        if exact: E=((c@Aex_t)**2).mean(dim=1,keepdim=True)
        else:
            Ar=torch.randint(0,2,(N,M),device=device,dtype=torch.float32)*2-1
            E=((c@Ar)**2).mean(dim=1,keepdim=True)
        loss=((c/torch.sqrt(E+1e-5)-target)**2).mean()
        grads.append(torch.autograd.grad(loss,c)[0].detach().cpu().numpy())
        vals.append(float(loss.detach().cpu()))
    G=np.stack(grads)
    return float(np.mean(vals)), float(np.mean(np.var(G,axis=0))), float(np.mean(np.linalg.norm(G-G.mean(0),axis=tuple(range(1,G.ndim)))**2))

grad_exact=grad_samples(True); grad_random=grad_samples(False)
# Make JSON serializable and include enough precision to audit.
result = {
 'seed': SEED, 'N':N, 'evaluations':M, 'device':str(device),
 'math': {'matrix_frobenius_residual':float(np.linalg.norm(residual)),
          'max_relative_energy_error':float(relative_error),
          'random_relative_rmse':rand_rel_rmse,
          'random_mean_relative_std':rand_rel_std_mean,
          'float16_input_float32_accum_max_relative_error':mixed_rel},
 'gradient_variance': {'exact_loss_mean':grad_exact[0], 'random_loss_mean':grad_random[0],
                       'exact_mean_element_variance':grad_exact[1], 'random_mean_element_variance':grad_random[1]},
 'training': {'exact':run(True), 'random':run(False)}
}
print(json.dumps(result, indent=2))