#################################################################################
# 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/"
#################################################################################
"""
mutable parameters for optimization:
    m.fs.system_recovery_target
"""
from pyomo.environ import (
    ConcreteModel,
    Objective,
    Expression,
    Constraint,
    TransformationFactory,
    value,
    Param,
    Var,
)
from pyomo.environ import units as pyunits
from pyomo.network import Arc

from idaes.core import FlowsheetBlock
from idaes.core.util.initialization import propagate_state
from idaes.core.util.scaling import calculate_scaling_factors

from watertap.examples.flowsheets.full_treatment_train.flowsheet_components import (
    desalination,
    feed_block,
    gypsum_saturation_index,
    translator_block,
    costing,
)
from watertap.examples.flowsheets.full_treatment_train.model_components import (
    property_models,
)
from watertap.examples.flowsheets.full_treatment_train.util import (
    solve_block,
    check_dof,
)


desal_kwargs = {
    "has_desal_feed": False,
    "is_twostage": False,
    "has_ERD": True,
    "RO_type": "0D",
    "RO_base": "TDS",
    "RO_level": "detailed",
}


def build_components(m, pretrt_type="NF", **kwargs):
    kwargs_desalination = {
        k: kwargs[k]
        for k in (
            "has_desal_feed",
            "is_twostage",
            "has_ERD",
            "RO_type",
            "RO_base",
            "RO_level",
        )
    }

    desal_port = desalination.build_desalination(m, **kwargs_desalination)
    m.fs.s_tb_desal = Arc(
        source=m.fs.tb_pretrt_to_desal.outlet, destination=desal_port["in"]
    )

    if pretrt_type == "softening":
        property_models.build_prop(m, base="eNRTL")
    gypsum_saturation_index.build(
        m, section="desalination", pretrt_type=pretrt_type, **kwargs
    )

    m.fs.RO.area.fix(80)
    m.fs.pump_RO.control_volume.properties_out[0].pressure.fix(60e5)

    if kwargs["is_twostage"]:
        m.fs.RO2.area.fix(20)
        m.fs.pump_RO2.control_volume.properties_out[0].pressure.fix(90e5)

    # touch some properties used in optimization
    if kwargs["is_twostage"]:
        product_water_sb = m.fs.mixer_permeate.mixed_state[0]
    else:
        product_water_sb = m.fs.RO.mixed_permeate[0]

    feed_flow_vol = 0.0009769808  # value of feed flowrate using the seawater property package with 1 kg/s mass flowrate
    m.fs.system_recovery = Expression(expr=product_water_sb.flow_vol / feed_flow_vol)

    # RO recovery
    m.fs.RO_recovery = Var(
        initialize=0.5,
        bounds=(0.01, 0.99),
        doc="Total volumetric water recovery for RO",
    )
    m.fs.eq_RO_recovery = Constraint(
        expr=m.fs.RO_recovery
        == product_water_sb.flow_vol
        / m.fs.tb_pretrt_to_desal.properties_out[0].flow_vol
    )

    # annual water production
    m.fs.treated_flow_vol = Expression(expr=product_water_sb.flow_vol)
    costing.build_costing(m, **kwargs)

    return desal_port


def build(m, **kwargs):
    """
    build an RO
    """
    assert not kwargs["has_desal_feed"]
    property_models.build_prop(m, base="ion")
    feed_block.build_feed(m, base="ion")

    property_models.build_prop(m, base=kwargs["RO_base"])
    translator_block.build_tb(
        m,
        base_inlet="ion",
        base_outlet=kwargs["RO_base"],
        name_str="tb_pretrt_to_desal",
    )
    m.fs.s_pretrt_tb = Arc(
        source=m.fs.feed.outlet, destination=m.fs.tb_pretrt_to_desal.inlet
    )

    property_models.build_prop(m, base="eNRTL")
    desal_port = build_components(m, **kwargs)


