# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
""" Provides object representation for the model that is conducive to code 
    generation using templates. """

from typing import Dict, List, Optional, Sequence
import string
import textwrap

from tflite_micro.codegen.operators import factory
from tflite_micro.codegen.operators import operator
from tflite_micro.codegen import tensor
from tflite_micro.codegen import utils
from tflite_micro.tensorflow.lite.python import schema_py_generated as schema_fb
from tflite_micro.tensorflow.lite.tools import visualize


class OpCode(object):

  def __init__(self, op_code: schema_fb.OperatorCodeT):
    self._op_code: schema_fb.OperatorCodeT = op_code

  @property
  def name(self) -> str:
    if self._op_code.customCode:
      return self._op_code.customCode
    return visualize.BuiltinCodeToName(self._op_code.builtinCode)

  @property
  def register_function(self) -> str:
    return "tflite::RegisterInference_{}".format(self.name)

  @property
  def enum_name(self) -> str:
    return "k{}".format(utils.to_pascal_case(self.name))

  @property
  def full_enum_name(self) -> str:
    return "OpCode::" + self.enum_name


class Subgraph(object):

  def __init__(self, model: schema_fb.ModelT, buffers: Sequence[tensor.Buffer],
               subgraph_idx: int, subgraph: schema_fb.SubGraphT):
    self._subgraph_idx: int = subgraph_idx
    self._subgraph: schema_fb.SubGraphT = subgraph
    self._op_codes: List[OpCode] = [
        OpCode(op_code) for op_code in model.operatorCodes
    ]
    self._tensors: List[Tensor] = []
    for t in subgraph.tensors:
      self._tensors.append(tensor.Tensor(buffers[t.buffer], t))

    self._operators: List[operator.Operator] = []
    for op in subgraph.operators:
      op_code = model.operatorCodes[op.opcodeIndex]
      self._operators.append(factory.create_operator(op_code, op))

  @property
  def index(self) -> int:
    return self._subgraph_idx

  @property
  def inputs(self) -> Sequence[int]:
    return self._subgraph.inputs

  @property
  def outputs(self) -> Sequence[int]:
    return self._subgraph.outputs

  @property
  def operators(self) -> Sequence[operator.Operator]:
    return self._operators

  @property
  def tensors(self) -> Sequence[tensor.Tensor]:
    return self._tensors

  @property
  def needs_zero_length_int_array(self) -> bool:
    return any(t.needs_zero_length_int_array for t in self.tensors)

  @property
  def invoke_fn_name(self) -> str:
    return f"InvokeSubgraph{self.index}"

  @property
  def inputs_array_name(self) -> str:
    return f"kSubgraph{self.index}Inputs"

  @property
  def outputs_array_name(self) -> str:
    return f"kSubgraph{self.index}Outputs"

  @property
  def nodes_array(self) -> str:
    return f"subgraph{self.index}_nodes_"

  def nodes_element(self, operator_idx: int) -> str:
    return self.nodes_array + f"[{operator_idx}]"

  def node_data_type(self, operator_idx: int) -> str:
    return f"Node{self.index}_{operator_idx}"

  def node_data_name(self, operator_idx: int) -> str:
    return f"node_{self.index}_{operator_idx}"

  def generate_c_node_data(self, indent: str) -> str:
    node_data_strs: List[str] = []
    for op_idx, op in enumerate(self.operators):
      type_name = self.node_data_type(op_idx)
      node_name = self.node_data_name(op_idx)
      node_data_strs.append(op.generate_c_node_data(type_name, node_name))
    return textwrap.indent("\n\n".join(node_data_strs), indent)

  def generate_c_node_init(self, indent: str) -> str:
    node_init_strs: List[str] = []
    for op_idx, op in enumerate(self.operators):
      tflite_node_name = self.nodes_element(op_idx)
      node_data_name = self.node_data_name(op_idx)
      node_init_strs.append(
          op.generate_c_node_init(tflite_node_name, node_data_name))
    return textwrap.indent("\n".join(node_init_strs), indent)

  def generate_c_invoke(self, indent: str) -> str:
    function_template = string.Template(
        "TfLiteStatus ${function_name}(TfLiteContext* context,\n"
        "                             tflite::Span<TfLiteNode> nodes) {\n"
        "  TFLITE_DCHECK(nodes.size() == ${num_nodes});\n"
        "${body}\n"
        "  return kTfLiteOk;\n"
        "}")

    body_template = string.Template(
        "  TF_LITE_ENSURE_OK(\n"
        "      context, op_table[${op_code}].invoke(context, &${node}));\n")
    invoke_strs: List[str] = []
    for op_idx, op in enumerate(self.operators):
      invoke_strs.append(
          body_template.substitute(
              op_code=self._op_codes[op.op_code_index].full_enum_name,
              node=f"nodes[{op_idx}]"))

    invoke = function_template.substitute(function_name=self.invoke_fn_name,
                                          num_nodes=len(self.operators),
                                          body="".join(invoke_strs))
    return textwrap.indent(invoke, indent)

  def generate_c_input_array(self, indent: str) -> str:
    return utils.generate_c_int_array(indent, "size_t", self.inputs_array_name,
                                      self.inputs)

  def generate_c_output_array(self, indent: str) -> str:
    return utils.generate_c_int_array(indent, "size_t",
                                      self.outputs_array_name, self.outputs)

  def generate_c_subgraph_init(self, indent: str) -> str:
    init_template = string.Template(
        "{.inputs = {&${input_array}[0], ${input_size}},\n"
        " .outputs = {&${output_array}[0], ${output_size}},\n"
        " .nodes = {&${node_array}[0], ${node_size}},\n"
        " .tensors = {&${tensor_array}[0], ${tensor_size}},\n"
        " .invoke = &${invoke}},")
    return textwrap.indent(
        init_template.substitute(input_array=self.inputs_array_name,
                                 input_size=len(self.inputs),
                                 output_array=self.outputs_array_name,
                                 output_size=len(self.outputs),
                                 node_array=self.nodes_array,
                                 node_size=len(self.operators),
                                 tensor_array=self.tensors_array,
                                 tensor_size=len(self.tensors),
                                 invoke=self.invoke_fn_name), indent)

  @property
  def tensors_array(self) -> str:
    return f"subgraph{self.index}_tensors_"

  def tensors_element(self, tensor_idx: int) -> str:
    return self.tensors_array + f"[{tensor_idx}]"

  def tensor_data_type(self, tensor_idx: int) -> str:
    return f"Tensor{self.index}_{tensor_idx}"

  def tensor_data_name(self, tensor_idx: int) -> str:
    return f"tensor{self.index}_{tensor_idx}"

  def generate_c_tensor_data(self, indent: str) -> str:
    tensor_dims_strs: List[str] = []
    for tensor_idx, tensor in enumerate(self.tensors):
      type_name = self.tensor_data_type(tensor_idx)
      tensor_name = self.tensor_data_name(tensor_idx)
      tensor_dims_strs.append(
          tensor.generate_c_tensor_dims(type_name, tensor_name))
    return textwrap.indent("\n\n".join(tensor_dims_strs), indent)

  def generate_c_tensor_init(self, indent: str) -> str:
    tensor_init_strs: List[str] = []
    for tensor_idx, tensor in enumerate(self.tensors):
      tflite_tensor_name = self.tensors_element(tensor_idx)
      tensor_data_name = self.tensor_data_name(tensor_idx)
      tensor_init_strs.append(
          tensor.generate_c_tensor_init(tflite_tensor_name, tensor_data_name))
    return textwrap.indent("\n".join(tensor_init_strs), indent)


