"""Log one of every rich telemetry kind to LaunchHelm.

Set LAUNCHHELM_URL, LAUNCHHELM_API_KEY, LAUNCHHELM_ORGANIZATION and
LAUNCHHELM_PROJECT in the environment or a .env file, then:

    python rich_values.py

Media uploads don't block the loop: ``image``/``media`` stage the bytes
locally, and ``finish()`` delivers every pending upload. A NumPy array
is logged as an image when NumPy is installed; otherwise a small PNG is
uploaded from bytes.
"""

import base64
import os

import launchhelm

try:
    import numpy
except ImportError:
    numpy = None

# A 1x1 transparent PNG, used when NumPy is unavailable.
PIXEL = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mIA"
    "AQAFAAH/q842iQAAAABJRU5ErkJggg=="
)

with launchhelm.init("quickstart", name="rich-values") as run:
    run.text("notes", "## epoch 3\nloss **0.25**, eval queued", format="markdown",
             caption="run notes", step=3)

    if numpy is not None:
        frame = numpy.zeros((32, 32, 3), dtype=numpy.uint8)
        frame[..., 0] = 200
        frame[8:24, 8:24] = [40, 120, 220]
        run.image("frames", frame, caption="confusion matrix", step=3)
    else:
        run.image("frames", PIXEL, media_type="image/png",
                  caption="confusion matrix", step=3)
    remote = run.media_url("https://media.example.test/waveform.png", "image/png")
    run.media("gallery", [remote], step=4)

    run.rows("sweep", [
        {"trial": "a", "lr": 0.001, "loss": 0.42},
        {"trial": "b", "lr": 0.003, "loss": 0.31},
    ], step=3)

    run.histogram("gradients", [0.0, 0.1, 0.4, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 1.0],
                  bins=4, step=3)
    run.histogram_counts("precomputed", [-1.0, 0.0, 1.0, 2.0], [2.0, 5.0, 3.0],
                         total=10.0, min=-1.5, max=1.75, mean=0.35, step=3)

print("run finished:", run.run_id)

# Read back bounded samples of each renderer's stream.
client = launchhelm.Client(
    os.environ["LAUNCHHELM_URL"],
    os.environ["LAUNCHHELM_API_KEY"],
    os.environ["LAUNCHHELM_ORGANIZATION"],
)
streams = {}
cursor = ""
project_id = run.project_id
while True:
    page = client.list_documents(
        kind="stream", project_id=project_id, run_id=run.run_id, cursor=cursor
    )
    for document in page.get("items", []):
        spec = document.get("spec") or {}
        name = spec.get("name")
        if name:
            streams[name] = (document.get("id"), (spec.get("presentation") or {}).get("renderer"))
    cursor = page.get("next_cursor") or ""
    if not cursor:
        break

for name, select in (
    ("notes", {"mode": "latest", "limit": 1}),
    ("frames", {"mode": "latest", "limit": 1}),
    ("sweep", {"mode": "index", "limit": 16}),
    ("gradients", {"mode": "index", "limit": 16}),
):
    stream_id, renderer = streams[name]
    result = client.read_samples(stream_id, select=select)
    points = result.get("points", [])
    print(f"{name}: {len(points)} points (renderer {result.get('renderer') or renderer},"
          f" complete {result.get('complete')})")