#################################################################################
# The Institute for the Design of Advanced Energy Systems Integrated Platform
# Framework (IDAES IP) was produced under the DOE Institute for the
# Design of Advanced Energy Systems (IDAES).
#
# Copyright (c) 2018-2024 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.md and LICENSE.md
# for full copyright and license information.
#################################################################################
"""
This module contains utility functions for use in testing IDAES models.
"""
# These are methods for testing, and break a few conventions
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring
# pylint: disable=protected-access

__author__ = "Andrew Lee"


import os
from typing import Callable, Union

from pyomo.environ import Constraint, Set, units, Var
from pyomo.common.config import ConfigBlock
from pyomo.common import Executable
from pyomo.common.dependencies import attempt_import

from idaes.core import (
    declare_process_block_class,
    PhysicalParameterBlock,
    StateBlock,
    StateBlockData,
    ReactionParameterBlock,
    ReactionBlockBase,
    ReactionBlockDataBase,
    MaterialFlowBasis,
    MaterialBalanceType,
    EnergyBalanceType,
    Component,
    Phase,
)

from idaes.core.util.model_statistics import (
    degrees_of_freedom,
    fixed_variables_set,
    activated_constraints_set,
)
import idaes.logger as idaeslog

_log = idaeslog.getLogger(__name__)


def initialization_tester(m, dof=0, unit=None, **init_kwargs):
    """
    A method to test initialization methods on IDAES models. This method is
    designed to be used as part of the tests for most models.

    This method checks that the initialization methods runs as expected
    and that the state of the model (active/inactive and fixed/unfixed)
    remains the same.

    This method also add some dummy constraints to the model and deactivates
    them to make sure that the initialization does not affect their status.

    Args:
        m: a Concrete model which contains a flowsheet and a model named unit
            (i.e. m.fs.unit) which will be initialized
        dof: expected degrees of freedom during initialization, default=0
        unit: unit object to test, if None assume m.fs.unit, default='None'
        init_kwargs: model specific arguments to pass to initialize method
                     (e.g. initial guesses for states)

    Returns:
        None

    Raises:
        AssertionErrors if an issue is found
    """
    if unit is None:
        unit = m.fs.unit
    # Add some extra constraints and deactivate them to make sure
    # they remain deactivated
    # Test both indexed and unindexed constraints
    unit.__dummy_var = Var()
    unit.__dummy_equality = Constraint(expr=unit.__dummy_var == 5)
    unit.__dummy_inequality = Constraint(expr=unit.__dummy_var <= 10)
    unit.__dummy_set = Set(initialize=[1])

    def deq_idx(b, i):
        return unit.__dummy_var == 5

    unit.__dummy_equality_idx = Constraint(unit.__dummy_set, rule=deq_idx)

    def dieq_idx(b, i):
        return unit.__dummy_var <= 10

    unit.__dummy_inequality_idx = Constraint(unit.__dummy_set, rule=dieq_idx)

    unit.__dummy_equality.deactivate()
    unit.__dummy_inequality.deactivate()
    unit.__dummy_equality_idx[1].deactivate()
    unit.__dummy_inequality_idx[1].deactivate()

    orig_fixed_vars = fixed_variables_set(m)
    orig_act_consts = activated_constraints_set(m)

    unit.initialize(**init_kwargs)

    assert degrees_of_freedom(m) == dof

    fin_fixed_vars = fixed_variables_set(m)
    fin_act_consts = activated_constraints_set(m)

    assert len(fin_act_consts) == len(orig_act_consts)
    assert len(fin_fixed_vars) == len(orig_fixed_vars)

    for c in fin_act_consts:
        assert c in orig_act_consts
    for v in fin_fixed_vars:
        assert v in orig_fixed_vars

    # Check dummy constraints and clean up
    assert not unit.__dummy_equality.active
    assert not unit.__dummy_inequality.active
    assert not unit.__dummy_equality_idx[1].active
    assert not unit.__dummy_inequality_idx[1].active

    unit.del_component(unit.__dummy_inequality)
    unit.del_component(unit.__dummy_equality)
    unit.del_component(unit.__dummy_inequality_idx)
    unit.del_component(unit.__dummy_equality_idx)
    unit.del_component(unit.__dummy_var)
    unit.del_component(unit.__dummy_set)


