"""Functions for formatting PyTensor compute graphs.

Author: Christof Angermueller <cangermueller@gmail.com>
"""

from functools import reduce
from pathlib import Path

import numpy as np

import pytensor
from pytensor.compile import Function, builders
from pytensor.graph.basic import Apply, Constant, Variable
from pytensor.graph.fg import FunctionGraph
from pytensor.graph.traversal import graph_inputs
from pytensor.printing import _try_pydot_import


class PyDotFormatter:
    """Create `pydot` graph object from PyTensor function.

    Parameters
    ----------
    compact : bool
        if True, will remove intermediate variables without name.

    Attributes
    ----------
    node_colors : dict
        Color table of node types.
    apply_colors : dict
        Color table of apply nodes.
    shapes : dict
        Shape table of node types.
    """

    def __init__(self, compact=True):
        """Construct PyDotFormatter object."""
        _try_pydot_import()

        self.compact = compact
        self.node_colors = {
            "input": "limegreen",
            "constant_input": "SpringGreen",
            "shared_input": "YellowGreen",
            "output": "dodgerblue",
            "unused": "lightgrey",
        }
        self.apply_colors = {
            "Scan": "yellow",
            "Shape": "cyan",
            "IfElse": "magenta",
            "Elemwise": "#FFAABB",  # dark pink
            "Subtensor": "#FFAAFF",  # purple
            "Alloc": "#FFAA22",
        }  # orange
        self.shapes = {"input": "box", "output": "box", "apply": "ellipse"}
        self.__node_prefix = "n"

    def __add_node(self, node):
        """Add new node to node list and return unique id.

        Parameters
        ----------
        node : PyTensor graph node
            Apply node, tensor variable, or shared variable in compute graph.

        Returns
        -------
        str
            Unique node id.
        """
        assert node not in self.__nodes
        _id = f"{self.__node_prefix}{len(self.__nodes) + 1}"
        self.__nodes[node] = _id
        return _id

    def __node_id(self, node):
        """Return unique node id.

        Parameters
        ----------
        node : PyTensor graph node
            Apply node, tensor variable, or shared variable in compute graph.

        Returns
        -------
        str
            Unique node id.
        """
        if node in self.__nodes:
            return self.__nodes[node]
        else:
            return self.__add_node(node)

    def __call__(self, fct, graph=None):
        """Create pydot graph from function.

        Parameters
        ----------
        fct : pytensor.compile.function.types.Function
            A compiled PyTensor function, variable, apply or a list of variables.
        graph: pydot.Dot
            `pydot` graph to which nodes are added. Creates new one if
            undefined.

        Returns
        -------
        pydot.Dot
            Pydot graph of `fct`
        """
        pd = _try_pydot_import()

        if graph is None:
            graph = pd.Dot()

        self.__nodes = {}

        profile = None

        if isinstance(fct, Function):
            profile = getattr(fct, "profile", None)
            fgraph = fct.maker.fgraph
        elif isinstance(fct, FunctionGraph):
            fgraph = fct
        else:
            if isinstance(fct, Variable):
                fct = [fct]
            elif isinstance(fct, Apply):
                fct = fct.outputs
            assert isinstance(fct, list | tuple)
            assert all(isinstance(v, Variable) for v in fct)
            fgraph = FunctionGraph(inputs=graph_inputs(fct), outputs=fct)

        outputs = fgraph.outputs
        topo = fgraph.toposort()
        outputs = list(outputs)

        # Loop over apply nodes
        for node in topo:
            nparams = {}
            __node_id = self.__node_id(node)
            nparams["name"] = __node_id
            nparams["label"] = apply_label(node)
            nparams["profile"] = apply_profile(fgraph, node, profile)
            nparams["node_type"] = "apply"
            nparams["apply_op"] = nparams["label"]
            nparams["shape"] = self.shapes["apply"]

            use_color = None
            for opName, color in self.apply_colors.items():
                if opName in node.op.__class__.__name__:
                    use_color = color
            if use_color:
                nparams["style"] = "filled"
                nparams["fillcolor"] = use_color
                nparams["type"] = "colored"

            pd_node = dict_to_pdnode(nparams)
            graph.add_node(pd_node)

            # Loop over input nodes
            for id, var in enumerate(node.inputs):
                var_id = self.__node_id(var.owner if var.owner else var)
                if var.owner is None:
                    vparams = {
                        "name": var_id,
                        "label": var_label(var),
                        "node_type": "input",
                    }
                    if isinstance(var, Constant):
                        vparams["node_type"] = "constant_input"
                    elif isinstance(
                        var, pytensor.tensor.sharedvar.TensorSharedVariable
                    ):
                        vparams["node_type"] = "shared_input"
                    vparams["dtype"] = type_to_str(var.type)
                    vparams["tag"] = var_tag(var)
                    vparams["style"] = "filled"
                    vparams["fillcolor"] = self.node_colors[vparams["node_type"]]
                    vparams["shape"] = self.shapes["input"]
                    pd_var = dict_to_pdnode(vparams)
                    graph.add_node(pd_var)

                edge_params = {}
                if node.op.view_map and id in reduce(
                    list.__add__, node.op.view_map.values(), []
                ):
                    edge_params["color"] = self.node_colors["output"]
                elif node.op.destroy_map and id in reduce(
                    list.__add__, node.op.destroy_map.values(), []
                ):
                    edge_params["color"] = "red"

                edge_label = vparams["dtype"]
                if len(node.inputs) > 1:
                    edge_label = str(id) + " " + edge_label
                pdedge = pd.Edge(var_id, __node_id, label=edge_label, **edge_params)
                graph.add_edge(pdedge)

            # Loop over output nodes
            for id, var in enumerate(node.outputs):
                var_id = self.__node_id(var)

                if var in outputs or len(fgraph.clients[var]) == 0:
                    vparams = {
                        "name": var_id,
                        "label": var_label(var),
                        "node_type": "output",
                        "dtype": type_to_str(var.type),
                        "tag": var_tag(var),
                        "style": "filled",
                    }
                    if len(fgraph.clients[var]) == 0:
                        vparams["fillcolor"] = self.node_colors["unused"]
                    else:
                        vparams["fillcolor"] = self.node_colors["output"]
                    vparams["shape"] = self.shapes["output"]
                    pd_var = dict_to_pdnode(vparams)
                    graph.add_node(pd_var)

                    graph.add_edge(pd.Edge(__node_id, var_id, label=vparams["dtype"]))
                elif var.name or not self.compact:
                    graph.add_edge(pd.Edge(__node_id, var_id, label=vparams["dtype"]))

            # Create sub-graph for OpFromGraph nodes
            if isinstance(node.op, builders.OpFromGraph):
                subgraph = pd.Cluster(__node_id)
                gf = PyDotFormatter()
                # Use different node prefix for sub-graphs
                gf.__node_prefix = __node_id
                gf(node.op.fn, subgraph)
                graph.add_subgraph(subgraph)
                pd_node.get_attributes()["subg"] = subgraph.get_name()

                def format_map(m):
                    return str([list(x) for x in m])

                # Inputs mapping
                ext_inputs = [self.__node_id(x) for x in node.inputs]
                int_inputs = [gf.__node_id(x) for x in node.op.inner_inputs]
                assert len(ext_inputs) == len(int_inputs)
                h = format_map(zip(ext_inputs, int_inputs, strict=True))
                pd_node.get_attributes()["subg_map_inputs"] = h

                # Outputs mapping
                ext_outputs = [self.__node_id(x) for x in node.outputs]
                int_outputs = [gf.__node_id(x) for x in node.op.inner_outputs]
                assert len(ext_outputs) == len(int_outputs)
                h = format_map(zip(int_outputs, ext_outputs, strict=True))
                pd_node.get_attributes()["subg_map_outputs"] = h

        return graph


