/* Copyright 2017 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.
==============================================================================*/
#ifndef TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_
#define TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_

#include <algorithm>
#include <cmath>
#include <iostream>
#include <limits>
#include <memory>
#include <string>
#include <vector>

#include "absl/strings/string_view.h"
#include "tensorflow/core/platform/logging.h"
#if TOCO_SUPPORT_PORTABLE_PROTOS
#include "third_party/protobuf/include/google/protobuf/text_format.h"
#endif  // TOCO_SUPPORT_PORTABLE_PROTOS
#include "tensorflow/contrib/lite/kernels/internal/types.h"
#include "tensorflow/contrib/lite/toco/model.h"
#include "tensorflow/contrib/lite/toco/model_flags.pb.h"
#include "tensorflow/contrib/lite/toco/runtime/types.h"
#include "tensorflow/contrib/lite/toco/toco_flags.pb.h"
#include "tensorflow/contrib/lite/toco/types.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"

// TODO(aselle): Replace with using a container specific hash override instead.
namespace std {
template <>
struct hash<toco::OperatorType> {
  size_t operator()(const toco::OperatorType& op) const {
    return std::hash<size_t>()(static_cast<size_t>(op));
  }
};
}  // namespace std

namespace toco {

constexpr int kLogLevelModelChanged = 1;
constexpr int kLogLevelModelUnchanged = 2;

absl::string_view FindLongestCommonPrefix(absl::string_view a,
                                          absl::string_view b);
string LogName(const Operator& op);

string ArrayDataTypeName(ArrayDataType data_type);

// Returns true if the given array is specified as a model input array.
bool IsInputArray(const Model& model, const string& array_name);
// Returns true if the given array is specified as a model output array.
bool IsOutputArray(const Model& model, const string& array_name);

bool IsArrayConsumed(const Model& model, const string& name);
int CountTrueOutputs(const Model& model, const Operator& op);

int CountOpsWithInput(const Model& model, const string& array_name);
bool DeleteArrayIfUnused(const string& array_name, Model* model);
bool DeleteArrayIfUsedOnce(const string& array_name, Model* model);

// Deletes the op and any of its input and output arrays if they are unused
// after the op has been deleted.
void DeleteOpAndArraysIfUnused(Model* model, Operator* op);

std::vector<std::unique_ptr<Operator>>::const_iterator FindOpWithOutput(
    const Model& model, const string& array_name);
Operator* GetOpWithOutput(const Model& model, const string& array_name);

std::vector<std::unique_ptr<Operator>>::iterator FindOpWithOutput(
    Model& model, const string& array_name);

std::vector<std::unique_ptr<Operator>>::const_iterator FindOpWithInput(
    const Model& model, const string& array_name);

std::vector<std::unique_ptr<Operator>>::iterator FindOpWithInput(
    Model& model, const string& array_name);

Operator* GetOpWithInput(const Model& model, const string& array_name);
Operator* GetFirstOpWithInput(const Model& model, const string& array_name);

// Replaces all uses of the |old_array_name| with the |new_array_name|.
void ReplaceArrayUsage(Model* model, const string& old_array_name,
                       const string& new_array_name);

std::vector<std::unique_ptr<Operator>>::const_iterator FindOp(
    const Model& model, const Operator* op);
std::vector<std::unique_ptr<Operator>>::iterator FindOp(Model& model,
                                                        const Operator* op);

const char* OperatorTypeName(OperatorType type);
string HelpfulOperatorTypeName(const Operator& op);

// Whether the operator can be fused with an activation function. Note that this
// will return false by default for new operators; fusing support is opt-in.
bool OperatorSupportsFusedActivation(OperatorType type);

void DumpGraphvizVideoFrame(const Model& model);
void LogDump(int log_level, const string& message, const Model& model);
void LogSummary(int log_level, const string& message, const Model& model);

// TODO(b/36075966): Clean up when dims superseded by array shape.
void ExtendShape(Shape* shape, int new_shape_size);

// TODO(b/36075966): Clean up when dims superseded by array shape.
void UnextendShape(Shape* shape, int new_shape_size);

// Checks that all dimensions of 'shape' are at least 1. Note that scalars,
// lacking dimensions, satisfy this condition and are considered non-empty.
bool IsNonEmpty(const Shape& shape);

// Given two shapes with potentially different dimensionality and dimension
// arrays d0 and d1. Without loss of generality, assume that shape0 may have
// higher dimensionality (length(d0) >= length(d1)). Then shape0 and shape1
// "agree up to broadcasting" if:
// - When walking the d0 and d1 from back to front with indices i0, i1,
//   d0[i0] == d1[i1] or d0[i0] == 1 or d1[i1] == 1, for each dimension until
//   i1 == 0 (inclusive).
bool ShapesAgreeUpToBroadcasting(const Shape& shape0, const Shape& shape1);

// A stricter constraint than ShapesAgreeUpToBroadcasting().
//
// Given two shapes with potentially different dimensionality and dimension
// arrays d0 and d1. Without loss of generality, assume that shape0 may have
// higher dimensionality (length(d0) >= length(d1)). Then shape0 and shape1
// "agree up to extending" if:
// - When walking the d0 and d1 from back to front with indices i0, i1,
//   d0[i0] == d1[i1] for each dimension until i1 == 0 (inclusive).
// - For the remaining indices [0..i0), d0[i0] == 1.
bool ShapesAgreeUpToExtending(const Shape& shape0, const Shape& shape1);

inline ::tflite::RuntimeShape ToRuntimeShape(const Shape& shape) {
  return ::tflite::RuntimeShape(shape.dimensions_count(), shape.dims().data());
}

bool IsArrayFullyConnectedWeights(const Model& model, const string& name);

// If there is a wildcard dimension (-1), this may return a negative value.
int RequiredBufferSizeForShape(const Shape& shape);

bool IsConstantParameterArray(const Model& model, const string& name);

// Compares two constant parameter arrays for exact equality.
bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array);

