#!/usr/bin/env python
"""Seed the ML-pipeline example: a batch flow, a dashboard, and a worker to try.

What it builds::

    demo_training   prepare -> train -> evaluate      a batch flow, run on demand
    demo_training   a dashboard watching a run as it happens

The flow is a small supervised fit — recover ``y = 3x + 2`` from noisy samples
— chosen because it is a real training loop with none of the setup. What it is
actually demonstrating is the shape of an ML pipeline on this engine:

* **Batch runs.** ``mode: batch``, so nothing is activated and nothing fires
  until a run asks. Its ``inputs`` are the run's parameters and its ``outputs``
  are the run's result.
* **Streaming outputs.** ``train`` is a generator: every ``yield`` publishes on
  the ``loss`` port the instant it happens, so the dashboard chart draws the
  curve while the training is still running. There is no logging call — the
  numbers are outputs, and the run keeps the whole series because ``loss``
  declares ``stream: true``.
* **``fluksio.emit``** for where a yield cannot reach: ``progress`` is
  published from inside a helper the loop calls, which is the shape a training
  framework's callback has.
* **Artifacts.** The dataset and the trained weights are far too big to be
  messages, so they travel as references into the content-addressed store and
  are opened on the other side — including when the other side is a different
  machine.
* **A GPU node beside CPU nodes.** ``train`` declares ``device: gpu``, so it
  runs on an attached worker carrying that label. Its policy is ``prefer``, so
  with no worker attached it runs on the engine instead — and it reports which
  machine and which numeric backend it used, so the routing is visible in the
  result rather than taken on trust.
* **Sweeps.** Runs of one flow are isolated from each other, so the same flow
  can be swept over parameters and the runs compared.

Run it against a stack that is already up::

    make -C app seed-demo

Then, from the dashboard or the API::

    POST /api/v1/runs/flows/demo_training  {"params": {"epochs": 40}}
    POST /api/v1/runs/flows/demo_training/sweep {"runs": [...]}

Environment (the Makefile passes these):
    API_URL, FIRST_SUPERUSER, FIRST_SUPERUSER_PASSWORD
"""

from __future__ import annotations

import os
import sys
from typing import Any

import httpx

API = os.environ.get("API_URL", "http://api.localhost")
EMAIL = os.environ.get("FIRST_SUPERUSER", "")
PASSWORD = os.environ.get("FIRST_SUPERUSER_PASSWORD", "")

FLOW = "demo_training"
PANEL = "demo_training"

PREPARE_SOURCE = '''"""Make the dataset, and hand it on as an artifact.

Runs on the engine: no device, so it goes to the local worker pool. Only the
standard library, because the engine's venv is its own and this node has no
business asking it for anything.
"""

import json
import random

import fluksio

#: The line the training is supposed to recover.
TRUE_SLOPE = 3.0
TRUE_INTERCEPT = 2.0


def process(seed, noise, samples):
    rng = random.Random(int(seed))
    rows = [
        [x, TRUE_SLOPE * x + TRUE_INTERCEPT + rng.gauss(0.0, float(noise))]
        for x in (rng.uniform(-1.0, 1.0) for _ in range(int(samples)))
    ]
    payload = json.dumps(
        {"rows": rows, "truth": {"slope": TRUE_SLOPE, "intercept": TRUE_INTERCEPT}}
    ).encode()
    # Data between stages is an artifact, not a message: the reference is
    # what travels, and it stays valid on whichever machine opens it.
    return {
        "dataset": fluksio.save_artifact(payload, "dataset.json", "application/json"),
        "sample_count": len(rows),
    }
'''