def var_label(var, precision=3):
    """Return label of variable node."""
    if var.name is not None:
        return var.name
    elif isinstance(var, Constant):
        h = np.asarray(var.data)
        is_const = False
        if h.ndim == 0:
            is_const = True
            h = np.array([h])
        dstr = np.array2string(h, precision=precision)
        if "\n" in dstr:
            dstr = dstr[: dstr.index("\n")]
        if is_const:
            dstr = dstr.replace("[", "").replace("]", "")
        return dstr
    else:
        return type_to_str(var.type)


def var_tag(var):
    """Parse tag attribute of variable node."""
    tag = var.tag
    if hasattr(tag, "trace") and len(tag.trace) and len(tag.trace[0]) == 4:
        if isinstance(tag.trace[0][0], tuple | list):
            path, line, _, src = tag.trace[0][-1]
        else:
            path, line, _, src = tag.trace[0]
        path = Path(path).name
        path = path.replace("<", "")
        path = path.replace(">", "")
        src = src.encode()
        return [path, line, src]
    else:
        return None


def apply_label(node):
    """Return label of apply node."""
    return node.op.__class__.__name__


def apply_profile(fgraph, node, profile):
    """Return apply profiling information."""
    if not profile or profile.fct_call_time == 0:
        return None
    time = profile.apply_time.get((fgraph, node), 0)
    call_time = profile.fct_call_time
    return [time, call_time]


def broadcastable_to_str(b):
    """Return string representation of broadcastable."""
    named_broadcastable = {
        (): "scalar",
        (False,): "vector",
        (False, True): "col",
        (True, False): "row",
        (False, False): "matrix",
    }
    if b in named_broadcastable:
        bcast = named_broadcastable[b]
    else:
        bcast = ""
    return bcast


def dtype_to_char(dtype):
    """Return character that represents data type."""
    dtype_char = {
        "complex64": "c",
        "complex128": "z",
        "float32": "f",
        "float64": "d",
        "int8": "b",
        "int16": "w",
        "int32": "i",
        "int64": "l",
    }
    if dtype in dtype_char:
        return dtype_char[dtype]
    else:
        return "X"


def type_to_str(t):
    """Return str of variable type."""
    if not hasattr(t, "broadcastable"):
        return str(t)
    s = broadcastable_to_str(t.broadcastable)
    if s == "":
        s = str(t.dtype)
    else:
        s = dtype_to_char(t.dtype) + s
    return s


def dict_to_pdnode(d):
    """Create pydot node from dict."""
    pd = _try_pydot_import()

    e = dict()
    for k, v in d.items():
        if v is not None:
            if isinstance(v, list):
                v = "\t".join(str(x) for x in v)
            else:
                v = str(v)
            v = str(v)
            v = v.replace('"', "'")
            e[k] = v
    pynode = pd.Node(**e)
    return pynode
