#################################################################################
# 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/"
#################################################################################
import pytest
import watertap.property_models.seawater_prop_pack as props
from idaes.models.properties.tests.test_harness import (
    PropertyTestHarness as PropertyTestHarness_idaes,
)
from watertap.property_models.tests.property_test_harness import (
    PropertyTestHarness,
    PropertyRegressionTest,
    PropertyCalculateStateTest,
)

# -----------------------------------------------------------------------------
class TestSeawaterProperty_idaes(PropertyTestHarness_idaes):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}
        self.prop_args = {}
        self.has_density_terms = False


class TestSeawaterProperty(PropertyTestHarness):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}
        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e2,
        }
        self.stateblock_statistics = {
            "number_variables": 26,
            "number_total_constraints": 22,
            "number_unused_variables": 0,
            "default_degrees_of_freedom": 4,
        }  # 4 state vars, but pressure is not active
        self.default_solution = {
            ("mass_frac_phase_comp", ("Liq", "H2O")): 0.965,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 0.035,
            ("dens_mass_phase", "Liq"): 1023.5,
            ("dens_mass_solvent", None): 996.9,
            ("flow_vol_phase", "Liq"): 9.770e-4,
            ("conc_mass_phase_comp", ("Liq", "H2O")): 987.7,
            ("conc_mass_phase_comp", ("Liq", "TDS")): 35.82,
            ("flow_mol_phase_comp", ("Liq", "H2O")): 53.57,
            ("flow_mol_phase_comp", ("Liq", "TDS")): 1.1145,
            ("mole_frac_phase_comp", ("Liq", "H2O")): 0.9796,
            ("mole_frac_phase_comp", ("Liq", "TDS")): 2.038e-2,
            ("molality_phase_comp", ("Liq", "TDS")): 1.155,
            ("visc_d_phase", "Liq"): 9.588e-4,
            ("osm_coeff", None): 0.9068,
            ("pressure_osm_phase", "Liq"): 2.588e6,
            ("enth_mass_phase", "Liq"): 9.896e4,
            ("pressure_sat", None): 3111,
            ("cp_mass_phase", "Liq"): 4001,
            ("therm_cond_phase", "Liq"): 0.6086,
            ("dh_vap_mass", None): 2.356e6,
            ("diffus_phase_comp", ("Liq", "TDS")): 1.471e-9,
            ("boiling_point_elevation_phase", "Liq"): 0.3093,
        }


@pytest.mark.component
class TestSeawaterPropertySolution_1(PropertyRegressionTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e2,
        }
        self.state_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 0.99,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 0.01,
            ("temperature", None): 273.15 + 50,
            ("pressure", None): 2e5,
        }
        self.regression_solution = {
            ("mass_frac_phase_comp", ("Liq", "H2O")): 0.99,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 0.01,
            ("dens_mass_phase", "Liq"): 995.4,
            ("dens_mass_solvent", None): 988.0,
            ("flow_vol_phase", "Liq"): 1.005e-3,
            ("conc_mass_phase_comp", ("Liq", "H2O")): 985.5,
            ("conc_mass_phase_comp", ("Liq", "TDS")): 9.954,
            ("flow_mol_phase_comp", ("Liq", "H2O")): 54.95,
            ("flow_mol_phase_comp", ("Liq", "TDS")): 0.3184,
            ("mole_frac_phase_comp", ("Liq", "H2O")): 0.9942,
            ("mole_frac_phase_comp", ("Liq", "TDS")): 5.761e-3,
            ("molality_phase_comp", ("Liq", "TDS")): 0.3216,
            ("visc_d_phase", "Liq"): 5.596e-4,
            ("osm_coeff", None): 0.9029,
            ("pressure_osm_phase", "Liq"): 7.710e5,
            ("enth_mass_phase", "Liq"): 2.057e5,
            ("pressure_sat", None): 1.229e4,
            ("cp_mass_phase", "Liq"): 4.130e3,
            ("therm_cond_phase", "Liq"): 0.6400,
            ("dh_vap_mass", None): 2.358e6,
            ("diffus_phase_comp", ("Liq", "TDS")): 1.493e-9,
            ("boiling_point_elevation_phase", "Liq"): 0.0989,
        }


@pytest.mark.component
class TestSeawaterPropertySolution_2(PropertyRegressionTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e2,
        }
        self.state_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 0.95,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 0.05,
            ("temperature", None): 273.15 + 10,
            ("pressure", None): 100e5,
        }
        self.regression_solution = {
            ("mass_frac_phase_comp", ("Liq", "H2O")): 0.95,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 0.05,
            ("dens_mass_phase", "Liq"): 1.039e3,
            ("dens_mass_solvent", None): 999.5,
            ("flow_vol_phase", "Liq"): 9.628e-4,
            ("conc_mass_phase_comp", ("Liq", "H2O")): 986.8,
            ("conc_mass_phase_comp", ("Liq", "TDS")): 51.93,
            ("flow_mol_phase_comp", ("Liq", "H2O")): 52.73,
            ("flow_mol_phase_comp", ("Liq", "TDS")): 1.592,
            ("mole_frac_phase_comp", ("Liq", "H2O")): 0.9707,
            ("mole_frac_phase_comp", ("Liq", "TDS")): 2.931e-2,
            ("molality_phase_comp", ("Liq", "TDS")): 1.676,
            ("visc_d_phase", "Liq"): 1.443e-3,
            ("osm_coeff", None): 0.9106,
            ("pressure_osm_phase", "Liq"): 3.591e6,
            ("enth_mass_phase", "Liq"): 4.783e4,
            ("pressure_sat", None): 1.194e3,
            ("cp_mass_phase", "Liq"): 3.916e3,
            ("therm_cond_phase", "Liq"): 0.5854,
            ("dh_vap_mass", None): 2.353e6,
            ("diffus_phase_comp", ("Liq", "TDS")): 1.471e-9,
            ("boiling_point_elevation_phase", "Liq"): 0.4069,
        }


