#################################################################################
# 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 miscellaneous utility functions for use in IDAES models.
"""
from enum import Enum
import sys

import pyomo.environ as pyo
from pyomo.common.config import ConfigBlock
from pyomo.core.expr.visitor import _ToStringVisitor
from pyomo.core.base.expression import ExpressionData

import idaes.logger as idaeslog

_log = idaeslog.getLogger(__name__)


# Author: Andrew Lee
def add_object_reference(self, local_name, remote_object):
    """
    Method to create a reference in the local model to a remote Pyomo object.
    This method should only be used where Pyomo Reference objects are not
    suitable (such as for referencing scalar Pyomo objects where the None
    index is undesirable).

    Args:
        local_name : name to use for local reference (str)
        remote_object : object to make a reference to

    Returns:
        None
    """
    try:
        object.__setattr__(self, local_name, remote_object)
    except AttributeError:
        raise AttributeError(
            "{} failed to construct reference to {} - remote "
            "object does not exist.".format(self.name, remote_object)
        )


# Author: Jaffer Ghouse
def extract_data(data_dict):
    """
    General method that returns a rule to extract data from a python
    dictionary. This method allows the param block to have a database for
    a parameter but extract a subset of this data to initialize a Pyomo
    param object.
    """

    def _rule_initialize(m, *args):
        if len(args) > 1:
            return data_dict[args]
        else:
            return data_dict[args[0]]

    return _rule_initialize


def set_param_from_config(b, param, config=None, index=None):
    """
    Utility method to set parameter value from a config block. This allows for
    converting units if required. This method directly sets the value of the
    parameter.

    This method supports three forms for defining the parameter value:

    1. a 2-tuple of the form (value, units) where units are the units that
    the value are defined in
    2. a float where the float is assumed to be the value of the parameter
    value in the base units of the property package
    3. a Python Class which has a get_parameter_value method which will
    return a 2-tuple of (value, units) based on a lookup of the parameter name

    Args:
        b - block on which parameter and config block are defined
        param - name of parameter as str. Used to find param and config arg
        units - units of param object (used if conversion required)
        config - (optional) config block to get parameter data from. If
                unset, assumes b.config.
        index - (optional) used for pure component properties where a single
                property may have multiple parameters associated with it.

    Returns:
        None
    """
    if config is None:
        try:
            config = b.config
        except AttributeError:
            raise AttributeError(
                "{} - set_param_from_config method was not provided with a "
                "config argument, but no default Config block exists. Please "
                "specify the Config block to use via the config argument.".format(
                    b.name
                )
            )

    # Check that config is an instance of a Config Block
    if not isinstance(config, ConfigBlock):
        raise TypeError(
            "{} - set_param_from_config - config argument provided is not an "
            "instance of a Config Block.".format(b.name)
        )

    if index is None:
        try:
            param_obj = getattr(b, param)
        except AttributeError:
            raise AttributeError(
                "{} - set_param_from_config method was provided with param "
                "argument {}, but no attribute of that name exists.".format(
                    b.name, param
                )
            )

        try:
            p_data = config.parameter_data[param]
        except (KeyError, AttributeError):
            raise KeyError(
                "{} - set_param_from_config method was provided with param "
                "argument {}, but the config block does not contain a "
                "value for this parameter.".format(b.name, param)
            )
    else:
        try:
            param_obj = getattr(b, param + "_" + index)
        except AttributeError:
            raise AttributeError(
                f"{b.name} - set_param_from_config method was provided with param and"
                f" index arguments {param} {index}, but no attribute with that "
                f"combination ({param}_{index}) exists."
            )

        try:
            p_data = config.parameter_data[param][index]
        except (KeyError, AttributeError):
            raise KeyError(
                "{} - set_param_from_config method was provided with param and"
                " index arguments {} {}, but the config block does not contain"
                " a value for this parameter and index.".format(b.name, param, index)
            )

    # Check to see if p_data is callable, and if so, try to call the
    # get_parameter_value method to get 2-tuple
    if hasattr(p_data, "get_parameter_value"):
        p_data = p_data.get_parameter_value(b.local_name, param)

    if isinstance(p_data, tuple):
        if p_data[1] is None:
            p_val = p_data[0] * pyo.units.dimensionless
        else:
            p_val = p_data[0] * p_data[1]
    else:
        _log.debug(
            "{} no units provided for parameter {} - assuming default "
            "units".format(b.name, param)
        )
        p_val = p_data

    param_obj.set_value(p_val)


class StrEnum(str, Enum):
    """
    Multiple inheritance string-Enum for representing Enums with string values
    """

    def __str__(self):
        return str(self.value)


class _ToExprStringVisitor(_ToStringVisitor):
    """
    Derived version of the Pyomo _ToStringVisitor class
    which checks for named Expressions in the expression tree
    and prints their name instead of expanding the expression tree.
    """

    def visiting_potential_leaf(self, node):
        # Check if node is a named Expression
        if isinstance(node, ExpressionData):
            # If it is, return the name of the Expression
            return True, node.name

        # Otherwise, continue descending as normal
        return super().visiting_potential_leaf(node)


def compact_expression_to_string(expr):
    """Return a compact string representation of an expression.

    Unlike the normal Pyomo string representations, this function will
    identify any Expressions that appear within the expression tree and
    print the name of these rather than expanding the expression tree.

    Args:
        expr: ExpressionBase. The root node of an expression tree.

    Returns:
        A string representation for the expression.
    """
    # Create and execute the visitor pattern
    #
    visitor = _ToExprStringVisitor(verbose=False, smap=None)
    return visitor.dfs_postorder_stack(expr)


def print_compact_form(expr, stream=None):
    """
    Writes a compact string representation of the given component or
    expression to the specified stream.

    Unlike the normal Pyomo string representations, this function will
    identify any Expressions that appear within the expression tree and
    print the name of these rather than expanding the expression tree.

    Args:
        expr: component or expression to print
        stream: StringIO object to write to. Default is stdout.

    Returns:
        None
    """
    if stream is None:
        stream = sys.stdout

    if hasattr(expr, "expr"):
        # We have a Constraint or Expression
        # We want to print the expression, not the object name
        expr = expr.expr

    stream.write(compact_expression_to_string(expr))
