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

"""Pretreatment flowsheet components"""

from pyomo.environ import ConcreteModel, TransformationFactory
from pyomo.network import Arc
from idaes.core import FlowsheetBlock
from idaes.models.unit_models import (
    Separator,
    Mixer,
)
from idaes.models.unit_models.separator import (
    SplittingType,
    EnergySplittingType,
)
from idaes.core.util.scaling import (
    calculate_scaling_factors,
    set_scaling_factor,
    get_scaling_factor,
    constraint_scaling_transform,
)
from idaes.core.util.initialization import propagate_state
from watertap.unit_models.pressure_changer import Pump
from watertap.examples.flowsheets.full_treatment_train.flowsheet_components import (
    feed_block,
)
from watertap.examples.flowsheets.full_treatment_train.model_components import (
    unit_separator,
    unit_ZONF,
    property_models,
)
from watertap.examples.flowsheets.full_treatment_train.util import (
    solve_block,
    check_dof,
)


def build_pretreatment_NF(m, has_bypass=True, NF_type="ZO", NF_base="ion"):
    """
    Builds NF pretreatment including specified feed and auxiliary equipment.
    Arguments:
        has_bypass: True or False, default = True
        NF_type: 'Sep' or 'ZO', default = 'ZO'
        NF_base: 'ion' or 'salt', default = 'ion'
    """
    pretrt_port = {}
    prop = property_models.get_prop(m, base=NF_base)

    # build feed
    feed_block.build_feed(m, base=NF_base)

    # build NF
    if NF_type == "Sep":
        unit_separator.build_SepNF(m, base=NF_base)
    elif NF_type == "ZO":
        unit_ZONF.build_ZONF(m, base=NF_base)
        m.fs.pump_NF = Pump(property_package=prop)
    else:
        raise ValueError(
            "Unexpected model type {NF_type} provided to build_NF_no_bypass"
            "".format(NF_type=NF_type)
        )

    if has_bypass:
        # build auxiliary units
        m.fs.splitter = Separator(
            property_package=prop,
            outlet_list=["pretreatment", "bypass"],
            split_basis=SplittingType.totalFlow,
            energy_split_basis=EnergySplittingType.equal_temperature,
        )
        m.fs.mixer = Mixer(property_package=prop, inlet_list=["pretreatment", "bypass"])

        # connect models
        m.fs.s_pretrt_feed_splitter = Arc(
            source=m.fs.feed.outlet, destination=m.fs.splitter.inlet
        )
        m.fs.s_pretrt_splitter_mixer = Arc(
            source=m.fs.splitter.bypass, destination=m.fs.mixer.bypass
        )
        if NF_type == "ZO":
            m.fs.s_pretrt_splitter_pumpNF = Arc(
                source=m.fs.splitter.pretreatment, destination=m.fs.pump_NF.inlet
            )
            m.fs.s_pretrt_pumpNF_NF = Arc(
                source=m.fs.pump_NF.outlet, destination=m.fs.NF.inlet
            )
        else:
            m.fs.s_pretrt_splitter_NF = Arc(
                source=m.fs.splitter.pretreatment, destination=m.fs.NF.inlet
            )
        m.fs.s_pretrt_NF_mixer = Arc(
            source=m.fs.NF.permeate, destination=m.fs.mixer.pretreatment
        )

        # specify (NF and feed is already specified, mixer has 0 DOF, splitter has 1 DOF, NF pump has 2 DOF)
        # splitter
        m.fs.splitter.split_fraction[0, "bypass"].fix(0.1)
        if NF_type == "ZO":
            m.fs.pump_NF.efficiency_pump.fix(0.80)
            m.fs.pump_NF.control_volume.properties_out[0].pressure.fix(4e5)

        # inlet/outlet ports for pretreatment
        pretrt_port["out"] = m.fs.mixer.outlet
        pretrt_port["waste"] = m.fs.NF.retentate

    else:  # no bypass
        # build auxiliary units (none)

        # connect models
        if NF_type == "ZO":
            m.fs.s_pretrt_feed_pumpNF = Arc(
                source=m.fs.feed.outlet, destination=m.fs.pump_NF.inlet
            )
            m.fs.s_pretrt_pumpNF_NF = Arc(
                source=m.fs.pump_NF.outlet, destination=m.fs.NF.inlet
            )
            # TODO: should source be m.fs.pump_NF.outlet? Double-check here and other arcs with pump_NF
        else:
            m.fs.s_pretrt_feed_NF = Arc(
                source=m.fs.feed.outlet, destination=m.fs.NF.inlet
            )

        # specify (NF and feed are already specified, NF pump has 2 DOF)
        if NF_type == "ZO":
            m.fs.pump_NF.efficiency_pump.fix(0.80)
            m.fs.pump_NF.control_volume.properties_out[0].pressure.fix(4e5)

        # inlet/outlet ports for pretreatment
        pretrt_port["out"] = m.fs.NF.permeate
        pretrt_port["waste"] = m.fs.NF.retentate

    return pretrt_port


