from __future__ import annotations

from pathlib import Path

import numpy as np

from openimpact.evaluation import evaluate_federated_setup
from openimpact.federated import FederatedTrainer
from openimpact.ledger import GovernanceLedger
from openimpact.model import TemporaryImpactModel
from openimpact.replay import DeterministicReplayEngine
from openimpact.synthetic import SyntheticMarketConfig, generate_synthetic_market


def test_temporary_impact_model_fits_synthetic_market():
    market = generate_synthetic_market(SyntheticMarketConfig(seed=11, venue_count=2, samples_per_venue=64))
    requests = market.requests
    impacts = market.impacts

    model = TemporaryImpactModel().fit(requests, impacts)
    predictions = model.predict(requests)

    assert predictions.shape == impacts.shape
    assert float(np.sqrt(np.mean((predictions - impacts) ** 2))) < 0.05


def test_federated_training_and_evaluation_are_deterministic():
    market = generate_synthetic_market(SyntheticMarketConfig(seed=3, venue_count=3, samples_per_venue=32))
    trainer = FederatedTrainer(seed=3)
    metrics_a, model_a, baseline_a = evaluate_federated_setup(market, trainer)
    metrics_b, model_b, baseline_b = evaluate_federated_setup(market, trainer)

    assert metrics_a == metrics_b
    assert np.allclose(model_a.coefficients_, model_b.coefficients_)
    assert np.allclose(baseline_a.coefficients_, baseline_b.coefficients_)
    assert metrics_a.rmse < 0.08


def test_replay_and_ledger_work_end_to_end(tmp_path: Path):
    market = generate_synthetic_market(SyntheticMarketConfig(seed=5, venue_count=2, samples_per_venue=16))
    trainer = FederatedTrainer(seed=5)
    result = trainer.fit(market.batches)
    replay = DeterministicReplayEngine(seed=5, base_latency_ms={"venue-1": 2.5, "venue-2": 3.5})
    steps = replay.replay(market.requests, result.global_model, market.impacts)

    ledger = GovernanceLedger(tmp_path / "ledger.sqlite3")
    for round_id, update in enumerate(result.updates, start=1):
        ledger.record_update(round_id=round_id, model_version="v1", update=update)

    entries = ledger.list_entries()
    assert len(entries) == 2
    assert len(steps) == len(market.requests)
    assert steps[0].delivery_ns > steps[0].timestamp_ns