TRAIN_SOURCE = '''"""Fit the dataset, publishing the loss as it goes.

A generator, so this node produces values over time: every ``yield`` is a dict
keyed by output port and is published the moment it happens, and what the
function returns at the end is the node's result. Nothing here logs anything —
the loss is an output of the graph, which is why a chart can bind to it.

Bound to ``device: gpu``. With a worker carrying that label attached it runs
there; without one it falls back to the engine, and ``trained_on`` says which
happened. numpy is used when the machine it landed on has it, which is how the
two environments tell themselves apart.
"""

import json
import platform
import time

import fluksio

try:
    import numpy as np
except ImportError:  # The engine's own venv has no numpy; a GPU box will.
    np = None


def process(dataset, learning_rate, epochs, pace):
    with open(fluksio.load_artifact(dataset)) as handle:
        data = json.load(handle)

    xs = [row[0] for row in data["rows"]]
    ys = [row[1] for row in data["rows"]]
    count = len(xs)
    total = int(epochs)
    rate = float(learning_rate)
    slope, intercept = 0.0, 0.0

    if np is not None:
        axis_x, axis_y = np.array(xs), np.array(ys)

    def report(epoch):
        # Published from inside a helper, where a yield cannot reach — the
        # shape a training framework's callback has. Same port, same checking.
        fluksio.emit(progress=round(100.0 * (epoch + 1) / total, 1))

    loss = 0.0
    for epoch in range(total):
        if np is not None:
            error = (slope * axis_x + intercept) - axis_y
            loss = float((error**2).mean())
            slope -= rate * float((2 * error * axis_x).mean())
            intercept -= rate * float((2 * error).mean())
        else:
            error = [slope * x + intercept - y for x, y in zip(xs, ys)]
            loss = sum(e * e for e in error) / count
            slope -= rate * sum(2 * e * x for e, x in zip(error, xs)) / count
            intercept -= rate * sum(2 * e for e in error) / count

        yield {"loss": round(loss, 6)}
        report(epoch)
        # Only so the curve is watchable; a real epoch takes as long as it takes.
        if float(pace) > 0:
            time.sleep(float(pace))

    weights = json.dumps({"slope": slope, "intercept": intercept}).encode()
    return {
        "weights": fluksio.save_artifact(weights, "weights.json", "application/json"),
        "final_loss": round(loss, 6),
        "trained_on": "{} ({})".format(
            platform.node(),
            "numpy " + np.__version__ if np is not None else "pure python",
        ),
    }
'''

EVALUATE_SOURCE = '''"""Score the fit, back on the engine.

Opens two artifacts: the dataset this run prepared and the weights the
training produced — which may have been written on another machine entirely.
A reference names content, so where it came from does not matter.
"""

import json

import fluksio


def process(weights, dataset):
    with open(fluksio.load_artifact(weights)) as handle:
        fit = json.load(handle)
    with open(fluksio.load_artifact(dataset)) as handle:
        data = json.load(handle)

    rows = data["rows"]
    truth = data["truth"]
    mean = sum(y for _x, y in rows) / len(rows)
    total = sum((y - mean) ** 2 for _x, y in rows)
    residual = sum((y - (fit["slope"] * x + fit["intercept"])) ** 2 for x, y in rows)
    r2 = 1.0 - residual / total if total else 0.0

    good = r2 > 0.9
    return {
        "accuracy": round(100.0 * max(0.0, r2), 2),
        "report": {
            "title": "Recovered {:.2f}x + {:.2f}".format(
                fit["slope"], fit["intercept"]
            ),
            "body": "Truth is {}x + {}. R² = {:.4f} over {} samples.".format(
                truth["slope"], truth["intercept"], r2, len(rows)
            ),
            "severity": "info" if good else "warning",
        },
    }
'''

