"""Which nodes make up which flow.

Membership is this list, not the file a function sits in: `fit` is declared in
`train.py` and used by both flows below, once as itself and once rewired and
reconfigured. Wiring is not membership either — nodes connect because one
provides a message another requires, never because one imported the other.

    fluksio sync examples/myresearch
    fluksio run train --lr 0.05 --wait
"""

from __future__ import annotations

from fluksio import Flow, Port, use
from myresearch.data import augment, prepare
from myresearch.evaluate import evaluate
from myresearch.train import fit

train = Flow(
    "train",
    title="Train",
    nodes=[prepare, fit, evaluate],
    inputs=[Port("lr", "float", initial=0.01)],
    outputs=["score", "final_loss"],
)

finetune = Flow(
    "finetune",
    title="Finetune",
    nodes=[
        prepare,
        augment,
        # The same function, reading `augmented` instead of `dataset` and with
        # a shorter schedule. `train` is unaffected.
        use(fit, wire={"dataset": "augmented"}, epochs=3),
        evaluate,
    ],
    inputs=[Port("lr", "float", initial=0.0001)],
    outputs=["score"],
)


if __name__ == "__main__":
    # Run it here, with no engine involved: the decorators changed nothing
    # about calling these functions.
    dataset = prepare(limit=64)
    losses = list(fit(dataset["dataset"], lr=0.05, epochs=5))
    print("losses:", [round(step["loss"], 4) for step in losses])