/*
 * Copyright (c) 2019 Samsung Electronics Co., Ltd. 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.
 */

#include "moco/Pass/Passes/FuseBinaryIntoPreceding.h"

#include <moco/Support/TFShapeInferenceHelper.h>

#include <moco/IR/TFDialect.h>
#include <moco/IR/Nodes/TFAdd.h>
#include <moco/IR/Nodes/TFBiasAdd.h>
#include <moco/IR/Nodes/TFConst.h>
#include <moco/IR/Nodes/TFConv2D.h>
#include <moco/IR/Nodes/TFDepthwiseConv2dNative.h>
#include <moco/IR/Nodes/TFMul.h>

#include <cassert>
#include <memory>

namespace
{

/**
 * @brief Fusable operation type
 */
enum class FuseType
{
  Conv2D,
  DepthwiseConv2D,
  // TODO Support FullyConnected
};

// TODO rename this method when there is a better name
bool is_only_one_valid(moco::TFConst *xc, moco::TFConst *yc)
{
  if (xc == nullptr && yc == nullptr)
    return false;
  if (xc != nullptr && yc != nullptr)
    return false;

  return true;
}

// TODO Put this in some common place
void copy_shape(const moco::TFConst *src, moco::TFConst *dst)
{
  assert(src != nullptr);
  assert(dst != nullptr);

  uint32_t rank = src->rank();
  dst->rank(rank);
  for (uint32_t index = 0; index < rank; ++index)
  {
    if (src->dim(index).known())
      dst->dim(index) = src->dim(index);
    else
      dst->dim(index).unset();
  }
}

/**
 * @brief return true if shape is identical
 */
bool shape_match(const moco::TFConst *c1, const moco::TFConst *c2)
{
  assert(c1 != nullptr);
  assert(c2 != nullptr);

  uint32_t rank = c1->rank();
  if (rank != c2->rank())
    return false;

  for (uint32_t index = 0; index < rank; ++index)
  {
    if (!c1->dim(index).known() || !c2->dim(index).known())
      return false;

    if (c1->dim(index).value() != c2->dim(index).value())
      return false;
  }
  return true;
}

template <FuseType FT>
moco::TFConst *create_kernel_from_fuse_mulparam(loco::Graph *graph, moco::TFConst *ker,
                                                moco::TFConst *mulparam);

template <>
moco::TFConst *create_kernel_from_fuse_mulparam<FuseType::Conv2D>(loco::Graph *graph,
                                                                  moco::TFConst *ker,
                                                                  moco::TFConst *mulparam)
{
  auto ker_shape_inf = moco::node_shape(ker);
  assert(ker_shape_inf.domain() != loco::Domain::Unknown);
  auto ker_shape = ker_shape_inf.as<loco::TensorShape>();

  auto mulparam_shape_inf = moco::node_shape(mulparam);
  assert(mulparam_shape_inf.domain() != loco::Domain::Unknown);
  auto mulparam_shape = mulparam_shape_inf.as<loco::TensorShape>();

  // create new ker_fused with same size of ker
  auto ker_fused = graph->nodes()->create<moco::TFConst>();

  assert(ker_shape.rank() == 4);
  assert(mulparam_shape.rank() == 1);
  assert(ker_shape.dim(3).value() == mulparam_shape.dim(0).value());

  ker_fused->dtype(loco::DataType::FLOAT32);
  copy_shape(ker, ker_fused);
  auto ker_num_elements = ker->size<loco::DataType::FLOAT32>();
  ker_fused->size<loco::DataType::FLOAT32>(ker_num_elements);

  // TensorFlow Conv2D Kernel has HWIO format
  // Broadcast Mul vector to Kernel tensor by the Output
  const uint32_t ker_height = ker_shape.dim(0).value();
  const uint32_t ker_width = ker_shape.dim(1).value();
  const uint32_t ker_input = ker_shape.dim(2).value();
  const uint32_t ker_output = ker_shape.dim(3).value();

  for (uint32_t ker_y = 0; ker_y < ker_height; ++ker_y)
  {
    for (uint32_t ker_x = 0; ker_x < ker_width; ++ker_x)
    {
      for (uint32_t in_ch = 0; in_ch < ker_input; ++in_ch)
      {
        uint32_t num_items = ((ker_y * ker_width + ker_x) * ker_input + in_ch) * ker_output;
        for (uint32_t out_ch = 0; out_ch < ker_output; ++out_ch)
        {
          auto mulparam_v = mulparam->at<loco::DataType::FLOAT32>(out_ch);
          auto ker_v = ker->at<loco::DataType::FLOAT32>(num_items + out_ch);
          ker_fused->at<loco::DataType::FLOAT32>(num_items + out_ch) = ker_v * mulparam_v;
        }
      }
    }
  }

  return ker_fused;
}

/**
 * @brief Create a kernel from fuse mulparam<FuseType::DepthwiseConv2D> object
 * @return Kernel of fused mulparam
 */
template <>
moco::TFConst *create_kernel_from_fuse_mulparam<FuseType::DepthwiseConv2D>(loco::Graph *graph,
                                                                           moco::TFConst *ker,
                                                                           moco::TFConst *mulparam)
{
  auto ker_shape_inf = moco::node_shape(ker);
  assert(ker_shape_inf.domain() != loco::Domain::Unknown);
  auto ker_shape = ker_shape_inf.as<loco::TensorShape>();

  auto mulparam_shape_inf = moco::node_shape(mulparam);
  assert(mulparam_shape_inf.domain() != loco::Domain::Unknown);
  auto mulparam_shape = mulparam_shape_inf.as<loco::TensorShape>();

  // create new ker_fused with same size of ker
  auto ker_fused = graph->nodes()->create<moco::TFConst>();

  assert(ker_shape.rank() == 4);
  assert(mulparam_shape.rank() == 1);
  assert(ker_shape.dim(2).value() * ker_shape.dim(3).value() == mulparam_shape.dim(0).value());

  ker_fused->dtype(loco::DataType::FLOAT32);
  copy_shape(ker, ker_fused);
  auto ker_num_elements = ker->size<loco::DataType::FLOAT32>();
  ker_fused->size<loco::DataType::FLOAT32>(ker_num_elements);

  // TensorFlow DepthwiseConv2DNative Kernel has HWIM format
  // Broadcast Mul vector to Kernel tensor by the Output
  const uint32_t ker_height = ker_shape.dim(0).value();
  const uint32_t ker_width = ker_shape.dim(1).value();
  const uint32_t ker_input = ker_shape.dim(2).value();
  const uint32_t ker_multiplier = ker_shape.dim(3).value();

  for (uint32_t ker_y = 0; ker_y < ker_height; ++ker_y)
  {
    for (uint32_t ker_x = 0; ker_x < ker_width; ++ker_x)
    {
      for (uint32_t in_ch = 0; in_ch < ker_input; ++in_ch)
      {
        uint32_t num_items = ((ker_y * ker_width + ker_x) * ker_input + in_ch) * ker_multiplier;
        for (uint32_t ker_ch = 0; ker_ch < ker_multiplier; ++ker_ch)
        {
          auto mulparam_v = mulparam->at<loco::DataType::FLOAT32>(in_ch + ker_ch * ker_input);
          auto ker_v = ker->at<loco::DataType::FLOAT32>(num_items + ker_ch);
          ker_fused->at<loco::DataType::FLOAT32>(num_items + ker_ch) = ker_v * mulparam_v;
        }
      }
    }
  }

  return ker_fused;
}

/**
 * @brief Create a fused convolution opertion from kernel of fused mulparam
 * @return Fused convolution operation
 */
template <FuseType FT, class T>
T *fused_conv_node(loco::Graph *graph, moco::TFConst *mulparam, T *conv_node)
{
  // LOGGER(l);

  // ker should be constant
  auto ker = dynamic_cast<moco::TFConst *>(conv_node->filter());
  if (ker == nullptr)
  {
    // Wait until ker is becomes TFConst: there are cases when it's Identity.
    // INFO(l) << "Mul fuse_to_preceding: precedingOp ker is not TFConst";
    return nullptr;
  }
  auto ifm = conv_node->input();
  assert(ifm != nullptr);

  // we need shape information, if not wait till it's ready
  auto ker_shape_inf = moco::node_shape(ker);
  if (ker_shape_inf.domain() == loco::Domain::Unknown)
  {
    // INFO(l) << "Mul fuse_to_preceding: precedingOp ker has no shape";
    return nullptr;
  }
  auto mulparam_shape_inf = moco::node_shape(mulparam);
  if (mulparam_shape_inf.domain() == loco::Domain::Unknown)
  {
    // INFO(l) << "Mul fuse_to_preceding: precedingOp mulparam has no shape";
    return nullptr;
  }
  // if MulParam rank is not 1 we cannot fuse, just skip
  auto mulparam_shape = mulparam_shape_inf.as<loco::TensorShape>();
  if (mulparam_shape.rank() != 1)
  {
    // INFO(l) << "Mul fuse_to_preceding: Mul rank is not 1";
    return nullptr;
  }

  auto ker_fused = create_kernel_from_fuse_mulparam<FT>(graph, ker, mulparam);
  auto conv_fused = graph->nodes()->create<T>();

  conv_fused->input(ifm);
  conv_fused->filter(ker_fused);
  conv_fused->padding(conv_node->padding());
  conv_fused->data_layout(conv_node->data_layout());
  conv_fused->strides(conv_node->strides());

  return conv_fused;
}

/**
 * @note This creates fused ker:2 from ker:1, 'mulparam' and
 *       new precedingOp:2 that uses ker:2 as the kernel.
 *       Then make C to use precedingOp:2 as new input.
 *
 * <Before>
 *                                   mulparam-\
 *          ker:1  --\                         \
 *          ifm  ----- precedingOp:1 ----------- Mul --- C
 *
 *
 * <After>
 *                                   mulparam-\
 *          ker:1  --\                         \
 *                   - precedingOp:1 ----------- Mul ---
 *                   /
 *          ifm  ----- precedingOp:2 ------------------- C
 *          ker:2 ---/
 *
 *
 * [Where]
 *     - precedingOp:1 can be one of TFConv2D, TFDepthwiseConv2dNative, FullyConnected
 *     - 'mulparam' and Mul will be disconnected from the Output.
 *     - ker:2 is added with fused values of ker:1 and mulparam
 *     - precedingOp:2 is added using ifm and ker:2 and other parameters
 *       same as precedingOp:1.
 *     - ker:1, precedingOp:1, 'mulparam' and Mul should be removed in
 *       RemoveDeadNodeTransform if not used.
 */
bool fuse_to_preceding(loco::Graph *graph, moco::TFMul *node)
{
  auto xc = dynamic_cast<moco::TFConst *>(node->x());
  auto yc = dynamic_cast<moco::TFConst *>(node->y());

  // Note: if both are constants, it should be done by constant-folding
  if (!(is_only_one_valid(xc, yc)))
    return false;

  moco::TFConst *mulparam = nullptr;
  moco::TFNode *precedingOp = nullptr;

  if (xc != nullptr)
  {
    mulparam = xc;
    precedingOp = dynamic_cast<moco::TFNode *>(node->y());
  }
  else // yc != nullptr
  {
    mulparam = yc;
    precedingOp = dynamic_cast<moco::TFNode *>(node->x());
  }

  assert(mulparam->dtype() == loco::DataType::FLOAT32);

  // TODO support FullyConnected
  moco::TFNode *fused_node = nullptr;
  if (auto conv2d = dynamic_cast<moco::TFConv2D *>(precedingOp))
    fused_node = fused_conv_node<FuseType::Conv2D, moco::TFConv2D>(graph, mulparam, conv2d);
  else if (auto dw_conv2d = dynamic_cast<moco::TFDepthwiseConv2dNative *>(precedingOp))
    fused_node = fused_conv_node<FuseType::DepthwiseConv2D, moco::TFDepthwiseConv2dNative>(
      graph, mulparam, dw_conv2d);

  // Not ready yet
  if (fused_node == nullptr)
    return false;

  // Replace TFMul node with new precedingOp with fused kernel
  // This will leave existing precedingOp as-is but can be removed if not used
  // from other transformations
  replace(node).with(fused_node);
  // TODO check if need to disconnect
  // node->x(nullptr);
  // node->y(nullptr);
  // fused_node->ifm(nullptr);
  // fused_node->ker(nullptr);

  return true;
}

/**
 * @brief Create zero-filled BiasAdd opertion and insert after precedingOp
 *        The plan is to fuse 'addparam' to TFBiasAdd bias
 * @return Zero-filled BiasAdd operation
 */
template <class T>
moco::TFBiasAdd *create_biasadd_node(loco::Graph *graph, moco::TFConst *addparam, T *precedingOp)
{
  auto dtype = addparam->dtype();
  assert(dtype == loco::DataType::FLOAT32);

  // Create TFConst(bias of TFBiasAdd) with same shape and dtype of 'addparam' but
  // with values 0.0
  auto biasadd_param = graph->nodes()->create<moco::TFConst>();
  biasadd_param->dtype(dtype);
  copy_shape(addparam, biasadd_param);
  auto biasadd_num_elements = addparam->size<loco::DataType::FLOAT32>();
  biasadd_param->size<loco::DataType::FLOAT32>(biasadd_num_elements);
  for (int32_t i = 0; i < biasadd_num_elements; i++)
  {
    biasadd_param->at<loco::DataType::FLOAT32>(i) = 0.0f;
  }

  // Create TFBiasAdd with same shape as TFAdd
  auto data_layout = precedingOp->data_layout();
  auto tf_biasadd = graph->nodes()->create<moco::TFBiasAdd>();
  tf_biasadd->data_layout(data_layout);

  loco::replace(precedingOp).with(tf_biasadd);
  tf_biasadd->value(precedingOp);
  tf_biasadd->bias(biasadd_param);

  return tf_biasadd;
}

/**
 * @note TFAdd will be fused to TFBiasAdd
 *
 * <Before>
 * If precedingOp is not TFBiasAdd, then insert TFConst:1 + TFBiasAdd that
 * TFConst:1 has zero values.
 *
 *                                      addparam --\
 *                                                  \
 *          precedingOp  ---------------------------- TFAdd ----- C
 *
 *
 * <Intermediate>
 * If it's TFBiasAdd and one of the input is TFConst type,
 * then we can fuse 'addparam' to the input TFConst:2 value of TFBiasAdd, where
 * TFConst:2 has added values from 'addparam'
 *
 *                                      addparam --\
 *          TFConst:1  --------\                    \
 *          precedingOp  ------- TFBiasAdd ---------- TFAdd ----- C
 *
 *
 * <After>
 *                                      addparam --\
 *          TFConst:2  --------\                    \
 *          precedingOp  ------- TFBiasAdd ---------- TFAdd -----
 *                                         \--------------------- C
 *
 *
 * [Where]
 *     - precedingOp can be TFConv2D, TFDepthwiseConv2dNative, FullyConnected,
 *       TFBiasAdd.
 *     - Intermediate is to insert TFBiasAdd + TFConst:1
 *     - After is to fuse 'addparam' of TFAdd into TFConst:1 + TFBiasAdd
 *       that becomes TFConst:2 + TFBiasAdd
 */
bool fuse_to_preceding(loco::Graph *graph, moco::TFAdd *node)
{
  // LOGGER(l);

  auto xc = dynamic_cast<moco::TFConst *>(node->x());
  auto yc = dynamic_cast<moco::TFConst *>(node->y());

  // Note: if both are constants, it should be done by constant-folding
  if (!(is_only_one_valid(xc, yc)))
    return false;

  moco::TFConst *addparam = nullptr;
  moco::TFNode *precedingOp = nullptr;

  if (xc != nullptr)
  {
    addparam = xc;
    precedingOp = dynamic_cast<moco::TFNode *>(node->y());
  }
  else // yc != nullptr
  {
    addparam = yc;
    precedingOp = dynamic_cast<moco::TFNode *>(node->x());
  }

  auto addparam_shape_inf = moco::node_shape(addparam);
  if (addparam_shape_inf.domain() == loco::Domain::Unknown)
  {
    // INFO(l) << "Add fuse_to_preceding: addparam has no shape";
    return false;
  }
  // if AddParam rank is not 0 or 1 we cannot fuse, just skip
  auto addparam_shape = addparam_shape_inf.as<loco::TensorShape>();
  if (addparam_shape.rank() > 1)
  {
    // INFO(l) << "Add fuse_to_preceding: Add rank is not 0 or 1";
    return false;
  }

  // TODO do something when rank() is 0
  if (addparam_shape.rank() == 0)
  {
    // Not supported yet
    return false;
  }
  assert(addparam_shape.rank() != 0);

  // TODO support FullyConnected
  moco::TFBiasAdd *biasadd = nullptr;
  if (auto conv2d = dynamic_cast<moco::TFConv2D *>(precedingOp))
    biasadd = create_biasadd_node<moco::TFConv2D>(graph, addparam, conv2d);
  else if (auto dw_conv2d = dynamic_cast<moco::TFDepthwiseConv2dNative *>(precedingOp))
    biasadd = create_biasadd_node<moco::TFDepthwiseConv2dNative>(graph, addparam, dw_conv2d);
  else if (auto old_bias_add = dynamic_cast<moco::TFBiasAdd *>(precedingOp))
    biasadd = old_bias_add;

  if (biasadd == nullptr)
  {
    // try next turn
    return false;
  }

  // Let's fuse addparam into biasadd bias
  auto biasadd_bias = loco::must_cast<moco::TFConst *>(biasadd->bias());
  if (!shape_match(biasadd_bias, addparam))
  {
    // INFO(l) << "TFBiasAdd bias and TFAdd input shape mismatch";
    return false;
  }
  auto add_num_elements = addparam->size<loco::DataType::FLOAT32>();
  assert(add_num_elements == biasadd_bias->size<loco::DataType::FLOAT32>());
  for (int32_t i = 0; i < add_num_elements; i++)
  {
    biasadd_bias->at<loco::DataType::FLOAT32>(i) += addparam->at<loco::DataType::FLOAT32>(i);
  }

  replace(node).with(biasadd);
  // TODO check if need to disconnect
  // node->x(nullptr);
  // node->y(nullptr);

  return true;
}

} // namespace

namespace moco
{

bool FuseBinaryIntoPreceding::run(loco::Graph *graph)
{
  bool changed = false;
  auto active_nodes = loco::active_nodes(loco::output_nodes(graph));

  for (auto node : active_nodes)
  {
    if (node->dialect() == moco::TFDialect::get())
    {
      {
        auto tf_node = dynamic_cast<moco::TFMul *>(node);
        if (tf_node != nullptr)
        {
          if (fuse_to_preceding(graph, tf_node))
            changed = true;
        }
      }
      {
        // TODO support Div
      }

      {
        auto tf_node = dynamic_cast<moco::TFAdd *>(node);
        if (tf_node != nullptr)
        {
          if (fuse_to_preceding(graph, tf_node))
            changed = true;
        }
      }
      {
        // TODO support Sub
      }
    }
  }

  return changed;
}

} // namespace moco