NODES = [
    {
        "id": "prepare",
        "type": "python",
        "title": "Make the dataset",
        "requires": [
            {"name": "seed", "dtype": "int"},
            {"name": "noise", "dtype": "float"},
            {"name": "samples", "dtype": "int"},
        ],
        "provides": [
            {"name": "dataset", "dtype": "artifact"},
            {"name": "sample_count", "dtype": "int"},
        ],
    },
    {
        "id": "train",
        "type": "python",
        "title": "Fit it",
        # Runs on a worker carrying this label; falls back here without one,
        # which is what makes the example work before a GPU box is attached.
        "device": "gpu",
        "device_policy": "prefer",
        # Generous, and an *idle* timeout: a node that keeps publishing keeps
        # its deadline reset, so this is how long it may go quiet.
        "timeout": 300,
        # A setting: a constant of this node's own code, reaching `process` as
        # an argument like the ports beside it. Only so a human can watch the
        # curve arrive — set it to 0 in the node panel for a sweep.
        "params": {"pace": 0.15},
        "requires": [
            {"name": "dataset", "dtype": "artifact"},
            {"name": "learning_rate", "dtype": "float"},
            {"name": "epochs", "dtype": "int"},
        ],
        "provides": [
            # The curve. `stream` says this port publishes repeatedly while the
            # node runs, and is what makes a run keep every value it takes.
            {"name": "loss", "dtype": "float", "stream": True},
            {"name": "progress", "dtype": "float", "stream": True},
            {"name": "weights", "dtype": "artifact"},
            {"name": "final_loss", "dtype": "float"},
            {"name": "trained_on", "dtype": "str"},
        ],
    },
    {
        "id": "evaluate",
        "type": "python",
        "title": "Score it",
        "requires": [
            {"name": "weights", "dtype": "artifact"},
            {"name": "dataset", "dtype": "artifact"},
        ],
        "provides": [
            {"name": "report", "dtype": "record"},
            {"name": "accuracy", "dtype": "float"},
        ],
    },
]

#: The run's parameters, with the values a run gets when it names none.
INPUTS = [
    {"spec": {"name": "seed", "dtype": "int"}, "initial": 1},
    {"spec": {"name": "noise", "dtype": "float"}, "initial": 0.25},
    {"spec": {"name": "samples", "dtype": "int"}, "initial": 200},
    {"spec": {"name": "learning_rate", "dtype": "float"}, "initial": 0.2},
    {"spec": {"name": "epochs", "dtype": "int"}, "initial": 40},
]

#: What a run reports as its result. Everything else the flow computed stays
#: in the run's own state and is dropped with it.
OUTPUTS = ["report", "accuracy", "final_loss", "trained_on", "sample_count"]

WIDGETS = [
    {
        "id": "curve",
        "type": "chart",
        "title": "Training loss",
        "layout": {"lg": {"x": 0, "y": 0, "w": 8, "h": 5}},
        "config": {
            # Bound to the port the node yields on — the same binding a
            # temperature would use. Nothing here knows what a metric is.
            "series": [{"message": f"{FLOW}.loss", "dtype": "float", "label": "loss"}],
            "history": {"points": 600},
        },
    },
    {
        "id": "how",
        "type": "markdown",
        "title": "",
        "layout": {"lg": {"x": 8, "y": 0, "w": 4, "h": 5}},
        "config": {
            "content": (
                "## Try it\n"
                "- Press Run on the flow, or POST /api/v1/runs/flows/demo_training\n"
                "- The curve fills in while the training runs. It is a port"
                " the node yields on, not a log.\n"
                "- Sweep it: POST .../sweep with a list of params. Runs are"
                " isolated, so they go in parallel.\n"
                "- Attach a worker labelled gpu and 'Trained on' changes"
                " machine.\n"
                "- The finished series lives on the run:"
                " GET /api/v1/runs/{id}/metrics?name=demo_training.loss"
            )
        },
    },
    {
        "id": "progress",
        "type": "gauge",
        "title": "Progress",
        "layout": {"lg": {"x": 0, "y": 5, "w": 3, "h": 3}},
        "config": {
            "message": f"{FLOW}.progress",
            "dtype": "float",
            "min": 0,
            "max": 100,
            "unit": "%",
        },
    },
    {
        "id": "accuracy",
        "type": "gauge",
        "title": "Fit quality (R²)",
        "layout": {"lg": {"x": 3, "y": 5, "w": 3, "h": 3}},
        "config": {
            "message": f"{FLOW}.accuracy",
            "dtype": "float",
            "min": 0,
            "max": 100,
            "unit": "%",
        },
    },
    {
        "id": "loss_now",
        "type": "stat",
        "title": "Final loss",
        "layout": {"lg": {"x": 6, "y": 5, "w": 3, "h": 3}},
        "config": {"message": f"{FLOW}.final_loss", "dtype": "float", "precision": 4},
    },
    {
        "id": "samples",
        "type": "stat",
        "title": "Samples",
        "layout": {"lg": {"x": 9, "y": 5, "w": 3, "h": 3}},
        "config": {"message": f"{FLOW}.sample_count", "dtype": "int"},
    },
    {
        "id": "where",
        "type": "stat",
        "title": "Trained on",
        "layout": {"lg": {"x": 0, "y": 8, "w": 5, "h": 3}},
        "config": {"message": f"{FLOW}.trained_on", "dtype": "str"},
    },
    {
        "id": "verdict",
        "type": "notification",
        "title": "Result",
        "layout": {"lg": {"x": 5, "y": 8, "w": 7, "h": 3}},
        "config": {"message": f"{FLOW}.report", "dtype": "record"},
    },
]


