"""
A collection of tests for pywr.domains.river

Specific additional functionality of the 'special' classes in the river domain
are tested here.
"""
from __future__ import print_function
import pywr.core
from pywr.core import Model, Input, Output, Catchment
from pywr.parameters import MonthlyProfileParameter
from pywr.domains import river
from numpy.testing import assert_allclose
import pytest

from helpers import assert_model, load_model

@pytest.fixture(params=[(10.0, 10.0, 10.0), (5.0, 5.0, 1.0)])
def simple_gauge_model(request):
    """
    Make a simple model with a single Input and Output and RiverGauge

    Input -> PiecewiseLink -> Output

    """
    in_flow, out_flow, benefit = request.param
    min_flow_req = 5.0

    model = pywr.core.Model()
    inpt = river.Catchment(model, name="Catchment", flow=in_flow)
    lnk = river.RiverGauge(model, name="Gauge", mrf=min_flow_req, mrf_cost=-1.0)
    inpt.connect(lnk)
    otpt = pywr.core.Output(model, name="Demand", max_flow=out_flow, cost=-benefit)
    lnk.connect(otpt)

    default = inpt.domain

    expected_sent = in_flow if benefit > 1.0 else out_flow

    expected_node_results = {
        "Catchment": expected_sent,
        "Gauge": expected_sent,
        "Gauge Sublink 0": min(min_flow_req, expected_sent),
        "Gauge Sublink 1": expected_sent - min(min_flow_req, expected_sent),
        "Demand": expected_sent,
    }
    return model, expected_node_results


@pytest.fixture
def simple_river_split_gauge_model():
    """
    Make a simple model with a single Input and Output and RiverGauge

    Input -> RiverSplit -> Output 1
                 \ --->--> Output 2

    """
    in_flow = 100.0
    min_flow_req = 40.0
    out_flow = 50.0
    model = pywr.core.Model()

    inpt = river.Catchment(model, name="Catchment", flow=in_flow)
    lnk = river.RiverSplitWithGauge(model, name="Gauge", mrf=min_flow_req, mrf_cost=-100,
                                    slot_names=("river", "abstraction"), factors=[3, 1])
    inpt.connect(lnk)
    estuary = pywr.core.Output(model, name="Estuary")
    lnk.connect(estuary, from_slot="river")
    otpt = pywr.core.Output(model, name="Demand", max_flow=out_flow, cost=-10)
    lnk.connect(otpt, from_slot="abstraction")

    net_flow_after_mrf = in_flow - min_flow_req
    expected_node_results = {
        "Catchment": in_flow,
        "Gauge": in_flow,
        "Gauge Sublink 0": min_flow_req,
        "Gauge Sublink 1": net_flow_after_mrf * 0.75,
        "Gauge Sublink 2": net_flow_after_mrf * 0.25,
        "Demand": min(out_flow, net_flow_after_mrf * 0.25),
    }
    return model, expected_node_results


def test_river_gauge():
    """
    Test loading a model with a RiverGauge from JSON, modifying it, then running it
    """
    model = load_model("river_mrf1.json")

    node = model.nodes["mrf"]
    demand = model.nodes["demand"]

    # test getting properties
    assert(isinstance(node.mrf, MonthlyProfileParameter))
    assert_allclose(node.mrf_cost, -1000)
    assert_allclose(node.cost, 0.0)

    # test setting properties
    node.mrf = 40
    node.mrf_cost = -999
    assert_allclose(node.mrf, 40)
    assert_allclose(node.mrf_cost, -999)

    # run the model and see if it works
    model.run()
    assert_allclose(node.flow, 40)
    assert_allclose(demand.flow, 60)


def test_piecewise_model(simple_gauge_model):
    assert_model(*simple_gauge_model)


def test_river_split_gauge(simple_river_split_gauge_model):
    assert_model(*simple_river_split_gauge_model)

def test_river_split_gauge_json():
    """As test_river_split_gauge, but model is defined in JSON"""

    model = load_model("river_split_with_gauge1.json")
    model.check()
    model.run()

    catchment_flow = 100.0
    mrf = 40.0
    expected_demand_flow = (catchment_flow - mrf) * 0.25

    catchment_node = model.nodes["Catchment"]
    demand_node = model.nodes["Demand"]
    assert_allclose(catchment_node.flow, catchment_flow)
    assert_allclose(demand_node.flow, expected_demand_flow)


def test_control_curve():
    """
    Use a simple model of a Reservoir to test that a control curve
    behaves as expected.

    The control curve should alter the cost of the Reservoir when it
    is above or below a particular threshold.

    (flow = 8.0)          (max_flow = 10.0)
    Catchment -> River -> DemandCentre
                     |        ^
    (max_flow = 2.0) v        | (max_flow = 2.0)
                    Reservoir

    """
    in_flow = 8

    model = pywr.core.Model()
    catchment = river.Catchment(model, name="Catchment", flow=in_flow)
    lnk = river.River(model, name="River")
    catchment.connect(lnk)
    demand = pywr.core.Output(model, name="Demand", cost=-10.0, max_flow=10)
    lnk.connect(demand)
    from pywr.parameters import ConstantParameter
    control_curve = ConstantParameter(model, 0.8)
    reservoir = river.Reservoir(model, name="Reservoir", max_volume=10, cost=-20, above_curve_cost=0.0,
                                control_curve=control_curve, initial_volume=10)
    reservoir.inputs[0].max_flow = 2.0
    reservoir.outputs[0].max_flow = 2.0
    lnk.connect(reservoir)
    reservoir.connect(demand)

    model.setup()

    model.step()
    # Reservoir is currently above control curve. 2 should be taken from the
    # reservoir
    assert(reservoir.volume == 8)
    assert(demand.flow == 10)
    # Reservoir is still at (therefore above) control curve. So 2 is still taken
    model.step()
    assert(reservoir.volume == 6)
    assert(demand.flow == 10)
    # Reservoir now below curve. Better to retain volume and divert some of the
    # inflow
    model.step()
    assert(reservoir.volume == 8)
    assert(demand.flow == 6)
    # Set the above_curve_cost function to keep filling
    from pywr.parameters.control_curves import ControlCurveParameter
    # We know what we're doing with the control_curve Parameter so unset its parent before overriding
    # the cost parameter.
    # We need to call setup() again because we're adding a parameter.
    reservoir.cost = ControlCurveParameter(model, reservoir, control_curve, [-20.0, -20.0])
    reservoir.initial_volume = 8
    model.setup()
    model.step()
    assert(reservoir.volume == 10)
    assert(demand.flow == 6)

def test_catchment_many_successors():
    """Test if node with fixed flow can have multiple successors. See #225"""
    model = Model()
    catchment = Catchment(model, "catchment", flow=100)
    out1 = Output(model, "out1", max_flow=10, cost=-100)
    out2 = Output(model, "out2", max_flow=15, cost=-50)
    out3 = Output(model, "out3")
    catchment.connect(out1)
    catchment.connect(out2)
    catchment.connect(out3)
    model.check()
    model.run()
    assert_allclose(out1.flow, 10)
    assert_allclose(out2.flow, 15)
    assert_allclose(out3.flow, 75)
