###############################################################################
# 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/"
#
###############################################################################
import pytest
from pyomo.environ import ConcreteModel, value
from idaes.core import FlowsheetBlock
from idaes.generic_models.properties.core.generic.generic_property import GenericParameterBlock
from idaes.core.util.scaling import calculate_scaling_factors, get_scaling_factor, constraint_scaling_transform
from idaes.core.util import get_solver
from watertap.examples.flowsheets.full_treatment_train.model_components.eNRTL import entrl_config_FTPx, entrl_config_FpcTP
from watertap.examples.flowsheets.full_treatment_train.util import check_scaling, solve_with_user_scaling

def simulate_enrtl_FTPx(state_var_args):
    m = ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": False})
    m.fs.params = GenericParameterBlock(default=entrl_config_FTPx.configuration)

    m.fs.state = m.fs.params.build_state_block(
        m.fs.time, default={"defined_state": True})

    for (v_name, ind), val in state_var_args.items():
        var = getattr(m.fs.state[0], v_name)
        var[ind].fix(val)
    m.fs.state[0].flow_mol_phase['Liq'].value = 1

    # scale model
    calculate_scaling_factors(m)

    # Regular solve
    solver = get_solver()
    results = solver.solve(m)

    ksp = 3.2e-9  # Gibbs energy gives 3.9e-8, but this fits expectations better
    saturation_index = value(m.fs.state[0].act_phase_comp["Liq", "Ca_2+"]
                             * m.fs.state[0].act_phase_comp["Liq", "SO4_2-"]
                             * m.fs.state[0].act_phase_comp["Liq", "H2O"] ** 2 /
                             ksp)
    return saturation_index


@pytest.mark.component
def test_enrtl_FTPx_0():
    # seawater concentration
    state_var_args = {('temperature', None): 298,
                      ('pressure', None): 101325,
                      ('flow_mol', None): 100,
                      ('mole_frac_comp', 'Na_+'): 0.008845,
                      ('mole_frac_comp', 'Ca_2+'): 0.000174,
                      ('mole_frac_comp', 'Mg_2+'): 0.001049,
                      ('mole_frac_comp', 'SO4_2-'): 0.000407,
                      ('mole_frac_comp', 'Cl_-'): 0.010479,
                      ('mole_frac_comp', 'H2O'): 0.979046}
    saturation_index = simulate_enrtl_FTPx(state_var_args)
    assert saturation_index == pytest.approx(0.2198, rel=1e-3)


@pytest.mark.component
def test_enrtl_FTPx_1():
    # 2 times seawater concentration
    state_var_args = {('temperature', None): 298,
                      ('pressure', None): 101325,
                      ('flow_mol', None): 100,
                      ('mole_frac_comp', 'Na_+'): 0.017327,
                      ('mole_frac_comp', 'Ca_2+'): 0.000341,
                      ('mole_frac_comp', 'Mg_2+'): 0.002054,
                      ('mole_frac_comp', 'SO4_2-'): 0.000796,
                      ('mole_frac_comp', 'Cl_-'): 0.020529,
                      ('mole_frac_comp', 'H2O'): 0.958952}
    saturation_index = simulate_enrtl_FTPx(state_var_args)
    assert saturation_index == pytest.approx(0.4333, rel=1e-3)


def simulate_enrtl_FpcTP(state_var_args):
    m = ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": False})
    m.fs.params = GenericParameterBlock(default=entrl_config_FpcTP.configuration)

    m.fs.state = m.fs.params.build_state_block(
        m.fs.time, default={"defined_state": True})

    for (v_name, ind), val in state_var_args.items():
        var = getattr(m.fs.state[0], v_name)
        var[ind].fix(val)

    # scale model
    calculate_scaling_factors(m)
    for (ind, c) in m.fs.state[0].true_to_appr_species.items():
        sf = get_scaling_factor(m.fs.state[0].flow_mol_phase_comp_apparent[ind])
        constraint_scaling_transform(c, sf)
    for (ind, c) in m.fs.state[0].appr_mole_frac_constraint.items():
        sf = get_scaling_factor(m.fs.state[0].mole_frac_phase_comp_apparent[ind])
        constraint_scaling_transform(c, sf)

    check_scaling(m)

    # User scaling
    solve_with_user_scaling(m)

    ksp = 3.2e-9  # Gibbs energy gives 3.9e-8, but this fits expectations better
    saturation_index = value(m.fs.state[0].act_phase_comp["Liq", "Ca_2+"]
                             * m.fs.state[0].act_phase_comp["Liq", "SO4_2-"]
                             * m.fs.state[0].act_phase_comp["Liq", "H2O"] ** 2 /
                             ksp)
    return saturation_index


@pytest.mark.component
def test_enrtl_FpcTP_1():
    # standard seawater concentration
    feed_flow_mass = 1  # kg/s
    feed_mass_frac_comp = {'Na_+': 11122e-6,
                           'Ca_2+': 382e-6,
                           'Mg_2+': 1394e-6,
                           'SO4_2-': 2136e-6,
                           'Cl_-': 20316.88e-6}
    feed_mass_frac_comp['H2O'] = 1 - sum(x for x in feed_mass_frac_comp.values())

    mw_comp = {'H2O': 18.015e-3,
               'Na_+': 22.990e-3,
               'Ca_2+': 40.078e-3,
               'Mg_2+': 24.305e-3,
               'SO4_2-': 96.06e-3,
               'Cl_-': 35.453e-3}

    feed_flow_mol_comp = {}
    for j in feed_mass_frac_comp:
        feed_flow_mol_comp[j] = feed_flow_mass * feed_mass_frac_comp[j] / mw_comp[j]

    state_var_args = {('temperature', None): 298,
                      ('pressure', None): 101325}
    for j in feed_flow_mol_comp:
        state_var_args[('flow_mol_phase_comp', ('Liq', j))] = feed_flow_mol_comp[j]

    saturation_index = simulate_enrtl_FpcTP(state_var_args)
    assert saturation_index == pytest.approx(0.2200, rel=1e-3)


@pytest.mark.component
def test_enrtl_FpcTP_2():
    # seawater concentration with 50% water removal
    feed_flow_mass = 1  # kg/s
    feed_mass_frac_comp = {'Na_+': 11122e-6,
                           'Ca_2+': 382e-6,
                           'Mg_2+': 1394e-6,
                           'SO4_2-': 2136e-6,
                           'Cl_-': 20316.88e-6}
    feed_mass_frac_comp['H2O'] = 1 - sum(x for x in feed_mass_frac_comp.values())

    mw_comp = {'H2O': 18.015e-3,
               'Na_+': 22.990e-3,
               'Ca_2+': 40.078e-3,
               'Mg_2+': 24.305e-3,
               'SO4_2-': 96.06e-3,
               'Cl_-': 35.453e-3}

    feed_flow_mol_comp = {}
    for j in feed_mass_frac_comp:
        feed_flow_mol_comp[j] = feed_flow_mass * feed_mass_frac_comp[j] / mw_comp[j]
        if j == 'H2O':
            feed_flow_mol_comp[j] = feed_flow_mol_comp[j] / 2

    state_var_args = {('temperature', None): 298,
                      ('pressure', None): 101325}
    for j in feed_flow_mol_comp:
        state_var_args[('flow_mol_phase_comp', ('Liq', j))] = feed_flow_mol_comp[j]

    saturation_index = simulate_enrtl_FpcTP(state_var_args)
    assert saturation_index == pytest.approx(0.4344, rel=1e-3)
