##############################################################################
# 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".
##############################################################################
"""
Tests for ControlVolume0D scaling.

Author: John Eslick
"""
import pytest
import pyomo.environ as pyo
from idaes.core import (
    ControlVolume1DBlock,
    MaterialBalanceType,
    EnergyBalanceType,
    MomentumBalanceType,
    FlowsheetBlock,
)
from idaes.core.util.testing import (
    PhysicalParameterTestBlock,
    ReactionParameterTestBlock,
)
from idaes.core.util import scaling as iscale


# -----------------------------------------------------------------------------
# Basic tests
@pytest.mark.unit
def test_basic_scaling():
    m = pyo.ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": False})
    m.fs.pp = PhysicalParameterTestBlock()
    # Set flag to include inherent reactions
    m.fs.pp._has_inherent_reactions = True

    m.fs.cv = ControlVolume1DBlock(default={
        "property_package": m.fs.pp,
        "transformation_method": "dae.finite_difference",
        "transformation_scheme": "BACKWARD",
        "finite_elements": 10})

    m.fs.cv.add_geometry()
    m.fs.cv.add_state_blocks(has_phase_equilibrium=False)

    m.fs.cv.add_material_balances(
        balance_type=MaterialBalanceType.componentTotal,
        has_phase_equilibrium=False)

    m.fs.cv.add_energy_balances(
        balance_type=EnergyBalanceType.enthalpyTotal)

    m.fs.cv.add_momentum_balances(
        balance_type=MomentumBalanceType.pressureTotal,
        has_pressure_change=True)

    m.fs.cv.apply_transformation()

    iscale.calculate_scaling_factors(m)

    # check scaling on select variables
    assert iscale.get_scaling_factor(m.fs.cv.area) == 1
    for (t, x), v in m.fs.cv.deltaP.items():
        assert iscale.get_scaling_factor(v) == 1040  # 10x the properties pressure scaling factor
    for t in m.fs.time:
        for x in m.fs.cv.length_domain:
            assert iscale.get_scaling_factor(
                m.fs.cv.properties[t, x].flow_vol) == 100

    # check scaling on mass, energy, and pressure balances.
    for c in m.fs.cv.material_balances.values():
        # this uses the minimum material flow term scale
        assert iscale.get_constraint_transform_applied_scaling_factor(c) == 112
    for c in m.fs.cv.enthalpy_balances.values():
        # this uses the minimum enthalpy flow term scale
        assert iscale.get_constraint_transform_applied_scaling_factor(c) == 110
    for c in m.fs.cv.pressure_balance.values():
        # This uses the inlet pressure scale
        assert iscale.get_constraint_transform_applied_scaling_factor(c) == 104

@pytest.mark.unit
def test_user_set_scaling():
    m = pyo.ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": False})
    m.fs.pp = PhysicalParameterTestBlock()
    m.fs.cv = ControlVolume1DBlock(default={
        "property_package": m.fs.pp,
        "transformation_method": "dae.finite_difference",
        "transformation_scheme": "BACKWARD",
        "finite_elements": 10})
    m.fs.cv.add_geometry()
    m.fs.cv.add_state_blocks(has_phase_equilibrium=False)
    m.fs.cv.add_material_balances(
        balance_type=MaterialBalanceType.componentTotal,
        has_phase_equilibrium=False)
    m.fs.cv.add_energy_balances(
        balance_type=EnergyBalanceType.enthalpyTotal,
        has_heat_transfer=True,
        has_work_transfer=True)
    # add momentum balance
    m.fs.cv.add_momentum_balances(
        balance_type=MomentumBalanceType.pressureTotal,
        has_pressure_change=True)

    m.fs.cv.apply_transformation()

    # The scaling factors used for this test were selected to be easy values to
    # test, they do not represent typical scaling factors.
    iscale.set_scaling_factor(m.fs.cv.heat, 11)
    iscale.set_scaling_factor(m.fs.cv.work, 12)
    iscale.set_scaling_factor(m.fs.cv.heat[0, 0], 17)
    iscale.calculate_scaling_factors(m)
    # Make sure the heat and work scaling factors are set and not overwritten
    # by the defaults in calculate_scaling_factors
    assert iscale.get_scaling_factor(m.fs.cv.heat) == 11
    assert iscale.get_scaling_factor(m.fs.cv.work) == 12
    # Make sure scaling factor is propagated but does not overwrite index specific factor
    for t in m.fs.time:
        for x in m.fs.cv.length_domain:
            if t == 0 and x == 0:
                assert iscale.get_scaling_factor(m.fs.cv.heat[t, x]) == 17
            else:
                assert iscale.get_scaling_factor(m.fs.cv.heat[t, x]) == 11
            assert iscale.get_scaling_factor(m.fs.cv.work[t, x]) == 12


