#!/usr/bin/env python3
"""Numerical checks for a small periodic 3-D KPZ evolution cell.
Run with: python3 kpz_experiment.py
"""
import json
from pathlib import Path
import numpy as np


def derivatives(h, dx=1.0):
    lap = np.zeros_like(h)
    grad2 = np.zeros_like(h)
    for ax in range(3):
        hp = np.roll(h, -1, axis=ax)
        hm = np.roll(h, 1, axis=ax)
        lap += (hp - 2.0 * h + hm) / dx**2
        grad = (hp - hm) / (2.0 * dx)
        grad2 += grad * grad
    return lap, grad2


def kpz_step(h, nu, lam, D, dt, dx, rng, noise=True):
    lap, grad2 = derivatives(h, dx)
    z = rng.standard_normal(h.shape) if noise and D > 0 else 0.0
    return h + dt * (nu * lap + 0.5 * lam * grad2) + np.sqrt(2 * D * dt / dx**3) * z


def noise_variance_check():
    # Prediction: Var(noise increment) = 2 D dt / dx^3.
    rng = np.random.default_rng(11)
    shape = (8, 8, 8)
    D, dt = 0.37, 0.021
    rows = []
    for dx in (0.5, 1.0, 2.0):
        vals = []
        for _ in range(1000):
            vals.append(kpz_step(np.zeros(shape), 0, 0, D, dt, dx, rng)[0, 0, 0])
        observed = float(np.var(vals, ddof=1))
        predicted = 2 * D * dt / dx**3
        rows.append({"dx": dx, "predicted": predicted, "observed": observed,
                     "relative_error": abs(observed-predicted)/predicted})
    return rows


def nonlinear_drift_check():
    # Prediction: for a fixed shape, mean height velocity is linear in lambda,
    # with slope 0.5 * spatial mean(|grad h|^2); lambda=0 has zero nonlinear drift.
    rng = np.random.default_rng(22)
    shape = (16, 16, 16)
    h = rng.normal(size=shape)
    _, grad2 = derivatives(h)
    q = float(grad2.mean())
    rows = []
    for lam in (-2.0, -1.0, 0.0, 1.0, 2.0):
        out = kpz_step(h, 0, lam, 0, 0.01, 1, rng, noise=False)
        observed_velocity = float((out.mean() - h.mean()) / 0.01)
        predicted_velocity = 0.5 * lam * q
        rows.append({"lambda": lam, "predicted_velocity": predicted_velocity,
                     "observed_velocity": observed_velocity,
                     "relative_error": abs(observed_velocity-predicted_velocity)/(abs(predicted_velocity)+1e-12)})
    return {"mean_grad2": q, "rows": rows}


def stability_check():
    # For the 7-point Laplacian, explicit diffusion is stable when
    # dt*nu*2*d/dx^2 <= 1 (d=3). A highest-frequency checkerboard mode
    # should decay below the boundary and grow above it.
    n = 12
    x = np.indices((n, n, n)).sum(axis=0)
    h0 = (-1.0) ** x
    nu, dx = 1.0, 1.0
    boundary = dx**2 / (2 * 3 * nu)
    rows = []
    for factor in (0.8, 1.0, 1.2):
        dt = factor * boundary
        h = h0.copy()
        amplitudes = []
        rng = np.random.default_rng(33)
        for _ in range(8):
            h = kpz_step(h, nu, 0, 0, dt, dx, rng, noise=False)
            amplitudes.append(float(np.max(np.abs(h))))
        # exact checkerboard amplification is |1-12 nu dt/dx^2|.
        predicted_amp = abs(1 - 12 * nu * dt / dx**2)
        observed_amp = amplitudes[0]
        rows.append({"factor_of_bound": factor, "dt": dt,
                     "predicted_one_step_amplification": predicted_amp,
                     "observed_one_step_amplification": observed_amp,
                     "final_amplitude": amplitudes[-1],
                     "stable_observed": amplitudes[-1] <= 1.0 + 1e-9})
    return {"predicted_dt_bound": boundary, "rows": rows}


def roughening_control():
    # Small same-setup rollout: KPZ has nonlinear mean growth whereas
    # diffusion-only preserves the mean exactly (periodic boundaries).
    rng_k, rng_d = np.random.default_rng(44), np.random.default_rng(44)
    shape = (16, 16, 16)
    h_k = np.zeros(shape); h_d = np.zeros(shape)
    nu, lam, D, dt, dx = 0.15, 2.0, 0.03, 0.01, 1.0
    kpz_means, diff_means, kpz_rms, diff_rms = [], [], [], []
    for _ in range(100):
        h_k = kpz_step(h_k, nu, lam, D, dt, dx, rng_k)
        h_d = kpz_step(h_d, nu, 0, D, dt, dx, rng_d)
        kpz_means.append(float(h_k.mean())); diff_means.append(float(h_d.mean()))
        kpz_rms.append(float(h_k.std())); diff_rms.append(float(h_d.std()))
    return {"kpz_final_mean": kpz_means[-1], "diffusion_final_mean": diff_means[-1],
            "kpz_final_rms": kpz_rms[-1], "diffusion_final_rms": diff_rms[-1],
            "kpz_mean_abs_change": abs(kpz_means[-1]),
            "diffusion_mean_abs_change": abs(diff_means[-1])}


def main():
    result = {"noise_variance": noise_variance_check(),
              "nonlinear_drift": nonlinear_drift_check(),
              "stability": stability_check(),
              "same_setup_control": roughening_control()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()

try:
    import torch
    from torch import nn

    class KPZCell(nn.Module):
        """Differentiable periodic KPZ transition for [B,C,X,Y,Z] latent fields."""
        def __init__(self, nu=0.1, lam=0.0, D=0.0, dt=0.01, dx=1.0,
                     stochastic=True):
            super().__init__()
            self.nu_raw = nn.Parameter(torch.tensor(float(np.log(np.expm1(nu)))))
            self.D_raw = nn.Parameter(torch.tensor(float(np.log(np.expm1(D))))) if D > 0 else nn.Parameter(torch.tensor(-20.0))
            self.lambda_ = nn.Parameter(torch.tensor(float(lam)))
            self.dt, self.dx, self.stochastic = float(dt), float(dx), stochastic

        def forward(self, h, noise=None):
            lap = torch.zeros_like(h)
            grad2 = torch.zeros_like(h)
            for axis in (2, 3, 4):
                hp = torch.roll(h, -1, dims=axis)
                hm = torch.roll(h, 1, dims=axis)
                lap = lap + (hp - 2 * h + hm) / self.dx ** 2
                grad2 = grad2 + ((hp - hm) / (2 * self.dx)) ** 2
            nu = torch.nn.functional.softplus(self.nu_raw)
            D = torch.nn.functional.softplus(self.D_raw)
            if noise is None:
                noise = torch.randn_like(h) if self.training and self.stochastic else torch.zeros_like(h)
            return h + self.dt * (nu * lap + 0.5 * self.lambda_ * grad2) + torch.sqrt(2 * D * self.dt / self.dx ** 3) * noise
except ImportError:
    KPZCell = None