text: ###############################################################################
  # 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/"
  #
  ###############################################################################
  
  """
      This test is to establish that the core chemistry packages in IDAES solve
      a simple water dissociation problem and return the correct pH value.
  """
  import enum
  
  # Importing testing libraries
  import pytest
  
  # Importing the object for units from pyomo
  from pyomo.environ import units as pyunits
  
  # Imports from idaes core
  from idaes.core import AqueousPhase
  from idaes.core.components import Solvent, Solute, Cation, Anion
  from idaes.core.phases import PhaseType as PT
  
  # Imports from idaes generic models
  import idaes.generic_models.properties.core.pure.Perrys as Perrys
  from idaes.generic_models.properties.core.state_definitions import FTPx
  from idaes.generic_models.properties.core.eos.ideal import Ideal
  
  # Importing the enum for concentration unit basis used in the 'get_concentration_term' function
  from idaes.generic_models.properties.core.generic.generic_reaction import (
      ConcentrationForm,
  )
  
  # Import the object/function for heat of reaction
  from idaes.generic_models.properties.core.reactions.dh_rxn import constant_dh_rxn
  
  # Import safe log power law equation
  from idaes.generic_models.properties.core.reactions.equilibrium_forms import (
      log_power_law_equil,
  )
  
  # Import k-value functions
  from idaes.generic_models.properties.core.reactions.equilibrium_constant import (
      gibbs_energy,
      van_t_hoff,
  )
  
  # Import specific pyomo objects
  from pyomo.environ import (
      ConcreteModel,
      SolverStatus,
      TerminationCondition,
      value,
      Suffix,
  )
  
  from idaes.core.util import scaling as iscale
  
  # Import pyomo methods to check the system units
  from pyomo.util.check_units import assert_units_consistent
  
  # Import idaes methods to check the model during construction
  from idaes.core.util import get_solver
  from idaes.core.util.model_statistics import (
      degrees_of_freedom,
      fixed_variables_set,
      activated_constraints_set,
      number_variables,
      number_total_constraints,
      number_unused_variables,
  )
  
  # Import the idaes objects for Generic Properties and Reactions
  from idaes.generic_models.properties.core.generic.generic_property import (
      GenericParameterBlock,
  )
  from idaes.generic_models.properties.core.generic.generic_reaction import (
      GenericReactionParameterBlock,
  )
  
  # Import the idaes object for the EquilibriumReactor unit model
  from idaes.generic_models.unit_models.equilibrium_reactor import EquilibriumReactor
  
  # Import the core idaes objects for Flowsheets and types of balances
  from idaes.core import FlowsheetBlock
  
  # Import log10 function from pyomo
  from pyomo.environ import log10
  
  
  __authors__ = [
      "Austin Ladshaw",
      "Ludovico Bianchi",
      "Dan Gunter",
  ]
  __author__ = __authors__[0]
  
  
  class Variant(str, enum.Enum):
      equilibrium = "equilibrium"
      inherent = "inherent"
  
      def __str__(self):
          return f'{self.__class__.__name__}.{self.name}'
  
      @property
      def is_equilibrium(self):
          return self is Variant.equilibrium
  
  
  @pytest.fixture(
      scope="module",
      params=[
          Variant.equilibrium,
          Variant.inherent
      ],
      ids=str
  )
  def variant(request) -> Variant:
      return request.param
  
  
  @pytest.fixture(scope="module")
  def base_units():
      return {
          "time": pyunits.s,
          "length": pyunits.m,
          "mass": pyunits.kg,
          "amount": pyunits.mol,
          "temperature": pyunits.K,
      }
  
  
  @pytest.fixture(scope="module")
  def solver():
      s = get_solver()
      s.options["max_iter"] = 200
      return s
  
  
  @pytest.fixture(scope="module")
  def thermo_config(base_units):
      # Configuration dictionary
      data = {
          "components": {
              "H2O": {
                  "type": Solvent,
                  # Define the methods used to calculate the following properties
                  "dens_mol_liq_comp": Perrys,
                  "enth_mol_liq_comp": Perrys,
                  "cp_mol_liq_comp": Perrys,
                  "entr_mol_liq_comp": Perrys,
                  # Parameter data is always associated with the methods defined above
                  "parameter_data": {
                      "mw": (18.0153, pyunits.g / pyunits.mol),
                      "pressure_crit": (220.64e5, pyunits.Pa),
                      "temperature_crit": (647, pyunits.K),
                      # Comes from Perry's Handbook:  p. 2-98
                      "dens_mol_liq_comp_coeff": {
                          "1": (5.459, pyunits.kmol * pyunits.m ** -3),
                          "2": (0.30542, pyunits.dimensionless),
                          "3": (647.13, pyunits.K),
                          "4": (0.081, pyunits.dimensionless),
                      },
                      "enth_mol_form_liq_comp_ref": (-285.830, pyunits.kJ / pyunits.mol),
                      "enth_mol_form_vap_comp_ref": (0, pyunits.kJ / pyunits.mol),
                      # Comes from Perry's Handbook:  p. 2-174
                      "cp_mol_liq_comp_coeff": {
                          "1": (2.7637e5, pyunits.J / pyunits.kmol / pyunits.K),
                          "2": (-2.0901e3, pyunits.J / pyunits.kmol / pyunits.K ** 2),
                          "3": (8.125, pyunits.J / pyunits.kmol / pyunits.K ** 3),
                          "4": (-1.4116e-2, pyunits.J / pyunits.kmol / pyunits.K ** 4),
                          "5": (9.3701e-6, pyunits.J / pyunits.kmol / pyunits.K ** 5),
                      },
                      "cp_mol_ig_comp_coeff": {
                          "A": (30.09200, pyunits.J / pyunits.mol / pyunits.K),
                          "B": (
                              6.832514,
                              pyunits.J
                              * pyunits.mol ** -1
                              * pyunits.K ** -1
                              * pyunits.kiloK ** -1,
                          ),
                          "C": (
                              6.793435,
                              pyunits.J
                              * pyunits.mol ** -1
                              * pyunits.K ** -1
                              * pyunits.kiloK ** -2,
                          ),
                          "D": (
                              -2.534480,
                              pyunits.J
                              * pyunits.mol ** -1
                              * pyunits.K ** -1
                              * pyunits.kiloK ** -3,
                          ),
                          "E": (
                              0.082139,
                              pyunits.J
                              * pyunits.mol ** -1
                              * pyunits.K ** -1
                              * pyunits.kiloK ** 2,
                          ),
                          "F": (-250.8810, pyunits.kJ / pyunits.mol),
                          "G": (223.3967, pyunits.J / pyunits.mol / pyunits.K),
                          "H": (0, pyunits.kJ / pyunits.mol),
                      },
                      "entr_mol_form_liq_comp_ref": (
                          69.95,
                          pyunits.J / pyunits.K / pyunits.mol,
                      ),
                      "pressure_sat_comp_coeff": {
                          "A": (4.6543, None),  # [1], temperature range 255.9 K - 373 K
                          "B": (1435.264, pyunits.K),
                          "C": (-64.848, pyunits.K),
                      },
                  },
                  # End parameter_data
              },
              "H_+": {
                  "type": Cation,
                  "charge": 1,
                  # Define the methods used to calculate the following properties
                  "dens_mol_liq_comp": Perrys,
                  "enth_mol_liq_comp": Perrys,
                  "cp_mol_liq_comp": Perrys,
                  "entr_mol_liq_comp": Perrys,
                  # Parameter data is always associated with the methods defined above
                  "parameter_data": {
                      "mw": (1.00784, pyunits.g / pyunits.mol),
                      "dens_mol_liq_comp_coeff": {
                          "1": (5.459, pyunits.kmol * pyunits.m ** -3),
                          "2": (0.30542, pyunits.dimensionless),
                          "3": (647.13, pyunits.K),
                          "4": (0.081, pyunits.dimensionless),
                      },
                      "enth_mol_form_liq_comp_ref": (-230.000, pyunits.kJ / pyunits.mol),
                      "cp_mol_liq_comp_coeff": {
                          "1": (2.7637e5, pyunits.J / pyunits.kmol / pyunits.K),
                          "2": (-2.0901e3, pyunits.J / pyunits.kmol / pyunits.K ** 2),
                          "3": (8.125, pyunits.J / pyunits.kmol / pyunits.K ** 3),
                          "4": (-1.4116e-2, pyunits.J / pyunits.kmol / pyunits.K ** 4),
                          "5": (9.3701e-6, pyunits.J / pyunits.kmol / pyunits.K ** 5),
                      },
                      "entr_mol_form_liq_comp_ref": (
                          -10.75,
                          pyunits.J / pyunits.K / pyunits.mol,
                      ),
                  },
                  # End parameter_data
              },
              "OH_-": {
                  "type": Anion,
                  "charge": -1,
                  # Define the methods used to calculate the following properties
                  "dens_mol_liq_comp": Perrys,
                  "enth_mol_liq_comp": Perrys,
                  "cp_mol_liq_comp": Perrys,
                  "entr_mol_liq_comp": Perrys,
                  # Parameter data is always associated with the methods defined above
                  "parameter_data": {
                      "mw": (17.008, pyunits.g / pyunits.mol),
                      "dens_mol_liq_comp_coeff": {
                          "1": (5.459, pyunits.kmol * pyunits.m ** -3),
                          "2": (0.30542, pyunits.dimensionless),
                          "3": (647.13, pyunits.K),
                          "4": (0.081, pyunits.dimensionless),
                      },
                      "enth_mol_form_liq_comp_ref": (-230.000, pyunits.kJ / pyunits.mol),
                      "cp_mol_liq_comp_coeff": {
                          "1": (2.7637e5, pyunits.J / pyunits.kmol / pyunits.K),
                          "2": (-2.0901e3, pyunits.J / pyunits.kmol / pyunits.K ** 2),
                          "3": (8.125, pyunits.J / pyunits.kmol / pyunits.K ** 3),
                          "4": (-1.4116e-2, pyunits.J / pyunits.kmol / pyunits.K ** 4),
                          "5": (9.3701e-6, pyunits.J / pyunits.kmol / pyunits.K ** 5),
                      },
                      "entr_mol_form_liq_comp_ref": (
                          -10.75,
                          pyunits.J / pyunits.K / pyunits.mol,
                      ),
                  },
                  # End parameter_data
              },
          },
          # End Component list
          "phases": {
              "Liq": {"type": AqueousPhase, "equation_of_state": Ideal},
          },
          "state_definition": FTPx,
          "state_bounds": {
              "flow_mol": (0, 50, 100),
              "temperature": (273.15, 300, 650),
              "pressure": (5e4, 1e5, 1e6),
          },
          "pressure_ref": 1e5,
          "temperature_ref": 300,
          "base_units": base_units,
          "inherent_reactions": {
              "H2O_Kw": {
                  "stoichiometry": {
                      ("Liq", "H2O"): -1,
                      ("Liq", "H_+"): 1,
                      ("Liq", "OH_-"): 1,
                  },
                  "heat_of_reaction": constant_dh_rxn,
                  "equilibrium_constant": van_t_hoff,
                  "equilibrium_form": log_power_law_equil,
                  "concentration_form": ConcentrationForm.moleFraction,
                  "parameter_data": {
                      "dh_rxn_ref": (55.830, pyunits.J / pyunits.mol),
                      "k_eq_ref": (10 ** -14 / 55.2 / 55.2, pyunits.dimensionless),
                      "T_eq_ref": (298, pyunits.K),
                      # By default, reaction orders follow stoichiometry
                      #    manually set reaction order here to override
                      "reaction_order": {
                          ("Liq", "H2O"): 0,
                          ("Liq", "H_+"): 1,
                          ("Liq", "OH_-"): 1,
                      },
                  }
                  # End parameter_data
              }
              # End R1
          }
      }
      return data
  
  
  # Define the reaction_config for water dissociation
  @pytest.fixture(scope="module")
  def water_reaction_config(base_units):
      return {
          "base_units": base_units,
          "equilibrium_reactions": {
              "H2O_Kw": {
                  "stoichiometry": {
                      ("Liq", "H2O"): -1,
                      ("Liq", "H_+"): 1,
                      ("Liq", "OH_-"): 1,
                  },
                  "heat_of_reaction": constant_dh_rxn,
                  "equilibrium_constant": van_t_hoff,
                  "equilibrium_form": log_power_law_equil,
                  "concentration_form": ConcentrationForm.moleFraction,
                  "parameter_data": {
                      "dh_rxn_ref": (55.830, pyunits.J / pyunits.mol),
                      "k_eq_ref": (10 ** -14 / 55.2 / 55.2, pyunits.dimensionless),
                      "T_eq_ref": (298, pyunits.K),
                      # By default, reaction orders follow stoichiometry
                      #    manually set reaction order here to override
                      "reaction_order": {
                          ("Liq", "H2O"): 0,
                          ("Liq", "H_+"): 1,
                          ("Liq", "OH_-"): 1,
                      },
                  }
                  # End parameter_data
              }
              # End R1
          }
          # End equilibrium_reactions
      }
      # End reaction_config definition
  
  def _get_without_inherent_reactions(config_dict: dict, key="inherent_reactions") -> dict:
      data = dict(config_dict)
      del data[key]
      return data
  
  class TestPureWater:
  
  
      @pytest.fixture
      def model(self, thermo_config, variant: Variant, water_reaction_config):
          if variant.is_equilibrium:
              thermo_config = _get_without_inherent_reactions(thermo_config)
          model = ConcreteModel()
          model.fs = FlowsheetBlock(default={"dynamic": False})
          model.fs.thermo_params = GenericParameterBlock(default=thermo_config)
          print(water_reaction_config)
          model.fs.rxn_params = GenericReactionParameterBlock(
              default={"property_package": model.fs.thermo_params, **water_reaction_config}
          )
          model.fs.unit = EquilibriumReactor(
              default={
                  "property_package": model.fs.thermo_params,
                  "reaction_package": model.fs.rxn_params,
                  "has_rate_reactions": False,
                  "has_equilibrium_reactions": variant.is_equilibrium,
                  "has_heat_transfer": False,
                  "has_heat_of_reaction": False,
                  "has_pressure_change": False,
              }
          )
  
          model.fs.unit.inlet.mole_frac_comp[0, "H_+"].fix(0.0)
          model.fs.unit.inlet.mole_frac_comp[0, "OH_-"].fix(0.0)
          model.fs.unit.inlet.mole_frac_comp[0, "H2O"].fix(1.0)
          model.fs.unit.inlet.pressure.fix(101325.0)
          model.fs.unit.inlet.temperature.fix(298.0)
          model.fs.unit.inlet.flow_mol.fix(10)
  
          return model
  
      @pytest.mark.unit
      def test_build_model(self, model):
  
          assert hasattr(model.fs.thermo_params, "component_list")
          assert len(model.fs.thermo_params.component_list) == 3
          assert "H2O" in model.fs.thermo_params.component_list
          assert isinstance(model.fs.thermo_params.H2O, Solvent)
          assert "H_+" in model.fs.thermo_params.component_list
          assert isinstance(model.fs.thermo_params.component("H_+"), Cation)
          assert "OH_-" in model.fs.thermo_params.component_list
          assert isinstance(model.fs.thermo_params.component("OH_-"), Anion)
  
          assert hasattr(model.fs.thermo_params, "phase_list")
          assert len(model.fs.thermo_params.phase_list) == 1
          assert isinstance(model.fs.thermo_params.Liq, AqueousPhase)
  
      @pytest.mark.unit
      def test_units_inherent(self, model):
          assert_units_consistent(model)
  
      @pytest.mark.unit
      def test_dof(self, model):
          assert degrees_of_freedom(model) == 0
  
      @pytest.fixture
      def expected_stats(self, variant):
          expected = {
              Variant.inherent: {
                  number_variables: 81,
                  number_total_constraints: 28,
                  number_unused_variables: 12,
              },
              Variant.equilibrium: {
                  number_variables: 75,
                  number_total_constraints: 28,
                  number_unused_variables: 6
              }
          }
          return expected[variant]
  
      @pytest.mark.unit
      def test_stats(self, model, expected_stats):
          for check_func, expected_result in expected_stats.items():
              check_result = check_func(model)
              assert check_result == expected_result
  
      def _do_inherent_reaction_scaling(self, model):
          for i in model.fs.unit.control_volume.inherent_reaction_extent_index:
              scale = value(
                  model.fs.unit.control_volume.properties_out[0.0].k_eq[i[1]].expr
              )
              iscale.set_scaling_factor(
                  model.fs.unit.control_volume.inherent_reaction_extent[0.0, i[1]],
                  10 / scale,
              )
              iscale.constraint_scaling_transform(
                  model.fs.unit.control_volume.properties_out[
                      0.0
                  ].inherent_equilibrium_constraint[i[1]],
                  0.1,
              )
          return model
  
      def _do_equilibrium_reaction_scaling(self, model):
          for i in model.fs.unit.control_volume.equilibrium_reaction_extent_index:
              scale = value(model.fs.unit.control_volume.reactions[0.0].k_eq[i[1]].expr)
              iscale.set_scaling_factor(
                  model.fs.unit.control_volume.equilibrium_reaction_extent[0.0, i[1]],
                  10 / scale,
              )
              iscale.constraint_scaling_transform(
                  model.fs.unit.control_volume.reactions[0.0].equilibrium_constraint[
                      i[1]
                  ],
                  0.1,
              )
          return model
  
      @pytest.fixture
      def model_scaling(self, model, variant: Variant):
          map_variant_reaction_scale_func = {
              Variant.equilibrium: self._do_equilibrium_reaction_scaling,
              Variant.inherent: self._do_inherent_reaction_scaling,
          }
  
          scale_func = map_variant_reaction_scale_func[variant]
          scale_func(model)
  
          # Next, try adding scaling for species
          min = 1e-3
          for i in model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp:
              # i[0] = phase, i[1] = species
              if model.fs.unit.inlet.mole_frac_comp[0, i[1]].value > min:
                  scale = model.fs.unit.inlet.mole_frac_comp[0, i[1]].value
              else:
                  scale = min
              iscale.set_scaling_factor(
                  model.fs.unit.control_volume.properties_out[0.0].mole_frac_comp[i[1]],
                  10 / scale,
              )
              iscale.set_scaling_factor(
                  model.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp[
                      i
                  ],
                  10 / scale,
              )
              iscale.set_scaling_factor(
                  model.fs.unit.control_volume.properties_out[0.0].flow_mol_phase_comp[i],
                  10 / scale,
              )
              iscale.constraint_scaling_transform(
                  model.fs.unit.control_volume.properties_out[
                      0.0
                  ].component_flow_balances[i[1]],
                  10 / scale,
              )
              iscale.constraint_scaling_transform(
                  model.fs.unit.control_volume.material_balances[0.0, i[1]], 10 / scale
              )
  
          iscale.calculate_scaling_factors(model.fs.unit)
          return model
  
      @pytest.mark.component
      def test_scaling(self, model_scaling, variant: Variant):
          model = model_scaling
  
          assert isinstance(model.fs.unit.control_volume.scaling_factor, Suffix)
  
          assert isinstance(
              model.fs.unit.control_volume.properties_out[0.0].scaling_factor, Suffix
          )
  
          assert isinstance(
              model.fs.unit.control_volume.properties_in[0.0].scaling_factor, Suffix
          )
  
          if variant.is_equilibrium:
              # When using equilibrium reactions, there is another set of scaling factors calculated
              assert isinstance(
                  model.fs.unit.control_volume.reactions[0.0].scaling_factor, Suffix
              )
  
      @pytest.fixture
      def state_args(self):
          return {
              "mole_frac_comp": {
                  "H2O": 1,
                  "H_+": 10 ** -7 / 55.6,
                  "OH_-": 10 ** -7 / 55.6,
              },
              "pressure": 101325,
              "temperature": 298,
              "flow_mol": 10,
          }
  
      @pytest.fixture
      def model_initialize(self, model, state_args, solver):
          def _collect_data_to_check(m):
              return {
                  fixed_variables_set: fixed_variables_set(m),
                  activated_constraints_set: activated_constraints_set(m)
              }
          data_before = _collect_data_to_check(model)
          model.fs.unit.initialize(state_args=state_args, optarg=solver.options)
  
          data_after = _collect_data_to_check(model)
          return model, data_before, data_after
  
      @pytest.mark.component
      def test_initialize(self, model_initialize):
          model, data_before, data_after = model_initialize
  
          assert degrees_of_freedom(model) == 0
  
          for key, value_before in data_before.items():
              value_after = data_after[key]
              assert len(value_before) == len(value_after)
  
      @pytest.fixture
      def model_solve(self, model, solver):
          results = solver.solve(model, tee=True)
          return model, results
  
      @pytest.mark.component
      def test_solve(self, model_solve):
          _, results = model_solve
          assert results.solver.termination_condition == TerminationCondition.optimal
          assert results.solver.status == SolverStatus.ok
  
      @pytest.mark.component
      def test_solution(self, model_solve):
          model, _ = model_solve
  
          assert pytest.approx(298, rel=1e-5) == value(
              model.fs.unit.outlet.temperature[0]
          )
          assert pytest.approx(10, rel=1e-5) == value(model.fs.unit.outlet.flow_mol[0])
          assert pytest.approx(101325, rel=1e-5) == value(
              model.fs.unit.outlet.pressure[0]
          )
  
          total_molar_density = (
              value(
                  model.fs.unit.control_volume.properties_out[0.0].dens_mol_phase["Liq"]
              )
              / 1000
          )
          assert pytest.approx(55.2336, rel=1e-5) == total_molar_density
  
          pH = -value(
              log10(model.fs.unit.outlet.mole_frac_comp[0, "H_+"] * total_molar_density)
          )
          pOH = -value(
              log10(model.fs.unit.outlet.mole_frac_comp[0, "OH_-"] * total_molar_density)
          )
          assert pytest.approx(6.9997414, rel=1e-5) == pH
          assert pytest.approx(6.9997414, rel=1e-5) == pOH
          assert pytest.approx(0.99999, rel=1e-5) == value(
              model.fs.unit.outlet.mole_frac_comp[0.0, "H2O"]
          )
  
  
  class TestPureWaterEDB(TestPureWater):
  
      @pytest.fixture
      def thermo_config(self, edb):
          base = edb.get_base("default_thermo")
          elements = ["H", "O"]
          components = []
          # Add the components
          for c in edb.get_components(element_names=elements):
              # Should remove these since we don't have a vapor phase
              # NOTE: The model still runs without removing these, but
              #       the 'stats' won't pass their checks due to the addition
              #       of new system variables that they inherently bring
              c.remove("enth_mol_ig_comp")
              c.remove("pressure_sat_comp")
              base.add(c)
              components.append(c.name)
          # Add the reactions
          for r in edb.get_reactions(component_names=components):
              # Set a custom reaction order
              r.set_reaction_order("Liq", {"H2O": 0, "H_+": 1, "OH_-": 1})
              r._data["type"] = "inherent"
              base.add(r)
          return base.idaes_config
  
  
      @pytest.fixture
      def water_reaction_config(self, edb):
          elements = ["H", "O"]
          components = [c.name for c in edb.get_components(element_names=elements)]
          base = edb.get_base("reaction")
          # Add the reactions
          for r in edb.get_reactions(component_names=components):
              # Set a custom reaction order
              r.set_reaction_order("Liq", {"H2O": 0, "H_+": 1, "OH_-": 1})
              # Need to remove this to avoid errors when using the generated config
              r.remove_parameter("ds_rxn_ref")
              base.add(r)
          return base.idaes_config
  