from __future__ import annotations

from typing import Any, Dict, List

from .base import BaseAgent
from ..config import TrainConfig
from ..ml.dataset import Dataset
from ..ml.logistic import LogisticRegression
from ..ml.metrics import evaluate


class EvaluatorAgent(BaseAgent):
    name = "EvaluatorAgent"

    def run(
        self,
        *,
        cfg: TrainConfig,
        dataset: Dataset,
        metrics: List[Dict[str, Any]],
        attempt: int,
    ) -> Dict[str, Any]:
        self.emit(
            receiver="MessageBus",
            kind="stage",
            payload={"stage": "eval", "attempt": attempt},
        )

        model_payload = self.ctx.store.read_json(f"model.attempt{attempt}.json")
        model = LogisticRegression(w=list(model_payload["w"]), b=float(model_payload["b"]))

        train_loss, train_acc = evaluate(model, dataset.x_train, dataset.y_train, eps=cfg.loss_eps)
        val_loss, val_acc = evaluate(model, dataset.x_val, dataset.y_val, eps=cfg.loss_eps)

        out = {
            "attempt": attempt,
            "train_loss": float(train_loss),
            "train_acc": float(train_acc),
            "val_loss": float(val_loss),
            "val_acc": float(val_acc),
            "epochs_ran": len(metrics),
            "cfg": cfg.to_dict(),
        }

        self.ctx.store.write_json(f"eval.attempt{attempt}.json", out)
        self.emit(receiver="Orchestrator", kind="eval_finished", payload=out)
        return out