import json, math, random
import numpy as np
import torch

SEED = 540
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device = torch.device('cpu')

# 2D nonequilibrium Langevin system: double well plus rotational current.
omega, T, dt, K = 1.25, 0.08, 0.05, 50
A = torch.tensor([-1.0, 0.0], device=device)
B = torch.tensor([1.0, 0.0], device=device)

def mobility(z):
    # Positive, state-dependent diagonal mobility (known ground truth).
    x, y = z[..., 0], z[..., 1]
    return torch.stack((0.35 + 0.85*torch.sigmoid(2.0*x),
                        0.55 + 0.45*torch.sigmoid(-1.5*y)), dim=-1)

def potential_grad(z):
    x, y = z[..., 0], z[..., 1]
    return torch.stack((x*(x*x-1.0), y), dim=-1)

def drift(z):
    g = potential_grad(z)
    # nonequilibrium rotational component, not a gradient or time reversal
    rot = torch.stack((-z[..., 1], z[..., 0]), dim=-1)
    return -g + omega*rot

def action(path):
    z0, z1 = path[:-1], path[1:]
    m = mobility(z0)
    v = (z1-z0)/dt
    residual = v-drift(z0)
    # a = 2 T M; FW action = .5 dt residual^T a^-1 residual
    return (0.25*dt/T * (residual*residual/m).sum(dim=-1)).sum()

def optimize_path(init, steps=1200, lr=0.025):
    inner = torch.nn.Parameter(init[1:-1].clone())
    opt = torch.optim.Adam([inner], lr=lr)
    history=[]
    for i in range(steps):
        path = torch.cat((A[None], inner, B[None]), dim=0)
        loss = action(path)
        opt.zero_grad(); loss.backward(); opt.step()
        if i % 100 == 0: history.append(float(loss.detach().cpu()))
    path = torch.cat((A[None], inner.detach(), B[None]), dim=0)
    return path, history

def check_sde(n=150000):
    # Verify both claims in the stochastic equation: conditional mean drift and variance 2TMdt.
    z = torch.tensor([0.35, -0.7], device=device)
    m = mobility(z); b = drift(z)
    eps = torch.randn(n, 2, device=device)
    inc = b*dt + torch.sqrt(2*T*m*dt)*eps
    mean_err = float(torch.max(torch.abs(inc.mean(0)/dt-b)).cpu())
    var_ratio = (inc.var(0)/(2*T*m*dt)).detach().cpu().numpy()
    return {'mean_abs_error_per_dt': mean_err, 'variance_ratio': var_ratio.tolist(),
            'expected_variance_ratio': [1.0,1.0]}

def stochastic_success(path, mode, n=1200, gain=7.0, seed=1234):
    gen = torch.Generator(device=device); gen.manual_seed(seed)
    z = A[None,:].repeat(n,1)
    min_dist = torch.full((n,), 99.0, device=device)
    for k in range(K):
        if mode == 'optimized':
            # local tube controller following the optimized nonreversible path
            target = path[k]
            r = (path[k+1]-path[k])/dt + gain*(target-z)
        elif mode == 'reverse_relaxation':
            # standard time-reversed-relaxation control, with no path optimization
            r = -drift(z)
        else:
            raise ValueError(mode)
        m = mobility(z)
        noise = torch.randn(z.shape, generator=gen, device=device)
        z = z + r*dt + torch.sqrt(2*T*m*dt)*noise
        min_dist = torch.minimum(min_dist, torch.linalg.vector_norm(z-B, dim=1))
    finite = torch.isfinite(z).all(dim=1)
    dist = torch.linalg.vector_norm(torch.nan_to_num(z, nan=1e6, posinf=1e6, neginf=-1e6)-B, dim=1)
    reach = finite & (dist < 0.35)
    finite_dist = dist[finite]
    mean_dist = float(finite_dist.mean().cpu()) if finite_dist.numel() else float('inf')
    finite_min = min_dist[torch.isfinite(min_dist)]
    mean_min = float(finite_min.mean().cpu()) if finite_min.numel() else float('inf')
    return float(reach.float().mean().cpu()), mean_dist, mean_min, int(finite.sum().cpu())

def main():
    # Core action sanity check: constant path velocity has the stated quadratic residual.
    straight = torch.linspace(0,1,K+1,device=device)[:,None]*B + torch.linspace(1,0,K+1,device=device)[:,None]*A
    straight_s = float(action(straight).cpu())
    # Add small deterministic perturbations and check action rises around this fixed path.
    perturb = straight.clone(); perturb[1:-1,1] += 0.08*torch.sin(torch.arange(1,K,device=device))
    pert_s = float(action(perturb).cpu())
    path, hist = optimize_path(straight + torch.cat((torch.zeros(1,2,device=device), 0.03*torch.randn(K-1,2,device=device), torch.zeros(1,2,device=device))))
    opt_s = float(action(path).cpu())
    check = check_sde()
    base = stochastic_success(path, 'reverse_relaxation')
    idea = stochastic_success(path, 'optimized')
    result = {'device':str(device), 'K':K, 'dt':dt, 'T':T, 'omega':omega,
      'math_check':check, 'action_check':{'straight':straight_s,'perturbed':pert_s,'optimized':opt_s,
                                           'optimization_history':hist},
      'baseline_reverse_relaxation':{'success_rate':base[0],'mean_final_distance':base[1],'mean_min_distance':base[2], 'finite_trajectories':base[3]},
      'idea_action_optimized':{'success_rate':idea[0],'mean_final_distance':idea[1],'mean_min_distance':idea[2], 'finite_trajectories':idea[3]}}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    np.savetxt('optimized_path.csv', path.cpu().numpy(), delimiter=',', header='x,y', comments='')
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()