# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
"""Fuction fusion unittest producer.

Takes in a full model, function keyword, and example inputs, produces unit model protos
that contains only a single node calling the target function proto.

- All initializers are lifted as model inputs.
- Example inputs and outputs are saved as test data for each unit model proto.
"""

from __future__ import annotations

import argparse
import itertools
import logging
import os
import sys

import numpy as np
import onnx
import onnx.inliner
import onnxruntime
from onnx import helper as onnx_helper
from onnx import numpy_helper

from onnxscript import _legacy_ir as ir
from onnxscript._legacy_ir import visitor
from onnxscript.utils import evaluation_utils, utils

logger = logging.getLogger(__name__)


# Copied from common.py from pytorch torchbench
def save_tensor_data(numpy_tensor, output_path: str):
    proto_tensor = numpy_helper.from_array(numpy_tensor)
    with open(output_path, "wb") as f:
        f.write(proto_tensor.SerializeToString())


class FunctionToKeepVisitor(visitor.ProtoVisitorCore):
    def __init__(self, function_keyword):
        self.function_keyword = function_keyword
        self.functions_to_keep = []
        self.in_target_function = False
        self._functions = {}
        super().__init__()

    def visit_function_node(self, node: onnx.NodeProto):
        prev_in_target_function = self.in_target_function
        function_id = ir.get_function_id_from_node(node)
        function = self._functions[function_id]
        if node.op_type.find(self.function_keyword) != -1:
            self.functions_to_keep.append(function_id)
            self.in_target_function = True
        elif prev_in_target_function:
            self.functions_to_keep.append(function_id)

        for subnode in function.node:
            self.visit_node(subnode)

        self.in_target_function = prev_in_target_function

    def process_node(self, node: onnx.NodeProto):
        if visitor.is_local_function_node(node, self._functions):
            return self.visit_function_node(node)
        return None

    def visit_model(self, model: onnx.ModelProto) -> None:
        for function in model.functions:
            self._functions[ir.get_function_id(function)] = function
        super().visit_model(model)


class TargetFunctionMetaVisitor(visitor.ProtoVisitorCore):
    def __init__(self, function_keyword):
        self.function_keyword = function_keyword
        # Map from (domain, name) to (actual_input_names, actual_output_names)
        self.function_meta: dict[tuple[str, str], tuple[list[str], list[str]]] = {}
        self._functions = {}
        super().__init__()

    def visit_function_node(self, node: onnx.NodeProto):
        function = self._functions[ir.get_function_id_from_node(node)]
        if node.op_type.find(self.function_keyword) != -1:
            self.function_meta[(function.domain, function.name)] = (
                node.input,
                node.output,
            )
        for subnode in function.node:
            self.visit_node(subnode)

    def process_node(self, node: onnx.NodeProto):
        if visitor.is_local_function_node(node, self._functions):
            return self.visit_function_node(node)
        return None

    def visit_model(self, model: onnx.ModelProto) -> None:
        for function in model.functions:
            self._functions[ir.get_function_id(function)] = function
        super().visit_model(model)


