##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2020, by the
# software owners: The Regents of the University of California, through
# Lawrence Berkeley National Laboratory,  National Technology & Engineering
# Solutions of Sandia, LLC, Carnegie Mellon University, West Virginia
# University Research Corporation, et al. All rights reserved.
#
# Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
# license information, respectively. Both files are also available online
# at the URL "https://github.com/IDAES/idaes-pse".
##############################################################################
"""
Methods for setting up FpcTP as the state variables in a generic property
package
"""
from pyomo.environ import Constraint, Expression, NonNegativeReals, Var, value

from idaes.core import (MaterialFlowBasis,
                        MaterialBalanceType,
                        EnergyBalanceType)


def set_metadata(b):
    # The default metadata should be correct in this case, so no need to update
    pass


def define_state(b):
    # FpcTP contains full information on the phase equilibrium, so flash
    # calculations re not always needed
    b.always_flash = False

    # Get bounds if provided
    try:
        f_bounds = b.params.config.state_bounds["flow_mol_phase_comp"]
    except (KeyError, TypeError):
        f_bounds = (None, None)

    try:
        t_bounds = b.params.config.state_bounds["temperature"]
    except (KeyError, TypeError):
        t_bounds = (None, None)

    try:
        p_bounds = b.params.config.state_bounds["pressure"]
    except (KeyError, TypeError):
        p_bounds = (None, None)

    # Set an initial value for each state var
    if f_bounds == (None, None):
        # No bounds, default to 1
        f_init = 1
    elif f_bounds[1] is None and f_bounds[0] is not None:
        # Only lower bound, use lower bound + 10
        f_init = f_bounds[0] + 10
    elif f_bounds[1] is not None and f_bounds[0] is None:
        # Only upper bound, use half upper bound
        f_init = f_bounds[1]/2
    else:
        # Both bounds, use mid point
        f_init = (f_bounds[0] + f_bounds[1])/2

    if t_bounds == (None, None):
        # No bounds, default to 298.15
        t_init = 298.15
    elif t_bounds[1] is None and t_bounds[0] is not None:
        # Only lower bound, use lower bound + 10
        t_init = t_bounds[0] + 10
    elif t_bounds[1] is not None and t_bounds[0] is None:
        # Only upper bound, use half upper bound
        t_init = t_bounds[1]/2
    else:
        # Both bounds, use mid point
        t_init = (t_bounds[0] + t_bounds[1])/2

    if p_bounds == (None, None):
        # No bounds, default to 101325
        p_init = 101325
    elif p_bounds[1] is None and p_bounds[0] is not None:
        # Only lower bound, use lower bound + 10
        p_init = p_bounds[0] + 10
    elif p_bounds[1] is not None and p_bounds[0] is None:
        # Only upper bound, use half upper bound
        p_init = p_bounds[1]/2
    else:
        # Both bounds, use mid point
        p_init = (p_bounds[0] + p_bounds[1])/2

    # Add state variables
    b.flow_mol_phase_comp = Var(b.params._phase_component_set,
                                initialize=f_init,
                                domain=NonNegativeReals,
                                bounds=f_bounds,
                                doc='Phase-component molar flowrate')
    b.pressure = Var(initialize=p_init,
                     domain=NonNegativeReals,
                     bounds=p_bounds,
                     doc='State pressure')
    b.temperature = Var(initialize=t_init,
                        domain=NonNegativeReals,
                        bounds=t_bounds,
                        doc='State temperature')

    # Add supporting variables
    b.flow_mol = Expression(
        expr=sum(b.flow_mol_phase_comp[i]
                 for i in b.params._phase_component_set),
        doc="Total molar flowrate")

    def flow_mol_phase(b, p):
        return sum(b.flow_mol_phase_comp[p, j]
                   for j in b.params.component_list
                   if (p, j) in b.params._phase_component_set)
    b.flow_mol_phase = Expression(b.params.phase_list,
                                  rule=flow_mol_phase,
                                  doc='Phase molar flow rates')

    def rule_flow_mol_comp(b, j):
        return sum(b.flow_mol_phase_comp[p, j]
                   for p in b.params.phase_list
                   if (p, j) in b.params._phase_component_set)
    b.flow_mol_comp = Expression(b.params.component_list,
                                 rule=rule_flow_mol_comp,
                                 doc='Component molar flow rates')

    def mole_frac_comp(b, j):
        return (sum(b.flow_mol_phase_comp[p, j]
                    for p in b.params.phase_list
                    if (p, j) in b.params._phase_component_set) / b.flow_mol)
    b.mole_frac_comp = Expression(b.params.component_list,
                                  rule=mole_frac_comp,
                                  doc='Mixture mole fractions [-]')

    b.mole_frac_phase_comp = Var(
            b.params._phase_component_set,
            initialize=1/len(b.params.component_list),
            doc='Phase mole fractions [-]')

    def rule_mole_frac_phase_comp(b, p, j):
        return b.mole_frac_phase_comp[p, j] * b.flow_mol_phase[p] == \
            b.flow_mol_phase_comp[p, j]
    b.mole_frac_phase_comp_eq = Constraint(
        b.params._phase_component_set, rule=rule_mole_frac_phase_comp)

    def rule_phase_frac(b, p):
        if len(b.params.phase_list) == 1:
            return 1
        else:
            return b.flow_mol_phase[p] / b.flow_mol
    b.phase_frac = Expression(
        b.params.phase_list,
        rule=rule_phase_frac,
        doc='Phase fractions')

    # -------------------------------------------------------------------------
    # General Methods
    def get_material_flow_terms_FpcTP(p, j):
        """Create material flow terms for control volume."""
        return b.flow_mol_phase_comp[p, j]
    b.get_material_flow_terms = get_material_flow_terms_FpcTP

    def get_enthalpy_flow_terms_FpcTP(p):
        """Create enthalpy flow terms."""
        return b.flow_mol_phase[p] * b.enth_mol_phase[p]
    b.get_enthalpy_flow_terms = get_enthalpy_flow_terms_FpcTP

    def get_material_density_terms_FpcTP(p, j):
        """Create material density terms."""
        if j in b.params.component_list:
            return b.dens_mol_phase[p] * b.mole_frac_phase_comp[p, j]
        else:
            return 0
    b.get_material_density_terms = get_material_density_terms_FpcTP

    def get_energy_density_terms_FpcTP(p):
        """Create energy density terms."""
        return b.dens_mol_phase[p] * b.enth_mol_phase[p]
    b.get_energy_density_terms = get_energy_density_terms_FpcTP

    def default_material_balance_type_FpcTP():
        return MaterialBalanceType.componentTotal
    b.default_material_balance_type = default_material_balance_type_FpcTP

    def default_energy_balance_type_FpcTP():
        return EnergyBalanceType.enthalpyTotal
    b.default_energy_balance_type = default_energy_balance_type_FpcTP

    def get_material_flow_basis_FpcTP():
        return MaterialFlowBasis.molar
    b.get_material_flow_basis = get_material_flow_basis_FpcTP

    def define_state_vars_FpcTP():
        """Define state vars."""
        return {"flow_mol_phase_comp": b.flow_mol_phase_comp,
                "temperature": b.temperature,
                "pressure": b.pressure}
    b.define_state_vars = define_state_vars_FpcTP

    def define_display_vars_FpcTP():
        """Define display vars."""
        return {"Molar Flowrate": b.flow_mol_phase_comp,
                "Temperature": b.temperature,
                "Pressure": b.pressure}
    b.define_display_vars = define_display_vars_FpcTP


def state_initialization(b):
    for i in b.params._phase_component_set:
        b.mole_frac_phase_comp[i].value = value(
            b.flow_mol_phase_comp[i] / b.flow_mol_phase[i[0]])


do_not_initialize = []
