#################################################################################
# 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/"
#################################################################################
"""
This module contains common methods for determining the electricity intensity
and demand for zero-order unit models.
"""

import idaes.logger as idaeslog

from pyomo.environ import Param, Var, units as pyunits

# Some more inforation about this module
__author__ = "Andrew Lee"

# Set up logger
_log = idaeslog.getLogger(__name__)


def _common(self):
    # Add electricity consumption to model
    self.electricity = Var(
        self.flowsheet().time,
        units=pyunits.kW,
        bounds=(0, None),
        doc="Electricity consumption of unit",
    )

    self._perf_var_dict["Electricity Demand"] = self.electricity


def constant_intensity(self):
    """
    Helper method for implementing electricity demand assuming constant
    intensity based on the inlet volumetric flow rate.

    E[t] = Q[t] * intensity

    Two variables are added to the model:
        * electricity (indexed by time)
        * energy_electric_flow_vol_inlet (unindexed)

    One constraint is added to the model:
        * electricity_consumption (indexed by time)
    """
    _common(self)

    self.energy_electric_flow_vol_inlet = Var(
        units=pyunits.kWh / pyunits.m**3,
        doc="Electricity intensity with respect to inlet flowrate of unit",
    )

    @self.Constraint(
        self.flowsheet().time,
        doc="Constraint for electricity consumption based on " "feed flowrate.",
    )
    def electricity_consumption(b, t):
        return b.electricity[t] == (
            b.energy_electric_flow_vol_inlet
            * pyunits.convert(
                b.get_inlet_flow(t), to_units=pyunits.m**3 / pyunits.hour
            )
        )

    self._fixed_perf_vars.append(self.energy_electric_flow_vol_inlet)
    self._perf_var_dict["Electricity Intensity"] = self.energy_electric_flow_vol_inlet


def pump_electricity(self, flow_rate):
    """
    Helper method for calculating electricity demand based on a
    pump flow equation.

    E[t] = (0.746 * Q[t] * H / (3960 * eta_pump * eta_motor))

    Here Q is the volumetric flowrate to be used to calculate electricity demand,
    H is the lift height for the pump and eta_pump and eta_motor are the
    efficiencies of the pump and motor respectively.

    Args:
        flow_rate - term to use for Q in the equation above.
    """
    _common(self)

    self.lift_height = Param(
        initialize=100, units=pyunits.feet, mutable=True, doc="Lift height for pump"
    )
    self.eta_pump = Param(
        initialize=0.9,
        units=pyunits.dimensionless,
        mutable=True,
        doc="Efficiency of pump",
    )
    self.eta_motor = Param(
        initialize=0.9,
        units=pyunits.dimensionless,
        mutable=True,
        doc="Efficiency of motor",
    )

    @self.Constraint(
        self.flowsheet().time,
        doc="Constraint for electricity consumption based on " "pump flowrate.",
    )
    def electricity_consumption(b, t):
        A = 3960 * pyunits.gallon * pyunits.foot / pyunits.minute / pyunits.horsepower
        return b.electricity[t] == pyunits.convert(
            flow_rate[t] * self.lift_height / (A * self.eta_pump * self.eta_motor),
            to_units=pyunits.kW,
        )
