"""Attach LaunchHelm to a PyTorch Lightning Trainer.

Requires ``pip install "launchhelm[lightning]"`` plus the LAUNCHHELM_*
environment. The callback runs the same safe-point protocol after each
optimizer step, on every rank.

    python lightning_trainer.py
"""

import torch
from lightning.pytorch import LightningModule, Trainer
from torch import nn
from torch.utils.data import DataLoader, Dataset

from launchhelm.torch.lightning import LaunchHelmCallback


class Module(LightningModule):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(4, 2)

    def training_step(
        self, batch: dict[str, torch.Tensor], batch_idx: int
    ) -> torch.Tensor:
        loss = (self.linear(batch["x"]) - batch["y"]).pow(2).mean()
        self.log("train_loss", loss, prog_bar=False)
        return loss

    def configure_optimizers(self) -> torch.optim.Optimizer:
        return torch.optim.SGD(self.parameters(), lr=0.1)


class Rows(Dataset):
    def __len__(self) -> int:
        return 16

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        return {"x": torch.ones(4), "y": torch.tensor([1.0, 0.0])}


trainer = Trainer(
    max_epochs=1,
    callbacks=[LaunchHelmCallback("quickstart", name="lightning-example")],
    enable_progress_bar=False,
    enable_model_summary=False,
    logger=False,
    default_root_dir="/tmp/lh-lightning-example",
)
trainer.fit(Module(), train_dataloaders=DataLoader(Rows(), batch_size=4))