class Api:
    def __init__(self) -> None:
        self.http = httpx.Client(base_url=f"{API}/api/v1", timeout=60)
        token = self.http.post(
            "/login/access-token",
            data={"username": EMAIL, "password": PASSWORD},
        ).json()["access_token"]
        self.http.headers["Authorization"] = f"Bearer {token}"

    def __call__(self, method: str, path: str, body: Any = None) -> Any:
        response = self.http.request(method, path, json=body)
        response.raise_for_status()
        return response.json() if response.content else None


def main() -> int:
    if not EMAIL or not PASSWORD:
        print("FIRST_SUPERUSER / FIRST_SUPERUSER_PASSWORD are unset.", file=sys.stderr)
        return 1

    api = Api()

    try:
        api("DELETE", f"/flows/{FLOW}")
    except httpx.HTTPStatusError:
        pass
    api(
        "PUT",
        f"/flows/{FLOW}",
        {
            "name": FLOW,
            "title": "Training run (example)",
            "mode": "batch",
            "outputs": OUTPUTS,
            "nodes": NODES,
            "inputs": INPUTS,
        },
    )
    for node_id, code in (
        ("prepare", PREPARE_SOURCE),
        ("train", TRAIN_SOURCE),
        ("evaluate", EVALUATE_SOURCE),
    ):
        api("PUT", f"/flows/{FLOW}/nodes/{node_id}/source", {"code": code})
    version = api("GET", f"/flows/{FLOW}?draft=true")["definition"]["version"]
    published = api("POST", f"/flows/{FLOW}/publish", {"version": version})
    print(f"  {FLOW}: {len(NODES)} nodes, published")
    for issue in published.get("issues") or []:
        print(f"    ! {issue['message']}")

    try:
        api("DELETE", f"/dashboards/{PANEL}")
    except httpx.HTTPStatusError:
        pass
    api("POST", f"/dashboards/{PANEL}", {"name": PANEL, "title": "Training"})
    current = api("GET", f"/dashboards/{PANEL}")
    api(
        "PUT",
        f"/dashboards/{PANEL}",
        {
            **current,
            "pages": [
                {
                    "id": "main",
                    "title": "Overview",
                    "sections": [{"id": "main", "widgets": WIDGETS}],
                }
            ],
        },
    )
    version = api("GET", f"/dashboards/{PANEL}?draft=true")["version"]
    api("POST", f"/dashboards/{PANEL}/publish", {"version": version})
    print(f"  dashboard '{PANEL}': published")

    print(
        f"\nOpen /view/{PANEL}, then start a run:\n"
        f'  POST {API}/api/v1/runs/flows/{FLOW} {{"params": {{}}}}\n'
        "The curve fills in as it trains. For a sweep, set the training\n"
        "node's 'pace' setting to 0."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())