void CheckNoMissingArray(const Model& model);
void CheckInvariants(const Model& model);

void CheckModelCounts(const Model& model);

void FixOperatorOrdering(Model* model);
void FixNoMissingArray(Model* model);
void FixNoOrphanedArray(Model* model);

// Fixes input/output arrays that may have issues during export or inference.
void FixEdgeArrays(Model* model);

// Finds and deduplicates large constant arrays in the model.
// After constant propagation runs it's possible to end up with several of the
// same large array (whether they be zeros or otherwise).
//
// |min_size| is used to adjust the minimum size in bytes of an array before
// it's considered for deduping. As deduping can make the graphs more difficult
// to read this helps prevent small arrays from spidering out.
void DedupeConstantArrays(Model* model, size_t min_size);

// Copies the contents of an array into another.
// Expects that the shape and data type match.
template <ArrayDataType A>
void CopyArrayBuffer(const Array& source_array, Array* target_array) {
  int source_buffer_size = RequiredBufferSizeForShape(source_array.shape());
  int target_buffer_size = RequiredBufferSizeForShape(target_array->shape());
  CHECK_EQ(source_buffer_size, target_buffer_size)
      << "Buffer sizes must match in element count";
  CHECK(source_array.data_type == target_array->data_type)
      << "Data types must match";
  if (source_array.buffer) {
    const auto& source_buffer = source_array.GetBuffer<A>();
    auto& target_buffer = target_array->GetMutableBuffer<A>();
    target_buffer.data = source_buffer.data;
  }
}

// Inserts a no-op reshape operator between the source array and the target
// array. This effectively just copies the data.
void InsertCopyOperator(Model* model, const string& source_array_name,
                        const string& target_array_name);

// Clones an array with all data and parameters.
void CloneArray(Model* model, const string& source_array_name,
                const string& target_array_name);

void ResolveModelFlags(const ModelFlags& model_flags, Model* model);

template <typename T>
T ConvertOperator(Operator* o, OperatorType type) {
  if (o != nullptr && o->type == type) {
    return static_cast<T>(o);
  }

  return nullptr;
}

void CheckIsReadyForQuantization(const Model& model);

bool ReshapeIsEquivalentToTranspose(const Model& model,
                                    const TensorFlowReshapeOperator* op,
                                    bool allow_extra_unary_dims);

inline int Offset(const Shape& shape, const std::vector<int>& indices) {
  DCHECK_EQ(shape.dimensions_count(), indices.size());
  const int dims_count = shape.dimensions_count();
  int offset = 0;
  for (int i = 0; i < dims_count; i++) {
    const int index = indices[i];
    DCHECK(index >= 0 && index < shape.dims(i));
    offset *= shape.dims(i);
    offset += index;
  }
  return offset;
}

inline std::vector<int> ReverseOffset(const Shape& shape, int index) {
  DCHECK_GE(index, 0);
  DCHECK_LT(index, RequiredBufferSizeForShape(shape));
  const int dims_count = shape.dimensions_count();
  std::vector<int> indices(dims_count);
  int residual = index;
  for (int i = dims_count - 1; i >= 0; i--) {
    indices[i] = residual % shape.dims(i);
    residual /= shape.dims(i);
  }
  return indices;
}

