from idea154_civicmesh_studio_federated.contracts import Asset, DualVariable, LocalModel, PlanDelta, SharedSignals
from idea154_civicmesh_studio_federated.solver import FederatedResilienceSolver, NeighborhoodProblem
from idea154_civicmesh_studio_federated.sync import DeltaSyncLog


def _problem(neighborhood_id: str, demand: float, capacity: float, energy_limit: float) -> NeighborhoodProblem:
    return NeighborhoodProblem(
        local_model=LocalModel(
            neighborhood_id=neighborhood_id,
            assets=[Asset(asset_id=f"{neighborhood_id}-asset", asset_type="microgrid", neighborhood_id=neighborhood_id, capacity_kw=capacity)],
            energy_limit_kw=energy_limit,
        ),
        signals=SharedSignals(
            neighborhood_id=neighborhood_id,
            source="weather",
            forecast_demand_kw=demand,
            available_capacity_kw=capacity,
            weather_risk=2.0,
        ),
        dual=DualVariable(neighborhood_id=neighborhood_id),
    )


def test_delta_sync_replays_deterministically():
    log = DeltaSyncLog()
    first = log.append(PlanDelta(neighborhood_id="a", iteration=1, actions={"load_shift_kw": 2.5}, notes="alpha"))
    second = log.append(PlanDelta(neighborhood_id="b", iteration=1, actions={"load_shift_kw": 1.5}, notes="beta"))
    assert first.sequence == 1
    assert second.sequence == 2
    assert log.verify()

    replayed = log.replay()
    assert len(replayed) == 2
    assert replayed[0].neighborhood_id == "a"


def test_solver_converges_under_district_limit():
    solver = FederatedResilienceSolver(
        problems=[
            _problem("north", 8.0, 2.0, 6.0),
            _problem("south", 7.0, 1.0, 5.0),
        ],
        district_energy_limit_kw=6.0,
        max_iterations=50,
        tolerance=1e-5,
        delay_steps=1,
    )

    result = solver.solve()

    assert result.iterations >= 1
    assert sum(delta.actions["load_shift_kw"] for delta in result.deltas) <= 6.000001
    assert result.primal_residual < 1.0