"""One training script, and everything it puts on the screen.

Every control, stream and inspectable tensor the run's workspace shows is
declared or logged by a numbered part of this file; the guide "From the
training script to the screen" walks through each one. Trains a small model
on synthetic data drawn from four sources, at a tenth of a second a step so
there is time to watch and steer it.

Requires ``pip install "launchhelm[torch,system]"`` and the same
LAUNCHHELM_* environment variables as log_run.py.

    python script_to_screen.py
"""

import time

import torch

import launchhelm.torch as lht
from launchhelm.models import Capability, CommandContext, opaque_schema

STEPS = 600
SOURCES = ["web", "code", "math", "books"]
mixture = {"web": 0.4, "code": 0.3, "math": 0.2, "books": 0.1}

torch.manual_seed(0)
# Each source is a regression target of its own.
targets = {name: torch.randn(16, 1) for name in SOURCES}
model = torch.nn.Sequential(
    torch.nn.Linear(16, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 64),
    torch.nn.ReLU(),
    torch.nn.Linear(64, 1),
)
layers = {
    index: module for index, module in enumerate(model) if isinstance(module, torch.nn.Linear)
}
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.1)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=STEPS)


def draw_batch(size: int = 64) -> tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
    """A batch drawn by the current mixture, and the share each source got."""
    weights = torch.tensor([mixture[name] for name in SOURCES])
    picks = torch.multinomial(weights, size, replacement=True)
    x = torch.randn(size, 16)
    y = torch.empty(size, 1)
    for index, name in enumerate(SOURCES):
        rows = picks == index
        y[rows] = x[rows] @ targets[name]
    shares = {name: float((picks == index).float().mean()) for index, name in enumerate(SOURCES)}
    return x, y, shares


def run_benchmarks() -> dict[str, float]:
    """Held-out loss on each source. The eval control runs it on request."""
    scores = {}
    with torch.no_grad():
        for name in SOURCES:
            x = torch.randn(256, 16)
            loss = torch.nn.functional.mse_loss(model(x), x @ targets[name])
            scores[f"loss/{name}"] = float(loss)
    return scores


# 1. The run.
tr = lht.init(
    "quickstart",
    name="script-to-screen",
    config={"lr": 3e-3, "weight_decay": 0.1, "batch_size": 64, "steps": STEPS},
    system_metrics=True,
)

# 2. Controls and their streams.
tr.watch(
    model,
    optimizer,
    scheduler=scheduler,
    controls=["lr", "weight_decay", "grad_clip", "checkpoint"],
    eval_fn=run_benchmarks,
    captures=True,  # 3. the model's tensors, and a control to capture them
)
tr.plan(step=STEPS)


# 4. Your own control: a typed command with a mixture of four weights.
def set_mixture(value: dict, context: CommandContext) -> dict:
    """Draw the next batches by these weights, scaled to sum to 1."""
    requested = value["weights"]["weights"]
    total = sum(float(share) for share in requested.values())
    if total > 0:
        mixture.update({name: float(share) / total for name, share in requested.items()})
    return {name: round(share, 4) for name, share in mixture.items()}


tr.run.on(
    "after_step",
    Capability(
        name="set_mixture",
        provider="script",
        version="1.0.0",
        input_schema={
            "revision": "1",
            "root": {
                "kind": "record",
                "renderer": "form",
                "fields": [
                    {
                        "name": "weights",
                        "type": {
                            "kind": "mixture",
                            "components": SOURCES,
                            "sampled": "data/mix/{component}",
                        },
                    },
                ],
            },
        },
        output_schema=opaque_schema(),
        safe_points=("after_step",),
        starting={"weights": {"weights": {name: str(share) for name, share in mixture.items()}}},
    ),
    set_mixture,
)

for step in range(STEPS):
    x, y, shares = draw_batch()
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    tr.clip_grad_norm_()  # uses the grad_clip control
    health = step % 10 == 0
    if health:
        before = {index: layer.weight.detach().clone() for index, layer in layers.items()}
    optimizer.step()
    scheduler.step()

    # 5. Many streams under one name pattern, every ten steps.
    if health:
        ratios = {
            f"health/layer{index}/update_ratio": float(
                (layer.weight.detach() - before[index]).norm() / before[index].norm()
            )
            for index, layer in layers.items()
        }
        drawn = {f"data/mix/{name}": share for name, share in shares.items()}
        tr.log(ratios | drawn, step=step)

    # 6. The core metrics, and the safe points.
    tr.step(step=step, loss=loss.item(), samples=len(x))
    tr.run.sync("after_step", step=step)
    time.sleep(0.1)

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