int ElementSize(ArrayDataType data_type);

void DropMinMax(Model* model, const string& array_name);

bool IsAllocatableTransientArray(const Model& model, const string& array_name);

void CreateOrCheckRnnStateArray(const string& name, int size, Model* model);

string AvailableArrayName(const Model& model, const string& name);

// Formats a shape as a string: [ dims(0), dims(1), ..., dims(num_dims-1) ].
string ShapeToString(const Shape& shape);

void PrintArrayShape(Model* model, const string& name);

void MakeArrayDims(int num_dims, int batch, int height, int width, int depth,
                   std::vector<int>* out_dims);

// Defines a constant int32 array with the provided values formatted for use
// as op parameters.
string CreateInt32Array(Model* model, const string& param_name,
                        const std::vector<int>& value);

bool EstimateArithmeticOpsCount(const Model& model, int64* result);

int AxesCount(AxesOrder axes_order);

// Returns the permutation of the dimensions based on the input axes order and
// output axes order.
void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order,
                     std::vector<int>* shuffle);

// Extend shuffle is designed to match ExtendShape, which pads the shape with
// unit dimensions at the beginning.
void ExtendShuffle(const std::vector<int>& input_shuffle, int newdim,
                   std::vector<int>* extended_shuffle);

void ShuffleDims(const Shape& input_shape, AxesOrder input_axes_order,
                 AxesOrder output_axes_order, Shape* output_shape);
void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order,
                  AxesOrder output_axes_order, const Shape& output_shape,
                  const float* input_data, float* output_data);
void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order,
                  AxesOrder output_axes_order, const Shape& output_shape,
                  const uint8* input_data, uint8* output_data);

// Returns true if it may be OK for any graph transformation to ever discard
// that array. The idea is that we can't ever discard arrays that are either
// an input or an output of the whole graph, or that appear in RNN back-edges,
// as that would undercut explicit flags that the user might pass.
bool IsDiscardableArray(const Model& model, const string& array_name);

void CheckFinalDataTypesSatisfied(const Model& model);

ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type);

// The process of building models varies according to the import format.
//
// (a) In some cases, such as model-proto format, the model should be fully
// specified. In these cases, no extra action should be taken by this function.
// (b) In other cases, such as TF graphdef format, the desired types of RNN
// arrays are not specified directly in the model, neither can they be inferred.
// However, we can set the types of RNN destination arrays to float. This breaks
// any cycles such as when resolution of the type of an RNN source array depends
// on the type of its destination array.
//
// This function is applied after the main import, after resolution of flags and
// after application of ArraysExtraInfo. It only defaults destination RNN arrays
// to float. If the model is subsequently quantized, it is assumed that the
// model contains sufficient information for that to be completed. If it is
// already quantized, then case (a) should hold.
void FinishBuildingRNNStates(Model* model);

void UseArraysExtraInfo(Model* model, bool quantize_output);

// Calculates the number of elements in tensor given a shape. Shape elements
// are assumed to be of type T, while the result total is of type U. If U
// doesn't have enough range to represent the sum of elements, an error is
// returned.
template <typename T, typename U>
tensorflow::Status NumElements(const std::vector<T>& shape, U* num_elements) {
  static_assert(
      std::numeric_limits<T>::max() <= std::numeric_limits<uint64_t>::max(),
      "vector type exceed capabilities of NumElements");

  *num_elements = 1;
  for (const T& dim : shape) {
    if (dim < 0) {
      // TensorFlow's shapes sometimes include -1 to represent an "unknown"
      // size but TOCO isn't able to create arrays of unknown sizes and will
      // crash in RequiredBufferSizeForShape().
      return tensorflow::errors::InvalidArgument(
          "Tensor shape should not include negative values");
    }
    if (static_cast<uint64_t>(dim) >
        std::numeric_limits<U>::max() / *num_elements) {
      *num_elements = 0;
      return tensorflow::errors::InvalidArgument("Tensor shape is too large");
    }
    *num_elements *= dim;
  }
  return tensorflow::Status::OK();
}

// A model file may have shuffled FC weights.
// When that happens, we want to de-shuffle them immediately on import,
// so that the rest of toco doesn't need to know about shuffled weights.
void UndoWeightsShuffling(Model* model);

// Copies minmax, quantization_params, and narrow_range.
void CopyMinMaxAndQuantizationRelatedFields(const Array& src, Array* dst);

}  // namespace toco

#endif  // TENSORFLOW_CONTRIB_LITE_TOCO_TOOLING_UTIL_H_