#################################################################################
# WaterTAP Copyright (c) 2020-2023, 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 to check reactions of water softening by adding lime (CaOH2) to calcium
   bicarbonate and magnesium bicarbonate with stoichiometric reactor and return
   correct hardness value.

   Kinetic Reactions:
        Ca(HCO3)2 +   Ca(OH)2 --> 2 CaCO3 +  2 H2O
        Mg(HCO3)2 + 2 Ca(OH)2 --> 2 CaCO3 + Mg(OH)2 +  2 H2O
"""
import pytest

# Importing the object for units from pyomo
from pyomo.environ import units as pyunits

from pyomo.environ import (
    ConcreteModel,
    TerminationCondition,
    SolverStatus,
    value,
    Suffix,
)

from pyomo.util.check_units import assert_units_consistent

from idaes.core import AqueousPhase, FlowsheetBlock, EnergyBalanceType

import idaes.logger as idaeslog

from idaes.models.unit_models.stoichiometric_reactor import (
    StoichiometricReactor,
)

from idaes.core.util import scaling as iscale

from idaes.core.solvers 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,
)

from idaes.core.base.components import Solvent, Solute
from idaes.core.base.phases import PhaseType as PT

# Import the idaes objects for Generic Properties and Reactions
from idaes.models.properties.modular_properties.base.generic_property import (
    GenericParameterBlock,
)
from idaes.models.properties.modular_properties.base.generic_reaction import (
    GenericReactionParameterBlock,
)

# Imports from idaes generic models
import idaes.models.properties.modular_properties.pure.Perrys as Perrys
from idaes.models.properties.modular_properties.pure.ConstantProperties import Constant
from idaes.models.properties.modular_properties.state_definitions import FTPx
from idaes.models.properties.modular_properties.eos.ideal import Ideal
from idaes.models.properties.modular_properties.reactions.rate_constant import arrhenius
from idaes.models.properties.modular_properties.reactions.rate_forms import (
    power_law_rate,
)

# Importing the enum for concentration unit basis used in the 'get_concentration_term' function
from idaes.models.properties.modular_properties.base.generic_reaction import (
    ConcentrationForm,
)

# Import the object/function for heat of reaction
from idaes.models.properties.modular_properties.reactions.dh_rxn import constant_dh_rxn

# Import safe log power law equation
from idaes.models.properties.modular_properties.reactions.equilibrium_forms import (
    log_power_law_equil,
)

__author__ = "Srikanth Allu, Austin Ladshaw"

thermo_config = {
    "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": {
                    "eqn_type": 1,
                    "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
        },
        "Ca(OH)2": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            # Define the methods used to calculate the following properties
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            # Parameter data is always associated with the methods defined above
            "parameter_data": {
                "mw": (74.093, pyunits.g / pyunits.mol),
                "dens_mol_liq_comp_coeff": (55, pyunits.kmol * pyunits.m**-3),
                "enth_mol_form_liq_comp_ref": (-945.53, pyunits.kJ / pyunits.mol),
                "cp_mol_liq_comp_coeff": (167039, pyunits.J / pyunits.kmol / pyunits.K),
                "entr_mol_form_liq_comp_ref": (
                    100,
                    pyunits.J / pyunits.K / pyunits.mol,
                ),
            },
            # End parameter_data
        },
        "CaCO3": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            # Define the methods used to calculate the following properties
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            # Parameter data is always associated with the methods defined above
            "parameter_data": {
                "dens_mol_liq_comp_coeff": (55, pyunits.kmol * pyunits.m**-3),
                "enth_mol_form_liq_comp_ref": (-945.53, pyunits.kJ / pyunits.mol),
                "cp_mol_liq_comp_coeff": (167039, pyunits.J / pyunits.kmol / pyunits.K),
                "entr_mol_form_liq_comp_ref": (
                    100,
                    pyunits.J / pyunits.K / pyunits.mol,
                ),
            },
            # End parameter_data
        },
        "Ca(HCO3)2": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            # Define the methods used to calculate the following properties
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            # Parameter data is always associated with the methods defined above
            "parameter_data": {
                "dens_mol_liq_comp_coeff": (55, pyunits.kmol * pyunits.m**-3),
                "enth_mol_form_liq_comp_ref": (-945.53, pyunits.kJ / pyunits.mol),
                "cp_mol_liq_comp_coeff": (167039, pyunits.J / pyunits.kmol / pyunits.K),
                "entr_mol_form_liq_comp_ref": (
                    100,
                    pyunits.J / pyunits.K / pyunits.mol,
                ),
            },
            # End parameter_data
        },
        "Mg(OH)2": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            # Define the methods used to calculate the following properties
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            # Parameter data is always associated with the methods defined above
            "parameter_data": {
                "mw": (74.093, pyunits.g / pyunits.mol),
                "dens_mol_liq_comp_coeff": (55, pyunits.kmol * pyunits.m**-3),
                "enth_mol_form_liq_comp_ref": (-945.53, pyunits.kJ / pyunits.mol),
                "cp_mol_liq_comp_coeff": (167039, pyunits.J / pyunits.kmol / pyunits.K),
                "entr_mol_form_liq_comp_ref": (
                    100,
                    pyunits.J / pyunits.K / pyunits.mol,
                ),
            },
            # End parameter_data
        },
        "Mg(HCO3)2": {
            "type": Solute,
            "valid_phase_types": PT.aqueousPhase,
            # Define the methods used to calculate the following properties
            "dens_mol_liq_comp": Constant,
            "enth_mol_liq_comp": Constant,
            "cp_mol_liq_comp": Constant,
            "entr_mol_liq_comp": Constant,
            # Parameter data is always associated with the methods defined above
            "parameter_data": {
                "dens_mol_liq_comp_coeff": (55, pyunits.kmol * pyunits.m**-3),
                "enth_mol_form_liq_comp_ref": (-945.53, pyunits.kJ / pyunits.mol),
                "cp_mol_liq_comp_coeff": (167039, pyunits.J / pyunits.kmol / pyunits.K),
                "entr_mol_form_liq_comp_ref": (
                    100,
                    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": {
        "time": pyunits.s,
        "length": pyunits.m,
        "mass": pyunits.kg,
        "amount": pyunits.mol,
        "temperature": pyunits.K,
    },
}
# End thermo_config definition


# This config is REQUIRED to use EquilibriumReactor even if we have no equilibrium reactions
reaction_config = {
    "base_units": {
        "time": pyunits.s,
        "length": pyunits.m,
        "mass": pyunits.kg,
        "amount": pyunits.mol,
        "temperature": pyunits.K,
    },
    "equilibrium_reactions": {
        "dummy": {
            "stoichiometry": {},
            "equilibrium_form": log_power_law_equil,
        }
    },
    # End equilibrium_reactions
    "rate_reactions": {
        "R1": {
            "stoichiometry": {
                ("Liq", "Ca(HCO3)2"): -1,
                ("Liq", "Ca(OH)2"): -1,
                ("Liq", "CaCO3"): 2,
                ("Liq", "H2O"): 2,
            },
            "heat_of_reaction": constant_dh_rxn,
            "rate_constant": arrhenius,
            "rate_form": power_law_rate,
            "concentration_form": ConcentrationForm.moleFraction,
            "parameter_data": {
                "arrhenius_const": (1, pyunits.mol / pyunits.m**3 / pyunits.s),
                "energy_activation": (0, pyunits.J / pyunits.mol),
                "dh_rxn_ref": (0, pyunits.J / pyunits.mol),
            },
        },
        "R2": {
            "stoichiometry": {
                ("Liq", "Mg(HCO3)2"): -1,
                ("Liq", "Ca(OH)2"): -2,
                ("Liq", "CaCO3"): 2,
                ("Liq", "Mg(OH)2"): 1,
                ("Liq", "H2O"): 2,
            },
            "heat_of_reaction": constant_dh_rxn,
            "rate_constant": arrhenius,
            "rate_form": power_law_rate,
            "concentration_form": ConcentrationForm.moleFraction,
            "parameter_data": {
                "arrhenius_const": (1, pyunits.mol / pyunits.m**3 / pyunits.s),
                "energy_activation": (0, pyunits.J / pyunits.mol),
                "dh_rxn_ref": (0, pyunits.J / pyunits.mol),
            },
        },
    },
}
# End reaction_config definition


# -----------------------------------------------------------------------------
# Get default solver for testing
solver = get_solver()

# -----------------------------------------------------------------------------
class TestWaterStoich(object):
    @pytest.fixture(scope="class")
    def water_stoich(self):
        m = ConcreteModel()
        m.fs = FlowsheetBlock(dynamic=False)

        m.fs.thermo_params = GenericParameterBlock(**thermo_config)
        m.fs.rxn_params = GenericReactionParameterBlock(
            property_package=m.fs.thermo_params, **reaction_config
        )

        m.fs.unit = StoichiometricReactor(
            property_package=m.fs.thermo_params,
            reaction_package=m.fs.rxn_params,
            has_heat_transfer=False,
            has_heat_of_reaction=False,
            energy_balance_type=EnergyBalanceType.none,
            has_pressure_change=False,
        )

        m.fs.unit.inlet.mole_frac_comp[0, "Mg(HCO3)2"].fix(0.00003)
        m.fs.unit.inlet.mole_frac_comp[0, "Mg(OH)2"].fix(0.0)
        m.fs.unit.inlet.mole_frac_comp[0, "Ca(HCO3)2"].fix(0.00003)
        m.fs.unit.inlet.mole_frac_comp[0, "CaCO3"].fix(0.0)
        m.fs.unit.inlet.mole_frac_comp[0, "H2O"].fix(0.99991)

        m.fs.unit.inlet.pressure.fix(101325.0)
        m.fs.unit.inlet.temperature.fix(298.0)
        m.fs.unit.inlet.flow_mol.fix(10)

        m.fs.unit.outlet.temperature.fix(298.0)

        m.fs.unit.outlet.mole_frac_comp[0, "Ca(HCO3)2"].fix(0.000015)
        m.fs.unit.outlet.mole_frac_comp[0, "Mg(HCO3)2"].fix(0.000015)
        m.fs.unit.outlet.mole_frac_comp[0, "Ca(OH)2"].fix(0.0000003)

        return m

    @pytest.mark.build
    @pytest.mark.unit
    def test_build(self, water_stoich):

        m = water_stoich
        assert hasattr(m.fs.thermo_params, "component_list")
        assert len(m.fs.thermo_params.component_list) == 6
        assert "H2O" in m.fs.thermo_params.component_list
        assert isinstance(m.fs.thermo_params.H2O, Solvent)
        assert "Ca(HCO3)2" in m.fs.thermo_params.component_list
        assert isinstance(m.fs.thermo_params.component("Ca(HCO3)2"), Solute)
        assert "Ca(OH)2" in m.fs.thermo_params.component_list
        assert isinstance(m.fs.thermo_params.component("Ca(OH)2"), Solute)
        assert "CaCO3" in m.fs.thermo_params.component_list
        assert isinstance(m.fs.thermo_params.component("CaCO3"), Solute)
        assert "Mg(HCO3)2" in m.fs.thermo_params.component_list
        assert isinstance(m.fs.thermo_params.component("Mg(HCO3)2"), Solute)

        assert hasattr(m.fs.thermo_params, "phase_list")
        assert len(m.fs.thermo_params.phase_list) == 1
        assert isinstance(m.fs.thermo_params.Liq, AqueousPhase)

    @pytest.mark.unit
    def test_units_stoich(self, water_stoich):
        m = water_stoich
        assert_units_consistent(m)

    @pytest.mark.unit
    def test_dof_stoich(self, water_stoich):
        m = water_stoich
        assert degrees_of_freedom(m) == 0

    @pytest.mark.unit
    def test_stats_stoich(self, water_stoich):
        m = water_stoich
        assert number_variables(m) == 123
        assert number_total_constraints(m) == 54
        assert number_unused_variables(m) == 59

    @pytest.mark.component
    def test_scaling_stoich(self, water_stoich):
        m = water_stoich

        # Next, try adding scaling for species
        min = 1e-3
        for i in m.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp:
            # i[0] = phase, i[1] = species
            if m.fs.unit.inlet.mole_frac_comp[0, i[1]].value > min:
                scale = m.fs.unit.inlet.mole_frac_comp[0, i[1]].value
            else:
                scale = min
            iscale.set_scaling_factor(
                m.fs.unit.control_volume.properties_out[0.0].mole_frac_comp[i[1]],
                10 / scale,
            )
            iscale.set_scaling_factor(
                m.fs.unit.control_volume.properties_out[0.0].mole_frac_phase_comp[i],
                10 / scale,
            )
            iscale.set_scaling_factor(
                m.fs.unit.control_volume.properties_out[0.0].flow_mol_phase_comp[i],
                10 / scale,
            )
            iscale.constraint_scaling_transform(
                m.fs.unit.control_volume.properties_out[0.0].component_flow_balances[
                    i[1]
                ],
                10 / scale,
            )
            iscale.constraint_scaling_transform(
                m.fs.unit.control_volume.material_balances[0.0, i[1]], 10 / scale
            )

        iscale.set_scaling_factor(
            m.fs.unit.control_volume.rate_reaction_extent[0.0, "R1"], 1
        )
        iscale.calculate_scaling_factors(m.fs.unit)

        assert isinstance(m.fs.unit.control_volume.scaling_factor, Suffix)

        assert isinstance(
            m.fs.unit.control_volume.properties_out[0.0].scaling_factor, Suffix
        )

        assert isinstance(
            m.fs.unit.control_volume.properties_in[0.0].scaling_factor, Suffix
        )

    @pytest.mark.component
    def test_initialize(self, water_stoich):
        m = water_stoich

        orig_fixed_vars = fixed_variables_set(m)
        orig_act_consts = activated_constraints_set(m)

        #   Manually fix the unknown inlet and force unfix the known outlets
        #       This customization is REQUIRED for StoichiometricReactor
        #       since IDAES assumes that inlets are knowns during a solve
        m.fs.unit.inlet.mole_frac_comp[0, "Ca(OH)2"].fix(0.0000006)
        m.fs.unit.outlet.mole_frac_comp[0, "Ca(OH)2"].unfix()
        m.fs.unit.outlet.mole_frac_comp[0, "Ca(HCO3)2"].unfix()
        m.fs.unit.outlet.mole_frac_comp[0, "Mg(HCO3)2"].unfix()

        m.fs.unit.initialize(optarg=solver.options, outlvl=idaeslog.DEBUG)

        #   Undo the fixing we just did and return all values to there
        #       original states
        m.fs.unit.outlet.mole_frac_comp[0, "Ca(OH)2"].fix(0.0000003)
        m.fs.unit.inlet.mole_frac_comp[0, "Ca(OH)2"].unfix()
        m.fs.unit.outlet.mole_frac_comp[0, "Ca(HCO3)2"].fix(0.000015)
        m.fs.unit.outlet.mole_frac_comp[0, "Mg(HCO3)2"].fix(0.000015)

        fin_fixed_vars = fixed_variables_set(m)
        fin_act_consts = activated_constraints_set(m)

        print(value(m.fs.unit.outlet.temperature[0]))
        assert degrees_of_freedom(m) == 0

        assert len(fin_act_consts) == len(orig_act_consts)
        assert len(fin_fixed_vars) == len(orig_fixed_vars)

    @pytest.mark.component
    def test_solve(self, water_stoich):
        m = water_stoich
        solver.options["max_iter"] = 4000
        results = solver.solve(m, tee=True)
        assert results.solver.termination_condition == TerminationCondition.optimal
        assert results.solver.status == SolverStatus.ok

    @pytest.mark.component
    def test_solution(self, water_stoich):
        m = water_stoich
        total_molar_density = (
            value(m.fs.unit.control_volume.properties_out[0.0].dens_mol_phase["Liq"])
            / 1000
        )
        total_hardness1 = (
            50000
            * 2
            * m.fs.unit.outlet.mole_frac_comp[0, "Ca(HCO3)2"].value
            * total_molar_density
        )
        total_hardness2 = (
            50000
            * 2
            * m.fs.unit.outlet.mole_frac_comp[0, "Mg(HCO3)2"].value
            * total_molar_density
        )
        assert pytest.approx(55.23359, rel=1e-5) == total_molar_density
        assert pytest.approx(165.70077, rel=1e-5) == total_hardness1 + total_hardness2
