#include <ATen/ATen.h>
#include <ATen/AccumulateType.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/Exceptions.h>
// Another possibility:
// #include <torch/all.h>

#include <assert.h>

#include "multi_tensor_apply.cuh"
#include "type_shim.h"

#define BLOCK_SIZE 512
#define ILP 4

using MATH_T = float;

// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
// It computes new parameter value.
template <typename T, typename UPD_T>
struct LAMBStage2Functor {
  __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<2>& tl,
                                             const float* per_tensor_param_norm, const float* per_tensor_update_norm,
                                             const float learning_rate, const float decay, bool use_nvlamb) {
    // I'd like this kernel to propagate infs/nans.
    // if(*noop_gmem == 1)
    //   return;

    int tensor_loc = tl.block_to_tensor[blockIdx.x];
    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
    int chunk_idx = tl.block_to_chunk[blockIdx.x];
    int n = tl.sizes[tensor_loc];

    MATH_T ratio = learning_rate;
    // nvlamb: apply adaptive learning rate to all parameters
    // otherwise, only apply to those with non-zero weight decay
    if (use_nvlamb || (decay != 0.0)) {
      float param_norm = per_tensor_param_norm[tensor_num];
      float update_norm = per_tensor_update_norm[tensor_num];
      ratio = (update_norm != 0.0f && param_norm != 0.0f) ? learning_rate * (param_norm / update_norm) : learning_rate;
    }

    T* p = (T*)tl.addresses[0][tensor_loc];
    p += chunk_idx * chunk_size;

    UPD_T* update = (UPD_T*)tl.addresses[1][tensor_loc];
    update += chunk_idx * chunk_size;

    n -= chunk_idx * chunk_size;

    for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) {
      T r_p[ILP];
      UPD_T r_update[ILP];
#pragma unroll
      for (int ii = 0; ii < ILP; ii++) {
        int i = i_start + threadIdx.x + ii * blockDim.x;
        if (i < n && i < chunk_size) {
          r_p[ii] = p[i];
          r_update[ii] = update[i];
        }
      }
#pragma unroll
      for (int ii = 0; ii < ILP; ii++) {
        r_p[ii] = r_p[ii] - (ratio * (T)r_update[ii]);
      }
#pragma unroll
      for (int ii = 0; ii < ILP; ii++) {
        int i = i_start + threadIdx.x + ii * blockDim.x;
        if (i < n && i < chunk_size) {
          p[i] = r_p[ii];
        }
      }
    }
  }
};

void multi_tensor_lamb_stage2_cuda(int chunk_size, at::Tensor noop_flag,
                                   std::vector<std::vector<at::Tensor>> tensor_lists, at::Tensor per_tensor_param_norm,
                                   at::Tensor per_tensor_update_norm, const float lr, const float weight_decay,
                                   at::optional<bool> use_nvlamb_python) {
  bool use_nvlamb = use_nvlamb_python.has_value() ? use_nvlamb_python.value() : false;

  using namespace at;

  DISPATCH_FLOAT_AND_HALF(
      tensor_lists[0][0].scalar_type(), 0, "lamb_stage_2",
      DISPATCH_FLOAT_AND_HALF(
          tensor_lists[1][0].scalar_type(), 1, "lamb_stage_2",
          multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
                                LAMBStage2Functor<scalar_t_0, scalar_t_1>(), per_tensor_param_norm.data_ptr<float>(),
                                per_tensor_update_norm.data_ptr<float>(), lr, weight_decay, use_nvlamb);))

  AT_CUDA_CHECK(cudaGetLastError());

  // AT_CUDA_CHECK(cudaDeviceSynchronize());
}