#################################################################################
# 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 (
    unit_separator,
    unit_0DRO,
    unit_1DRO,
    unit_ZONF,
)


@pytest.mark.component
def test_unit_separator_SepRO():
    m = unit_separator.solve_SepRO(base="TDS")
    assert value(
        m.fs.RO.permeate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.4825, rel=1e-3)
    assert value(
        m.fs.RO.permeate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(3.5e-4, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.4825, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(3.465e-2, rel=1e-3)


@pytest.mark.component
def test_unit_separator_SepNF_ion():
    m = unit_separator.solve_SepNF(base="ion")
    assert value(
        m.fs.NF.mixed_state[0].mass_frac_phase_comp["Liq", "Na"]
    ) == pytest.approx(11122e-6, rel=1e-3)
    assert value(
        m.fs.NF.mixed_state[0].mass_frac_phase_comp["Liq", "Cl"]
    ) == pytest.approx(20317e-6, rel=1e-3)
    assert value(m.fs.NF.split_fraction[0, "permeate", "Cl"]) == pytest.approx(
        0.7753, rel=1e-3
    )
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.8682, rel=1e-3)
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "Na"]
    ) == pytest.approx(1.000e-2, rel=1e-3)
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "Cl"]
    ) == pytest.approx(1.575e-2, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "Ca"]
    ) == pytest.approx(3.438e-4, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "Mg"]
    ) == pytest.approx(1.255e-3, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "SO4"]
    ) == pytest.approx(1.922e-3, rel=1e-3)


@pytest.mark.component
def test_unit_separator_SepNF_salt():
    m = unit_separator.solve_SepNF(base="salt")
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.8682, rel=1e-3)
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "NaCl"]
    ) == pytest.approx(2.544e-2, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "CaSO4"]
    ) == pytest.approx(1.168e-3, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "MgSO4"]
    ) == pytest.approx(1.376e-3, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "MgCl2"]
    ) == pytest.approx(3.401e-3, rel=1e-3)
    assert value(
        m.fs.NF.retentate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(9.645e-2, rel=1e-3)


@pytest.mark.component
def test_unit_0DRO_simple():
    m = unit_0DRO.solve_RO(base="TDS", level="simple")
    assert value(
        m.fs.RO.permeate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.3389, rel=1e-3)
    assert value(
        m.fs.RO.permeate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(7.875e-5, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.6261, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(3.492e-2, rel=1e-3)


@pytest.mark.component
def test_unit_0DRO_detailed():
    m = unit_0DRO.solve_RO(base="TDS", level="detailed")
    assert value(
        m.fs.RO.permeate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.2653, rel=1e-3)
    assert value(
        m.fs.RO.permeate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(8.476e-5, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.6997, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(3.492e-2, rel=1e-3)


@pytest.mark.component
def test_unit_1DRO_detailed():
    m = unit_1DRO.solve_RO(base="TDS", level="detailed")
    assert value(
        m.fs.RO.mixed_permeate[0].flow_mass_phase_comp["Liq", "H2O"]
    ) == pytest.approx(0.2541, rel=1e-3)
    assert value(
        m.fs.RO.mixed_permeate[0].flow_mass_phase_comp["Liq", "TDS"]
    ) == pytest.approx(8.577e-5, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.7109, rel=1e-3)
    assert value(
        m.fs.RO.retentate.flow_mass_phase_comp[0, "Liq", "TDS"]
    ) == pytest.approx(3.491e-2, rel=1e-3)


@pytest.mark.component
def test_unit_ZONF():
    m = unit_ZONF.solve_ZONF(base="ion")
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "H2O"]
    ) == pytest.approx(0.8350, rel=1e-3)
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "Ca"]
    ) == pytest.approx(6.897e-5, rel=1e-3)
    assert value(
        m.fs.NF.permeate.flow_mass_phase_comp[0, "Liq", "Cl"]
    ) == pytest.approx(1.475e-2, rel=1e-3)
