##############################################################################
# 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".
##############################################################################
try:
    from pyomo.core.expr.sympy_tools import (
        _prod, _sum, _functionMap, _operatorMap, _pyomo_operator_map)
except ImportError:
    from pyomo.core.base.symbolic import (
        _prod, _sum, _functionMap, _operatorMap, _pyomo_operator_map)

from pyomo.environ import ExternalFunction, Var, Expression, value
from pyomo.core.base.constraint import _ConstraintData, Constraint
from pyomo.core.base.expression import _ExpressionData
from pyomo.core.base.block import _BlockData
from pyomo.core.expr.visitor import StreamBasedExpressionVisitor
from pyomo.core.expr.numeric_expr import (ExternalFunctionExpression,
                                          ExpressionBase)
from pyomo.core.expr import current as EXPR, native_types
from pyomo.core.kernel.component_map import ComponentMap

import sympy
from IPython.display import display, Markdown

import logging
import re

_log = logging.getLogger(__name__)

#TODO<jce> Look into things like sum operator and template expressions

def _add_latex_subscripts(x, s):
    if not "_" in x:
        return "{}_{{ {} }}".format(x, s)
    else:
        return re.sub(r"^(.+_)({.+}|.)", "\\1{{ \\2 ,{} }}".format(s), x)

def deduplicate_symbol(x, v, used):
    """
    Check if x is a duplicated LaTeX symbol if so add incrementing Di subscript

    Args:
        x: symbol string
        v: pyomo object
        used: dictionary of pyomo objects with symbols as keys

    Returns:
        Returns a unique symbol.  If x was not in used keys, returns x,
        otherwise adds exponents to make it unique.
    """
    y = x
    k = 1
    while x in used and id(used[x]) != id(v):
        x = _add_latex_subscripts(y, "D{}".format(k))
        k += 1
        if k > 1000:
            # Either the symbol string is not updating or there are lots of dupes
            break
    used[x] = v
    return x

class PyomoSympyBimap(object):
    """
    This is based on the class of the same name in pyomo.core.base.symbolic, but
    it adds mapping latex symbols to the sympy symbols. This will get you pretty
    equations when using sympy's LaTeX writer.
    """
    def __init__(self):
        self.pyomo2sympy = ComponentMap()
        self.parent_symbol = ComponentMap()
        self.sympy2pyomo = {}
        self.sympy2latex = {}
        self.used = {}
        self.i_var = 0
        self.i_expr = 0
        self.i_func = 0
        self.i = 0

    def getPyomoSymbol(self, sympy_object, default=None):
        return self.sympy2pyomo.get(sympy_object, default)

    def getSympySymbol(self, pyomo_object):
        if pyomo_object in self.pyomo2sympy:
            return self.pyomo2sympy[pyomo_object]
        else:
            return self._add_sympy(pyomo_object)

    def sympyVars(self):
        return self.sympy2pyomo.keys()

    def _add_sympy(self, pyomo_object):
        parent_object = pyomo_object.parent_component()
        if isinstance(parent_object, Var):
            sympy_class = sympy.Symbol
            base_name = "x"
            i = self.i_var
            self.i_var += 1
        elif isinstance(parent_object, Expression):
            sympy_class = sympy.Symbol
            base_name = "u"
            i = self.i_expr
            self.i_expr += 1
        elif isinstance(parent_object, ExternalFunction):
            sympy_class = sympy.Function
            base_name = "func"
            i = self.i_func
            self.i_func += 1
        else:
            raise Exception("Should be Var, Exression, or ExternalFunction")

        if parent_object.is_indexed() and parent_object in self.parent_symbol:
            x = self.parent_symbol[parent_object][0]
            latex_symbol = self.parent_symbol[parent_object][1]
        else:
            x = "{}_{}".format(base_name, i)
            base_latex = getattr(parent_object, "latex_symbol", None)
            if base_latex is not None:
                latex_symbol = deduplicate_symbol(base_latex, parent_object,
                                                  self.used)
            else:
                latex_symbol = x
        if parent_object.is_indexed():
            if parent_object not in self.parent_symbol:
                self.parent_symbol[parent_object] = (x, latex_symbol)
            idx = pyomo_object.index()
            idxl = idx
            if isinstance(idx, tuple):
                idxl = ",".join(idx)
                idx = "|".join(idx)
            x = "{}[{}]".format(x, idx)
            latex_symbol = _add_latex_subscripts(latex_symbol, idxl)
        sympy_obj = sympy_class(x)
        self.pyomo2sympy[pyomo_object] = sympy_obj
        self.sympy2pyomo[sympy_obj] = pyomo_object
        self.sympy2latex[sympy_obj] = latex_symbol
        return sympy_obj


class Pyomo2SympyVisitor(StreamBasedExpressionVisitor):
    """
    This is based on the class of the same name in pyomo.core.base.symbolic, but
    it catches ExternalFunctions and does not decend into named expressions.
    """
    def __init__(self, object_map):
        super(Pyomo2SympyVisitor, self).__init__()
        self.object_map = object_map

    def exitNode(self, node, values):
        if isinstance(node, ExternalFunctionExpression):
            # catch ExternalFunction
            _op = self.object_map.getSympySymbol(node._fcn)
        else:
            if node.__class__ is EXPR.UnaryFunctionExpression:
                return _functionMap[node._name](values[0])
            _op = _pyomo_operator_map.get(node.__class__, None)
        if _op is None:
            return node._apply_operation(values)
        else:
            return _op(*tuple(values))

    def beforeChild(self, node, child):
        # Don't replace native or sympy types
        if type(child) in native_types:
            return False, child
        # We will not descend into named expressions...
        if child.is_expression_type():
            if child.is_named_expression_type():
                # To keep expressions from becoming too crazy complicated
                # treat names expressions like variables.
                return False, self.object_map.getSympySymbol(child)
            else:
                return True, None
        # Replace pyomo variables with sympy variables
        if child.is_potentially_variable():
            return False, self.object_map.getSympySymbol(child)
        # Everything else is a constant...
        return False, value(child)


