#################################################################################
# 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/"
#################################################################################

"""eNRTL property blocks to calculate and constrain gypsum saturation index"""

from pyomo.environ import Block, Var, Expression, Constraint, units as pyunits


def build(m, section="desalination", pretrt_type="NF", **kwargs):
    if section == "desalination":
        m.fs.desal_saturation = Block()
        m.fs.desal_saturation.properties = m.fs.prop_eNRTL.build_state_block([0])
        sb_eNRTL = m.fs.desal_saturation.properties[0]
    elif section == "pretreatment":
        m.fs.pretrt_saturation = Block()
        m.fs.pretrt_saturation.properties = m.fs.prop_eNRTL.build_state_block([0])
        sb_eNRTL = m.fs.pretrt_saturation.properties[0]
    else:
        raise ValueError(
            "{section} is not an expected section for building the saturation index"
            "".format(section=section)
        )

    # populate initial values
    populate_eNRTL_state_vars(sb_eNRTL, base="FpcTP")

    # ksp = 3.9e-9  # Gibbs energy gives 3.9e-8, but this fits expectations better
    ksp = 3.2e-9  # This fits expectations even better

    # constraints
    if section == "desalination":
        sb_dilute = m.fs.tb_pretrt_to_desal.properties_in[0]
        if kwargs["is_twostage"]:
            sb_perm = m.fs.mixer_permeate.mixed_state[0]
            sb_conc = m.fs.RO2.feed_side.properties[0, 1]
            sb_conc_inter = m.fs.RO2.feed_side.properties_interface[0, 1]
        else:
            sb_perm = m.fs.RO.mixed_permeate[0]
            sb_conc = m.fs.RO.feed_side.properties[0, 1]
            sb_conc_inter = m.fs.RO.feed_side.properties_interface[0, 1]

        m.fs.desal_saturation.cp_modulus = Expression(
            expr=sb_conc_inter.conc_mass_phase_comp["Liq", "TDS"]
            / sb_conc.conc_mass_phase_comp["Liq", "TDS"]
        )

        # constraints
        m.fs.desal_saturation.eq_temperature = Constraint(
            expr=sb_eNRTL.temperature == sb_conc.temperature
        )
        m.fs.desal_saturation.eq_pressure = Constraint(
            expr=sb_eNRTL.pressure == sb_conc.pressure
        )

        if pretrt_type == "NF":
            # assumes pretreatment uses the ion property basis
            comp_match_dict = {
                "Na_+": "Na",
                "Ca_2+": "Ca",
                "Mg_2+": "Mg",
                "SO4_2-": "SO4",
                "Cl_-": "Cl",
                "H2O": "H2O",
            }

            @m.fs.desal_saturation.Constraint(comp_match_dict.keys())
            def eq_flow_mol_balance(b, j):
                if j in ["Cl_-", "Na_+"]:
                    bulk_flow = sb_dilute.flow_mol_phase_comp[
                        "Liq", comp_match_dict[j]
                    ] - sb_perm.flow_mass_phase_comp["Liq", "TDS"] / (
                        58.44e-3 * pyunits.kg / pyunits.mol
                    )
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == bulk_flow * m.fs.desal_saturation.cp_modulus
                    )
                elif j in ["Ca_2+", "Mg_2+", "SO4_2-"]:
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == sb_dilute.flow_mol_phase_comp["Liq", comp_match_dict[j]]
                        * m.fs.desal_saturation.cp_modulus
                    )
                elif j == "H2O":
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == sb_conc.flow_mol_phase_comp["Liq", "H2O"]
                    )

        elif pretrt_type == "softening":
            # assumes pretreatment uses the softening property basis
            comp_match_dict = {
                "Na_+": "NaCl",
                "Ca_2+": "Ca(HCO3)2",
                "Mg_2+": "Mg(HCO3)2",
                "SO4_2-": "SO4_2-",
                "Cl_-": "Cl_-",
                "H2O": "H2O",
            }

            @m.fs.desal_saturation.Constraint(comp_match_dict.keys())
            def eq_flow_mol_balance(b, j):
                if j == "Na_+":
                    bulk_flow = sb_dilute.flow_mol_phase_comp[
                        "Liq", "NaCl"
                    ] - sb_perm.flow_mass_phase_comp["Liq", "TDS"] / (
                        58.44e-3 * pyunits.kg / pyunits.mol
                    )
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == bulk_flow * m.fs.desal_saturation.cp_modulus
                    )
                if j in "Cl_-":
                    bulk_flow = (
                        sb_dilute.flow_mol_phase_comp["Liq", "Cl_-"]
                        + sb_dilute.flow_mol_phase_comp["Liq", "NaCl"]
                        - sb_perm.flow_mass_phase_comp["Liq", "TDS"]
                        / (58.44e-3 * pyunits.kg / pyunits.mol)
                    )
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == bulk_flow * m.fs.desal_saturation.cp_modulus
                    )
                elif j in ["Ca_2+", "Mg_2+", "SO4_2-"]:
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == sb_dilute.flow_mol_phase_comp["Liq", comp_match_dict[j]]
                        * m.fs.desal_saturation.cp_modulus
                    )
                elif j == "H2O":
                    return (
                        sb_eNRTL.flow_mol_phase_comp["Liq", j]
                        == sb_conc.flow_mol_phase_comp["Liq", "H2O"]
                    )

        m.fs.desal_saturation.saturation_index = Var(
            initialize=0.5,
            bounds=(0.0, 10),
            units=pyunits.dimensionless,
            doc="Gypsum saturation index",
        )

        m.fs.desal_saturation.eq_saturation_index = Constraint(
            expr=m.fs.desal_saturation.saturation_index
            == sb_eNRTL.act_phase_comp["Liq", "Ca_2+"]
            * sb_eNRTL.act_phase_comp["Liq", "SO4_2-"]
            * sb_eNRTL.act_phase_comp["Liq", "H2O"] ** 2
            / ksp
        )

    elif section == "pretreatment":
        comp_match_dict = {
            "Na_+": "Na",
            "Ca_2+": "Ca",
            "Mg_2+": "Mg",
            "SO4_2-": "SO4",
            "Cl_-": "Cl",
            "H2O": "H2O",
        }

        sb_conc = m.fs.NF.feed_side.properties_out[0]

        m.fs.pretrt_saturation.eq_temperature = Constraint(
            expr=sb_eNRTL.temperature == sb_conc.temperature
        )
        m.fs.pretrt_saturation.eq_pressure = Constraint(
            expr=sb_eNRTL.pressure == sb_conc.pressure
        )

        @m.fs.pretrt_saturation.Constraint(comp_match_dict.keys())
        def eq_flow_mol_balance(b, j):
            return (
                sb_eNRTL.flow_mol_phase_comp["Liq", j]
                == sb_conc.flow_mol_phase_comp["Liq", comp_match_dict[j]]
            )

        m.fs.pretrt_saturation.saturation_index = Var(
            initialize=0.5,
            bounds=(0.0, 1e6),
            units=pyunits.dimensionless,
            doc="Gypsum saturation index",
        )

        m.fs.pretrt_saturation.eq_saturation_index = Constraint(
            expr=m.fs.pretrt_saturation.saturation_index
            == sb_eNRTL.act_phase_comp["Liq", "Ca_2+"]
            * sb_eNRTL.act_phase_comp["Liq", "SO4_2-"]
            * sb_eNRTL.act_phase_comp["Liq", "H2O"] ** 2
            / ksp
        )


def populate_eNRTL_state_vars(blk, base="FpcTP"):
    blk.temperature = 298
    blk.pressure = 101325

    if base == "FpcTP":
        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,
        }

        for j in feed_mass_frac_comp:
            blk.flow_mol_phase_comp["Liq", j] = (
                feed_flow_mass * feed_mass_frac_comp[j] / mw_comp[j]
            )
            if j == "H2O":
                blk.flow_mol_phase_comp["Liq", j] /= 2
    elif base == "FTPx":
        blk.mole_frac_comp["Na_+"] = 0.017327
        blk.mole_frac_comp["Ca_2+"] = 0.000341
        blk.mole_frac_comp["Mg_2+"] = 0.002054
        blk.mole_frac_comp["SO4_2-"] = 0.000796
        blk.mole_frac_comp["Cl_-"] = 0.020529
        blk.mole_frac_comp["H2O"] = 0.958952
