/*
 * Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.
 *
 * This source code and/or documentation ("Licensed Deliverables") are
 * subject to NVIDIA intellectual property rights under U.S. and
 * international Copyright laws.
 */

#include <cugraph-ops/operators/common.hpp>

#include <cmath>
#include <cstddef>
#include <cstring>
#include <memory>

namespace cugraph::ops::mha_simple_csc {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_kernel(DataT* output,
                          DataT* sm_scores,
                          const DataT* key_feat /* key_src_feat */,
                          const DataT* query_feat /* query_dst_feat */,
                          const DataT* value_feat /* key_src_feat */,
                          const DataT* edge_feat,
                          const DataT* score_bias,
                          const DataT* score_weight,
                          bool norm_by_dim,
                          const IdxT* offsets,
                          const IdxT* indices,
                          const IdxT* map_csc_to_coo,
                          IdxT n_dst_nodes,
                          IdxT n_indices,
                          int dim_node,
                          const mha_params& params)
{
  // if we add own node, sm scores come after the edge sm sores of all edges
  auto n_edges          = n_indices;
  int dim_head          = dim_node / params.num_heads;
  DataT norm_head       = DataT{1} / static_cast<DataT>(params.num_heads);
  bool add_bias         = score_bias != nullptr;
  bool use_score_weight = score_weight != nullptr;
  DataT scalar_inv_norm = DataT{1} / ::sqrt(static_cast<DataT>(dim_head));

  if (params.concat_heads) {
    std::memset(output, 0, n_dst_nodes * dim_node * sizeof(DataT));
  } else {
    std::memset(output, 0, n_dst_nodes * dim_head * sizeof(DataT));
  }
  std::memset(sm_scores, 0, 2 * n_edges * params.num_heads * sizeof(DataT));
  static constexpr DataT ZERO = DataT{0};

  for (auto dst_node_id = 0; dst_node_id < n_dst_nodes; ++dst_node_id) {
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];
    auto n_neighbors = off_end - off_start;
    auto scores      = std::make_unique<DataT[]>(n_neighbors);

    for (int h = 0; h < params.num_heads; ++h) {
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        scores[edge] = ZERO;
      }
      auto sm_off_pre  = h * n_edges;
      auto sm_off_post = (params.num_heads + h) * n_edges;
      auto sm_off_pre_mul =
        (params.num_heads * 2 + h) * n_edges;  // only needed if score is weighted
      // we calculate softmax manually here. first get all the scores
      for (int i = 0; i < dim_head; ++i) {
        auto h_i = h * dim_head + i;
        for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
          auto eid       = off_start + edge;
          auto neigh_off = indices[eid];
          auto d = query_feat[dst_node_id * dim_node + h_i] * key_feat[neigh_off * dim_node + h_i];
          if (edge_feat != nullptr) {
            IdxT eid     = off_start + edge;
            IdxT eid_idx = (map_csc_to_coo != nullptr) ? map_csc_to_coo[eid] : eid;
            d += query_feat[dst_node_id * dim_node + h_i] * edge_feat[eid_idx * dim_node + h_i];
          }
          scores[edge] += d;
        }
      }

