##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2020, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
#
# Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
# license information, respectively. Both files are also available online
# at the URL "https://github.com/IDAES/idaes-pse".
##############################################################################
"""
Tests for math util methods.
"""

import pytest
from pyomo.environ import ConcreteModel, Param, Var, value
from idaes.core.util.math import (smooth_abs, smooth_minmax,
                                  smooth_min, smooth_max)

__author__ = "Andrew Lee"


@pytest.fixture(scope="module")
def simple_model():
    """Build a simple model for testing.
    """
    m = ConcreteModel()
    m.a = Var(initialize=4.0)
    m.b = Var(initialize=-4.0)
    m.e = Param(default=1e-4)

    return m


@pytest.mark.unit
def test_smooth_abs_maths():
    # Test basic smooth_abs functionalliy
    assert smooth_abs(4, 0) == 4.0
    assert smooth_abs(-4, 0) == 4.0
    assert smooth_abs(10.0, 0.0) == 10.0
    assert smooth_abs(-10.0, 0.0) == 10.0

    assert smooth_abs(2, 1e-4) == pytest.approx(2.0, abs=1e-4)
    assert smooth_abs(-2, 1e-4) == pytest.approx(2.0, abs=1e-4)
    assert smooth_abs(10) == pytest.approx(10.0, abs=1e-4)
    assert smooth_abs(-10) == pytest.approx(10.0, abs=1e-4)


@pytest.mark.unit
def test_smooth_abs_expr(simple_model):
    # Test that smooth_abs works with Pyomo components
    assert value(smooth_abs(simple_model.a, 0)) == 4.0
    assert value(smooth_abs(simple_model.b, 0)) == 4.0

    assert (value(smooth_abs(simple_model.a, simple_model.e)) ==
            pytest.approx(4.0, abs=1e-4))
    assert (value(smooth_abs(simple_model.b, simple_model.e)) ==
            pytest.approx(4.0, abs=1e-4))


@pytest.mark.unit
def test_smooth_abs_a_errors():
    # Test that smooth_abs returns meaningful errors when given invalid arg
    with pytest.raises(TypeError):
        smooth_abs("foo")
    with pytest.raises(TypeError):
        smooth_abs([1, 2, 3])


@pytest.mark.unit
def test_smooth_abs_eps_errors():
    # Test that smooth_abs returns meaningful errors when given invalid eps
    with pytest.raises(TypeError):
        smooth_abs(1.0, "a")
    with pytest.raises(TypeError):
        smooth_abs(1.0, [1, 2, 3])


@pytest.mark.unit
def test_smooth_minmax_maths():
    # Test basic smooth_minmax functionality
    assert smooth_minmax(1, 2, 0, sense='max') == 2
    assert smooth_minmax(1, 2, 0, sense='min') == 1
    assert smooth_minmax(5.0, 3, 0.0, sense='max') == 5
    assert smooth_minmax(5.0, 3, 0.0, sense='min') == 3

    assert (smooth_minmax(2.0, 12.0, 1e-4, 'max') ==
            pytest.approx(12.0, abs=1e-4))
    assert (smooth_minmax(2.0, 12.0, 1e-4, 'min') ==
            pytest.approx(2.0, abs=1e-4))
    assert (smooth_minmax(32.0, 12.0, sense='max') ==
            pytest.approx(32.0, abs=1e-4))
    assert (smooth_minmax(32.0, 12.0, sense='min') ==
            pytest.approx(12.0, abs=1e-4))


@pytest.mark.unit
def test_smooth_minmax_default_sense():
    # Test that smooth_minmax defaults to maximise
    assert smooth_minmax(1, 2, 0,) == 2


@pytest.mark.unit
def test_smooth_minmax_expr(simple_model):
    # Test that smooth_minmax works with Pyomo components
    assert value(smooth_minmax(simple_model.a,
                               simple_model.b,
                               0,
                               sense='max')) == 4.0
    assert value(smooth_minmax(simple_model.a,
                               simple_model.b,
                               0,
                               sense='min')) == -4.0

    assert value(smooth_minmax(simple_model.a,
                               simple_model.b,
                               sense='max')) == pytest.approx(4.0, abs=1e-4)
    assert value(smooth_minmax(simple_model.a,
                               simple_model.b,
                               sense='min')) == pytest.approx(-4.0, abs=1e-4)

    assert value(smooth_minmax(simple_model.a,
                               simple_model.b,
                               simple_model.e,
                               sense='max')) == pytest.approx(4.0, abs=1e-4)
    assert value(smooth_minmax(simple_model.a,
                               simple_model.b,
                               simple_model.e,
                               sense='min')) == pytest.approx(-4.0, abs=1e-4)


@pytest.mark.unit
def test_smooth_abs_ab_errors():
    # Test that smooth_abs returns meaningful errors when given invalid args
    with pytest.raises(TypeError):
        smooth_abs("foo", 1)
    with pytest.raises(TypeError):
        smooth_abs(3, [1, 2, 3])


@pytest.mark.unit
def test_smooth_minmax_eps_errors():
    # Test that smooth_abs returns meaningful errors when given invalid eps
    with pytest.raises(TypeError):
        smooth_minmax(1.0, 1.0, "foo")
    with pytest.raises(TypeError):
        smooth_minmax(1.0, 1.0, [1, 2, 3])


@pytest.mark.unit
def test_smooth_minmax_sense_errors():
    # Test that smooth_abs returns meaningful errors when given invalid sense
    with pytest.raises(ValueError):
        smooth_minmax(1.0, 1.0, sense="foo")
    with pytest.raises(ValueError):
        smooth_minmax(1.0, 1.0, sense=1.0)
    with pytest.raises(ValueError):
        smooth_minmax(1.0, 1.0, sense=[1.0])


@pytest.mark.unit
def test_smooth_max(simple_model):
    # Test that smooth_max gives correct values
    assert smooth_max(3.0, 12.0) == pytest.approx(12.0, abs=1e-4)
    assert value(smooth_max(simple_model.a,
                            simple_model.b,
                            simple_model.e)) == pytest.approx(4.0, abs=1e-4)


@pytest.mark.unit
def test_smooth_min(simple_model):
    # Test that smooth_min gives correct values
    assert smooth_min(3.0, 12.0) == pytest.approx(3.0, abs=1e-4)
    assert value(smooth_min(simple_model.a,
                            simple_model.b,
                            simple_model.e)) == pytest.approx(-4.0, abs=1e-4)
