##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2019, 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 flowsheet_model.

Author: Andrew Lee
"""
import pytest
import inspect
from pyomo.environ import ConcreteModel, Constraint, Set, Var
from pyomo.common.config import ConfigBlock
from idaes.core import (declare_process_block_class, ReactionParameterBlock,
                        ReactionBlockBase, ReactionBlockDataBase,
                        PhysicalParameterBlock, StateBlock,
                        StateBlockData)
from idaes.core.util.exceptions import (PropertyPackageError,
                                        PropertyNotSupportedError)


# -----------------------------------------------------------------------------
# Test ParameterBlock
@declare_process_block_class("PropertyParameterBlock")
class _PropertyParameterBlock(PhysicalParameterBlock):
    def build(self):
        super(_PropertyParameterBlock, self).build()

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'prop1': {'method': None, 'units': 'm'},
                            'prop3': {'method': False}})
        obj.add_default_units({'time': 's',
                               'length': 'm',
                               'mass': 'g',
                               'amount': 'mol',
                               'temperature': 'K',
                               'energy': 'J',
                               'holdup': 'mol'})


@declare_process_block_class("ReactionParameterTestBlock")
class _ReactionParameterBlock(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()


def test_config_block():
    # Test that PhysicalParameterBlock gets module information
    m = ConcreteModel()
    m.r = ReactionParameterTestBlock()

    assert len(m.r.config) == 2
    assert isinstance(m.r.config.default_arguments, ConfigBlock)
    assert len(m.r.config.default_arguments) == 0


def test_ReactionParameter_NotImplementedErrors():
    # Test that class methods return NotImplementedError
    m = ConcreteModel()
    m.r = ReactionParameterTestBlock()

    with pytest.raises(NotImplementedError):
        m.r.get_metadata()


@declare_process_block_class("ReactionParameterBlock2")
class _ReactionParameterBlock2(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()

    @classmethod
    def get_required_properties(self):
        return ['prop1']

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'rxn1': {'method': None}})
        obj.add_default_units({'time': 'hr',
                               'length': 'm',
                               'mass': 'g',
                               'amount': 'mol',
                               'temperature': 'K',
                               'energy': 'J',
                               'holdup': 'mol'})

def test_validate_state_block_invalid_units():
    # Test validation of associated PropertyParameterBlock
    m = ConcreteModel()
    m.p = PropertyParameterBlock()
    m.r = ReactionParameterBlock2(default={"property_package": m.p})

    with pytest.raises(PropertyPackageError):
        m.r._validate_property_parameter_units()


@declare_process_block_class("ReactionParameterBlock3")
class _ReactionParameterBlock3(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'rxn1': {'method': None}})
        obj.add_default_units({'time': 's',
                               'length': 'm',
                               'mass': 'g',
                               'amount': 'mol',
                               'temperature': 'K',
                               'energy': 'J',
                               'holdup': 'mol'})
        obj.add_required_properties({'prop2': 'some'})


def test_validate_state_block_unsupported_prop():
    # Test validation of associated PropertyParameterBlock
    m = ConcreteModel()
    m.p = PropertyParameterBlock()
    m.r = ReactionParameterBlock3(default={"property_package": m.p})

    with pytest.raises(PropertyPackageError):
        m.r._validate_property_parameter_properties()


@declare_process_block_class("ReactionParameterBlock4")
class _ReactionParameterBlock4(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'rxn1': {'method': None}})
        obj.add_default_units({'time': 's',
                               'length': 'm',
                               'mass': 'g',
                               'amount': 'mol',
                               'temperature': 'K',
                               'energy': 'J',
                               'holdup': 'mol'})
        obj.add_required_properties({'prop3': 'some'})


def test_validate_state_block_unsupported_prop_False():
    # Test validation of associated PropertyParameterBlock
    m = ConcreteModel()
    m.p = PropertyParameterBlock()
    m.r = ReactionParameterBlock4(default={"property_package": m.p})

    with pytest.raises(PropertyPackageError):
        m.r._validate_property_parameter_properties()


@declare_process_block_class("ReactionParameterBlock5")
class _ReactionParameterBlock5(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'rxn1': {'method': None}})
        obj.add_default_units({'time': 's',
                               'length': 'm',
                               'mass': 'g',
                               'amount': 'mol',
                               'temperature': 'K',
                               'energy': 'J',
                               'holdup': 'mol'})
        obj.add_required_properties({'prop1': 'km'})


def test_validate_state_block_req_prop_wrong_units():
    # Test validation of associated PropertyParameterBlock
    m = ConcreteModel()
    m.p = PropertyParameterBlock()
    m.r = ReactionParameterBlock5(default={"property_package": m.p})

    with pytest.raises(PropertyPackageError):
        m.r._validate_property_parameter_properties()


@declare_process_block_class("ReactionParameterBlock6")
class _ReactionParameterBlock6(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'rxn1': {'method': None}})
        obj.add_default_units({'time': 's',
                               'length': 'm',
                               'mass': 'g',
                               'amount': 'mol',
                               'temperature': 'K',
                               'energy': 'J',
                               'holdup': 'mol'})
        obj.add_required_properties({'prop1': 'm'})


def test_ReactionParameterBase_build():
    # Test that ReactionParameterBlock builds correctly
    m = ConcreteModel()
    m.p = PropertyParameterBlock()
    m.r = ReactionParameterBlock6(default={"property_package": m.p})
    super(_ReactionParameterBlock6, m.r).build()


# -----------------------------------------------------------------------------
# Test ReactionBlockBase
@declare_process_block_class("ReactionBlock",
                             block_class=ReactionBlockBase)
class ReactionBlockData(ReactionBlockDataBase):
    def build(self):
        super(ReactionBlockDataBase, self).build()


def test_ReactionBlockBase_initialize():
    # Test that ReactionBlockBase initialize method raises NotImplementedError
    m = ConcreteModel()
    # Need to index block so that it does not do multiple inheritance
    m.s = Set(initialize=[1, 2])
    m.r = ReactionBlock(m.s)

    with pytest.raises(NotImplementedError):
        m.r.initialize()


def test_ReactionBlockBase_report():
    # Test that ReactionBlockBase report method raises NotImplementedError
    m = ConcreteModel()
    # Need to index block so that it does not do multiple inheritance
    m.s = Set(initialize=[1, 2])
    m.r = ReactionBlock(m.s)

    with pytest.raises(NotImplementedError):
        m.r.report()


# -----------------------------------------------------------------------------
# Test ReactionBlockDataBase
def test_StateBlock_config():
    # Test that ReactionBlockDataBase config has correct arguments
    m = ConcreteModel()
    m.p = ReactionBlock()

    assert len(m.p.config) == 3
    assert hasattr(m.p.config, "parameters")
    assert hasattr(m.p.config, "state_block")
    assert hasattr(m.p.config, "has_equilibrium")

    m.p.config.has_equilibrium = True
    m.p.config.has_equilibrium = False
    with pytest.raises(ValueError):
        m.p.config.has_equilibrium = 'foo'
    with pytest.raises(ValueError):
        m.p.config.has_equilibrium = 10


@declare_process_block_class("TestStateBlock", block_class=StateBlock)
class StateTestBlockData(StateBlockData):
    def build(self):
        super(StateBlockData, self).build()


def test_validate_state_block_fail():
    # Test that ReactionBlockDataBase validate_state_block returns error
    m = ConcreteModel()
    m.p = PropertyParameterBlock()
    m.p2 = PropertyParameterBlock()

    m.pb = TestStateBlock(default={"parameters": m.p2})

    m.r = ReactionParameterBlock6(default={"property_package": m.p})
    super(_ReactionParameterBlock6, m.r).build()

    m.rb = ReactionBlock(default={"parameters": m.r,
                                  "state_block": m.pb})

    with pytest.raises(PropertyPackageError):
        m.rb._validate_state_block()


@declare_process_block_class("ReactionBlock2",
                             block_class=ReactionBlockBase)
class ReactionBlockData2(ReactionBlockDataBase):
    def build(self):
        super(ReactionBlockData2, self).build()


def test_build():
    # Test that ReactionBlockDataBase builds correctly with good argumnets
    m = ConcreteModel()
    m.p = PropertyParameterBlock()

    m.pb = TestStateBlock(default={"parameters": m.p})

    m.r = ReactionParameterBlock6(default={"property_package": m.p})
    super(_ReactionParameterBlock6, m.r).build()

    m.rb = ReactionBlock2(default={"parameters": m.r,
                                   "state_block": m.pb})

    assert hasattr(m.rb, "state_ref")
    assert m.rb._params == m.rb.config.parameters


# -----------------------------------------------------------------------------
# Test reaction __getattr__ method
@declare_process_block_class("Parameters")
class _Parameters(ReactionParameterBlock):
    def build(self):
        super(ReactionParameterBlock, self).build()
        frm = inspect.stack()[1]
        self.__package_module = inspect.getmodule(frm[0])

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties({'a': {'method': 'a_method'},
                            'recursion1': {'method': '_recursion1'},
                            'recursion2': {'method': '_recursion2'},
                            'not_callable': {'method': 'test_obj'},
                            'raise_exception': {'method': '_raise_exception'},
                            'not_supported': {'method': False},
                            'does_not_create_component': {
                                    'method': '_does_not_create_component'}})


@declare_process_block_class("Reaction", block_class=ReactionBlockBase)
class _Reaction(ReactionBlockDataBase):
    def build(self):
        super(ReactionBlockDataBase, self).build()

        self.test_obj = 1

    def a_method(self):
        self.a = Var(initialize=1)

    def _recursion1(self):
        self.recursive_cons1 = Constraint(expr=self.recursion2 == 1)

    def _recursion2(self):
        self.recursive_cons2 = Constraint(expr=self.recursion1 == 1)

    def _raise_exception(self):
        raise Exception()

    def _does_not_create_component(self):
        pass


@pytest.fixture()
def m():
    m = ConcreteModel()
    m.pb = Parameters()
    m.p = Reaction(default={"parameters": m.pb})

    return m


def test_getattr_add_var(m):
    assert isinstance(m.p.a, Var)
    assert m.p.a.value == 1


def test_getattr_protected(m):
    with pytest.raises(PropertyPackageError):
        # Call a protected component that does not exist
        m.p.cons = Constraint(expr=m.p._foo == 1)


def test_getattr_recursion(m):
    with pytest.raises(PropertyPackageError):
        # Call a component that triggers a recursive loop of calls
        m.p.cons = Constraint(expr=m.p.recursion1 == 1)


def test_getattr_does_not_exist(m):
    with pytest.raises(PropertyNotSupportedError):
        m.p.cons = Constraint(expr=m.p.does_not_exist == 1)


def test_getattr_not_callable(m):
    with pytest.raises(PropertyPackageError):
        m.p.cons = Constraint(expr=m.p.not_callable == 1)


def test_getattr_not_supported(m):
    with pytest.raises(PropertyNotSupportedError):
        m.p.cons = Constraint(expr=m.p.not_supported == 1)


def test_getattr_raise_exception(m):
    with pytest.raises(Exception):
        m.p.cons = Constraint(expr=m.p.raise_exception == 1)


# TODO : Need a test for cases where method does not create property
#def test_getattr_does_not_create_component(m):
#    with pytest.raises(PropertyPackageError):
#        m.p.cons = Constraint(expr=m.p.does_not_create_component == 1)
