###############################################################################
# WaterTAP Copyright (c) 2021, The Regents of the University of California,
# through Lawrence Berkeley National Laboratory, Oak Ridge National
# Laboratory, National Renewable Energy Laboratory, and National Energy
# Technology Laboratory (subject to receipt of any required approvals from
# the U.S. Dept. of Energy). All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and license
# information, respectively. These files are also available online at the URL
# "https://github.com/watertap-org/watertap/"
#
###############################################################################
"""
Tests for general zero-order property package
"""
import pytest
from io import StringIO
import os

from idaes.core import declare_process_block_class, FlowsheetBlock
from idaes.core.util.model_statistics import degrees_of_freedom
from idaes.core.util.testing import initialization_tester
from idaes.core.util import get_solver
import idaes.core.util.scaling as iscale
from pyomo.environ import (check_optimal_termination,
                           ConcreteModel,
                           Constraint,
                           Param,
                           Set,
                           value,
                           Var)
from pyomo.network import Port
from pyomo.util.check_units import assert_units_consistent


from watertap.core import (
    Database, WaterParameterBlock, WaterStateBlock, ZeroOrderBaseData)
from watertap.core.zero_order_sido_reactive import (
    build_sido_reactive,
    initialize_sidor,
    calculate_scaling_factors_sidor,
    _get_Q_sidor)

solver = get_solver()

local_path = os.path.dirname(os.path.abspath(__file__))


@declare_process_block_class("DerivedSIDOR")
class DerivedSIDORData(ZeroOrderBaseData):
    def build(self):
        super().build()

        self._tech_type = "test_sidor_data"

        build_sido_reactive(self)


