# %% [markdown]
# # Calibrate a decay model with Gauss–Newton
# Recover amplitude A, rate k, and offset b from y(t) = A exp(−kt) + b.
# These noiseless synthetic observations provide a checkable inverse problem.
# Install: python -m pip install --pre "omnibias-torch==0.5.0rc1" matplotlib
# Gauss–Newton uses residual-Jacobian curvature. It is distinct from the full loss Hessian.
# %%
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.optim import GaussNewton

torch.set_num_threads(1); torch.manual_seed(31)
dtype = torch.float64
started = time.perf_counter()
t = torch.linspace(0,5,60,dtype=dtype)
truth = torch.tensor([2.,.7,.2],dtype=dtype)
observed = truth[0]*torch.exp(-truth[1]*t)+truth[2]
def physical(theta):
    return torch.stack([theta[0].exp(),theta[1].exp(),theta[2]])
def model(theta,time):
    a,k,b = physical(theta).unbind()
    return a*torch.exp(-k*time)+b
def residual(theta):
    return model(theta,t)-observed
theta = torch.tensor([1.2,.5,.1],dtype=dtype)
theta[:2] = theta[:2].log()
initial = model(theta,t).detach()

# %% [markdown]
# ## Check parameter sensitivities and take functional steps
# step returns a new flat vector; it does not mutate an nn.Module.
# Log parameterization keeps amplitude and decay rate positive. The offset is unconstrained.
# %%
jacobian = torch.func.jacrev(residual)(theta)
a,k,b = physical(theta).unbind()
expected = torch.stack([a*torch.exp(-k*t),-a*k*t*torch.exp(-k*t),torch.ones_like(t)],1)
torch.testing.assert_close(jacobian,expected,rtol=1e-10,atol=1e-12)
probe = theta.clone().requires_grad_()
gradient = torch.autograd.grad(residual(probe).square().mean(),probe)[0]
assert torch.isfinite(gradient).all() and bool((gradient.abs()>0).all())
optimizer = GaussNewton(solver="qr",damping=1e-3)
history = [float(residual(theta).square().mean())]
for _ in range(12):
    theta,info = optimizer.step(residual,theta)
    history.append(float(residual(theta).square().mean().detach()))
recovered = physical(theta).detach()
torch.testing.assert_close(recovered,truth,rtol=1e-6,atol=1e-7)
heldout = torch.linspace(.025,4.975,101,dtype=dtype)
heldout_rmse = float((model(theta,heldout)-(truth[0]*torch.exp(-truth[1]*heldout)+truth[2])).square().mean().sqrt().detach())
assert heldout_rmse < 1e-6
assert history[-1] < history[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":"decay-calibration","seed":31,"dtype":"float64","device":"CPU","threads":1,
        "packages":{"torch":torch.__version__,"omnibias-torch":version("omnibias-torch")},
        "seconds":seconds,"steps":12,"metrics":{"heldout_rmse":heldout_rmse,"max_parameter_error":float((recovered-truth).abs().max())},
        "recovered_parameters":recovered.tolist(),"x":t.tolist(),"reference":observed.tolist(),"initial":initial.tolist(),"prediction":model(theta,t).detach().tolist(),"loss":history}
(out/"decay-calibration.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="Synthetic observations"); axes[0].plot(data["x"],data["initial"],"--",label="Initial model"); axes[0].plot(data["x"],data["prediction"],label="Calibrated model"); axes[0].legend(); axes[0].set(xlabel="Time",ylabel="Response",title="Parameter recovery")
axes[1].semilogy(history); axes[1].set(xlabel="Gauss–Newton step",ylabel="Residual MSE",title="Convergence")
fig.savefig(out/"decay-calibration.svg"); plt.close(fig)
print(json.dumps({"metrics":data["metrics"],"seconds":seconds,"recovered":recovered.tolist()},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/"decay-calibration.svg")))