def scale_pretreatment_NF(m, **kwargs):
    calculate_scaling_factors(m.fs.feed)
    calculate_scaling_factors(m.fs.NF)

    if kwargs["has_bypass"]:
        calculate_scaling_factors(m.fs.splitter)
        set_scaling_factor(
            m.fs.splitter.split_fraction, 1
        )  # TODO: should have an IDAES default
        constraint_scaling_transform(
            m.fs.splitter.sum_split_frac[0], 1
        )  # TODO: should have an IDAES default
        calculate_scaling_factors(m.fs.mixer)
        set_scaling_factor(
            m.fs.mixer.minimum_pressure,
            get_scaling_factor(m.fs.mixer.mixed_state[0].pressure),
        )  # TODO: IDAES should have a default and link to the constraint
        for c in [
            m.fs.mixer.minimum_pressure_constraint[0, 1],
            m.fs.mixer.minimum_pressure_constraint[0, 2],
            m.fs.mixer.mixture_pressure[0.0],
        ]:
            constraint_scaling_transform(
                c, get_scaling_factor(m.fs.mixer.minimum_pressure)
            )

    if kwargs["NF_type"] == "ZO":
        set_scaling_factor(m.fs.pump_NF.control_volume.work, 1e-3)
        calculate_scaling_factors(m.fs.pump_NF)
        set_scaling_factor(
            m.fs.pump_NF.ratioP, 1
        )  # TODO: IDAES should have a default and link to the constraint


def initialize_pretreatment_NF(m, **kwargs):
    optarg = {"nlp_scaling_method": "user-scaling"}

    if kwargs["has_bypass"]:
        m.fs.feed.initialize(optarg=optarg)
        propagate_state(m.fs.s_pretrt_feed_splitter)
        m.fs.splitter.initialize(optarg=optarg)
        propagate_state(m.fs.s_pretrt_splitter_mixer)
        if kwargs["NF_type"] == "ZO":
            propagate_state(m.fs.s_pretrt_splitter_pumpNF)
            m.fs.pump_NF.initialize(optarg=optarg)
            propagate_state(m.fs.s_pretrt_pumpNF_NF)
            m.fs.NF.initialize(optarg=optarg)
        else:  # NF_type == 'Sep'
            propagate_state(m.fs.s_pretrt_splitter_NF)
            # m.fs.NF.initialize(optarg=optarg)  # IDAES error when NF is a separator TODO: address in IDAES
        propagate_state(m.fs.s_pretrt_NF_mixer)
        m.fs.mixer.initialize(optarg=optarg)

    else:  # no bypass
        m.fs.feed.initialize(optarg=optarg)
        if kwargs["NF_type"] == "ZO":
            propagate_state(m.fs.s_pretrt_feed_pumpNF)
            m.fs.pump_NF.initialize(optarg=optarg)
            propagate_state(m.fs.s_pretrt_pumpNF_NF)
            m.fs.NF.initialize(optarg=optarg)
        else:  # NF_type == 'Sep'
            propagate_state(m.fs.s_pretrt_feed_NF)
            # m.fs.NF.initialize(optarg=optarg)  # IDAES error when NF is a separator TODO: address in IDAES


def display_pretreatment_NF(m, **kwargs):

    m.fs.feed.report()

    if kwargs["has_bypass"]:
        m.fs.splitter.report()
        if kwargs["NF_type"] == "ZO":
            m.fs.pump_NF.report()
        m.fs.NF.report()
        m.fs.mixer.report()
    else:  # no bypass
        if kwargs["NF_type"] == "ZO":
            m.fs.pump_NF.report()
        m.fs.NF.report()


def solve_pretreatment_NF(**kwargs):
    m = ConcreteModel()
    m.fs = FlowsheetBlock(dynamic=False)

    property_models.build_prop(m, base=kwargs["NF_base"])
    build_pretreatment_NF(m, **kwargs)
    TransformationFactory("network.expand_arcs").apply_to(m)

    scale_pretreatment_NF(m, **kwargs)

    initialize_pretreatment_NF(m, **kwargs)

    check_dof(m)
    solve_block(m, tee=True, fail_flag=True)

    display_pretreatment_NF(m, **kwargs)


if __name__ == "__main__":
    solve_pretreatment_NF(has_bypass=True, NF_type="ZO", NF_base="ion")
