"""The part that would be on the GPU."""

from __future__ import annotations

import json
import math

import fluksio
from fluksio import Port, node


def train_curve(lr: float, epochs: int) -> list[float]:
    """The loss per epoch, with no Fluksio in it."""
    return [
        math.exp(-lr * epoch * 10) * (1 + 0.05 * (epoch % 3)) / (1 + lr)
        for epoch in range(epochs)
    ]


@node(
    requires=["dataset", Port("lr", "float")],
    provides=[
        Port("loss", "float", stream=True),
        Port("weights", "artifact"),
        Port("final_loss", "float"),
    ],
    device="gpu",
    device_policy="prefer",
    timeout=600,
    resources={"cpus": 2},
)
def fit(dataset, lr, epochs=25):
    """Train, reporting the loss as it goes.

    Yielding is the reporting: each one publishes on the `loss` port the
    instant it happens, and the run keeps every value as a series — which is
    why there is no `log_metric()` to call. `device="gpu"` with
    `device_policy="prefer"` sends this to a worker carrying that label when
    one is attached, and runs it here when none is.

    `resources` is what a real training node needs and the rest of a flow does
    not: two cores held for the length of the call, and a worker started with
    thread limits saying so — otherwise a numerical library sizes itself to
    every core on the machine, and a sweep of these starves the engine. A GPU
    one would say `{"gpus": 1}` and hold the card exclusively.
    """
    rows = json.loads(open(fluksio.load_artifact(dataset)).read())["rows"]
    loss = 1.0
    for loss in train_curve(lr, epochs):
        yield {"loss": loss}
    weights = json.dumps({"lr": lr, "epochs": epochs, "n": len(rows)}).encode()
    return {
        "weights": fluksio.save_artifact(weights, "weights.json", "application/json"),
        "final_loss": loss,
    }