# %% [markdown]
# # A Poisson PINN with a direct Laplacian
# Solve Δu = −2π² sin(πx)sin(πy) on the unit square, with u = 0 on the boundary.
# The known solution sin(πx)sin(πy) lets us evaluate unseen coordinates independently.
# Install: python -m pip install --pre "omnibias-torch==0.5.0rc1" matplotlib
# CPU float64. Direct contraction avoids forming a full Hessian in the training residual.
# %%
import json
import os
import time
from pathlib import Path
from importlib.metadata import version
import matplotlib.pyplot as plt
import torch
from omnibias.torch.laplacian import deep_field_laplacian

torch.set_num_threads(1)
torch.manual_seed(17)
dtype = torch.float64
started = time.perf_counter()
weights = [torch.nn.Parameter(torch.randn(24,2,dtype=dtype)),
           torch.nn.Parameter(torch.randn(24,24,dtype=dtype)/5),
           torch.nn.Parameter(torch.randn(1,24,dtype=dtype)/5)]
biases = [torch.nn.Parameter(torch.zeros(n,dtype=dtype)) for n in (24,24,1)]
parameters = [*weights,*biases]
layers = [(weights[0],biases[0],"tanh"),(weights[1],biases[1],"tanh"),(weights[2],biases[2],None)]
def model(x):
    return torch.tanh(torch.tanh(x@weights[0].T+biases[0])@weights[1].T+biases[1])@weights[2].T+biases[2]
def exact(x):
    return torch.sin(torch.pi*x[:,0:1])*torch.sin(torch.pi*x[:,1:2])
def grid(axis):
    return torch.stack(torch.meshgrid(axis,axis,indexing="ij"),-1).reshape(-1,2)

# %% [markdown]
# ## Check the contraction against an independent Hessian
# Only three probe points use a dense Hessian. The residual uses deep_field_laplacian.
# This test validates the operator for this architecture and dtype, not every possible model.
# %%
probe = torch.tensor([[.17,.29],[.53,.71],[.83,.37]],dtype=dtype)
dense = torch.stack([torch.func.hessian(lambda z:model(z[None]).squeeze())(p).trace() for p in probe])
torch.testing.assert_close(deep_field_laplacian(probe,layers).flatten(),dense,rtol=1e-9,atol=1e-10)
interior = grid(torch.linspace(.025,.975,15,dtype=dtype))
edge = torch.linspace(0,1,25,dtype=dtype)
boundary = torch.cat([torch.stack([edge,torch.zeros_like(edge)],1),torch.stack([edge,torch.ones_like(edge)],1),
                      torch.stack([torch.zeros_like(edge),edge],1),torch.stack([torch.ones_like(edge),edge],1)])
forcing = -2*torch.pi**2*exact(interior)
optimizer = torch.optim.LBFGS(parameters,lr=1,max_iter=600,tolerance_grad=1e-10,
                             tolerance_change=1e-13,line_search_fn="strong_wolfe")
history = []
def closure():
    optimizer.zero_grad()
    residual = (deep_field_laplacian(interior,layers)-forcing)/(2*torch.pi**2)
    loss = residual.square().mean()+5*model(boundary).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]
# ## Validate interior error and boundary error separately
# Boundary conditions are penalized, not structurally exact. We report their maximum error.
# No end-to-end speedup is claimed: this notebook establishes a complete, trainable workflow.
# %%
axis = torch.linspace(0,1,41,dtype=dtype)
points = grid(axis)
prediction = model(points).detach(); truth = exact(points)
relative_l2 = float((prediction-truth).norm()/truth.norm())
boundary_max = float(model(boundary).detach().abs().max())
residual_rms = float(((deep_field_laplacian(points,layers).detach()+2*torch.pi**2*truth)/(2*torch.pi**2)).square().mean().sqrt())
assert relative_l2 < .02, relative_l2
assert boundary_max < .02, boundary_max
assert residual_rms < .03, residual_rms
seconds = time.perf_counter()-started
out = Path(os.environ.get("OMNIBIAS_EXAMPLE_OUTPUT","omnibias-results")); out.mkdir(parents=True,exist_ok=True)
data = {"id":"poisson-pinn","seed":17,"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,"boundary_max":boundary_max,"normalized_residual_rms":residual_rms},
        "axis":axis.tolist(),"prediction":prediction.reshape(41,41).tolist(),"error":(prediction-truth).reshape(41,41).tolist(),"loss":history}
(out/"poisson-pinn.json").write_text(json.dumps(data,indent=2)+"\n")
fig,axes = plt.subplots(1,2,figsize=(10,3.5),layout="constrained")
for ax,values,title in [(axes[0],data["prediction"],"PINN solution"),(axes[1],data["error"],"Error against analytic solution")]:
    im=ax.imshow(values,origin="lower",extent=[0,1,0,1]); ax.set(xlabel="y",ylabel="x",title=title); fig.colorbar(im,ax=ax)
fig.savefig(out/"poisson-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/"poisson-pinn.svg")))