"""Prepare data, train and evaluate on Labqoat. Run with `lab run app.py`."""

import json
from dataclasses import dataclass

from labqoat import Artifact, Dataset, RunContext, compute

from config import Config
from data import Sample, part, samples
from evaluate import score
from model import Model

app = compute.App("__NAME__", image=compute.Image.debian_slim())


@dataclass(frozen=True)
class Trained:
    """The weights a training run kept and how they did on validation rows."""

    weights: Artifact
    val_loss: float


@dataclass(frozen=True)
class Outcome:
    """What one full pass produced."""

    weights: Artifact
    val_loss: float
    test_loss: float


@app.function(memory=512, timeout=300)
def prepare(context: RunContext) -> Dataset[Sample]:
    """Publish the rows once, with their splits."""
    return context.datasets.publish("__NAME__/samples", row_type=Sample, rows=samples(), key="samples")


@app.function(memory=512, timeout=300)
def train(context: RunContext, dataset: Dataset[Sample], steps: int = 16, learning_rate: float = 0.5) -> Trained:
    """Fit the model on the train split, logging train and validation loss at every step."""
    config = Config(steps=steps, learning_rate=learning_rate)
    rows = list(dataset.rows(context))
    training, validation = part(rows, "train"), part(rows, "validation")
    model = Model()
    track = context.training()
    for step in range(1, config.steps + 1):
        model = model.step(training, config.learning_rate)
        track.step(step, loss=model.loss(training))
        track.validation(step, {"loss": model.loss(validation)})
        context.progress(step, total=config.steps)
    path = context.scratch / "weights.json"
    path.write_text(json.dumps({"weight": model.weight}))
    weights = context.assets.publish(path, media_type="application/json", name="weights", step=config.steps)
    return Trained(weights, model.loss(validation))


@app.function(memory=512, timeout=300)
def evaluate(context: RunContext, weights: Artifact, dataset: Dataset[Sample]) -> float:
    """Score kept weights on the test split."""
    with context.assets.local_file(weights) as path:
        model = Model(float(json.loads(path.read_text())["weight"]))
    test_loss = score(model, part(list(dataset.rows(context)), "test"))
    context.training().test({"loss": test_loss})
    return test_loss


@app.function(memory=512, retries=1)
async def main(context: RunContext, steps: int = 16, learning_rate: float = 0.5) -> Outcome:
    """Prepare, train, then evaluate; each stage starts when its input exists."""
    prepared = await context.runs.with_options(key="prepare").submit_async(prepare)
    dataset = await prepared.result_async()
    training = await context.runs.with_options(key="train").submit_async(
        train, dataset=dataset, steps=steps, learning_rate=learning_rate
    )
    trained = await training.result_async()
    evaluation = await context.runs.with_options(key="evaluate").submit_async(
        evaluate, weights=trained.weights, dataset=dataset
    )
    return Outcome(trained.weights, trained.val_loss, await evaluation.result_async())
