/*
 * 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/activation.hpp>
#include <cugraph-ops/operators/common.hpp>
#include <cugraph-ops/test/activation_cpu.hpp>

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

namespace cugraph::ops::mha_gat_csc {

template <typename DataT, typename IdxT>
void naive_cpu_fwd_kernel(DataT* output,
                          DataT* sm_scores,
                          const DataT* src_feat,
                          const DataT* dst_feat,
                          const DataT* edge_feat,
                          const DataT* attn_weights,
                          const IdxT* offsets,
                          const IdxT* indices,
                          const IdxT* map_csc_to_coo,
                          IdxT n_dst_nodes,
                          IdxT n_indices,
                          int dim_node,
                          int dim_edge,
                          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;
  int dim_head_edge = dim_edge / params.num_heads;
  DataT norm        = DataT{1} / static_cast<DataT>(params.num_heads);

  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));
  }
  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;
      // 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         = attn_weights[h_i] * src_feat[neigh_off * dim_node + h_i];
          d += attn_weights[dim_node + h_i] * dst_feat[dst_node_id * dim_node + h_i];
          scores[edge] += d;
        }
      }
      if (edge_feat != nullptr) {
        for (int i = 0; i < dim_head_edge; ++i) {
          auto h_i = h * dim_head_edge + i;
          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];
            auto d       = attn_weights[2 * dim_node + h_i] * edge_feat[eid_idx * dim_edge + h_i];
            scores[edge] += d;
          }
        }
      }
      // apply activation
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        scores[edge] = activation_fwd_cpu(params.activation, scores[edge]);
        // write back pre-sm (i.e. scores after activation)
        auto eid                        = off_start + edge;
        auto eid_idx                    = (map_csc_to_coo == nullptr) ? eid : map_csc_to_coo[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 eid       = off_start + edge;
          auto neigh_off = indices[eid];
          acc += scores[edge] * src_feat[neigh_off * 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;
        }
      }
    }
  }
}

template <typename DataT, typename IdxT>
void naive_cpu_bwd_kernel(DataT* d_src_feat,
                          DataT* d_dst_feat,
                          DataT* d_edge_feat,
                          DataT* d_attn_weights,
                          DataT* d_sm_scores,
                          const DataT* d_out,
                          const DataT* src_feat,
                          const DataT* dst_feat,
                          const DataT* edge_feat,
                          const DataT* attn_weights,
                          const DataT* sm_scores,
                          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,
                          int dim_edge,
                          const mha_params& params)
{
  auto n_edges                = n_indices;
  static constexpr DataT ZERO = DataT{0};
  int dim_head                = dim_node / params.num_heads;
  int dim_head_edge           = dim_edge / params.num_heads;
  DataT norm                  = DataT{1} / static_cast<DataT>(params.num_heads);

  auto dim_node_s   = static_cast<size_t>(dim_node);
  auto dim_edge_s   = static_cast<size_t>(dim_edge);
  auto dim_weight_s = (2 * dim_node_s + dim_edge_s);
  if (d_src_feat != nullptr)
    std::memset(d_src_feat, 0, static_cast<size_t>(n_src_nodes) * dim_node_s * sizeof(DataT));
  if (d_dst_feat != nullptr)
    std::memset(d_dst_feat, 0, static_cast<size_t>(n_dst_nodes) * dim_node_s * sizeof(DataT));
  if (d_attn_weights != nullptr) std::memset(d_attn_weights, 0, dim_weight_s * sizeof(DataT));
  if (d_edge_feat != nullptr) std::memset(d_edge_feat, 0, n_indices * dim_edge_s * 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;
      // first get all the post SM gradients
      for (IdxT edge = IdxT{0}; edge < n_neighbors; ++edge) {
        d_post[edge]   = 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];
        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] * src_feat[neigh_off * dim_node + h_i];
          } else {
            d_post[edge] +=
              (d_out[dst_node_id * dim_head + i] * norm) * src_feat[neigh_off * dim_node + h_i];
          }
        }
        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 eid       = off_start + edge;
        auto neigh_off = indices[eid];
        auto d_pre     = ZERO;
        auto eid_idx   = (map_csc_to_coo == nullptr) ? eid : map_csc_to_coo[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;

        // backward of activation
        auto sm_score_j = sm_scores[sm_off_pre + eid_idx];
        d_pre *= activation_bwd_cpu(params.activation, sm_score_j);
        // accumulate gradients based on gradients before softmax
        for (int i = 0; i < dim_head; ++i) {
          auto h_i = h * dim_head + i;
          if (d_src_feat != nullptr) {
            if (params.concat_heads) {
              d_src_feat[neigh_off * dim_node + h_i] +=
                d_out[dst_node_id * dim_node + h_i] * sm_score;
            } else {
              d_src_feat[neigh_off * dim_node + h_i] +=
                (d_out[dst_node_id * dim_head + i] * norm) * sm_score;
            }
          }
          if (d_src_feat != nullptr) {
            d_src_feat[neigh_off * dim_node + h_i] += d_pre * attn_weights[h_i];
          }
          if (d_dst_feat != nullptr) {
            d_dst_feat[dst_node_id * dim_node + h_i] += d_pre * attn_weights[dim_node + h_i];
          }
          if (d_attn_weights != nullptr) {
            d_attn_weights[h_i] += d_pre * src_feat[neigh_off * dim_node + h_i];
            d_attn_weights[dim_node + h_i] += d_pre * dst_feat[dst_node_id * dim_node + h_i];
          }
        }
        if (edge_feat != nullptr) {
          for (int i = 0; i < dim_head_edge; ++i) {
            auto h_i = h * dim_head_edge + i;
            if (d_edge_feat != nullptr) {
              d_edge_feat[eid_idx * dim_edge + h_i] += d_pre * attn_weights[2 * dim_node + h_i];
            }
            if (d_attn_weights != nullptr) {
              d_attn_weights[2 * dim_node + h_i] += d_pre * edge_feat[eid_idx * dim_edge + h_i];
            }
          }
        }
      }
    }
  }
}

}  // namespace cugraph::ops::mha_gat_csc
