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

SEED = 2378
C = 0.587
SQRT2PI = math.sqrt(2.0 * math.pi)
random.seed(SEED); np.random.seed(SEED)

def normal_cdf(x):
    return 0.5 * (1.0 + np.vectorize(math.erf)(x / math.sqrt(2.0)))

def run_math_check():
    rng = np.random.default_rng(SEED); rows = []
    for d in (1, 2, 8):
        x = rng.normal(size=(200000, d))
        jac = 2.0 / SQRT2PI * np.exp(-0.5 * x*x)
        empirical = float(np.max(jac)); directional = C * 2.0
        global_bound = C * 2.0 * math.sqrt(d)
        rows.append({'d': d, 'empirical_max_jacobian': empirical,
                     'exact_supremum': math.sqrt(2.0/math.pi),
                     'directional_bound': directional, 'global_bound': global_bound,
                     'directional_holds': empirical <= directional,
                     'global_holds': empirical <= global_bound})
    return rows

def train_comparison(d, steps=400, batch=128):
    try:
        import torch
        torch.manual_seed(SEED+d)
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        try: _ = torch.zeros(1, device=device) + 1
        except Exception: device = 'cpu'
    except Exception as e: return {'d': d, 'error': str(e)}
    rng = np.random.default_rng(SEED+d)
    x_np = rng.normal(size=(4096,d)).astype('float32')
    target_np = (2*normal_cdf(x_np)-1).astype('float32')
    x = torch.tensor(x_np, device=device); target = torch.tensor(target_np, device=device)
    L = C*2*math.sqrt(d); out = {'d': d, 'device': device, 'L_star': L, 'initial_scale': 4.0}
    for name, capped in [('baseline_uncapped',False), ('certificate_projected',True)]:
        torch.manual_seed(SEED+d+(1000 if capped else 0))
        scale = torch.nn.Parameter(torch.tensor(4.0,device=device)); opt = torch.optim.Adam([scale],lr=.08)
        post_peak = 0.; violations = 0; losses=[]
        for _ in range(steps):
            idx=torch.randint(0,x.shape[0],(batch,),device=device)
            loss=((scale*torch.tanh(x[idx])-target[idx])**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
            if capped:
                with torch.no_grad(): scale.clamp_(0.,L)
            with torch.no_grad():
                a=float(scale.cpu()); post_peak=max(post_peak,a); violations += int(a>L+1e-6)
                losses.append(float(loss.cpu()))
        with torch.no_grad():
            a=float(scale.cpu()); mse=float(((a*torch.tanh(x)-target)**2).mean().cpu())
        out[name]={'final_scale':a,'post_update_peak_jacobian':post_peak,
                   'cap_violations_after_updates':violations,'final_mse':mse,
                   'mean_last50_mse':float(np.mean(losses[-50:]))}
    return out

def main():
    out={'math_check':run_math_check(), 'experiments':[train_comparison(d) for d in (1,2,8)]}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()