"""Tests for the mock flow solver."""
import pytest

from pipeline_editor.calc import flow_solver
from pipeline_editor.model.document import DiagramDocument
from pipeline_editor.model.edge import EdgeModel
from pipeline_editor.model.node_library import default_library


@pytest.fixture
def lib():
    return default_library()


def _add(doc, lib, key, x=0, y=0):
    return doc.add_node(lib.get(key).instantiate(doc.next_id("n"), x, y))


def test_linear_network_flow_direction(qapp, lib):
    doc = DiagramDocument()
    src = _add(doc, lib, "source")          # supply 120
    pump = _add(doc, lib, "pump", 200)
    cons = _add(doc, lib, "consumer", 400)  # demand 40
    e1 = doc.add_edge(EdgeModel(doc.next_id("e"), src.node_id, "out", pump.node_id, "in"))
    e2 = doc.add_edge(EdgeModel(doc.next_id("e"), pump.node_id, "out", cons.node_id, "in"))

    result = flow_solver.apply_to_document(doc)
    # 40 units flow from source to consumer along both edges, in source->target dir
    assert result.edge_flow[e1.edge_id] == pytest.approx(40)
    assert result.edge_flow[e2.edge_id] == pytest.approx(40)
    assert e1.flow_direction == 1 and e2.flow_direction == 1
    assert e1.flow == pytest.approx(40)


def test_reversed_edge_gives_negative_direction(qapp, lib):
    doc = DiagramDocument()
    src = _add(doc, lib, "source")
    cons = _add(doc, lib, "consumer", 400)
    # edge authored consumer(source-of-edge) -> source(target-of-edge)
    e = doc.add_edge(EdgeModel(doc.next_id("e"), cons.node_id, "in", src.node_id, "out"))
    flow_solver.apply_to_document(doc)
    # physical flow is source->consumer, i.e. against edge authoring direction
    assert e.flow_direction == -1
    assert e.flow == pytest.approx(40)


def test_demand_capped_by_supply(qapp, lib):
    doc = DiagramDocument()
    src = _add(doc, lib, "source")     # supply 120
    c1 = _add(doc, lib, "consumer", 200)
    c2 = _add(doc, lib, "consumer", 400)
    c1.properties.set_value("demand", 100)
    c2.properties.set_value("demand", 100)  # total demand 200 > supply 120
    doc.add_edge(EdgeModel(doc.next_id("e"), src.node_id, "out", c1.node_id, "in"))
    doc.add_edge(EdgeModel(doc.next_id("e"), src.node_id, "out", c2.node_id, "in"))
    result = flow_solver.apply_to_document(doc)
    served = sum(abs(f) for f in result.edge_flow.values())
    assert served == pytest.approx(120)  # cannot exceed supply


def test_no_sources_uses_nominal_flow(qapp, lib):
    doc = DiagramDocument()
    a = _add(doc, lib, "junction")
    b = _add(doc, lib, "junction", 200)
    e = doc.add_edge(EdgeModel(doc.next_id("e"), a.node_id, "e", b.node_id, "w"))
    flow_solver.apply_to_document(doc)
    assert e.flow == pytest.approx(1.0)
    assert e.flow_direction == 1


def test_empty_document(qapp):
    doc = DiagramDocument()
    result = flow_solver.solve(doc)
    assert result.edge_flow == {}