/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 <tvm/ffi/cast.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/attrs/distributed.h>
#include <tvm/relax/attrs/linear_algebra.h>
#include <tvm/relax/attrs/manipulate.h>
#include <tvm/relax/attrs/nn.h>
#include <tvm/relax/attrs/statistical.h>
#include <tvm/relax/distributed/axis_group_graph.h>
#include <tvm/relax/expr.h>

#include <numeric>

namespace tvm {
namespace tirx {
Var GetShardingVarFromIndex(PrimExpr index, ffi::Map<Var, Range> var_range,
                            const arith::Analyzer& analyzer) {
  if (auto prim_var = index.as<PrimVar>()) {
    return prim_var.value();
  }
  ffi::Map<PrimVar, Range> primitive_var_range;
  for (const auto& [var, range] : var_range) {
    primitive_var_range.Set(var.as_or_throw<PrimVar>(), range);
  }
  arith::IterSumExpr iter_sum = arith::NormalizeToIterSum(index, primitive_var_range, analyzer);
  if (!is_zero(iter_sum->base)) {
    return Var();
  }
  if (iter_sum->args.empty()) {
    return Var();
  }
  // floormod(floordiv(source, lower_factor), extent) * scale
  arith::IterSplitExpr highest_iter_split = iter_sum->args[0];
  auto source_var = highest_iter_split->source->source.as<PrimVar>();
  if (!source_var) {
    return Var();
  }
  Var var = source_var.value();
  // the floormod must take no effect
  if (!analyzer->CanProve(floordiv(var_range[var]->extent, highest_iter_split->lower_factor) <=
                          highest_iter_split->extent)) {
    return Var();
  }
  return var;
}
}  // namespace tirx
}  // namespace tvm