# -----------------------------------------------------------------------------
# Define some generic PhysicalBlock and ReactionBlock classes for testing
@declare_process_block_class("PhysicalParameterTestBlock")
class _PhysicalParameterBlock(PhysicalParameterBlock):
    def build(self):
        super(_PhysicalParameterBlock, self).build()

        self.p1 = Phase()
        self.p2 = Phase()

        self.c1 = Component()
        self.c2 = Component()

        self.phase_equilibrium_idx = Set(initialize=["e1", "e2"])
        self.element_list = Set(initialize=["H", "He", "Li"])
        self.element_comp = {
            "c1": {"H": 1, "He": 2, "Li": 3},
            "c2": {"H": 4, "He": 5, "Li": 6},
        }

        self.phase_equilibrium_list = {
            "e1": ["c1", ("p1", "p2")],
            "e2": ["c2", ("p1", "p2")],
        }

        # Add inherent reactions for use when needed
        self.inherent_reaction_idx = Set(initialize=["i1", "i2"])
        self.inherent_reaction_stoichiometry = {
            ("i1", "p1", "c1"): 1,
            ("i1", "p1", "c2"): 1,
            ("i1", "p2", "c1"): 1,
            ("i1", "p2", "c2"): 1,
            ("i2", "p1", "c1"): 1,
            ("i2", "p1", "c2"): 1,
            ("i2", "p2", "c1"): 1,
            ("i2", "p2", "c2"): 1,
        }

        # Attribute to switch flow basis for testing
        self.basis_switch = 1
        self.default_balance_switch = 1

        self._state_block_class = StateBlockForTesting

        self.set_default_scaling("flow_vol", 100)
        self.set_default_scaling("flow_mol", 101)
        self.set_default_scaling("flow_mol_phase_comp", 102)
        self.set_default_scaling("test_var", 103)
        self.set_default_scaling("pressure", 104)
        self.set_default_scaling("temperature", 105)
        self.set_default_scaling("enth_mol", 106)
        self.set_default_scaling("gibbs_mol_phase_comp", 107)
        self.set_default_scaling("entr_mol", 108)
        self.set_default_scaling("mole_frac_phase_comp", 109)
        self.set_default_scaling("enthalpy_flow", 110)
        self.set_default_scaling("energy_dens", 111)
        self.set_default_scaling("material_flow_mol", 112)
        self.set_default_scaling("material_dens_mol", 113)
        self.set_default_scaling("material_flow_mass", 114)
        self.set_default_scaling("material_dens_mass", 115)
        self.set_default_scaling("material_flow_dimensionless", 116)
        self.set_default_scaling("material_dens_dimensionless", 117)

    @classmethod
    def define_metadata(cls, obj):
        obj.add_properties(
            {
                "enth_mol": {"method": None},
            }
        )
        obj.add_default_units(
            {
                "time": units.s,
                "length": units.m,
                "mass": units.kg,
                "amount": units.mol,
                "temperature": units.K,
            }
        )


class SBlockBase(StateBlock):
    def initialize(
        blk,
        outlvl=idaeslog.NOTSET,
        optarg=None,
        solver=None,
        hold_state=False,
        state_args=None,
    ):
        for k in blk.values():
            k.init_test = True
            k.hold_state = hold_state

    def release_state(blk, flags=None, outlvl=idaeslog.NOTSET):
        for k in blk.values():
            k.hold_state = not k.hold_state


@declare_process_block_class("StateBlockForTesting", block_class=SBlockBase)
class StateTestBlockData(StateBlockData):
    CONFIG = ConfigBlock(implicit=True)

    def build(self):
        super(StateTestBlockData, self).build()

        self.flow_vol = Var(initialize=20, units=units.m**3 / units.s)
        self.flow_mol = Var(
            initialize=14,
            units=units.mol / units.s,
        )
        self.flow_mol_phase_comp = Var(
            self.params.phase_list,
            self.params.component_list,
            initialize=2,
            units=units.mol / units.s,
        )
        self.test_var = Var(initialize=1)
        self.enthalpy_flow = Var(initialize=1, units=units.J / units.s)
        self.energy_dens = Var(initialize=1, units=units.J / units.m**3)
        self.material_flow_mol = Var(initialize=1, units=units.mol / units.s)
        self.material_dens_mol = Var(initialize=1, units=units.mol / units.m**3)
        self.material_flow_mass = Var(initialize=1, units=units.kg / units.s)
        self.material_dens_mass = Var(initialize=1, units=units.kg / units.m**3)
        self.material_flow_dimensionless = Var(initialize=1, units=units.dimensionless)
        self.material_dens_dimensionless = Var(initialize=1, units=units.dimensionless)
        self.pressure = Var(initialize=1e5, units=units.Pa)
        self.temperature = Var(initialize=300, units=units.K)

        self.cp_mol = Var(initialize=100, units=units.J / units.mol / units.K)
        self.enth_mol = Var(initialize=10000, units=units.J / units.mol)

        self.gibbs_mol_phase_comp = Var(
            self.params.phase_list,
            self.params.component_list,
            initialize=50,
            units=units.J / units.mol,
        )
        self.entr_mol = Var(initialize=1000, units=units.J / units.mol / units.K)

        self.mole_frac_phase_comp = Var(
            self.params.phase_list, self.params.component_list, initialize=0.5
        )

    def get_material_flow_terms(b, p, j):
        if b.config.parameters.basis_switch == 2:
            return b.material_flow_mass
        elif b.config.parameters.basis_switch == 1:
            return b.material_flow_mol
        else:
            return b.material_flow_dimensionless

    def get_material_density_terms(b, p, j):
        if b.config.parameters.basis_switch == 2:
            return b.material_dens_mass
        elif b.config.parameters.basis_switch == 1:
            return b.material_dens_mol
        else:
            return b.material_dens_dimensionless

    def get_enthalpy_flow_terms(b, p):
        return b.enthalpy_flow

    def get_energy_density_terms(b, p):
        return b.energy_dens

    def model_check(self):
        self.check = True

    def get_material_flow_basis(b):
        if b.config.parameters.basis_switch == 1:
            return MaterialFlowBasis.molar
        elif b.config.parameters.basis_switch == 2:
            return MaterialFlowBasis.mass
        else:
            return MaterialFlowBasis.other

    def default_material_balance_type(self):
        if self.params.default_balance_switch == 1:
            return MaterialBalanceType.componentPhase
        else:
            raise NotImplementedError

    def default_energy_balance_type(self):
        if self.params.default_balance_switch == 1:
            return EnergyBalanceType.enthalpyTotal
        else:
            raise NotImplementedError

    def define_state_vars(self):
        return {
            "component_flow_phase": self.flow_mol_phase_comp,
            "temperature": self.temperature,
            "pressure": self.pressure,
        }