class Graph(object):

  def __init__(self, model: schema_fb.ModelT):
    buffers: List[tensor.Buffer] = [
        tensor.Buffer("buffer_{}".format(idx), buffer)
        for idx, buffer in enumerate(model.buffers)
    ]
    self._subgraphs: List[SubGraph] = [
        Subgraph(model, buffers, idx, subgraph)
        for idx, subgraph in enumerate(model.subgraphs)
    ]

  @property
  def subgraphs(self) -> Sequence[Subgraph]:
    return self._subgraphs

  @property
  def buffers(self) -> Sequence[tensor.Buffer]:
    buffers: List[tensor.Buffer] = []
    for subgraph in self.subgraphs:
      for t in subgraph.tensors:
        buffers.append(t.buffer)
    return buffers

  @property
  def needs_zero_length_int_array(self) -> bool:
    return any(subgraph.needs_zero_length_int_array
               for subgraph in self.subgraphs)


class OpCodeTable(object):

  def __init__(self, models: Sequence[schema_fb.ModelT]):
    op_codes = []
    for model in models:
      for op_code in model.operatorCodes:
        op_codes.append(OpCode(op_code))

    self._op_codes: List([OpCode]) = list(set(op_codes))

  @property
  def op_codes(self) -> Sequence[OpCode]:
    return self._op_codes