/**
 * (C) Copyright 2020, 2021 IBM. All Rights Reserved.
 *
 * This code is licensed under the Apache License, Version 2.0. You may
 * obtain a copy of this license in the LICENSE.txt file in the root directory
 * of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
 *
 * Any modifications or derivative works of this code must retain this
 * copyright notice, and modified files need to carry a notice indicating
 * that they have been altered from the originals.
 */

#include "noise_manager.h"

#include <chrono>
#include <cmath>
#include <iostream>
#include <memory>
#include <random>

#include "cuda_math_util.h"
#include "cuda_util.h"
#include <cub/cub.cuh>

#include "io_iterator.h"

namespace RPU {

template <typename T>
__global__ void kernelAbsMaxNPSum(
    float *scale_values,
    const int m_batch,
    const float *amax_values,
    const T *psum_values,
    const T *nsum_values,
    const T out_bound,
    const T assumed_wmax,
    const T bm_max // io.max_bm_res/io.inp_res

) {
  int tid = blockDim.x * blockIdx.x + threadIdx.x;

  if (tid < m_batch) {
    // w_max*MAX(psum,nsum)/out_bound < scale
    T w_max = assumed_wmax;
    T amax = amax_values[tid];
    T psum = psum_values[tid];
    T nsum = -nsum_values[tid];
    T sum = MAX(psum, nsum);
    scale_values[tid] = MAX(amax, MIN(sum * w_max / out_bound, amax * bm_max));
    ;
  }
}

template <typename InputIteratorT, typename T>
__global__ void kernelNPSumBatchTrans(
    InputIteratorT input,
    const int total_size_in,
    const int m_batch_in,
    T *psum_values,
    T *nsum_values,
    T *psum_values0,
    T *nsum_values0) {

  // -- only use this version if m_batch < blockDim.x !!!
  // -- probably: strided version would be faster...

  int tid = blockDim.x * blockIdx.x + threadIdx.x;
  // assumes that shared is of size 2*nthreads*sizeof(T) !!!!!!!!!!
  extern __shared__ __align__(sizeof(double)) unsigned char rpu_smem_nm[];
  T *block_sum_values = reinterpret_cast<T *>(rpu_smem_nm);

  T *block_psum_values = &block_sum_values[0];
  T *block_nsum_values = &block_sum_values[blockDim.x];

  const int size = total_size_in;
  const int m_batch = m_batch_in;

  block_psum_values[threadIdx.x] = (T)0.0;
  block_nsum_values[threadIdx.x] = (T)0.0;

  __syncthreads();

  if (tid < m_batch) {
    psum_values0[tid] = (T)0.0;
    nsum_values0[tid] = (T)0.0;
  }

  if (tid < size) {

    T value = input[tid];
    int midx = tid % m_batch;

    if (value >= 0) {
      atomicAdd(&(block_psum_values[midx]), value);
    } else {
      atomicAdd(&(block_nsum_values[midx]), value);
    }
  }
  __syncthreads();

  int bidx = threadIdx.x;
  if (bidx < m_batch) {
    atomicAdd(&(psum_values[bidx]), block_psum_values[bidx]);
    atomicAdd(&(nsum_values[bidx]), block_nsum_values[bidx]);
  }
}

template <typename InputIteratorT, typename T>
__global__ void kernelNPSumBatchTrans_LargeBatch(
    InputIteratorT input,
    const int total_size_in,
    const int m_batch_in,
    T *psum_values,
    T *nsum_values,
    T *psum_values0,
    T *nsum_values0) {

  // -- use this version if m_batch >= blockDim.x
  // -- just uses atomic on global memory

  int tid = blockDim.x * blockIdx.x + threadIdx.x;

  const int size = total_size_in;
  const int m_batch = m_batch_in;

  if (tid < m_batch) {
    psum_values0[tid] = (T)0.0;
    nsum_values0[tid] = (T)0.0;
  }

  if (tid < size) {

    T value = input[tid];
    int midx = tid % m_batch;

    if (value >= 0) {
      atomicAdd(&psum_values[midx], value);
    } else {
      atomicAdd(&nsum_values[midx], value);
    }
  }
}

template <typename T>
__global__ void kernelAverageAbsMaxSetScales(
    float *scales, float *ravg, const float *sum, const int m_batch_in, T decay_rate_in) {
  int tid = blockDim.x * blockIdx.x + threadIdx.x;
  const int m_batch = m_batch_in;
  T decay_rate = decay_rate_in;
  T max_avg = (*sum) / m_batch;
  T run_avg = *ravg;

  if (tid < m_batch) {
    scales[tid] = (float)(run_avg * (1.0 - decay_rate) + decay_rate * max_avg);
  }
  if (tid == m_batch) {
    *ravg = (float)(run_avg * (1.0 - decay_rate) + decay_rate * max_avg);
  }
}

template <typename T>
__global__ void
kernelAverageAbsMaxSingleMomentum(float *ravg, const float *sum, const int m_batch, T decay_rate) {
  // just single block!
  int tid = blockDim.x * blockIdx.x + threadIdx.x;
  if (tid == 0) {
    T max_avg = (*sum) / m_batch;
    *ravg = (float)(*ravg * (1.0 - decay_rate) + decay_rate * max_avg);
  }
}

/****************************************************************************************************************/
/* NoiseManager */
/******************************************************************************************************************/
#define LAUNCH_NPSUM_KERNEL(KNAME, SHARED_MEM, ARGS)                                               \
  KNAME<InputIteratorT, T><<<nblocks, nthreads, SHARED_MEM, s>>> ARGS;

template <typename T>
NoiseManager<T>::NoiseManager(CudaContext *c, int size)
    : size_(size), context_(c), buffer_m_batch_(0), const_set_if_(false) {
  // initialize for m_batch=1
  dev_scale_values_ = RPU::make_unique<CudaArray<float>>(context_, 1);
  dev_psum_values_ = RPU::make_unique<CudaArray<T>>(context_, 1);
  dev_nsum_values_ = RPU::make_unique<CudaArray<T>>(context_, 1);
  dev_ravg_scale_value_ = RPU::make_unique<CudaArray<float>>(context_, 1);
  dev_ravg_scale_value_->setConst(1.0);

  amaximizer_ = RPU::make_unique<Maximizer<T>>(context_, size, true);
  maximizer_ = RPU::make_unique<Maximizer<T>>(context_, size, false);

  size_t temp_storage_bytes = 0;
  RPU::cub::DeviceReduce::Reduce(
      nullptr, temp_storage_bytes, dev_psum_values_->getData(), dev_psum_values_->getData(), size_,
      nsum_op_, 0, context_->getStream());

  dev_v_temp_storage_ = RPU::make_unique<CudaArray<char>>(context_, temp_storage_bytes);
  context_->synchronize();
}

template <typename T> void NoiseManager<T>::initializeBatchBuffer(int m_batch) {
  // this inits all the buffers needed for PMSum only !!

  if ((m_batch > 1) && (buffer_m_batch_ != m_batch)) {
    buffer_m_batch_ = m_batch;

    dev_psum_values_ = RPU::make_unique<CudaArray<T>>(context_, m_batch);
    dev_psum_values0_ = RPU::make_unique<CudaArray<T>>(context_, m_batch);
    dev_psum_values0_->setConst((T)0.0);

    dev_nsum_values_ = RPU::make_unique<CudaArray<T>>(context_, m_batch);
    dev_nsum_values0_ = RPU::make_unique<CudaArray<T>>(context_, m_batch);
    dev_nsum_values0_->setConst((T)0.0);

    int *offsets = new int[m_batch + 1];

    // not trans
    for (int i = 0; i <= m_batch; i++) {
      offsets[i] = i * size_;
    }

    dev_offsets_ = RPU::make_unique<CudaArray<int>>(context_, m_batch + 1, offsets);

    size_t temp_storage_bytes = 0;
    RPU::cub::DeviceSegmentedReduce::Reduce(
        nullptr, temp_storage_bytes, dev_psum_values_->getData(), dev_psum_values_->getData(),
        m_batch, dev_offsets_->getData(), dev_offsets_->getData() + 1, psum_op_, 0,
        context_->getStream());
    dev_m_temp_storage_ = RPU::make_unique<CudaArray<char>>(context_, temp_storage_bytes);

    const_set_if_ = false;
    context_->synchronize();
    delete[] offsets;
  }
}

template <typename T>
template <typename InputIteratorT>
void NoiseManager<T>::computeNPSum(InputIteratorT dev_input, int m_batch, bool trans) {
  cudaStream_t s = context_->getStream();

  if (m_batch == 1) {
    size_t ssz = dev_v_temp_storage_->getSize();
    RPU::cub::DeviceReduce::Reduce(
        (void *)dev_v_temp_storage_->getData(), ssz, dev_input, dev_psum_values_->getData(), size_,
        psum_op_, (T)0, s);
    RPU::cub::DeviceReduce::Reduce(
        (void *)dev_v_temp_storage_->getData(), ssz, dev_input, dev_nsum_values_->getData(), size_,
        nsum_op_, (T)0, s);

  } else {

    if (trans) {

      if (buffer_m_batch_ < m_batch) {
        this->initializeBatchBuffer(m_batch);
      }

      std::swap(dev_psum_values_, dev_psum_values0_);
      std::swap(dev_nsum_values_, dev_nsum_values0_);

      int nthreads = context_->getNThreads();
      int n = size_ * m_batch;
      int nblocks = context_->getNBlocks(n, nthreads);
      if (m_batch <= nthreads) {
        int shared_mem = 2 * nthreads * sizeof(T);

        LAUNCH_NPSUM_KERNEL(
            kernelNPSumBatchTrans, shared_mem,
            (dev_input, n, m_batch, dev_psum_values_->getData(), dev_nsum_values_->getData(),
             dev_psum_values0_->getData(), dev_nsum_values0_->getData()));

      } else {
        // simple atomic global memory version
        LAUNCH_NPSUM_KERNEL(
            kernelNPSumBatchTrans_LargeBatch, 0,
            (dev_input, n, m_batch, dev_psum_values_->getData(), dev_nsum_values_->getData(),
             dev_psum_values0_->getData(), dev_nsum_values0_->getData()));
      }

    } else {
      if (buffer_m_batch_ != m_batch) { // !! need to reinitilize offsets when batch changes !
        this->initializeBatchBuffer(m_batch);
      }

      // Fast Segmented reduction
      size_t ssz = dev_m_temp_storage_->getSize();
      RPU::cub::DeviceSegmentedReduce::Reduce(
          (void *)dev_m_temp_storage_->getData(), ssz, dev_input, dev_psum_values_->getData(),
          m_batch, dev_offsets_->getData(), dev_offsets_->getData() + 1, psum_op_, (T)0.0, s);
      RPU::cub::DeviceSegmentedReduce::Reduce(
          (void *)dev_m_temp_storage_->getData(), ssz, dev_input, dev_nsum_values_->getData(),
          m_batch, dev_offsets_->getData(), dev_offsets_->getData() + 1, nsum_op_, (T)0.0, s);
    }
  }
}

template <typename T> void NoiseManager<T>::setAverageAbsMax(float value) {
  dev_ravg_scale_value_->setConst(value);
  dev_scale_values_->setConst(value);
  ravg_initialized_ = true;
  context_->synchronize();
}

template <typename T> float NoiseManager<T>::getAverageAbsMax() const {
  float tmp;
  dev_ravg_scale_value_->copyTo(&tmp);
  return tmp;
};

template <typename T>
template <typename InputIteratorT>
void NoiseManager<T>::compute(
    InputIteratorT dev_input,
    const NoiseManagementType &nm_type,
    const IOMetaParameter<T> &io,
    int m_batch,
    bool trans,
    bool is_test) {

  // does not check for positive m_batch!
  nm_type_ = nm_type;

  switch (nm_type_) {

  case NoiseManagementType::None: {
    return;
  }
  case NoiseManagementType::Constant: {
    if (m_batch > dev_scale_values_->getSize()) {
      dev_scale_values_ = RPU::make_unique<CudaArray<float>>(context_, m_batch);
      const_set_if_ = false;
    }
    if (!const_set_if_) {
      dev_scale_values_->setConst(io.nm_thres > 0 ? (float)io.nm_thres : (float)1.0);
      const_set_if_ = true;
    }
    return;
  }

  case NoiseManagementType::Max: {
    this->maximizer_->compute(dev_input, m_batch, trans);
    if (io.nm_thres > 0) {
      this->maximizer_->saturateAbove(io.nm_thres);
    }

    return;
  }

  case NoiseManagementType::AbsMax: {
    this->amaximizer_->compute(dev_input, m_batch, trans);
    if (io.nm_thres > 0) {
      this->amaximizer_->saturateAbove(io.nm_thres);
    }

    return;
  }
  case NoiseManagementType::AbsMaxNPSum: {
    if (m_batch > dev_scale_values_->getSize()) {
      dev_scale_values_ = RPU::make_unique<CudaArray<float>>(context_, m_batch);
    }

    // get amax and npsum
    this->amaximizer_->compute(dev_input, m_batch, trans);
    if (io.nm_thres > 0) {
      this->amaximizer_->saturateAbove(io.nm_thres);
    }

    this->computeNPSum(dev_input, m_batch, trans);

    // combine
    int nthreads = context_->getNThreads();
    int nblocks = context_->getNBlocks(m_batch, nthreads);
    cudaStream_t s = context_->getStream();

    kernelAbsMaxNPSum<T><<<nblocks, nthreads, 0, s>>>(
        dev_scale_values_->getData(), m_batch, this->amaximizer_->getMaxValues(),
        dev_psum_values_->getDataConst(), dev_nsum_values_->getDataConst(), io.out_bound,
        io.nm_assumed_wmax, io.inp_res > 0 ? io.max_bm_res / io.inp_res : 1.0);
    return;
  }

  case NoiseManagementType::AverageAbsMax:
  case NoiseManagementType::AverageAbsMaxSingleValue: {
    // CAUTION: the running average will not be saved for checkpointing... so there might be a
    // glitch when continueing training from checkpoint...
    // ALSO: average max is computed across
    // mbatch whereas for CPU it is based running average of single mat-vecs

    if ((nm_type_ == NoiseManagementType::AverageAbsMax) &&
        (m_batch > dev_scale_values_->getSize())) {

      dev_scale_values_ = RPU::make_unique<CudaArray<float>>(context_, m_batch);

      cudaStream_t s = context_->getStream();
      int nthreads = context_->getNThreads();

      // set scales to ravg [first time, could be set from outside]
      kernelAverageAbsMaxSetScales<T>
          <<<context_->getNBlocks(m_batch + 1, nthreads), nthreads, 0, s>>>(
              dev_scale_values_->getData(), dev_ravg_scale_value_->getData(),
              dev_ravg_scale_value_->getDataConst(), m_batch, (T)0.0);
    }

    if (!is_test) {

      this->amaximizer_->compute(dev_input, m_batch, trans);

      cudaStream_t s = context_->getStream();
      int nthreads = context_->getNThreads();
      int nblocks = context_->getNBlocks(m_batch, nthreads);

      if (m_batch > 1) {
        // first compute the average of the max over batch
        if (!dev_a_temp_storage_) {
          dev_avgmax_value_ = RPU::make_unique<CudaArray<float>>(context_, 1);

          size_t temp_storage_bytes = 0;
          RPU::cub::DeviceReduce::Sum(
              nullptr, temp_storage_bytes, amaximizer_->getMaxValues(),
              dev_avgmax_value_->getData(), m_batch, s);
          dev_a_temp_storage_ = RPU::make_unique<CudaArray<char>>(context_, temp_storage_bytes);
        }

        size_t ssz = dev_v_temp_storage_->getSize();
        RPU::cub::DeviceReduce::Sum(
            (void *)dev_v_temp_storage_->getData(), ssz, amaximizer_->getMaxValues(),
            dev_avgmax_value_->getData(), m_batch, s);
      }

      if (nm_type_ == NoiseManagementType::AverageAbsMax) {
        // now update the running scale and set the current scales constant for all m_batch
        kernelAverageAbsMaxSetScales<T>
            <<<context_->getNBlocks(m_batch + 1, nthreads), nthreads, 0, s>>>(
                dev_scale_values_->getData(), dev_ravg_scale_value_->getData(),
                m_batch > 1 ? dev_avgmax_value_->getData() : amaximizer_->getMaxValues(),
                dev_scale_values_->getSize(),
                ravg_initialized_ ? MIN(io.nm_decay * m_batch, 1.) : 1.0);

      } else {
        // just update the running avg value as only single output requested
        kernelAverageAbsMaxSingleMomentum<T><<<1, 1, 0, s>>>(
            dev_ravg_scale_value_->getData(),
            m_batch > 1 ? dev_avgmax_value_->getData() : amaximizer_->getMaxValues(), m_batch,
            ravg_initialized_ ? MIN(io.nm_decay * m_batch, 1.0) : 1.0);
      }
      ravg_initialized_ = true;
    }
    return;
  }

  default:
    RPU_FATAL("Noise management type not implemented.");
  }
}

template <typename T> float *NoiseManager<T>::getScaleValues() const {
  switch (nm_type_) {
  case NoiseManagementType::None:
    return nullptr;
  case NoiseManagementType::AbsMaxNPSum:
  case NoiseManagementType::Constant:
  case NoiseManagementType::AverageAbsMax:
    return dev_scale_values_->getData();
  case NoiseManagementType::AbsMax:
    return amaximizer_->getMaxValues();
  case NoiseManagementType::Max:
    return maximizer_->getMaxValues();
  case NoiseManagementType::AverageAbsMaxSingleValue:
    return dev_ravg_scale_value_->getData();
  default:
    RPU_FATAL("Noise management type not implemented.");
  }
};

#define ARGS1(NUM_T) , const NoiseManagementType &, const IOMetaParameter<NUM_T> &, int, bool, bool
#define ARGS2 , int, bool

template class NoiseManager<float>;
RPU_GEN_IITER_TEMPLATES(float, void, NoiseManager<float>::compute, ARGS1(float));
RPU_GEN_IITER_TEMPLATES(float, void, NoiseManager<float>::computeNPSum, ARGS2);

#ifdef RPU_USE_DOUBLE
template class NoiseManager<double>;
RPU_GEN_IITER_TEMPLATES(double, void, NoiseManager<double>::compute, ARGS1(double));
RPU_GEN_IITER_TEMPLATES(double, void, NoiseManager<double>::computeNPSum, ARGS2);
#endif

#undef ARGS1
#undef ARGS2
#undef LAUNCH_NPSUM_KERNEL

} // namespace RPU