class FunctionProtoProducerWithData(visitor.ProtoVisitor):
    """Fuction fusion unittest producer.

    Creates unit model proto for selected function, as well as example inputs and outputs.

    Utilizes ORT fetch feature.

    Steps as follows:

    - Identify the target function, and all functions called within.
    - Call onnx.inliner to inline all other functions.
    - Identity inputs and outputs to target function calls, construct ort fetch.
    - Run the model with ort fetch to receive example inputs and outputs.
    - For each target function call, construct a unit model proto with example inputs and outputs from previous step.
    """

    def __init__(self, function_keyword: str, model_path: str, output_dir: str):
        self.function_keyword = function_keyword
        self.model_path = model_path
        self.output_dir = output_dir
        self.output_model_basename = function_keyword
        self._functions: dict[ir.FunctionId, onnx.FunctionProto] = {}
        self._unit_model_protos: list[onnx.ModelProto] = []
        self._unit_model_inputs = []  # type: ignore[var-annotated]
        self._unit_model_outputs = []  # type: ignore[var-annotated]
        # Example intermediate data values
        self._named_values: dict[str, np.ndarray] = {}
        super().__init__()

    @property
    def unit_model_protos(self) -> list[onnx.ModelProto]:
        return self._unit_model_protos

    @property
    def unit_model_inputs(self):
        return self._unit_model_inputs

    @property
    def unit_model_outputs(self):
        return self._unit_model_outputs

    def find_all_called_function_protos(
        self, function: onnx.FunctionProto
    ) -> list[onnx.FunctionProto]:
        result: dict[ir.FunctionId, onnx.FunctionProto] = {
            ir.get_function_id(function): function
        }
        for node in function.node:
            if visitor.is_local_function_node(node, self._functions):
                sub_function = self._functions[ir.get_function_id_from_node(node)]
                result.update(
                    {
                        ir.get_function_id(func): func
                        for func in self.find_all_called_function_protos(sub_function)
                    }
                )
        return result.values()  # type: ignore[return-value]

    def _generate_value_info_for_function_value(
        self, value: str, function: onnx.FunctionProto
    ) -> onnx.ValueInfoProto | None:
        value_ir = self.function_shape_env.lookup(function, value)
        if value_ir is None:
            return None
        return self.function_shape_env.save_to_value_info(
            value_ir, *ir.get_function_id(function)
        )

    def _generate_value_info_for_function_values(
        self, function: onnx.FunctionProto
    ) -> list[onnx.ValueInfoProto]:
        value_infos = []
        values = {
            *function.input,
            *function.output,
            *itertools.chain((*node.input, *node.output) for node in function.node),
        }

        for value in values:
            value_info = self._generate_value_info_for_function_value(value, function)
            if value_info is not None:
                value_infos.append(value_info)
        return value_infos

    def create_unit_model_proto(
        self,
        function_proto: onnx.FunctionProto,
        actual_input_value_infos: list[ir.Value | None],
        actual_output_value_infos: list[ir.Value | None],
    ) -> onnx.ModelProto | None:
        unit_model_proto = onnx.ModelProto()
        unit_model_proto.ir_version = self._model_proto.ir_version
        unit_model_proto.producer_name = self._model_proto.producer_name
        unit_model_proto.producer_version = self._model_proto.producer_version
        unit_model_proto.domain = self._model_proto.domain
        unit_model_proto.model_version = self._model_proto.model_version
        unit_model_proto.opset_import.extend(self._model_proto.opset_import)
        graph_proto = unit_model_proto.graph

        for actual_input_value_info, formal_input in zip(
            actual_input_value_infos, function_proto.input
        ):
            if actual_input_value_info is None:
                logger.error(
                    "Value info for input %s is not found. Skip model proto creation for function %s::%s",
                    formal_input,
                    function_proto.domain,
                    function_proto.name,
                )
                return None
            if actual_input_value_info.type is None:
                logger.error(
                    "Value info for input %s has no type. Skip model proto creation for function %s::%s",
                    formal_input,
                    function_proto.domain,
                    function_proto.name,
                )

            value_info = onnx.ValueInfoProto()
            value_info.name = actual_input_value_info.name
            value_info.type.CopyFrom(actual_input_value_info.type)
            graph_proto.input.append(value_info)

        for actual_output_value_info, formal_output in zip(
            actual_output_value_infos, function_proto.output
        ):
            if actual_output_value_info is None:
                logger.error(
                    "Value info for output %s is not found. Skip model proto creation for function %s::%s",
                    formal_output,
                    function_proto.domain,
                    function_proto.name,
                )
                return None
            if actual_output_value_info.type is None:
                logger.error(
                    "Value info for output %s has no type. Skip model proto creation for function %s::%s",
                    formal_output,
                    function_proto.domain,
                    function_proto.name,
                )

            value_info = onnx.ValueInfoProto()
            value_info.name = actual_output_value_info.name
            value_info.type.CopyFrom(actual_output_value_info.type)
            graph_proto.output.append(value_info)

        new_function_node = onnx.NodeProto()
        new_function_node.op_type = function_proto.name
        new_function_node.domain = function_proto.domain
        new_function_node.input.extend([input.name for input in actual_input_value_infos])  # type: ignore[union-attr]
        new_function_node.output.extend([output.name for output in actual_output_value_infos])  # type: ignore[union-attr]
        # TODO: Producing function node attribute is not supported yet.

        graph_proto.node.append(new_function_node)
        called_function_protos = self.find_all_called_function_protos(function_proto)
        for called_function_proto in called_function_protos:
            graph_proto.value_info.extend(
                self._generate_value_info_for_function_values(called_function_proto)
            )
        unit_model_proto.functions.extend(called_function_protos)
        return unit_model_proto

    def process_initializer(self, init: onnx.TensorProto):
        self.bind(
            init.name,
            ir.Value(name=init.name, type=utils.get_initializer_type(init)),
        )

    def lookup(self, name: str) -> ir.Value | None:
        """Override unit model proto inputs & outputs value infos with value info derived from actual example data.

        This step is required because onnx FunctionProto does not contain value info.
        The experimental solution from exporter writes value infos under root GraphProto, and associate them with
        FunctionProto by name mangling. This is lost during onnx.inliner because of the structural and value name
        changes.

        This step is not necessary once value info is natively supported in FunctionProto.

        This step by design cannot support dynamic shape.
        """
        if name in self._named_values:
            return ir.Value(
                name=name,
                type=onnx_helper.make_tensor_type_proto(
                    onnx_helper.np_dtype_to_tensor_dtype(self._named_values[name].dtype),
                    self._named_values[name].shape,
                ),
            )
        return super().lookup(name)

    def visit_model(self, model: onnx.ModelProto):
        functions_to_keep_visitor = FunctionToKeepVisitor(self.function_keyword)
        functions_to_keep_visitor.visit_model(model)
        functions_to_keep = functions_to_keep_visitor.functions_to_keep
        # TODO: bug report: IsScalar function inside if subgraph is not part of functions_to_keep.
        # Yet it is also not inlined. But its function_proto is removed by inliner.
        # To unblock us, we manually add it to functions_to_keep.
        functions_to_keep.append(("pkg.onnxscript.torch_lib.common", "IsScalar"))
        # TODO: Post ONNX 1.16, overload will be introduced.
        functions_to_keep = [function_id[:2] for function_id in functions_to_keep]
        inlined_model_proto = onnx.inliner.inline_selected_functions(
            model, functions_to_keep, exclude=True
        )
        target_function_meta_visitor = TargetFunctionMetaVisitor(self.function_keyword)
        target_function_meta_visitor.visit_model(inlined_model_proto)
        target_function_meta = target_function_meta_visitor.function_meta

        fetch_outputs = []  # type: ignore[var-annotated]
        for inputs, outputs in target_function_meta.values():
            fetch_outputs.extend((*inputs, *outputs))

        fetch_output_value_infos = []
        for fetch_output in fetch_outputs:
            value_info = onnx.ValueInfoProto()
            value_info.name = fetch_output
            fetch_output_value_infos.append(value_info)

        inlined_model_proto.graph.output.extend(fetch_output_value_infos)
        inlined_model_proto = onnx.shape_inference.infer_shapes(inlined_model_proto)

        self._model_proto = inlined_model_proto

        model_path = self.model_path
        model_dir = os.path.dirname(model_path)
        inputs, _ = evaluation_utils.load_test_data(  # type: ignore[assignment]
            model_dir, [i.name for i in model.graph.input]
        )
        tmp_model_path = f"{model_dir}/tmp_model.onnx"
        onnx.save(inlined_model_proto, tmp_model_path)

        sess = onnxruntime.InferenceSession(
            tmp_model_path, providers=["CUDAExecutionProvider"]
        )
        outputs = sess.run(fetch_outputs, inputs)
        assert len(outputs) == len(fetch_outputs), (
            f"Number of outputs mismatch. outputs: {len(outputs)}, fetch_outputs: {len(fetch_outputs)}"
        )

        self._named_values = dict(zip(fetch_outputs, outputs))  # type: ignore[arg-type]
        for inputs, outputs in target_function_meta.values():
            named_inputs = [(i, self._named_values[i]) for i in inputs]
            named_outputs = [(o, self._named_values[o]) for o in outputs]
            self._unit_model_inputs.append(named_inputs)
            self._unit_model_outputs.append(named_outputs)

        for function in inlined_model_proto.functions:
            self._functions[ir.get_function_id(function)] = function

        super().visit_model(inlined_model_proto)

    def process_function(self, function: onnx.FunctionProto):
        if function.name.find(self.function_keyword) == -1:
            return

        try:
            actual_input_value_infos = [self.lookup(input) for input in function.input]
            actual_output_value_infos = [self.lookup(output) for output in function.output]
        except ValueError as e:
            raise ValueError(
                "Cannot create ModelProto unittest for function. "
                f"Failed to find value info for function {function.domain}::{function.name}"
            ) from e
        unit_model_proto = self.create_unit_model_proto(
            function, actual_input_value_infos, actual_output_value_infos
        )
        if unit_model_proto is not None:
            self._unit_model_protos.append(unit_model_proto)


