##############################################################################
# Institute for the Design of Advanced Energy Systems Process Systems
# Engineering Framework (IDAES PSE Framework) Copyright (c) 2018-2019, 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".
##############################################################################
"""
Test for Caprese's module for NMPC.
"""

import pytest
from pyomo.environ import (Block, ConcreteModel,  Constraint, Expression,
                           Set, SolverFactory, Var, value, 
                           TransformationFactory, TerminationCondition)
from pyomo.network import Arc
from pyomo.common.collections import ComponentSet
from pyomo.dae.flatten import flatten_dae_components

from idaes.core import (FlowsheetBlock, MaterialBalanceType, EnergyBalanceType,
        MomentumBalanceType)
from idaes.core.util.model_statistics import (degrees_of_freedom, 
        activated_equalities_generator)
from idaes.apps.caprese.categorize import (
        categorize_dae_variables,
        categorize_dae_variables_and_constraints,
        )
from idaes.apps.caprese.common.config import VariableCategory as VC
import idaes.logger as idaeslog
from idaes.apps.caprese.examples.cstr_model import make_model

__author__ = "Robert Parker"


@pytest.mark.unit
def test_categorize_1():
    """
    This test categorizes the enzyme cstr "as intended."
    Volume is used as a measurement, and solvent flow rates are
    fixed, but otherwise equations are as expected.
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_accumulation[0, 'aq', :]),
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            model.fs.cstr.control_volume.volume[0],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_meas_set = ComponentSet(init_meas_list)
    # Solvent holdup is not a measurement; we measure volume instead
    init_meas_set.remove(
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'Solvent'])

    init_alg_list = [
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0


@pytest.mark.unit
def test_categorize_2():
    """
    In this instance, temperature is "measured" (used as an initial condition)
    instead of energy_holdup, conc[P] is measured instead of holdup[P],
    and accumulation[C] is measured instead of holdup[C].
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time
    t0 = time.first()

    material_holdup = model.fs.cstr.control_volume.material_holdup
    material_accumulation = model.fs.cstr.control_volume.material_accumulation
    energy_holdup = model.fs.cstr.control_volume.energy_holdup
    energy_accumulation = model.fs.cstr.control_volume.energy_accumulation
    conc_mol = model.fs.cstr.outlet.conc_mol

    # Specify temperature instead of energy holdup
    energy_holdup[t0, 'aq'].unfix()
    model.fs.cstr.outlet.temperature[t0].fix(300)

    # Specify C_P instead of holdup
    material_holdup[t0, 'aq', 'P'].unfix()
    conc_mol[t0, 'P'].fix(0.)

    # Specify accumulation of C instead of holdup
    # You might want to do this if you want to start at steady state,
    # but don't know the value of every variable at the steady state
    # you want to start at...
    material_holdup[t0, 'aq', 'C'].unfix()
    material_accumulation[t0, 'aq', 'C'].fix(0.)

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_accumulation[0, 'aq', :]),
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.outlet.temperature[0],
            model.fs.cstr.control_volume.volume[0],
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'S'],
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'E'],
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'S'],
            model.fs.cstr.control_volume.material_accumulation[0, 'aq', 'C'],
            model.fs.cstr.outlet.conc_mol[0,'P'],
            ]
    init_meas_set = ComponentSet(init_meas_list)
    # No need to remove solvent holdup here as it was not added to this list.

    init_alg_list = [
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0


@pytest.mark.unit
def test_categorize_3():
    """
    In this test we fix one of the differential variables.
    This is the case where somebody wants to run an isothermal
    CSTR.
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    CV = model.fs.cstr.control_volume
    CV.energy_holdup.fix(300)
    CV.energy_accumulation_disc_eq.deactivate()

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            *list(model.fs.cstr.control_volume.material_accumulation[0, 'aq', :]),
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            # Energy holdup has been fixed
            model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.control_volume.volume[0],
            *list(model.fs.cstr.control_volume.material_holdup[0, 'aq', :]),
            ]
    init_meas_set = ComponentSet(init_meas_list)
    # Solvent holdup is not a measurement; we measure volume instead
    init_meas_set.remove(
            model.fs.cstr.control_volume.material_holdup[0, 'aq', 'Solvent'])

    init_alg_list = [
        # Since energy_holdup is fixed, energy_accumulation is not
        # considered a derivative. Instead it is algebraic.
        model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0


@pytest.mark.unit
def test_categorize_4():
    """
    This tests categorization when a psuedo-steady state
    approximation is used. Energy accumulation and accumulation of E
    are fixed, the corresponding initial conditions unfixed, and
    the corresponding discretization equations deactivated.
    """
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    CV = model.fs.cstr.control_volume
    CV.energy_accumulation[:, 'aq'].fix(0.)
    CV.material_accumulation[:, 'aq', 'E'].fix(0.)
    CV.energy_holdup[0, 'aq'].unfix()
    CV.material_holdup[0, 'aq', 'E'].unfix()
    CV.energy_accumulation_disc_eq.deactivate()
    CV.material_accumulation_disc_eq.deactivate()

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            ]
    init_input_set = ComponentSet(init_input_list)

    init_deriv_list = [
            CV.material_accumulation[0, 'aq', 'C'],
            CV.material_accumulation[0, 'aq', 'S'],
            CV.material_accumulation[0, 'aq', 'P'],
            CV.material_accumulation[0, 'aq', 'Solvent'],
            ]
    init_deriv_set = ComponentSet(init_deriv_list)

    init_diff_list = [
            CV.material_holdup[0, 'aq', 'C'],
            CV.material_holdup[0, 'aq', 'S'],
            CV.material_holdup[0, 'aq', 'P'],
            CV.material_holdup[0, 'aq', 'Solvent'],
            ]
    init_diff_set = ComponentSet(init_diff_list)

    init_fixed_list = [
            model.fs.cstr.control_volume.energy_accumulation[0, 'aq'],
            CV.material_accumulation[0, 'aq', 'E'],
            model.fs.mixer.E_inlet.temperature[0],
            model.fs.mixer.S_inlet.temperature[0],
            *list(model.fs.mixer.E_inlet.conc_mol[0, :]),
            *list(model.fs.mixer.S_inlet.conc_mol[0, :]),
            model.fs.cstr.outlet.conc_mol[0, 'Solvent'],
            model.fs.mixer.outlet.conc_mol[0, 'Solvent'],
            ]
    init_fixed_set = ComponentSet(init_fixed_list)

    init_meas_list = [
            model.fs.cstr.control_volume.volume[0],
            CV.material_holdup[0,'aq','P'],
            CV.material_holdup[0,'aq','C'],
            CV.material_holdup[0,'aq','S'],
            ]
    init_meas_set = ComponentSet(init_meas_list)

    init_alg_list = [
        model.fs.cstr.control_volume.energy_holdup[0, 'aq'],
        CV.material_holdup[0, 'aq', 'E'],
        model.fs.cstr.outlet.flow_vol[0],
        model.fs.cstr.outlet.temperature[0],
        model.fs.cstr.inlet.flow_vol[0],
        model.fs.cstr.inlet.temperature[0],
        model.fs.mixer.outlet.flow_vol[0],
        model.fs.mixer.outlet.temperature[0],
        model.fs.cstr.control_volume.volume[0],
        *list(model.fs.cstr.control_volume.properties_out[0].flow_mol_comp[:]),
        *list(model.fs.cstr.inlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.properties_in[0].flow_mol_comp[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_generation[0, 'aq', :]),
        *list(model.fs.mixer.mixed_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.E_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.mixer.S_inlet_state[0].flow_mol_comp[:]),
        *list(model.fs.cstr.outlet.conc_mol[0, :]),
        *list(model.fs.mixer.outlet.conc_mol[0, :]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_coef[:]),
        *list(model.fs.cstr.control_volume.reactions[0].reaction_rate[:]),
        *list(model.fs.cstr.control_volume.rate_reaction_extent[0, :]),
        ]
    init_alg_set = ComponentSet(init_alg_list)
    # solvent outlet concentrations are not algebraic variables as
    # it is fixed.
    init_alg_set.remove(model.fs.cstr.outlet.conc_mol[0, 'Solvent'])
    init_alg_set.remove(model.fs.mixer.outlet.conc_mol[0, 'Solvent'])

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(dae_vars, time, init_input_list)

    input_vars = category_dict[VC.INPUT]
    diff_vars = category_dict[VC.DIFFERENTIAL]
    deriv_vars = category_dict[VC.DERIVATIVE]
    fixed_vars = category_dict[VC.FIXED]
    alg_vars = category_dict[VC.ALGEBRAIC]
    meas_vars = category_dict[VC.MEASUREMENT]

    assert len(input_vars) == len(init_input_set)
    for v in input_vars:
        assert v[0] in init_input_set

    assert len(deriv_vars) == len(init_deriv_set)
    for v in deriv_vars:
        assert v[0] in init_deriv_set

    assert len(diff_vars) == len(init_deriv_set)
    for v in diff_vars:
        assert v[0] in init_diff_set

    assert len(fixed_vars) == len(init_fixed_set)
    for v in fixed_vars:
        assert v[0] in init_fixed_set

    assert len(alg_vars) == len(init_alg_set)
    for v in alg_vars:
        assert v[0] in init_alg_set

    assert len(meas_vars) == len(init_meas_set)
    for v in meas_vars:
        assert v[0] in init_meas_set

    assert len(scalar_vars) == 0


@pytest.mark.unit
def test_categorize_error():
    model = make_model(horizon=1, ntfe=5, ntcp=2)
    time = model.fs.time

    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)

    # Add a dummy var to treat as an input.
    # This var is not in `dae_vars`, so it will not be located during
    # categorization, which should fail.
    model.dummy_var = Var(time)

    init_input_list = [
            model.fs.mixer.S_inlet.flow_vol[0],
            model.fs.mixer.E_inlet.flow_vol[0],
            model.dummy_var[0],
            ]

    with pytest.raises(RuntimeError, match=r'Not all inputs could be found'):
        category_dict = categorize_dae_variables(
                dae_vars,
                time,
                init_input_list,
                )

    # Re-run flattener. Now `dummy_var` should be included in `dae_vars`.
    scalar_vars, dae_vars = flatten_dae_components(model, time, ctype=Var)
    category_dict = categorize_dae_variables(
            dae_vars,
            time,
            init_input_list,
            )


if __name__ == '__main__':
    test_categorize_1()
    test_categorize_2()
    test_categorize_3()
    test_categorize_4()
    test_categorize_error()
