"""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, make_rows, prepare
from myresearch.evaluate import evaluate, score
from myresearch.train import fit, train_curve

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 the research here, with no engine involved. Not the node bodies: they
    # save and load artifacts, and there is nothing to save to out here. The
    # functions they wrap are ordinary Python, which is why they are worth
    # keeping separate — the arithmetic stays yours to run by hand.
    data = make_rows(limit=64)
    print("rows:", len(data["rows"]))
    print("losses:", [round(loss, 4) for loss in train_curve(0.05, 5)])
    print("score:", score(0.05, 5))