"""Save a checkpoint, then resume the same run in a new trainer.

``save_checkpoint`` uploads the model/optimizer/RNG state as an artifact.
``init(resume=<run id>)`` reopens the run and ``load_checkpoint("latest")``
restores the newest checkpoint. After a crash the same flow continues the
run from the last save; steps re-executed since that save appear again in
the series.

    python torch_checkpoint.py
"""

import torch

import launchhelm.torch as lht


def build() -> tuple[torch.nn.Module, torch.optim.Optimizer]:
    torch.manual_seed(0)
    model = torch.nn.Linear(8, 1)
    return model, torch.optim.SGD(model.parameters(), lr=0.01)


model, optimizer = build()
x = torch.randn(64, 8)
y = x.sum(dim=1, keepdim=True)

tr = lht.init("quickstart", name="checkpoint-demo")
tr.watch(model, optimizer)
for step in range(5):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()
    tr.step(step=step, loss=loss.item())
artifact_id = tr.save_checkpoint()
run_id = tr.run.run_id
print("checkpoint artifact:", artifact_id)
# close_run=False ends this session but leaves the run open, modelling a
# crashed process; a finished run cannot be resumed.
tr.finish(close_run=False)

# A relaunched process resumes the run and restores the latest checkpoint.
model2, optimizer2 = build()
tr2 = lht.init("quickstart", resume=run_id)
tr2.watch(model2, optimizer2)
extra = tr2.load_checkpoint("latest")
print("restored at step:", tr2.global_step)
print("extra payload:", extra)
for step in range(5, 10):
    optimizer2.zero_grad()
    loss = torch.nn.functional.mse_loss(model2(x), y)
    loss.backward()
    optimizer2.step()
    tr2.step(step=step, loss=loss.item())
tr2.finish()