from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .contracts import flatten_batches
from .federated import FederatedTrainer
from .model import TemporaryImpactModel
from .replay import DeterministicReplayEngine
from .synthetic import SyntheticMarket


@dataclass(frozen=True)
class EvaluationMetrics:
    rmse: float
    latency_ms: float
    privacy_budget: float
    leakage_bound: float


def _rmse(prediction: np.ndarray, target: np.ndarray) -> float:
    return float(np.sqrt(np.mean((prediction - target) ** 2)))


def evaluate_federated_setup(market: SyntheticMarket, trainer: FederatedTrainer | None = None) -> tuple[EvaluationMetrics, TemporaryImpactModel, TemporaryImpactModel]:
    trainer = trainer or FederatedTrainer()
    flat_requests, flat_impacts = flatten_batches(market.batches)
    baseline = TemporaryImpactModel().fit(flat_requests, flat_impacts)

    result = trainer.fit(market.batches)
    replay = DeterministicReplayEngine(seed=trainer.seed)
    steps = replay.replay(flat_requests, result.global_model, flat_impacts)
    latency_ms = float(np.mean([(step.delivery_ns - step.timestamp_ns) / 1_000_000.0 for step in steps]))

    predictions = result.global_model.predict(flat_requests)
    rmse = _rmse(predictions, flat_impacts)
    privacy_budget = float(sum(update.privacy_budget for update in result.updates))
    leakage_bound = float(np.max(np.abs(predictions - baseline.predict(flat_requests))))
    return EvaluationMetrics(rmse=rmse, latency_ms=latency_ms, privacy_budget=privacy_budget, leakage_bound=leakage_bound), result.global_model, baseline