class TestSIDOR:
    @pytest.fixture(scope="module")
    def model(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        m.fs.unit = DerivedSIDOR(
            default={"property_package": m.fs.water_props,
                     "database": m.db})

        m.fs.unit.inlet.flow_mass_comp[0, "H2O"].fix(1000)
        m.fs.unit.inlet.flow_mass_comp[0, "A"].fix(10)
        m.fs.unit.inlet.flow_mass_comp[0, "B"].fix(20)
        m.fs.unit.inlet.flow_mass_comp[0, "C"].fix(30)

        m.fs.unit.load_parameters_from_database(use_default_removal=True)

        return m

    @pytest.mark.unit
    def test_private_attributes(self, model):
        assert model.fs.unit._has_recovery_removal is True
        assert model.fs.unit._fixed_perf_vars == []
        assert model.fs.unit._initialize is initialize_sidor
        assert model.fs.unit._scaling is calculate_scaling_factors_sidor
        assert model.fs.unit._get_Q is _get_Q_sidor
        assert model.fs.unit._stream_table_dict == {
            "Inlet": model.fs.unit.inlet,
            "Treated": model.fs.unit.treated,
            "Byproduct": model.fs.unit.byproduct}
        assert model.fs.unit._perf_var_dict == {
            "Water Recovery": model.fs.unit.recovery_frac_mass_H2O,
            "Solute Removal": model.fs.unit.removal_frac_mass_solute,
            "Reaction Extent": model.fs.unit.extent_of_reaction}

    @pytest.mark.unit
    def test_build(self, model):
        assert isinstance(model.fs.unit.properties_in, WaterStateBlock)
        assert isinstance(model.fs.unit.properties_treated, WaterStateBlock)
        assert isinstance(model.fs.unit.properties_byproduct, WaterStateBlock)

        assert isinstance(model.fs.unit.inlet, Port)
        assert isinstance(model.fs.unit.treated, Port)
        assert isinstance(model.fs.unit.byproduct, Port)

        assert isinstance(model.fs.unit.recovery_frac_mass_H2O, Var)
        assert len(model.fs.unit.recovery_frac_mass_H2O) == 1
        assert isinstance(model.fs.unit.removal_frac_mass_solute, Var)
        assert len(model.fs.unit.removal_frac_mass_solute) == 3

        assert isinstance(model.fs.unit.water_recovery_equation, Constraint)
        assert len(model.fs.unit.water_recovery_equation) == 1
        assert isinstance(model.fs.unit.water_balance, Constraint)
        assert len(model.fs.unit.water_balance) == 1
        assert isinstance(model.fs.unit.solute_removal_equation, Constraint)
        assert len(model.fs.unit.solute_removal_equation) == 3
        assert isinstance(model.fs.unit.solute_treated_equation, Constraint)
        assert len(model.fs.unit.solute_treated_equation) == 3

        assert isinstance(model.fs.unit.reaction_set, Set)
        assert isinstance(model.fs.unit.generation_ratio, Param)

        assert isinstance(model.fs.unit.reaction_conversion, Var)
        assert isinstance(model.fs.unit.extent_of_reaction, Var)

        assert isinstance(model.fs.unit.reaction_extent_equation, Constraint)

        for r in model.fs.unit.reaction_set:
            assert r in ["Rxn1", "Rxn2"]
            assert (0, r) in model.fs.unit.reaction_conversion
            assert (0, r) in model.fs.unit.extent_of_reaction
            assert (0, r) in model.fs.unit.reaction_extent_equation
            for j in model.fs.water_props.component_list:
                assert (r, j) in model.fs.unit.generation_ratio

    @pytest.mark.unit
    def test_key_components(self, model):
        assert (str(model.fs.unit.properties_in[0].flow_mass_comp["A"]) in
                str(model.fs.unit.reaction_extent_equation[0, "Rxn1"].body))
        assert (str(model.fs.unit.properties_in[0].flow_mass_comp["B"]) in
                str(model.fs.unit.reaction_extent_equation[0, "Rxn2"].body))

    @pytest.mark.unit
    def test_loading_data(self, model):
        # RXn 1 is all conversion ratios
        # Rxn 2 is all stocihiometry (incl. water)
        cfactor = {"Rxn1": {"A": -1, "B": 1, "C": 0, "H2O": 0},
                   "Rxn2": {"A": 0, "B": -1, "C": 22*2/20, "H2O": -1*18/20}}

        for (r, j), p in model.fs.unit.generation_ratio.items():
            assert value(p) == cfactor[r][j]

        assert model.fs.unit.reaction_conversion[0, "Rxn1"].value == 0.8
        assert model.fs.unit.reaction_conversion[0, "Rxn2"].value == 0.1

        assert model.fs.unit.recovery_frac_mass_H2O[0].value == 0.85
        assert model.fs.unit.removal_frac_mass_solute[0, "A"].value == 0.5
        assert model.fs.unit.removal_frac_mass_solute[0, "B"].value == 0.4
        assert model.fs.unit.removal_frac_mass_solute[0, "C"].value == 0.0

    @pytest.mark.unit
    def test_degrees_of_freedom(self, model):
        assert degrees_of_freedom(model) == 0

    @pytest.mark.component
    def test_unit_consistency(self, model):
        assert_units_consistent(model)

    @pytest.mark.component
    def test_scaling(self, model):
        iscale.calculate_scaling_factors(model)

        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.water_recovery_equation[0]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.water_balance[0]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.solute_removal_equation[0, "A"]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.solute_removal_equation[0, "B"]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.solute_removal_equation[0, "C"]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.solute_treated_equation[0, "A"]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.solute_treated_equation[0, "B"]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.solute_treated_equation[0, "C"]) == 1e5

        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.reaction_extent_equation[0, "Rxn1"]) == 1e5
        assert iscale.get_constraint_transform_applied_scaling_factor(
            model.fs.unit.reaction_extent_equation[0, "Rxn2"]) == 1e5

    @pytest.mark.component
    def test_initialization(self, model):
        initialization_tester(model)

    @pytest.mark.component
    def test_solve(self, model):
        results = solver.solve(model)

        # Check for optimal solution
        assert check_optimal_termination(results)

    @pytest.mark.component
    def test_solution(self, model):
        model.fs.unit.treated.display()
        model.fs.unit.byproduct.display()
        assert (pytest.approx((1000-0.1*20*18/20)*0.85, rel=1e-5) ==
                value(model.fs.unit.treated.flow_mass_comp[0, "H2O"]))
        assert (pytest.approx((1000-0.1*20*18/20)*0.15, rel=1e-5) ==
                value(model.fs.unit.byproduct.flow_mass_comp[0, "H2O"]))

        assert (pytest.approx((10*0.2)*(1-0.5), rel=1e-5) ==
                value(model.fs.unit.treated.flow_mass_comp[0, "A"]))
        assert (pytest.approx((10*0.2)*0.5, rel=1e-5) ==
                value(model.fs.unit.byproduct.flow_mass_comp[0, "A"]))
        assert (pytest.approx((20+0.8*10-0.1*20)*(1-0.4), rel=1e-5) ==
                value(model.fs.unit.treated.flow_mass_comp[0, "B"]))
        assert (pytest.approx((20+0.8*10-0.1*20)*0.4, rel=1e-5) ==
                value(model.fs.unit.byproduct.flow_mass_comp[0, "B"]))
        assert (pytest.approx(30+0.1*20*2*22/20, rel=1e-5) ==
                value(model.fs.unit.treated.flow_mass_comp[0, "C"]))
        assert (pytest.approx(0, rel=1e-5) ==
                value(model.fs.unit.byproduct.flow_mass_comp[0, "C"]))

    # Conservation would be similar to above, as need to calculate generation

    @pytest.mark.component
    def test_report(self, model):
        stream = StringIO()
        model.fs.unit.report(ostream=stream)

        output = """
====================================================================================
Unit : fs.unit                                                             Time: 0.0
------------------------------------------------------------------------------------
    Unit Performance

    Variables: 

    Key                    : Value   : Fixed : Bounds
    Reaction Extent [Rxn1] :  8.0000 : False : (None, None)
    Reaction Extent [Rxn2] :  2.0000 : False : (None, None)
        Solute Removal [A] : 0.50000 :  True : (0, None)
        Solute Removal [B] : 0.40000 :  True : (0, None)
        Solute Removal [C] :  0.0000 :  True : (0, None)
            Water Recovery : 0.85000 :  True : (1e-08, 1.0000001)

------------------------------------------------------------------------------------
    Stream Table
                            Inlet  Treated  Byproduct
    Volumetric Flowrate    1.0600 0.89947     0.16113
    Mass Concentration H2O 943.40  943.30      929.25
    Mass Concentration A   9.4340  1.1118      6.2062
    Mass Concentration B   18.868  17.344      64.544
    Mass Concentration C   28.302  38.245  6.2062e-13
====================================================================================
"""

        assert output == stream.getvalue()