@pytest.mark.unit
def test_full_auto_scaling():
    m = pyo.ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": False})
    m.fs.pp = PhysicalParameterTestBlock()
    m.fs.rp = ReactionParameterTestBlock(default={"property_package": m.fs.pp})
    m.fs.cv = ControlVolume1DBlock(default={
        "property_package": m.fs.pp,
        "reaction_package": m.fs.rp,
        "transformation_method": "dae.finite_difference",
        "transformation_scheme": "BACKWARD",
        "finite_elements": 10})
    m.fs.cv.add_geometry()
    m.fs.cv.add_state_blocks(has_phase_equilibrium=True)
    m.fs.cv.add_reaction_blocks(has_equilibrium=True)

    m.fs.cv.add_material_balances(
        balance_type=MaterialBalanceType.componentTotal,
        has_rate_reactions=True,
        has_equilibrium_reactions=True,
        has_phase_equilibrium=True,
        has_mass_transfer=True)

    m.fs.cv.add_energy_balances(
        balance_type=EnergyBalanceType.enthalpyTotal,
        has_heat_of_reaction=True,
        has_heat_transfer=True,
        has_work_transfer=True,
        has_enthalpy_transfer=True)

    m.fs.cv.add_momentum_balances(
        balance_type=MomentumBalanceType.pressureTotal,
        has_pressure_change=True)

    m.fs.cv.apply_transformation()

    iscale.calculate_scaling_factors(m)

    # check that all variables have scaling factors
    unscaled_var_list = list(iscale.unscaled_variables_generator(m))
    # Unscaled variables are:
    # rate_reaction_extent (2 reactions, 11 spatial points)
    # equilibrium_reaction_extent  (2 reactions, 11 spatial points)
    assert len(unscaled_var_list) == 44
    # check that all constraints have been scaled
    unscaled_constraint_list = list(iscale.unscaled_constraints_generator(m))
    assert len(unscaled_constraint_list) == 0


@pytest.mark.unit
def test_full_auto_scaling_dynamic():
    m = pyo.ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": True, "time_units": pyo.units.s})
    m.fs.pp = PhysicalParameterTestBlock()
    m.fs.rp = ReactionParameterTestBlock(default={"property_package": m.fs.pp})
    m.fs.cv = ControlVolume1DBlock(default={
        "property_package": m.fs.pp,
        "reaction_package": m.fs.rp,
        "transformation_method": "dae.finite_difference",
        "transformation_scheme": "BACKWARD",
        "finite_elements": 10})
    m.fs.cv.add_geometry()
    m.fs.cv.add_state_blocks(has_phase_equilibrium=True)
    m.fs.cv.add_reaction_blocks(has_equilibrium=True)

    m.fs.cv.add_material_balances(
        balance_type=MaterialBalanceType.componentTotal,
        has_rate_reactions=True,
        has_equilibrium_reactions=True,
        has_phase_equilibrium=True,
        has_mass_transfer=True)

    m.fs.cv.add_energy_balances(
        balance_type=EnergyBalanceType.enthalpyTotal,
        has_heat_of_reaction=True,
        has_heat_transfer=True,
        has_work_transfer=True,
        has_enthalpy_transfer=True)

    m.fs.cv.add_momentum_balances(
        balance_type=MomentumBalanceType.pressureTotal,
        has_pressure_change=True)

    m.fs.cv.apply_transformation()
    m.discretizer = pyo.TransformationFactory('dae.finite_difference')
    m.discretizer.apply_to(m, nfe=3, wrt=m.fs.time, scheme="BACKWARD")

    iscale.calculate_scaling_factors(m)

    # check that all variables have scaling factors
    unscaled_var_list = list(iscale.unscaled_variables_generator(m))
    # Unscaled variables are:
    # rate_reaction_extent (2 reactions, 44 time & space points)
    # equilibrium_reaction_extent  (2 reactions, 44 time & space points)
    assert len(unscaled_var_list) == 176
    # check that all constraints have been scaled
    unscaled_constraint_list = list(iscale.unscaled_constraints_generator(m))
    assert len(unscaled_constraint_list) == 0


