# %% [markdown]
# # A fourth-order beam PINN, from equation to solution
# Solve u''''(x) = q(x) on [0, 1], with u = u' = 0 at both ends.
# This is a nondimensional clamped Euler–Bernoulli beam with EI = 1.
# We manufacture a nonuniform load from u*(x) = x²(1−x)²[1 + 0.3 sin(πx)].
# Install: python -m pip install --pre "omnibias-torch==0.5.0rc1" matplotlib
# The example uses CPU float64. It demonstrates one smooth problem, not general PDE convergence.
# %%
import json
import os
import time
from importlib.metadata import version
from pathlib import Path
import matplotlib.pyplot as plt
import torch
from omnibias.torch.jet import mlp_jet, jet_to_tower

torch.set_num_threads(1)
torch.manual_seed(42)
dtype = torch.float64
started = time.perf_counter()
weights = [torch.nn.Parameter(torch.randn(16, 1, dtype=dtype)),
           torch.nn.Parameter(torch.randn(16, 16, dtype=dtype) / 4),
           torch.nn.Parameter(torch.randn(1, 16, dtype=dtype) / 4)]
biases = [torch.nn.Parameter(torch.zeros(n, dtype=dtype)) for n in (16, 16, 1)]
parameters = [*weights, *biases]
layers = [(weights[0], biases[0], "tanh"), (weights[1], biases[1], "tanh"),
          (weights[2], biases[2], None)]

def network(x):
    return torch.tanh(torch.tanh(x @ weights[0].T + biases[0]) @ weights[1].T + biases[1]) @ weights[2].T + biases[2]

def envelope(x):
    return [x.square() * (1-x).square(), 2*x-6*x.square()+4*x**3,
            2-12*x+12*x.square(), -12+24*x, torch.full_like(x, 24)]

def exact(x):
    return envelope(x)[0] * (1 + .3 * torch.sin(torch.pi*x))

def load(x):
    p = envelope(x)
    g = [1+.3*torch.sin(torch.pi*x), .3*torch.pi*torch.cos(torch.pi*x),
         -.3*torch.pi**2*torch.sin(torch.pi*x), -.3*torch.pi**3*torch.cos(torch.pi*x),
         .3*torch.pi**4*torch.sin(torch.pi*x)]
    return sum(c*p[k]*g[4-k] for k,c in enumerate([1,4,6,4,1]))

def derivatives(x):
    g = jet_to_tower(mlp_jet(x, torch.ones_like(x), layers, order=4))
    assert g.shape == (5, len(x), 1)
    p = envelope(x)
    u = p[0]*g[0]
    fourth = sum(c*p[k]*g[4-k] for k,c in enumerate([1,4,6,4,1]))
    return u, fourth

# %% [markdown]
# ## Independently check the operator, then train
# The polynomial envelope enforces all four boundary conditions for every parameter value.
# omnibias supplies derivatives of the neural factor; the product rule supplies u''''.
# Taylor coefficients must be converted with jet_to_tower before using ordinary derivatives.
# %%
probe = torch.linspace(.11, .89, 5, dtype=dtype).reshape(-1,1).requires_grad_()
reference = envelope(probe)[0]*network(probe)
for _ in range(4):
    reference = torch.autograd.grad(reference.sum(), probe, create_graph=True)[0]
torch.testing.assert_close(derivatives(probe)[1], reference, rtol=1e-8, atol=1e-9)
x_train = torch.linspace(.01, .99, 80, dtype=dtype).reshape(-1,1)
forcing = load(x_train)
optimizer = torch.optim.LBFGS(parameters, lr=1, max_iter=300, tolerance_grad=1e-10,
                             tolerance_change=1e-13, line_search_fn="strong_wolfe")
history = []
def closure():
    optimizer.zero_grad()
    residual = (derivatives(x_train)[1] - forcing) / 24
    loss = residual.square().mean()
    loss.backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in parameters)
    history.append(float(loss.detach()))
    return loss
closure()
optimizer.step(closure)

# %% [markdown]
# ## Evaluate unseen coordinates and the physical constraints
# Require relative L2 error below 1%, normalized residual RMS below 1%, and exact boundary values.
# These are sampled checks against a known solution, not a certified global error bound.
# %%
x = torch.linspace(0, 1, 241, dtype=dtype).reshape(-1,1)
prediction, fourth = derivatives(x)
truth = exact(x)
relative_l2 = float((prediction-truth).norm().detach()/truth.norm())
residual_rms = float(((fourth-load(x))/24).square().mean().sqrt().detach())
ends = torch.tensor([[0.], [1.]], dtype=dtype, requires_grad=True)
boundary = envelope(ends)[0]*network(ends)
slopes = torch.autograd.grad(boundary.sum(), ends)[0]
assert relative_l2 < .01, relative_l2
assert residual_rms < .01, residual_rms
torch.testing.assert_close(boundary, torch.zeros_like(boundary), atol=1e-14, rtol=0)
torch.testing.assert_close(slopes, torch.zeros_like(slopes), atol=1e-14, rtol=0)
seconds = time.perf_counter()-started
out = Path(os.environ.get("OMNIBIAS_EXAMPLE_OUTPUT", "omnibias-results")); out.mkdir(parents=True, exist_ok=True)
data = {"id":"beam-pinn", "seed":42, "dtype":"float64", "device":"CPU", "threads":1,
        "packages":{"torch":torch.__version__, "omnibias-torch":version("omnibias-torch")},
        "seconds":seconds, "closure_evaluations":len(history),
        "metrics":{"relative_l2":relative_l2,"normalized_residual_rms":residual_rms,"boundary_max":float(boundary.abs().max().detach())},
        "x":x.flatten().tolist(), "reference":truth.flatten().tolist(), "prediction":prediction.detach().flatten().tolist(), "loss":history}
(out/"beam-pinn.json").write_text(json.dumps(data, indent=2)+"\n")
fig, axes = plt.subplots(1,2,figsize=(10,3.5),layout="constrained")
axes[0].plot(data["x"],data["reference"],label="Analytic solution"); axes[0].plot(data["x"],data["prediction"],"--",label="Trained PINN")
axes[0].set(xlabel="x",ylabel="Deflection u",title="Clamped beam"); axes[0].legend()
axes[1].semilogy(history); axes[1].set(xlabel="Closure evaluation",ylabel="Normalized residual MSE",title="Training trace")
fig.savefig(out/"beam-pinn.svg"); plt.close(fig)
print(json.dumps({"metrics":data["metrics"],"seconds":seconds},indent=2))

# Display the saved scientific figure when running as a Jupyter notebook.
if "get_ipython" in globals():
    from IPython.display import SVG, display
    display(SVG(filename=str(out/"beam-pinn.svg")))