class TestSIDORErrors:
    @pytest.mark.unit
    def test_reaction_list(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files["test_sidor_data"]["default"] = {}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not contain a list of "
                "reactions for this technology."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_missing_conversion(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not "
                "contain an entry for conversion for reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_missing_key_reactant(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not "
                "contain an entry for key_reactant for reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_invlaid_key_reactant(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "foo"

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                ValueError,
                match="fs.unit - key_reactant foo for reaction Rxn3 "
                "is not in the component list used by the assigned property "
                "package."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_missing_stoichiometry(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "A"

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not "
                "contain an entry for stoichiometry for reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_ratio_and_order(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "A"
        R3["stoichiometry"] = {"A": {"order": 1, "conversion_ratio": 1}}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                RuntimeError,
                match="fs.unit - database provides entries for both "
                "conversion_ratio and reaction order in reaction Rxn3. "
                "Please provide only one or the other."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_no_ratio_or_order(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "A"
        R3["stoichiometry"] = {"A": {}}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                RuntimeError,
                match="fs.unit - database provided does not "
                "contain any information for conversion_ratio or reaction "
                "order w.r.t. species A in reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_order_no_mw(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "B"
        R3["stoichiometry"] = {"A": {"order": 1},
                               "B": {"order": 1, "molecular_weight": 1}}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not "
                "contain an entry for molecular_weight w.r.t. "
                "species A in reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_order_no_key_order(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "B"
        R3["stoichiometry"] = {"A": {"order": 1, "molecular_weight": 1}}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not "
                "contain an entry for order w.r.t. species "
                "B in reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})

    @pytest.mark.unit
    def test_order_no_key_mw(self):
        m = ConcreteModel()

        m.db = Database(dbpath=local_path)

        # Load data from YAML so that we can modify it
        m.db._get_technology("test_sidor_data")
        m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"] = {}
        R3 = m.db._cached_files[
            "test_sidor_data"]["default"]["reactions"]["Rxn3"]
        R3["conversion"] = 0.5
        R3["key_reactant"] = "B"
        R3["stoichiometry"] = {"A": {"order": 1, "molecular_weight": 1},
                               "B": {"order": 1}}

        m.fs = FlowsheetBlock(default={"dynamic": False})

        m.fs.water_props = WaterParameterBlock(
            default={"solute_list": ["A", "B", "C"]})

        with pytest.raises(
                KeyError,
                match="fs.unit - database provided does not "
                "contain an entry for molecular_weight w.r.t. "
                "species B in reaction Rxn3."):
            m.fs.unit = DerivedSIDOR(
                default={"property_package": m.fs.water_props,
                         "database": m.db})
