# Copyright (c) ONNX Project Contributors
# SPDX-License-Identifier: Apache-2.0
"""Lift constants to initializers."""

from __future__ import annotations

__all__ = [
    "AddInitializersToInputsPass",
    "LiftConstantsToInitializersPass",
    "LiftSubgraphInitializersToMainGraphPass",
    "RemoveInitializersFromInputsPass",
]

import logging

import numpy as np

import onnx_ir as ir

logger = logging.getLogger(__name__)


class LiftConstantsToInitializersPass(ir.passes.InPlacePass):
    """Lift constants to initializers.

    Attributes:
        lift_all_constants: Whether to lift all Constant nodes, including those that does not contain a tensor attribute (e.g. with value_ints etc.)
            Default to False, where only Constants with the ``value`` attribute are lifted.
        size_limit: The minimum size of the tensor to be lifted. If the tensor contains
            number of elements less than size_limit, it will not be lifted. Default is 16.
    """

    def __init__(self, lift_all_constants: bool = False, size_limit: int = 16):
        super().__init__()
        self.lift_all_constants = lift_all_constants
        self.size_limit = size_limit

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        count = 0
        for node in ir.traversal.RecursiveGraphIterator(model.graph):
            assert node.graph is not None
            if node.op_type != "Constant" or node.domain not in ("", "onnx.ai"):
                continue
            if node.outputs[0].is_graph_output():
                logger.debug(
                    "Constant node '%s' is used as output, so it can't be lifted.", node.name
                )
                continue
            constant_node_attribute = set(node.attributes.keys())
            if len(constant_node_attribute) != 1:
                logger.debug(
                    "Invalid constant node '%s' has more than one attribute", node.name
                )
                continue

            attr_name, attr_value = next(iter(node.attributes.items()))
            initializer_name = node.outputs[0].name
            assert initializer_name is not None
            assert isinstance(attr_value, ir.Attr)
            tensor = self._constant_node_attribute_to_tensor(
                node, attr_name, attr_value, initializer_name
            )
            if tensor is None:
                # The reason of None is logged in _constant_node_attribute_to_tensor
                continue
            # Register an initializer with the tensor value
            initializer = ir.Value(
                name=initializer_name,
                shape=tensor.shape,  # type: ignore[arg-type]
                type=ir.TensorType(tensor.dtype),
                const_value=tensor,
                # Preserve metadata from Constant value into the onnx model
                metadata_props=node.outputs[0].metadata_props.copy(),
            )
            # Preserve value meta from the Constant output for intermediate analysis
            initializer.meta.update(node.outputs[0].meta)

            assert node.graph is not None
            node.graph.register_initializer(initializer)
            # Replace the constant node with the initializer
            node.outputs[0].replace_all_uses_with(initializer)
            node.graph.remove(node, safe=True)
            count += 1
            logger.debug(
                "Converted constant node '%s' to initializer '%s'", node.name, initializer_name
            )
        if count:
            logger.debug("Lifted %s constants to initializers", count)
        return ir.passes.PassResult(model, modified=bool(count))

    def _constant_node_attribute_to_tensor(
        self, node, attr_name: str, attr_value: ir.Attr, initializer_name: str
    ) -> ir.TensorProtocol | None:
        """Convert constant node attribute to tensor."""
        if not self.lift_all_constants and attr_name != "value":
            logger.debug(
                "Constant node '%s' has non-tensor attribute '%s'", node.name, attr_name
            )
            return None

        tensor: ir.TensorProtocol
        if attr_name == "value":
            tensor = attr_value.as_tensor()
        elif attr_name == "value_int":
            tensor = ir.tensor(
                attr_value.as_int(), dtype=ir.DataType.INT64, name=initializer_name
            )
        elif attr_name == "value_ints":
            tensor = ir.tensor(
                attr_value.as_ints(), dtype=ir.DataType.INT64, name=initializer_name
            )
        elif attr_name == "value_float":
            tensor = ir.tensor(
                attr_value.as_float(), dtype=ir.DataType.FLOAT, name=initializer_name
            )
        elif attr_name == "value_floats":
            tensor = ir.tensor(
                attr_value.as_floats(), dtype=ir.DataType.FLOAT, name=initializer_name
            )
        elif attr_name in ("value_string", "value_strings"):
            tensor = ir.StringTensor(
                np.array(attr_value.value, dtype=np.bytes_), name=initializer_name
            )
        else:
            raise ValueError(
                f"Unsupported constant node '{node.name}' attribute '{attr_name}'"
            )

        if tensor.size < self.size_limit:
            logger.debug(
                "Tensor from node '%s' has less than %s elements",
                node.name,
                self.size_limit,
            )
            return None
        return tensor


