#################################################################################
# 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
from pyomo.environ import value
from watertap.examples.flowsheets.full_treatment_train.model_components import (
    property_models,
)


@pytest.mark.component
def test_feed_TDS():
    m = property_models.solve_specify_feed(base="TDS")
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "TDS"]) == pytest.approx(
        0.035, rel=1e-3
    )


@pytest.mark.component
def test_feed_ion():
    m = property_models.solve_specify_feed(base="ion")
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "Ca"]) == pytest.approx(
        3.82e-4, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "Cl"]) == pytest.approx(
        2.032e-2, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "H2O"]) == pytest.approx(
        0.9646, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "Mg"]) == pytest.approx(
        1.394e-3, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "Na"]) == pytest.approx(
        1.112e-2, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "SO4"]) == pytest.approx(
        2.136e-3, rel=1e-3
    )


@pytest.mark.component
def test_feed_salt():
    m = property_models.solve_specify_feed(base="salt")
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "CaSO4"]) == pytest.approx(
        1.298e-3, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "H2O"]) == pytest.approx(
        0.9646, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "MgCl2"]) == pytest.approx(
        4.251e-3, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "MgSO4"]) == pytest.approx(
        1.529e-3, rel=1e-3
    )
    assert value(m.fs.stream[0].mass_frac_phase_comp["Liq", "NaCl"]) == pytest.approx(
        2.827e-2, rel=1e-3
    )