@pytest.mark.unit
def test_full_auto_scaling_mbtype_phase():
    m = pyo.ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": False})
    m.fs.pp = PhysicalParameterTestBlock()
    m.fs.rp = ReactionParameterTestBlock(default={"property_package": m.fs.pp})
    m.fs.cv = ControlVolume1DBlock(default={
        "property_package": m.fs.pp,
        "reaction_package": m.fs.rp,
        "transformation_method": "dae.finite_difference",
        "transformation_scheme": "BACKWARD",
        "finite_elements": 10})
    m.fs.cv.add_geometry()
    m.fs.cv.add_state_blocks(has_phase_equilibrium=True)
    m.fs.cv.add_reaction_blocks(has_equilibrium=True)

    m.fs.cv.add_material_balances(
        balance_type=MaterialBalanceType.componentPhase,
        has_rate_reactions=True,
        has_equilibrium_reactions=True,
        has_phase_equilibrium=True,
        has_mass_transfer=True)

    m.fs.cv.add_energy_balances(
        balance_type=EnergyBalanceType.enthalpyTotal,
        has_heat_of_reaction=True,
        has_heat_transfer=True,
        has_work_transfer=True,
        has_enthalpy_transfer=True)

    m.fs.cv.add_momentum_balances(
        balance_type=MomentumBalanceType.pressureTotal,
        has_pressure_change=True)

    m.fs.cv.apply_transformation()

    iscale.calculate_scaling_factors(m)

    # check that all variables have scaling factors
    unscaled_var_list = list(iscale.unscaled_variables_generator(m))
    # Unscaled variables are:
    # rate_reaction_extent (2 reactions, 11 spatial points)
    # equilibrium_reaction_extent  (2 reactions, 11 spatial points)
    # phase_equilibrium_generation  (2 reactions, 11 spatial points)
    assert len(unscaled_var_list) == 66
    # check that all constraints have been scaled
    unscaled_constraint_list = list(iscale.unscaled_constraints_generator(m))
    assert len(unscaled_constraint_list) == 0


@pytest.mark.unit
def test_full_auto_scaling_mbtype_element():
    m = pyo.ConcreteModel()
    m.fs = FlowsheetBlock(default={"dynamic": True, "time_units": pyo.units.s})
    m.fs.pp = PhysicalParameterTestBlock()
    m.fs.rp = ReactionParameterTestBlock(default={"property_package": m.fs.pp})
    m.fs.cv = ControlVolume1DBlock(default={
        "property_package": m.fs.pp,
        "reaction_package": m.fs.rp,
        "transformation_method": "dae.finite_difference",
        "transformation_scheme": "BACKWARD",
        "finite_elements": 10})
    m.fs.cv.add_geometry()
    m.fs.cv.add_state_blocks(has_phase_equilibrium=False)
    m.fs.cv.add_reaction_blocks(has_equilibrium=False)

    m.fs.cv.add_total_element_balances(has_mass_transfer=True)

    m.fs.cv.apply_transformation()
    m.discretizer = pyo.TransformationFactory('dae.finite_difference')
    m.discretizer.apply_to(m, nfe=3, wrt=m.fs.time, scheme="BACKWARD")

    iscale.calculate_scaling_factors(m)

    # check that all variables have scaling factors
    unscaled_var_list = list(iscale.unscaled_variables_generator(m))
    assert len(unscaled_var_list) == 0
    # check that all constraints have been scaled
    unscaled_constraint_list = list(iscale.unscaled_constraints_generator(m))
    assert len(unscaled_constraint_list) == 0