class LiftSubgraphInitializersToMainGraphPass(ir.passes.InPlacePass):
    """Lift subgraph initializers to main graph.

    This pass lifts the initializers of a subgraph to the main graph.
    It is used to ensure that the initializers are available in the main graph
    for further processing or optimization.

    Initializers that are also graph inputs will not be lifted.

    .. versionchanged:: 0.2.1
        Lifted initializers will be renamed if necessary to avoid name
        collisions with existing initializers, node outputs, or graph inputs in
        the main graph.
    """

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        count = 0
        registered_initializer_names: dict[str, int] = {}
        # Collect main graph output and input names to avoid collisions
        main_graph_output_names = {
            output.name for node in model.graph for output in node.outputs if output.name
        }
        main_graph_input_names = {input.name for input in model.graph.inputs if input.name}
        for graph in model.graphs():
            if graph is model.graph:
                continue
            for name in tuple(graph.initializers):
                assert name is not None
                initializer = graph.initializers[name]
                if initializer.is_graph_input():
                    # Skip the ones that are also graph inputs
                    logger.debug(
                        "Initializer '%s' is also a graph input, so it can't be lifted",
                        initializer.name,
                    )
                    continue
                if initializer.is_graph_output():
                    logger.debug(
                        "Initializer '%s' is used as output, so it can't be lifted",
                        initializer.name,
                    )
                    continue
                # Remove the initializer from the subgraph
                graph.initializers.pop(name)
                # To avoid name conflicts, we need to rename the initializer
                # to a unique name in the main graph that doesn't conflict with
                # existing initializers, node outputs, or graph inputs
                new_name = name
                while (
                    new_name in model.graph.initializers
                    or new_name in main_graph_output_names
                    or new_name in main_graph_input_names
                ):
                    if name in registered_initializer_names:
                        registered_initializer_names[name] += 1
                    else:
                        registered_initializer_names[name] = 1
                    new_name = f"{name}_{registered_initializer_names[name]}"
                initializer.name = new_name
                model.graph.register_initializer(initializer)
                count += 1
                logger.debug(
                    "Lifted initializer '%s' from subgraph '%s' to main graph",
                    initializer.name,
                    graph.name,
                )
        return ir.passes.PassResult(model, modified=bool(count))


class RemoveInitializersFromInputsPass(ir.passes.InPlacePass):
    """Remove initializers from inputs.

    This pass finds all graph inputs that have a const_value and removes them from the graph.inputs list.
    """

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        count = 0
        for graph in model.graphs():
            initializers = set(graph.initializers.values())
            new_inputs = []
            for input_value in graph.inputs:
                if input_value in initializers:
                    count += 1
                else:
                    new_inputs.append(input_value)
            graph.inputs.clear()
            graph.inputs.extend(new_inputs)
        logger.info("Removed %s initializers from graph inputs", count)
        return ir.passes.PassResult(model, modified=bool(count))


class AddInitializersToInputsPass(ir.passes.InPlacePass):
    """Add initializers to inputs.

    This pass finds all initializers and adds them to the graph.inputs list if they are not already present.
    """

    def call(self, model: ir.Model) -> ir.passes.PassResult:
        count = 0
        for graph in model.graphs():
            inputs_set = set(graph.inputs)
            for initializer in graph.initializers.values():
                if initializer not in inputs_set:
                    graph.inputs.append(initializer)
                    count += 1
        logger.info("Added %s initializers to graph inputs", count)
        return ir.passes.PassResult(model, modified=bool(count))