namespace tvm {

namespace relax {
namespace distributed {

const TensorTypeNode* GetTensorType(Expr tensor) {
  const auto* tensor_ty = GetTypeAs<TensorTypeNode>(tensor);
  if (tensor_ty) {
    return tensor_ty;
  }
  const auto* dtensor_ty = GetTypeAs<DTensorTypeNode>(tensor);
  if (dtensor_ty) {
    return dtensor_ty->tensor_ty.get();
  }
  TVM_FFI_THROW(InternalError) << tensor << " must be either Tensor or DTesor";
  throw;
}

void UnaryOpHelper(ffi::Array<Expr> tensor_list, distributed::AxisGroupGraph* axis_group_graph) {
  int n_dim = GetTensorType(tensor_list[0])->ndim;
  for (const auto& tensor : tensor_list) {
    TVM_FFI_ICHECK(GetTensorType(tensor)->ndim == n_dim);
  }
  for (int i = 0; i < n_dim; i++) {
    TVM_FFI_ICHECK(tensor_list.size() <= 2);
    for (int j = 0; j < static_cast<int>(tensor_list.size()) - 1; j++) {
      axis_group_graph->JoinAxis({tensor_list[j].get(), i}, {tensor_list[j + 1].get(), i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    }
  }
}

void BuildAxisGraphUnary(const Var& output_var, const Call& call,
                         distributed::AxisGroupGraph* axis_group_graph) {
  ffi::Array<Expr> tensor_list;  // vars in param and output
  if (call->args[0]->IsInstance<VarNode>()) {
    tensor_list.push_back(call->args[0]);
  }
  tensor_list.push_back(output_var);
  UnaryOpHelper(tensor_list, axis_group_graph);
}

void BuildAxisGraphBinary(const Var& output_var, const Call& call,
                          distributed::AxisGroupGraph* axis_group_graph) {
  ffi::Array<Expr> tensor_list;  // vars in param and output
  if (call->args[0]->ty.as<TensorTypeNode>() || call->args[0]->ty.as<DTensorTypeNode>()) {
    tensor_list.push_back(call->args[0]);
  }
  if (call->args[1]->ty.as<TensorTypeNode>() || call->args[1]->ty.as<DTensorTypeNode>()) {
    tensor_list.push_back(call->args[1]);
  }
  tensor_list.push_back(output_var);
  if (tensor_list.size() <= 2) {
    UnaryOpHelper(tensor_list, axis_group_graph);
    return;
  }
  const auto* x1_ty = GetTensorType(tensor_list[0]);
  const auto* x2_ty = GetTensorType(tensor_list[1]);
  int x1_ndim = x1_ty->ndim;
  int x2_ndim = x2_ty->ndim;
  const auto* x1_shape = x1_ty->shape.as<ShapeExprNode>();
  const auto* x2_shape = x2_ty->shape.as<ShapeExprNode>();
  TVM_FFI_ICHECK(x1_shape && x2_shape);
  arith::Analyzer analyzer;
  for (int i = 1; i <= std::min(x1_ndim, x2_ndim); ++i) {
    const PrimExpr& dim0 = x1_shape->values[x1_ndim - i];
    const PrimExpr& dim1 = x2_shape->values[x2_ndim - i];
    if (analyzer->CanProveEqual(dim0, dim1)) {
      // join batch dim
      axis_group_graph->JoinAxis({tensor_list[0].get(), x1_ndim - i},
                                 {tensor_list[2].get(), std::max(x1_ndim, x2_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
      axis_group_graph->JoinAxis({tensor_list[1].get(), x2_ndim - i},
                                 {tensor_list[2].get(), std::max(x1_ndim, x2_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    } else if (analyzer->CanProveEqual(dim0, 1)) {
      axis_group_graph->JoinAxis({tensor_list[1].get(), x2_ndim - i},
                                 {tensor_list[2].get(), std::max(x1_ndim, x2_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    } else if (analyzer->CanProveEqual(dim1, 1)) {
      axis_group_graph->JoinAxis({tensor_list[0].get(), x1_ndim - i},
                                 {tensor_list[2].get(), std::max(x1_ndim, x2_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    } else {
      TVM_FFI_THROW(InternalError) << "Invalid broadcast, dim0: " << dim0 << ", dim1: " << dim1;
    }
  }
  if (x1_ndim > x2_ndim) {
    for (int i = 0; i < x1_ndim - x2_ndim; i++) {
      axis_group_graph->JoinAxis({tensor_list[0].get(), i}, {tensor_list[2].get(), i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    }
  } else if (x1_ndim < x2_ndim) {
    for (int i = 0; i < x2_ndim - x1_ndim; i++) {
      axis_group_graph->JoinAxis({tensor_list[1].get(), i}, {tensor_list[2].get(), i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    }
  }
}

void BuildAxisGraphReduce(const Var& output_var, const Call& call,
                          distributed::AxisGroupGraph* axis_group_graph) {
  Expr input_tensor = call->args[0];
  ffi::Array<int64_t> axes;
  bool keepdims;
  if (const auto* attrs = call->attrs.as<StatisticalAttrs>()) {
    if (attrs->axis.has_value()) {
      axes = attrs->axis.value();
    }
    keepdims = attrs->keepdims;
  } else if (const auto* attrs = call->attrs.as<SoftmaxAttrs>()) {
    axes = {attrs->axis};
    keepdims = true;
  } else {
    TVM_FFI_THROW(InternalError) << "Unsupported reduce op: " << call->op;
  }

  int ndim = GetTensorType(input_tensor)->ndim;

  std::unordered_set<int> normalized_axes;
  for (int64_t i : axes) {
    int val = static_cast<int>(i);
    TVM_FFI_ICHECK(val < ndim && val >= -ndim);
    if (val < 0) {
      val = ndim + val;
    }
    normalized_axes.insert(val);
  }
  if (keepdims) {
    for (int i = 0; i < ndim; i++) {
      if (!normalized_axes.count(i)) {
        axis_group_graph->JoinAxis({input_tensor.get(), i}, {output_var.get(), i},
                                   distributed::AxisGroupGraph::EdgeType::kDescend);
      }
    }
  } else {
    for (int i = 0, j = 0; i < ndim; i++) {
      if (!normalized_axes.count(i)) {
        axis_group_graph->JoinAxis({input_tensor.get(), i}, {output_var.get(), j},
                                   distributed::AxisGroupGraph::EdgeType::kDescend);
        j++;
      }
    }
  }
}

void BuildAxisGraphMatmul(const Var& output_var, const Call& call,
                          distributed::AxisGroupGraph* axis_group_graph) {
  Expr x1 = call->args[0];
  Expr x2 = call->args[1];
  Var x3 = output_var;
  const auto* x1_ty = GetTensorType(x1);
  const auto* x2_ty = GetTensorType(x2);
  int x1_ndim = x1_ty->ndim;
  int x2_ndim = x2_ty->ndim;
  TVM_FFI_ICHECK(x1_ndim > 0 && x2_ndim > 0);
  int x1_prepended = 0;
  int x2_appended = 0;
  if (x1_ndim == 1) {
    x1_ndim = 2;
    x1_prepended = 1;
  }
  if (x2_ndim == 1) {
    x2_ndim = 2;
    x2_appended = 1;
  }
  const auto* x1_shape = x1_ty->shape.as<ShapeExprNode>();
  const auto* x2_shape = x2_ty->shape.as<ShapeExprNode>();
  TVM_FFI_ICHECK(x1_shape && x2_shape);
  ffi::Array<PrimExpr> x1_shape_prefix{x1_shape->values.begin(),
                                       x1_shape->values.end() - 2 + x1_prepended};
  ffi::Array<PrimExpr> x2_shape_prefix{x2_shape->values.begin(),
                                       x2_shape->values.end() - 2 + x2_appended};

  int x1_prefix_ndim = x1_shape_prefix.size();
  int x2_prefix_ndim = x2_shape_prefix.size();
  arith::Analyzer analyzer;
  for (int i = 1; i <= std::min(x1_prefix_ndim, x2_prefix_ndim); ++i) {
    const PrimExpr& dim0 = x1_shape_prefix[x1_prefix_ndim - i];
    const PrimExpr& dim1 = x2_shape_prefix[x2_prefix_ndim - i];
    // join batch dim
    if (analyzer->CanProveEqual(dim0, dim1)) {
      axis_group_graph->JoinAxis({x1.get(), x1_prefix_ndim - i},
                                 {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
      axis_group_graph->JoinAxis({x2.get(), x2_prefix_ndim - i},
                                 {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    } else if (analyzer->CanProveEqual(dim0, 1)) {
      axis_group_graph->JoinAxis({x2.get(), x2_prefix_ndim - i},
                                 {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    } else if (analyzer->CanProveEqual(dim1, 1)) {
      axis_group_graph->JoinAxis({x1.get(), x1_prefix_ndim - i},
                                 {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim) - i},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    } else {
      TVM_FFI_THROW(InternalError) << "Cannot broadcast " << dim0 << " and " << dim1;
    }
  }
  // join reduction dim
  axis_group_graph->JoinAxis({x1.get(), x1_ty->ndim - 1}, {x2.get(), x2_ndim - 2},
                             distributed::AxisGroupGraph::EdgeType::kSimbling);
  // join lhs_spatial dim and rhs_spatial dim
  if (!x1_prepended) {
    axis_group_graph->JoinAxis({x1.get(), x1_ndim - 2},
                               {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim)},
                               distributed::AxisGroupGraph::EdgeType::kDescend);
    if (!x2_appended) {
      axis_group_graph->JoinAxis({x2.get(), x2_ndim - 1},
                                 {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim) + 1},
                                 distributed::AxisGroupGraph::EdgeType::kDescend);
    }
  } else if (!x2_appended) {
    axis_group_graph->JoinAxis({x2.get(), x2_ndim - 1},
                               {x3.get(), std::max(x1_prefix_ndim, x2_prefix_ndim)},
                               distributed::AxisGroupGraph::EdgeType::kDescend);
  }
}

void BuildAxisGraphPermuteDims(const Var& output_var, const Call& call,
                               distributed::AxisGroupGraph* axis_group_graph) {
  Expr input_tensor = call->args[0];
  const auto* attrs = call->attrs.as<PermuteDimsAttrs>();
  TVM_FFI_ICHECK(attrs);
  int ndim = GetTensorType(input_tensor)->ndim;
  std::vector<int> normalized_axes;
  if (attrs->axes.has_value()) {
    for (int64_t i : attrs->axes.value()) {
      int val = static_cast<int>(i);
      TVM_FFI_ICHECK(val < ndim && val >= -ndim);
      if (val < 0) {
        val = ndim + val;
      }
      normalized_axes.push_back(val);
    }
  } else {
    normalized_axes.resize(ndim);
    std::iota(normalized_axes.rbegin(), normalized_axes.rend(), 0);
  }
  for (int i = 0; i < ndim; i++) {
    axis_group_graph->JoinAxis({input_tensor.get(), normalized_axes[i]}, {output_var.get(), i},
                               distributed::AxisGroupGraph::EdgeType::kDescend);
  }
}
void BuildAxisGraphReshape(const Var& output_var, const Call& call,
                           distributed::AxisGroupGraph* axis_group_graph) {
  Expr input_tensor = call->args[0];
  const auto* tensor_ty = GetTensorType(input_tensor);
  const auto* new_shape_ty = GetTypeAs<ShapeTypeNode>(call->args[1]);
  const auto* old_shape_ty = GetTypeAs<ShapeTypeNode>(tensor_ty->shape.value());
  TVM_FFI_ICHECK_NOTNULL(old_shape_ty);
  ffi::Array<PrimExpr> old_shape_values = old_shape_ty->values.value();
  ffi::Array<PrimExpr> new_shape_values = new_shape_ty->values.value();
  int i = old_shape_values.size();
  int j = new_shape_values.size();
  PrimExpr old_shape_product = 1, new_shape_product = 1;
  arith::Analyzer analyzer_;
  while (i > 0 && j > 0) {
    if (analyzer_->CanProve(new_shape_product > old_shape_product)) {
      i--;
      old_shape_product *= old_shape_values[i];
    } else if (analyzer_->CanProve(new_shape_product < old_shape_product)) {
      j--;
      new_shape_product *= new_shape_values[j];
    } else {
      if (i != static_cast<int>(old_shape_values.size())) {
        axis_group_graph->JoinAxis({input_tensor.get(), i}, {output_var.get(), j},
                                   distributed::AxisGroupGraph::EdgeType::kDescend);
      }
      i--;
      j--;
      old_shape_product *= old_shape_values[i];
      new_shape_product *= new_shape_values[j];
    }
  }
}

inline int GetNumOutput(Call call) {
  Type output_ty = call->ty_args[0];
  if (const auto* tuple_ty = output_ty.as<TupleTypeNode>()) {
    return tuple_ty->fields.size();
  } else {
    return 1;
  }
}

void BuildAxisGraphCallTIR(const Var& output_var, const Call& call, const tirx::PrimFunc& func,
                           distributed::AxisGroupGraph* axis_group_graph) {
  auto tir_var_axis_group_list = tirx::BufferAxisGraphExtractor::GetTIRVarAxisGraph(func);
  ffi::Map<tirx::Var, Expr> input_var_to_relax_expr;
  ffi::Array<Expr> input_list = call->args[1].as_or_throw<Tuple>()->fields;
  input_list.push_back(output_var);
  for (int i = 0; i < static_cast<int>(input_list.size()); i++) {
    if (func->params[i]->ty.as<tirx::BufferTypeNode>()) {
      input_var_to_relax_expr.Set(func->params[i], input_list[i]);
    }
  }
  int num_params = func->params.size();
  int num_outputs = GetNumOutput(call);
  for (const auto& var_axis_group : tir_var_axis_group_list) {
    std::unordered_map<int, int> output_tensor_indices;
    for (int i = 0; i < static_cast<int>(var_axis_group.size()); i++) {
      for (int j = num_params - num_outputs; j < num_params; j++) {
        if (func->params[j].same_as(var_axis_group[i].first)) {
          output_tensor_indices[i] = j - num_params + num_outputs;
          break;
        }
      }
    }
    if (output_tensor_indices.empty()) {
      for (int i = 1; i < static_cast<int>(var_axis_group.size()); i++) {
        axis_group_graph->JoinAxis(
            {input_var_to_relax_expr[var_axis_group[i].first].get(), var_axis_group[i].second},
            {input_var_to_relax_expr[var_axis_group[0].first].get(), var_axis_group[0].second},
            distributed::AxisGroupGraph::EdgeType::kSimbling);
      }
    } else {
      for (const auto& pr : output_tensor_indices) {
        for (int i = 0; i < static_cast<int>(var_axis_group.size()); i++) {
          if (!output_tensor_indices.count(i)) {
            axis_group_graph->JoinAxis(
                {input_var_to_relax_expr[var_axis_group[i].first].get(), var_axis_group[i].second},
                {output_var.get(), var_axis_group[pr.first].second, pr.second},
                distributed::AxisGroupGraph::EdgeType::kDescend);
          }
        }
      }
    }
  }
}
}  // namespace distributed
}  // namespace relax
}  // namespace tvm
