"""Tests for activation function parameters."""
from pywr.parameters import load_parameter
from pywr.recorders import assert_rec
from fixtures import simple_linear_model
import numpy as np
import pytest


class TestBinaryStepParameter:
    @pytest.mark.parametrize(['internal_value', 'output', ], [
        (-1.0, 1.0, ),
        (0.0, 1.0, ),
        (0.5, 1.0, ),
    ])
    def test_binary_step(self, simple_linear_model, internal_value, output):
        """Test the binary step profile parameter."""

        m = simple_linear_model
        m.timestepper.start = '2015-01-01'
        m.timestepper.end = '2015-12-31'

        if internal_value <= 0.0:
            expected_values = np.zeros(365)
        else:
            expected_values = np.ones(365) * output

        data = {
            'type': 'binarystep',
            'value': internal_value,
            'output': output,
        }

        p = load_parameter(m, data)

        @assert_rec(m, p)
        def expected_func(timestep, _scenario_index):
            return expected_values[timestep.index]

        m.run()

    @pytest.mark.parametrize(['lb', 'ub', ], [
        (-1.0, 1.0, ),
        (-10.0, 5.0, )
    ])
    def test_bounds(self, simple_linear_model, lb, ub):

        m = simple_linear_model

        data = {
            'type': 'binarystep',
            'value': 1.0,
            'output': 1.0,
            'lower_bounds': lb,
            'upper_bounds': ub
        }

        p = load_parameter(m, data)
        np.testing.assert_allclose(p.get_double_lower_bounds(), [lb])
        np.testing.assert_allclose(p.get_double_upper_bounds(), [ub])


class TestRectifierParameter:
    @pytest.mark.parametrize(['internal_value', 'max_output', 'upper_bounds'], [
        (-1.0, 1.0, 1.0),
        (0.0, 1.0, 1.0),
        (0.5, 1.0, 1.0),
        (0.5, 10.0, 1.0),
        (0.5, 10.0, 5.0),
        (1.0, 10.0, 1.0),
    ])
    def test_binary_step(self, simple_linear_model, internal_value, max_output, upper_bounds):
        """Test the binary step profile parameter."""

        m = simple_linear_model
        m.timestepper.start = '2015-01-01'
        m.timestepper.end = '2015-12-31'

        if internal_value <= 0.0:
            expected_values = np.zeros(365)
        else:
            expected_values = np.ones(365) * max_output * internal_value / upper_bounds

        data = {
            'type': 'rectifier',
            'value': internal_value,
            'max_output': max_output,
            'upper_bounds': upper_bounds
        }

        p = load_parameter(m, data)

        @assert_rec(m, p)
        def expected_func(timestep, _scenario_index):
            return expected_values[timestep.index]

        m.run()

    @pytest.mark.parametrize(['lb', 'ub', ], [
        (-1.0, 1.0, ),
        (-10.0, 5.0, )
    ])
    def test_bounds(self, simple_linear_model, lb, ub):

        m = simple_linear_model

        data = {
            'type': 'rectifier',
            'value': 1.0,
            'max_output': 1.0,
            'lower_bounds': lb,
            'upper_bounds': ub
        }

        p = load_parameter(m, data)
        np.testing.assert_allclose(p.get_double_lower_bounds(), [lb])
        np.testing.assert_allclose(p.get_double_upper_bounds(), [ub])

class TestLogisticParameter:
    @pytest.mark.parametrize(['internal_value', 'max_output', 'growth_rate', ], [
        (-1.0, 1.0, 1.0),
        (0.0, 1.0, 1.0),
        (0.5, 1.0, 1.0),
        (1.0, 1.0, 1.0),
    ])
    def test_logistic(self, simple_linear_model, internal_value, max_output, growth_rate):
        """Test the binary step profile parameter."""

        m = simple_linear_model
        m.timestepper.start = '2015-01-01'
        m.timestepper.end = '2015-12-31'

        expected_values = np.ones(365) * max_output / (1.0 + np.exp(-growth_rate * internal_value))

        data = {
            'type': 'logistic',
            'value': internal_value,
            'max_output': max_output,
        }

        p = load_parameter(m, data)

        @assert_rec(m, p)
        def expected_func(timestep, _scenario_index):
            return expected_values[timestep.index]

        m.run()

    @pytest.mark.parametrize(['lb', 'ub', ], [
        (-1.0, 1.0, ),
        (-10.0, 5.0, )
    ])
    def test_bounds(self, simple_linear_model, lb, ub):

        m = simple_linear_model

        data = {
            'type': 'logistic',
            'value': 1.0,
            'max_output': 1.0,
            'lower_bounds': lb,
            'upper_bounds': ub
        }

        p = load_parameter(m, data)
        np.testing.assert_allclose(p.get_double_lower_bounds(), [lb])
        np.testing.assert_allclose(p.get_double_upper_bounds(), [ub])
