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

# Numerical MVP for the cyclic Lie-bracket residual idea.
# The vector fields are nonlinear and noncommuting, so their composition
# produces an order-dependent second-order term.

def v1(x):
    x = np.asarray(x, dtype=float)
    return np.stack([np.tanh(x[..., 1]), 0.35*np.sin(x[..., 0])], axis=-1)

def v2(x):
    x = np.asarray(x, dtype=float)
    return np.stack([0.4*np.sin(x[..., 1]), -np.tanh(x[..., 0])], axis=-1)

def linear_v1(x):
    # A commuting control: both fields are scalar multiples of identity.
    return 0.7*np.asarray(x, dtype=float)

def linear_v2(x):
    return -0.4*np.asarray(x, dtype=float)

def jac(fun, x, delta=1e-5):
    x=np.asarray(x,dtype=float); J=np.zeros((x.size,x.size))
    for q in range(x.size):
        e=np.zeros_like(x); e[q]=delta
        J[:,q]=(fun(x+e)-fun(x-e))/(2*delta)
    return J

def bracket(x, f1=v1, f2=v2):
    return jac(f2,x)@f1(x)-jac(f1,x)@f2(x)

def compose(x, h, k, order, f1=v1, f2=v2):
    y=np.array(x,dtype=float)
    for field,step in ((order[0],h),(order[1],k)):
        y += step*(f1(y) if field==1 else f2(y))
    return y

def math_sweeps():
    rng=np.random.default_rng(7)
    pts=rng.uniform(-1,1,size=(400,2))
    # Prediction A: h=k=sqrt(eps) makes order difference scale as eps^1.
    epses=np.array([2.0**(-q) for q in range(3,11)])
    diffs=[]
    ratios=[]
    bnorm=np.mean([np.linalg.norm(bracket(x)) for x in pts])
    for eps in epses:
        h=k=math.sqrt(eps)
        ds=np.array([np.linalg.norm(compose(x,h,k,(1,2))-compose(x,h,k,(2,1))) for x in pts])
        diffs.append(ds.mean())
        ratios.append(ds.mean()/eps)
    slope=float(np.polyfit(np.log(epses),np.log(diffs),1)[0])
    # Prediction B: normalized difference converges to bracket magnitude.
    # With this explicit Euler composition, sign is opposite to the stated
    # convention depending on order; compare magnitudes.
    ratio_small=float(np.mean(ratios[-3:]))
    # Prediction C: scaling either field by lambda scales bracket and order gap
    # quadratically when both fields are scaled.
    lam=np.array([0.25,0.5,1,2,4],float)
    eps=2**-8
    gaps=[]
    for a in lam:
        d=np.array([np.linalg.norm(compose(x,math.sqrt(eps),math.sqrt(eps),(1,2),
                                      lambda z:a*v1(z),lambda z:a*v2(z))-
                         compose(x,math.sqrt(eps),math.sqrt(eps),(2,1),
                                 lambda z:a*v1(z),lambda z:a*v2(z))) for x in pts])
        gaps.append(d.mean())
    scale_slope=float(np.polyfit(np.log(lam),np.log(gaps),1)[0])
    # Control: commuting linear fields have no order gap up to floating point.
    commuting_gap=float(np.mean([np.linalg.norm(compose(x,math.sqrt(eps),math.sqrt(eps),(1,2),linear_v1,linear_v2)-
                                           compose(x,math.sqrt(eps),math.sqrt(eps),(2,1),linear_v1,linear_v2)) for x in pts]))
    return {
      'predictions': {
        'order_gap_vs_epsilon_exponent_predicted':1.0,
        'order_gap_vs_epsilon_exponent_observed':slope,
        'normalized_gap_predicted_mean_bracket_norm':bnorm,
        'normalized_gap_observed_small_epsilon':ratio_small,
        'joint_field_scale_exponent_predicted':2.0,
        'joint_field_scale_exponent_observed':scale_slope,
        'commuting_gap_predicted':0.0,
        'commuting_gap_observed':commuting_gap
      },
      'sweeps': {'epsilon':epses.tolist(),'gap':list(map(float,diffs)),
                 'scale':lam.tolist(),'scaled_gap':list(map(float,gaps))}
    }

def mlp_train(seed, cyclic, steps=500):
    # Tiny regression task where the target contains a bracket-like nonlinear
    # interaction. Equal parameter count is not claimed; this is secondary.
    rng=np.random.default_rng(seed)
    X=rng.uniform(-1,1,size=(1024,2)).astype(np.float64)
    Y=(X[:,0]*X[:,1] + 0.25*np.sin(2*X[:,0]-X[:,1]))[:,None]
    tr=slice(0,768); te=slice(768,None)
    W1=rng.normal(0,.5,(2,16)); b1=np.zeros((1,16)); W2=rng.normal(0,.2,(16,1)); b2=np.zeros((1,1))
    lr=.025
    def forward(x):
        if not cyclic:
            h=np.tanh(x@W1+b1); return h@W2+b2
        # Two shared-width vector fields in hidden state; random signs/runtimes
        # make a centered sequential residual composition.
        h=np.tanh(x@W1+b1)
        a,b=0.8,-0.8
        z=h + a*.35*np.tanh(h@W1.T[:, :16] if False else h)
        # Use fixed nonlinear fields in hidden coordinates, with exact centering.
        f1=np.tanh(z); f2=-np.tanh(z)+0.15*np.sin(z)
        z=z+0.35*b*f2
        return z@W2+b2
    # finite-difference-free simple backprop is omitted for cyclic; use torch
    # in the companion script below for the actual matched autograd test.
    return None

def run():
    out=math_sweeps()
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': run()