# @pytest.mark.component
# class TestSeawaterPropertySolution_3(PropertyRegressionTest):
#     def configure(self):
#         self.prop_pack = props.SeawaterParameterBlock
#         self.param_args = {}
#
#         self.solver = "ipopt"
#         self.optarg = {"nlp_scaling_method": "user-scaling"}
#
#         self.scaling_args = {
#             ("flow_vol_phase", "Liq"): 1,
#         }
#         self.state_args = {
#             ("flow_vol_phase", "Liq"): 10,
#             ("mass_frac_phase_comp", ("Liq", "TDS")): 0.04,
#             ("temperature", None): 273.15 + 10,
#             ("pressure", None): 2e6,
#         }
#         self.regression_solution = {
#             ("mass_frac_phase_comp", ("Liq", "H2O")): 0.96,
#             ("mass_frac_phase_comp", ("Liq", "TDS")): 0.04,
#             ("dens_mass_phase", "Liq"): 1.039e3,
#             ("dens_mass_solvent", None): 999.5,
#             ("flow_vol_phase", "Liq"): 9.628e-4,
#             ("conc_mass_phase_comp", ("Liq", "H2O")): 986.8,
#             ("conc_mass_phase_comp", ("Liq", "TDS")): 51.93,
#             ("flow_mol_phase_comp", ("Liq", "H2O")): 52.73,
#             ("flow_mol_phase_comp", ("Liq", "TDS")): 1.592,
#             ("mole_frac_phase_comp", ("Liq", "H2O")): 0.9707,
#             ("mole_frac_phase_comp", ("Liq", "TDS")): 2.931e-2,
#             ("molality_phase_comp", ("Liq", "TDS")): 1.676,
#             ("visc_d_phase", "Liq"): 1.443e-3,
#             ("osm_coeff", None): 0.9106,
#             ("pressure_osm_phase", "Liq"): 3.591e6,
#             ("enth_mass_phase", "Liq"): 4.783e4,
#             ("pressure_sat", None): 1.194e3,
#             ("cp_mass_phase", "Liq"): 3.916e3,
#             ("therm_cond_phase", "Liq"): 0.5854,
#             ("dh_vap_mass", None): 2.353e6,
#             ("diffus_phase_comp", ("Liq", "TDS")): 1.471e-9,
#             ("boiling_point_elevation_phase", "Liq"): 0.4069,
#         }


@pytest.mark.component
class TestSeawaterCalculateState_1(PropertyCalculateStateTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1e-1,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e1,
        }
        self.var_args = {
            ("flow_vol_phase", "Liq"): 2e-2,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 0.05,
            ("temperature", None): 273.15 + 25,
            ("pressure", None): 5e5,
        }
        self.state_solution = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 19.66,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1.035,
        }


@pytest.mark.component
class TestNaClCalculateState_2(PropertyCalculateStateTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e2,
        }
        self.var_args = {
            ("flow_vol_phase", "Liq"): 1e-3,
            ("pressure_osm_phase", "Liq"): 100e5,
            ("temperature", None): 273.15 + 25,
            ("pressure", None): 5e5,
        }
        self.state_solution = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 0.9604,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 0.1231,
        }


@pytest.mark.component
class TestNaClCalculateState_3(PropertyCalculateStateTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e2,
        }
        self.var_args = {
            ("flow_vol_phase", "Liq"): 1e-3,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 0.06,
            ("pressure_sat", None): 3e4,
            ("pressure", None): 5e5,
        }
        self.state_solution = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 0.9605,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 0.0613,
            ("temperature", None): 343.05,
        }


@pytest.mark.component
class TestSeawaterCalculateState_4(PropertyCalculateStateTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1e-4,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e1,
            ("enth_mass_phase", ("Liq", "TDS")): 1e-4,
        }
        self.var_args = {
            ("flow_vol_phase", "Liq"): 10,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 1e-5,
            ("temperature", None): 273.15 + 10,
            ("pressure", None): 2e6,
        }
        self.state_solution = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 9995.07,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 0.09995,
            ("enth_mass_phase", "Liq"): 4.3945e4,
        }


@pytest.mark.component
class TestSeawaterCalculateState_5(PropertyCalculateStateTest):
    def configure(self):
        self.prop_pack = props.SeawaterParameterBlock
        self.param_args = {}

        self.solver = "ipopt"
        self.optarg = {"nlp_scaling_method": "user-scaling"}

        self.scaling_args = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 1e-3,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1e-3,
            ("enth_mass_phase", ("Liq", "TDS")): 1e-5,
        }
        self.var_args = {
            ("flow_vol_phase", "Liq"): 10,
            ("mass_frac_phase_comp", ("Liq", "TDS")): 0.12,
            ("temperature", None): 273.15 + 120,
            ("pressure", None): 1.2e7,
        }
        self.state_solution = {
            ("flow_mass_phase_comp", ("Liq", "H2O")): 9091.08,
            ("flow_mass_phase_comp", ("Liq", "TDS")): 1239.69,
            ("enth_mass_phase", "Liq"): 3.8562e5,
        }
