"""Two commands share one safe point; ``sync`` services whichever waits.

set_schedule replaces the learning-rate plan: hold the current value, then
ramp linearly to a target. duplicate_layer inserts a copy of one layer.
Needs the LAUNCHHELM_* environment like log_run.py.

    python experiment.py
"""

import launchhelm
from launchhelm import CommandContext


class Experiment:
    """Local stand-in for the values a real trainer would mutate."""

    def __init__(self) -> None:
        self.lr = 0.0003
        self.layers = ["layer0", "layer1", "layer2"]
        self._hold_until = 0
        self._ramp_until = 0
        self._ramp_from = self.lr
        self._target = self.lr

    def set_schedule(
        self,
        hold_steps: int,
        ramp_steps: int,
        target_lr: float,
        context: CommandContext,
    ) -> None:
        """Hold, then ramp from the step where the command is applied."""
        step = int(context.coordinates[0]["value"])
        self._hold_until = step + hold_steps
        self._ramp_until = self._hold_until + ramp_steps
        self._ramp_from = self.lr
        self._target = target_lr

    def duplicate_layer(self, index: int) -> None:
        """Insert a copy immediately after ``index``."""
        self.layers.insert(index + 1, f"{self.layers[index]}-copy")

    def advance(self, step: int) -> None:
        """Apply the active plan at ``step``."""
        if step < self._hold_until:
            return
        span = self._ramp_until - self._hold_until
        if span <= 0 or step >= self._ramp_until:
            self.lr = self._target
            return
        done = step - self._hold_until
        self.lr = self._ramp_from + (self._target - self._ramp_from) * done / span


experiment = Experiment()
with launchhelm.init("quickstart", config={"lr": experiment.lr}) as run:
    run.command("after_step")(experiment.set_schedule)
    run.command("after_step")(experiment.duplicate_layer)
    for step in range(8):
        applied = run.sync("after_step", step=step)
        experiment.advance(step)
        run.log(
            {
                "loss": 2.0 / (step + 1),
                "lr": experiment.lr,
                "layers": len(experiment.layers),
            },
            step=step,
        )
        del applied

print("run finished:", run.run_id)