def scale(m, **kwargs):
    calculate_scaling_factors(m.fs.s_tb_desal)
    desalination.scale_desalination(m, **kwargs)


def initialize(m, **kwargs):
    propagate_state(m.fs.s_tb_desal)
    desalination.initialize_desalination(m, **kwargs)
    m.fs.desal_saturation.properties.initialize()


def report(m, **kwargs):
    m.fs.tb_pretrt_to_desal.report()
    desalination.display_desalination(m, **kwargs)
    print(
        "desalination solubility index:", value(m.fs.desal_saturation.saturation_index)
    )
    print("water recovery:", value(m.fs.system_recovery))
    costing.display_costing(m)


def set_optimization_components(m, system_recovery, **kwargs):
    """
    adds max_saturation_index and system_recovery_target
    """
    m.fs.pump_RO.control_volume.properties_out[0].pressure.unfix()
    m.fs.pump_RO.control_volume.properties_out[0].pressure.setlb(20e5)
    m.fs.pump_RO.control_volume.properties_out[0].pressure.setub(75e5)

    m.fs.RO.area.unfix()
    m.fs.RO.area.setlb(10)
    m.fs.RO.area.setub(300)

    m.fs.RO.feed_side.N_Re[0, 0].unfix()

    # Set lower bound for water flux at the RO outlet, based on a minimum net driving pressure, NDPmin
    m.fs.RO.NDPmin = Param(initialize=1e5, mutable=True, units=pyunits.Pa)
    m.fs.RO.flux_mass_phase_comp[0, 1, "Liq", "H2O"].setlb(
        value(m.fs.RO.A_comp[0, "H2O"] * m.fs.RO.dens_solvent * m.fs.RO.NDPmin)
    )

    # saturation index
    m.fs.max_saturation_index = Param(initialize=1.0, mutable=True)
    m.fs.eq_max_saturation_index_desal = Constraint(
        expr=m.fs.desal_saturation.saturation_index <= m.fs.max_saturation_index
    )

    m.fs.system_recovery_target = Param(initialize=system_recovery, mutable=True)
    m.fs.system_recovery_tol = Param(initialize=5e-3, mutable=True)
    m.fs.eq_system_recovery = Constraint(
        expr=(
            m.fs.system_recovery_target,
            m.fs.system_recovery,
            m.fs.system_recovery_target + m.fs.system_recovery_tol,
        )
    )

    # set objective
    m.fs.objective = Objective(expr=m.fs.costing.LCOW)


def set_up_optimization(m, system_recovery=0.50, **kwargs):
    set_optimization_components(m, system_recovery, **kwargs)
    check_dof(m, 3)


def optimize(m, check_termination=True):
    return solve_block(m, tee=False, fail_flag=check_termination)


def solve_flowsheet(**desal_kwargs):
    m = ConcreteModel()
    m.fs = FlowsheetBlock(dynamic=False)
    build(m, **desal_kwargs)
    TransformationFactory("network.expand_arcs").apply_to(m)

    # scale
    calculate_scaling_factors(m.fs.tb_pretrt_to_desal)
    scale(m, **desal_kwargs)
    calculate_scaling_factors(m)

    # initialize
    m.fs.feed.initialize()
    propagate_state(m.fs.s_pretrt_tb)
    optarg = {"nlp_scaling_method": "user-scaling"}
    m.fs.tb_pretrt_to_desal.initialize(optarg=optarg)
    initialize(m, **desal_kwargs)

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

    # report
    report(m, **desal_kwargs)

    return m


def optimize_flowsheet(system_recovery=0.50, **kwargs):
    m = solve_flowsheet(**kwargs)
    set_up_optimization(m, system_recovery=system_recovery, **kwargs)
    optimize(m)

    report(m, **kwargs)

    return m


if __name__ == "__main__":
    import sys

    if len(sys.argv) == 1:
        m = solve_flowsheet(**desal_kwargs)
    else:
        m = optimize_flowsheet(system_recovery=float(sys.argv[1]), **desal_kwargs)