def sympify_expression(expr):
    """
    Converts Pyomo expressions to sympy expressions.
    This is based on the function of the same name in pyomo.core.base.symbolic.
    The difference between this and the Pymomo is that this one checks if the
    expr argument is a named expression and expands it anyway. This allows named
    expressions to only be expanded if they are the top level object.
    """
    #
    # Create the visitor and call it.
    #
    object_map = PyomoSympyBimap()
    visitor = Pyomo2SympyVisitor(object_map)
    is_expr, ans = visitor.beforeChild(None, expr)
    try: # If I explicitly ask for a named expression then descend into it.
        if expr.is_named_expression_type():
            is_expr = True
    except:
        pass
    if not is_expr: # and not expr.is_named_expression_type():
        return object_map, ans

    return object_map, visitor.walk_expression(expr)

def _add_docs(object_map, docs, typ, head):
    """
    Adds documentation for a set of pyomo components to a markdown table

    Args:
        object_map (PyomoSympyBimap): Pyomo, sympy, LaTeX mapping
        docs: string containing a mardown table
        typ: the class of objects to document (Var, Expression, ExternalFunction)
        head: a string to used in the sybol table heading for this class of objects
    Returns:
        A new string markdown table with added doc rows.
    """
    docked = set() # components already documented, mainly for indexed compoents
    whead = True  # write heading before adding first item
    for i, sc in enumerate(object_map.sympyVars()):
        c = object_map.getPyomoSymbol(sc)
        c = c.parent_component() # Document the parent for indexed comps
        if not isinstance(c, typ): continue
        if whead: # add heading if is first entry in this section
            docs += "**{}** | **Doc** | **Path**\n".format(head)
            whead = False
        if id(c) not in docked:
            docked.add(id(c)) # make sure we don't get a line for every index
            try: #just document the parent of indexed vars
                s = object_map.parent_symbol[c][1]
            except KeyError: # non-indexed vars
                s = object_map.sympy2latex[sc]
            docs += "${}$|{}|{}\n".format(s, c.doc, c)
    return docs

def to_latex(expr):
    """Return a sympy expression for the given Pyomo expression

    Args:
        expr (Expression): Pyomo expression
    Returns:
        (dict): keys: sympy_expr, a sympy expression; where, markdown string
            with documentation table; latex_expr, a LaTeX string representation
            of the expression.
    """
    object_map, sympy_expr = sympify_expression(expr)

    # This next bit documents the expression, could use a lot of work, but
    # for now it generates markdown tables that are resonably readable in a
    # jupyter notebook.
    docs = "\nSymbol | Doc | Path\n ---: | --- | ---\n"
    docs = _add_docs(object_map, docs, Var, "Variable")
    docs = _add_docs(object_map, docs, Expression, "Expression")
    docs = _add_docs(object_map, docs, ExternalFunction, "Function")

    #probably should break this up, but this will do for now.
    return {"sympy_expr": sympy_expr,
            "where":docs,
            "latex_expr":sympy.latex(sympy_expr, symbol_names=object_map.sympy2latex)}

def document_constraints(comp, doc=True, descend_into=True):
    """
    Provides nicely formatted constraint documetntation in markdown format,
    assuming the $$latex math$$ and $latex math$ syntax is supported.

    Args:
        comp: A Pyomo component to document in {_ConstraintData, _ExpressionData,
                _BlockData}.
        doc: True adds a documentation table for each constraint or expression.
                Due to the way symbols are semi-automatiaclly generated, the
                exact symbol definitions may be unique to each constraint or
                expression, if unique LaTeX symbols were not provided
                everywhere in a block.
        descend_into: If True, look in subblocks for constraints.

    Returns:
        A string in markdown format with equations in LaTeX form.
    """
    s = None
    if isinstance(comp, _ExpressionData):
        d = to_latex(comp)
        if doc:
            s = "$${}$$\n{}".format(d["latex_expr"], d["where"])
        else:
            s = "$${}$$".format(d["latex_expr"])
    elif isinstance(comp, _ConstraintData):
        d = to_latex(comp.body)
        if doc:
            s = "$${} \le {}\le {}$$\n{}".format(
                comp.lower, d["latex_expr"], comp.upper, d["where"])
        else:
            s = "$${} \le {}\le {}$$".format(
                comp.lower, d["latex_expr"], comp.upper)
    elif isinstance(comp, _BlockData):
        cs = []
        for c in comp.component_data_objects(
            Constraint, descend_into=descend_into):
            if not c.active:
                continue
            cs.append("## {}".format(c))
            cs.append(document_constraints(c, doc))
        s = "\n".join(cs)
    return s

def ipython_document_constraints(comp, doc=True, descend_into=True):
    """
    See document_constraints, this just directly displays the markdown instead
    of returning a string.
    """
    s = document_constraints(comp, doc, descend_into)
    display(Markdown(s))