      // norm and bias
      // fwd sequence is score_weight, norm_by_dim, add_bias
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        auto eid     = off_start + edge;
        auto eid_idx = (map_csc_to_coo == nullptr) ? eid : map_csc_to_coo[eid];
        if (use_score_weight) {
          sm_scores[sm_off_pre_mul + eid_idx] = scores[edge];
          scores[edge]                        = scores[edge] * score_weight[h * n_indices + eid];
        }
        if (norm_by_dim) scores[edge] = scores[edge] * scalar_inv_norm;
        if (add_bias) {
          IdxT eid = off_start + edge;
          scores[edge] += score_bias[h * n_indices + eid];
        }
        sm_scores[sm_off_pre + eid_idx] = scores[edge];
      }
      // calculate the softmax (find max first, then subtract etc.)
      auto max_score = std::numeric_limits<DataT>::lowest();
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        if (scores[edge] > max_score) max_score = scores[edge];
      }
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        scores[edge] = std::exp(scores[edge] - max_score);
      }
      DataT sum_score = ZERO;
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        sum_score += scores[edge];
      }
      sum_score = DataT{1} / sum_score;
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        scores[edge] *= sum_score;
        auto eid                         = off_start + edge;
        auto eid_idx                     = (map_csc_to_coo == nullptr) ? eid : map_csc_to_coo[eid];
        sm_scores[sm_off_post + eid_idx] = scores[edge];
      }
      // at this point we have the softmax scores
      // now we simply want to weight the source features using these
      for (int i = 0; i < dim_head; ++i) {
        auto h_i  = h * dim_head + i;
        DataT acc = ZERO;
        for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
          auto neigh_off = indices[off_start + edge];
          acc += scores[edge] * value_feat[neigh_off * dim_node + h_i];
          if (edge_feat != nullptr) {
            IdxT eid     = off_start + edge;
            IdxT eid_idx = (map_csc_to_coo != nullptr) ? map_csc_to_coo[eid] : eid;
            acc += scores[edge] * edge_feat[eid_idx * dim_node + h_i];
          }
        }
        if (params.concat_heads) {
          output[dst_node_id * dim_node + h_i] += acc;
        } else {
          output[dst_node_id * dim_head + i] += acc * norm_head;
        }
      }
    }
  }
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_kernel(DataT* d_key_feat /* d_key_src_feat */,
                          DataT* d_query_feat /* d_query_dst_feat */,
                          DataT* d_value_feat /* d_value_src_feat */,
                          DataT* d_edge_feat,
                          DataT* d_sm_scores,
                          DataT* d_score_bias,
                          DataT* d_score_weight,
                          const DataT* d_out,
                          const DataT* key_feat /* key_src_feat */,
                          const DataT* query_feat /* query_dst_feat */,
                          const DataT* value_feat /* value_src_feat */,
                          const DataT* edge_feat,
                          const DataT* score_weight,
                          const DataT* sm_scores,
                          bool norm_by_dim,
                          const IdxT* offsets,
                          const IdxT* indices,
                          const IdxT* map_csc_to_coo,
                          IdxT n_src_nodes,
                          IdxT n_dst_nodes,
                          IdxT n_indices,
                          int dim_node,
                          const mha_params& params)
{
  auto n_edges                = n_indices;
  static constexpr DataT ZERO = DataT{0};
  int dim_head                = dim_node / params.num_heads;
  DataT norm_head             = DataT{1} / static_cast<DataT>(params.num_heads);
  bool add_bias               = (d_score_bias != nullptr);
  bool use_score_weight       = score_weight != nullptr;
  DataT scalar_inv_norm       = DataT{1} / ::sqrt(static_cast<DataT>(dim_head));

  if (d_key_feat != nullptr) std::memset(d_key_feat, 0, n_src_nodes * dim_node * sizeof(DataT));
  if (d_query_feat != nullptr) std::memset(d_query_feat, 0, n_dst_nodes * dim_node * sizeof(DataT));
  if (d_value_feat != nullptr) std::memset(d_value_feat, 0, n_src_nodes * dim_node * sizeof(DataT));
  if (d_edge_feat != nullptr) std::memset(d_edge_feat, 0, n_indices * dim_node * sizeof(DataT));

  for (auto dst_node_id = 0; dst_node_id < n_dst_nodes; ++dst_node_id) {
    auto off_start = offsets[dst_node_id], off_end = offsets[dst_node_id + 1];
    auto n_neighbors = off_end - off_start;
    auto d_post      = std::make_unique<DataT[]>(n_neighbors);

    for (int h = 0; h < params.num_heads; ++h) {
      auto sm_off_pre     = h * n_edges;
      auto sm_off_post    = (params.num_heads + h) * n_edges;
      auto sm_off_pre_mul = (params.num_heads * 2 + h) * n_edges;
      // first get all the post SM gradients
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        d_post[edge]   = ZERO;
        auto neigh_off = indices[off_start + edge];
        for (int i = 0; i < dim_head; ++i) {
          auto h_i = h * dim_head + i;
          if (params.concat_heads) {
            d_post[edge] +=
              d_out[dst_node_id * dim_node + h_i] * value_feat[neigh_off * dim_node + h_i];
            if (edge_feat != nullptr) {
              IdxT eid     = off_start + edge;
              IdxT eid_idx = (map_csc_to_coo != nullptr) ? map_csc_to_coo[eid] : eid;
              d_post[edge] +=
                d_out[dst_node_id * dim_node + h_i] * edge_feat[eid_idx * dim_node + h_i];
            }
          } else {
            IdxT eid     = off_start + edge;
            IdxT eid_idx = (map_csc_to_coo != nullptr) ? map_csc_to_coo[eid] : eid;
            d_post[edge] += (d_out[dst_node_id * dim_head + i] * norm_head) *
                            value_feat[neigh_off * dim_node + h_i];
            if (edge_feat != nullptr) {
              d_post[edge] += (d_out[dst_node_id * dim_head + i] * norm_head) *
                              edge_feat[eid_idx * dim_node + h_i];
            }
          }
        }
        auto eid     = off_start + edge;
        auto eid_idx = (map_csc_to_coo == nullptr) ? eid : map_csc_to_coo[eid];
        d_sm_scores[sm_off_post + eid_idx] = d_post[edge];
      }
      // calculate gradients before softmax (given out/attn node and head)
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        auto d_pre     = ZERO;
        auto eid       = off_start + edge;
        auto eid_idx   = (map_csc_to_coo == nullptr) ? eid : map_csc_to_coo[eid];
        auto neigh_off = indices[eid];
        auto sm_score  = sm_scores[sm_off_post + eid_idx];

        for (IdxT edge_j = IdxT{0}; edge_j < n_neighbors; ++edge_j) {
          auto eid_j      = off_start + edge_j;
          auto eid_idx_j  = map_csc_to_coo == nullptr ? eid_j : map_csc_to_coo[eid_j];
          auto sm_score_j = sm_scores[sm_off_post + eid_idx_j];
          auto d          = edge == edge_j ? DataT{1} : ZERO;
          auto d_sm_local = sm_score * (d - sm_score_j);
          d_pre += d_sm_local * d_post[edge_j];
        }

        d_sm_scores[sm_off_pre + eid_idx] = d_pre;

        // gradient on bias and backprop through norm
        // fwd sequence is use_score_weight, norm_by_dim, add_bias
        // so backward sequence is simply reversed starting with add_bias
        if (add_bias) { d_score_bias[h * n_indices + eid] = d_pre; }
        if (norm_by_dim) { d_pre *= scalar_inv_norm; }
        if (use_score_weight) {
          if (d_score_weight != nullptr) {
            d_score_weight[h * n_indices + eid] = d_pre * sm_scores[sm_off_pre_mul + eid_idx];
          }
          d_pre *= score_weight[h * n_indices + eid];
          d_sm_scores[sm_off_pre_mul + eid_idx] = d_pre;
        }

        // accumulate gradients based on gradients before softmax
        for (int i = 0; i < dim_head; ++i) {
          auto h_i = h * dim_head + i;
          if (params.concat_heads) {
            if (d_value_feat != nullptr) {
              d_value_feat[neigh_off * dim_node + h_i] +=
                d_out[dst_node_id * dim_node + h_i] * sm_score;
            }
            if (d_edge_feat != nullptr) {
              d_edge_feat[eid_idx * dim_node + h_i] +=
                d_out[dst_node_id * dim_node + h_i] * sm_score;
            }
          } else {
            if (d_value_feat != nullptr) {
              d_value_feat[neigh_off * dim_node + h_i] +=
                (d_out[dst_node_id * dim_head + i] * norm_head) * sm_score;
            }
            if (d_edge_feat != nullptr) {
              d_edge_feat[eid_idx * dim_node + h_i] +=
                (d_out[dst_node_id * dim_head + i] * norm_head) * sm_score;
            }
          }
          if (d_key_feat != nullptr) {
            d_key_feat[neigh_off * dim_node + h_i] +=
              d_pre * query_feat[dst_node_id * dim_node + h_i];
          }
          if (d_query_feat != nullptr) {
            d_query_feat[dst_node_id * dim_node + h_i] +=
              d_pre * key_feat[neigh_off * dim_node + h_i];
          }
          if (edge_feat != nullptr) {
            if (d_query_feat != nullptr) {
              d_query_feat[dst_node_id * dim_node + h_i] +=
                d_pre * edge_feat[eid_idx * dim_node + h_i];
            }
            if (d_edge_feat != nullptr) {
              d_edge_feat[eid_idx * dim_node + h_i] +=
                d_pre * query_feat[dst_node_id * dim_node + h_i];
            }
          }
        }
      }
    }
  }
}

}  // namespace cugraph::ops::mha_simple_csc