def produce_function_proto_unittest(
    model_path: str,
    function_keyword: str,
    output_dir: str,
) -> tuple[
    list[onnx.ModelProto],
    list[list[tuple[str, np.ndarray]]],
    list[list[tuple[str, np.ndarray]]],
]:
    model_proto = onnx.load(model_path, load_external_data=False)

    # model_proto = optimizer.optimize(model_proto, onnx_shape_inference=False)

    producer = FunctionProtoProducerWithData(
        function_keyword,
        model_path,
        output_dir,
    )

    producer.visit_model(model_proto)
    return (
        producer.unit_model_protos,
        producer.unit_model_inputs,
        producer.unit_model_outputs,
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", "--model_path", type=str)
    parser.add_argument("--function", type=str)
    parser.add_argument("--output-dir", "--output_dir", type=str)
    parser.add_argument("--max-outputs", "--max_outputs", type=int, default=sys.maxsize)
    parser.add_argument("--name", type=str)

    args = parser.parse_args()
    model_path = args.model_path
    function = args.function
    output_dir = args.output_dir
    max_outputs = args.max_outputs
    name = args.name

    (
        unit_model_protos,
        named_inputs_list,
        named_outputs_list,
    ) = produce_function_proto_unittest(model_path, function, output_dir)

    for i, unit_model_proto in enumerate(unit_model_protos[:max_outputs]):
        if logger.level <= logging.DEBUG:
            logger.debug("unit model proto %d:", i)
            # logger.debug(onnx.printer.to_text(unit_model_proto))
        output_model_dir = f"{output_dir}/{name}_{i}/"
        os.makedirs(output_model_dir, exist_ok=True)
        onnx.save(unit_model_proto, f"{output_model_dir}/{name}_{i}.onnx")
        # save test data
        test_data_dir = f"{output_model_dir}/test_data_set_0/"
        os.makedirs(test_data_dir, exist_ok=True)
        named_inputs = named_inputs_list[i]
        for j, (_, input) in enumerate(named_inputs):
            save_tensor_data(input, f"{test_data_dir}/input_{j}.pb")
        named_outputs = named_outputs_list[i]
        for j, (_, output) in enumerate(named_outputs):
            save_tensor_data(output, f"{test_data_dir}/output_{j}.pb")

    print(
        f"{len(unit_model_protos[:max_outputs])} unit model protos and test data are saved to {output_dir}."
    )


if __name__ == "__main__":
    # python tools/function_rewriter_testing/function_unittest_producer.py \
    #    --model_path tools/ort_rewriter_profiling/onnx_models/stable_diffusion_unet/dynamo/stable_diffusion_unet_dynamo.onnx \
    #    --function GEGLU --output-dir testdata/unittest_models/ --max_outputs 4 --name geglu_stable_diffusion_unet
    main()