"""Where the data comes from."""

from __future__ import annotations

import json

import fluksio
from fluksio import Port, node


@node(provides=[Port("dataset", "artifact"), Port("rows", "int")])
def prepare(source: str = "builtin", limit: int = 512):
    """Make the training set and store it, returning a reference to it.

    `source` and `limit` have defaults and are not ports, so they become the
    node's settings: the canvas can tune them without touching this file.
    """
    rows = [{"x": index / limit, "y": (index % 7) / 7} for index in range(limit)]
    payload = json.dumps({"source": source, "rows": rows}).encode()
    return {
        "dataset": fluksio.save_artifact(payload, "dataset.json", "application/json"),
        "rows": len(rows),
    }


@node(requires=["dataset"], provides=[Port("augmented", "artifact")])
def augment(dataset, factor: int = 2):
    """Grow the training set, so `finetune` has something of its own."""
    payload = json.loads(open(fluksio.load_artifact(dataset)).read())
    payload["rows"] = payload["rows"] * factor
    blob = json.dumps(payload).encode()
    return {
        "augmented": fluksio.save_artifact(blob, "augmented.json", "application/json")
    }