@declare_process_block_class("ReactionParameterTestBlock")
class _ReactionParameterBlock(ReactionParameterBlock):
    def build(self):
        super(_ReactionParameterBlock, self).build()

        self.rate_reaction_idx = Set(initialize=["r1", "r2"])
        self.equilibrium_reaction_idx = Set(initialize=["e1", "e2"])

        self.rate_reaction_stoichiometry = {
            ("r1", "p1", "c1"): 1,
            ("r1", "p1", "c2"): 1,
            ("r1", "p2", "c1"): 1,
            ("r1", "p2", "c2"): 1,
            ("r2", "p1", "c1"): 1,
            ("r2", "p1", "c2"): 1,
            ("r2", "p2", "c1"): 1,
            ("r2", "p2", "c2"): 1,
        }
        self.equilibrium_reaction_stoichiometry = {
            ("e1", "p1", "c1"): 1,
            ("e1", "p1", "c2"): 1,
            ("e1", "p2", "c1"): 1,
            ("e1", "p2", "c2"): 1,
            ("e2", "p1", "c1"): 1,
            ("e2", "p1", "c2"): 1,
            ("e2", "p2", "c1"): 1,
            ("e2", "p2", "c2"): 1,
        }

        self._reaction_block_class = ReactionBlock

        # Attribute to switch flow basis for testing
        self.basis_switch = 1

        self.set_default_scaling("reaction_rate", 101, "r1")
        self.set_default_scaling("reaction_rate", 102, "r2")

    @classmethod
    def define_metadata(cls, obj):
        obj.add_default_units(
            {
                "time": units.s,
                "length": units.m,
                "mass": units.kg,
                "amount": units.mol,
                "temperature": units.K,
            }
        )

    @classmethod
    def get_required_properties(cls):
        return {}


class RBlockBase(ReactionBlockBase):
    def initialize(
        blk, outlvl=idaeslog.NOTSET, optarg=None, solver=None, state_vars_fixed=False
    ):
        for k in blk.values():
            k.init_test = True


@declare_process_block_class("ReactionBlock", block_class=RBlockBase)
class ReactionBlockData(ReactionBlockDataBase):
    CONFIG = ConfigBlock(implicit=True)

    def build(self):
        super(ReactionBlockData, self).build()

        self.reaction_rate = Var(["r1", "r2"], units=units.mol / units.m**3 / units.s)

        self.dh_rxn = {
            "r1": 10 * units.J / units.mol,
            "r2": 20 * units.J / units.mol,
            "e1": 30 * units.J / units.mol,
            "e2": 40 * units.J / units.mol,
        }

    def model_check(self):
        self.check = True

    def get_reaction_rate_basis(b):
        if b.config.parameters.basis_switch == 1:
            return MaterialFlowBasis.molar
        elif b.config.parameters.basis_switch == 2:
            return MaterialFlowBasis.mass
        else:
            return MaterialFlowBasis.other


def _enable_scip_solver_for_testing(
    name: str = "scip",
) -> Union[Callable[[], None], None]:
    ampl_module_scip, is_available = attempt_import("ampl_module_scip")
    if not is_available:
        _log.warning(
            "ampl_module_scip must be installed to enable SCIP solver for testing"
        )
        return None

    new_path_entry = ampl_module_scip.bin_dir
    # TODO prepending new_path_entry to PATH means that the "testing" SCIP will "shadow"
    # existing executable directories already on PATH
    # this behavior would match the (implied) semantics of this function, i.e.
    # "enable SCIP solver used for testing (at the expense of other non-testing SCIP executables)"
    os.environ["PATH"] = os.pathsep.join([new_path_entry, os.environ["PATH"]])
    Executable.rehash()

    def remove_from_path():
        path_as_list = os.environ["PATH"].split(os.pathsep)

        try:
            path_as_list.remove(new_path_entry)
        except ValueError:
            # not in PATH
            pass
        else:
            os.environ["PATH"] = os.pathsep.join(path_as_list)
            Executable.rehash()

    return remove_from_path
