#include <cuda_fp16.h>
#include "rep_pen.cuh"
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>
#include "../util.h"
#include "../util.cuh"
#include <limits>

#define BLOCK_VOCAB_SPAN 4096
#define NUM_THREADS 1024
#define CLAMP(x, min, max) ((x) < (min) ? (min) : ((x) > (max) ? (max) : (x)))

__device__ __forceinline__
float shmemAtomicMaxF(float* addr, float val)
{
    if (val >= 0.0f) {
        unsigned int old = atomicMax(reinterpret_cast<unsigned int*>(addr), __float_as_uint(val));
        return __uint_as_float(old);
    } else {
        unsigned int old = atomicMin(reinterpret_cast<unsigned int*>(addr), __float_as_uint(val));
        return __uint_as_float(old);
    }
}

template <bool input_fp16>
__global__ __launch_bounds__(NUM_THREADS)
void apply_rep_pens_kernel
(
    const void* __restrict__ in_logits,
    float* __restrict__ out_logits,
    const uint64_t* __restrict__ past_ids,
    const int past_len,
    const int vocab_size,
    const float rep_p,
    const int sustain_range,
    const int decay_range
)
{
    // Each block processes a range of the logits
    int range_min = blockIdx.x * BLOCK_VOCAB_SPAN;
    int range_max = MIN(range_min + BLOCK_VOCAB_SPAN, vocab_size);

    __shared__ float factors[BLOCK_VOCAB_SPAN];
    for (int i = threadIdx.x; i < BLOCK_VOCAB_SPAN; i += NUM_THREADS)
        factors[i] = 0.0f;
    __syncthreads();

    // Record which tokens from the range appear in past_ids
    for (int i = threadIdx.x; i < past_len; i += NUM_THREADS)
    {
        if (i <= past_len - sustain_range - decay_range)
            continue;

        int tid = (int) past_ids[i];
        if (tid < range_min || tid >= range_max)
            continue;

        int dist = past_len - i;
        if (dist <= sustain_range)
            factors[tid - range_min] = 1.0f;
        else
        {
            float distf = (float)(past_len - i);
            float sustain_rangef = (float)sustain_range;
            float decay_rangef = (float)decay_range;
            float factor = decay_range > 0 ? 1.0f - (distf - sustain_rangef) / decay_rangef : 1.0f;
            factor = CLAMP(factor, 0.0f, 1.0f);
            shmemAtomicMaxF(factors + tid - range_min, factor);
        }
    }
    __syncthreads();

    // Apply penalties to range
    for (int i = threadIdx.x; i < BLOCK_VOCAB_SPAN; i += NUM_THREADS)
    {
        if (i + range_min >= vocab_size)
            break;

        float v;
        if constexpr (input_fp16)
            v = __half2float(((half*) in_logits)[i + range_min]);
        else
            v = ((float*) in_logits)[i + range_min];

        float w = v > 0.0f ? v / rep_p : v * rep_p;
        float f = factors[i] + 1e-30;
        float f1 = (1.0f - f) + 1e-30;
        float o = v * f1 + w * f;
        out_logits[i + range_min] = o;
    }
}

void apply_rep_pens
(
    const at::Tensor& in_logits,
    const at::Tensor& out_logits,
    const at::Tensor& past_ids,
    float rep_p,
    int sustain_range,
    int decay_range
)
{
    const at::cuda::OptionalCUDAGuard device_guard(in_logits.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    TORCH_CHECK_DTYPE(out_logits, kFloat);
    TORCH_CHECK_DTYPE(past_ids, kLong);
    TORCH_CHECK(in_logits.size(0) == 1, "rep. penalties only implemented for bsz 1");  // TODO
    TORCH_CHECK_SHAPES(past_ids, 0, in_logits, 0, 1);
    TORCH_CHECK_SHAPES_FULL(in_logits, out_logits);

    int past_len = past_ids.size(1);
    int vocab_size = in_logits.size(1);
    int num_blocks = CEIL_DIVIDE(vocab_size, BLOCK_VOCAB_SPAN);

    #define kernel_args \
        (const void*) in_logits.data_ptr(), \
        (float*) out_logits.data_ptr(), \
        (const uint64_t*) past_ids.data_ptr(), \
        past_len, \
        vocab_size, \
        rep_p, \
        sustain_range, \
        decay_range

    if (in_logits.dtype() == at::kHalf)
        apply_rep_pens_kernel<true><<<num_blocks, NUM_THREADS, 0, stream>>>(kernel_args);
    else
        apply_rep_pens_kernel<false><<<num_blocks, NUM_THREADS, 0, stream>>>(kernel_args);

    #undef kernel_args

    cuda_check(cudaPeekAtLastError());
}


template <bool input_fp16>
__global__ __launch_bounds__(NUM_THREADS)
void apply_pres_freq_pens_kernel
(
    const void* __restrict__ in_logits,
    float* __restrict__ out_logits,
    const uint64_t* __restrict__ past_ids,
    const int past_len,
    const int vocab_size,
    const float pres_p,
    const float freq_p,
    const int sustain_range,
    const int decay_range
)
{
    // Each block processes a range of the logits
    int range_min = blockIdx.x * BLOCK_VOCAB_SPAN;
    int range_max = MIN(range_min + BLOCK_VOCAB_SPAN, vocab_size);

    __shared__ float frequency[BLOCK_VOCAB_SPAN];
    __shared__ float presence[BLOCK_VOCAB_SPAN];
    for (int i = threadIdx.x; i < BLOCK_VOCAB_SPAN; i += NUM_THREADS)
    {
        frequency[i] = 0.0f;
        presence[i] = 0.0f;
    }
    __syncthreads();

    // Record which tokens from the range appear in past_ids
    for (int i = threadIdx.x; i < past_len; i += NUM_THREADS)
    {
        if (i <= past_len - sustain_range - decay_range)
            continue;

        int tid = (int) past_ids[i];
        if (tid < range_min || tid >= range_max)
            continue;

        float distf = (float)(past_len - i);
        float sustain_rangef = (float)sustain_range;
        float decay_rangef = (float)decay_range;
        float factor = decay_range > 0 ? 1.0f - (distf - sustain_rangef) / decay_rangef : 1.0f;
        factor = CLAMP(factor, 0.0f, 1.0f);
        atomicAdd(frequency + tid - range_min, factor * freq_p);
        shmemAtomicMaxF(presence + tid - range_min, factor * pres_p);
    }
    __syncthreads();

    // Apply frequency to range
    for (int i = threadIdx.x; i < BLOCK_VOCAB_SPAN; i += NUM_THREADS)
    {
        if (i + range_min >= vocab_size)
            break;

        float v;
        if constexpr (input_fp16)
            v = __half2float(((half*) in_logits)[i + range_min]);
        else
            v = ((float*) in_logits)[i + range_min];

        v -= frequency[i];
        v -= presence[i];
        out_logits[i + range_min] = v;
    }
}

void apply_pres_freq_pens
(
    const at::Tensor& in_logits,
    const at::Tensor& out_logits,
    const at::Tensor& past_ids,
    float pres_p,
    float freq_p,
    int sustain_range,
    int decay_range
)
{
    const at::cuda::OptionalCUDAGuard device_guard(in_logits.device());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    TORCH_CHECK_DTYPE(out_logits, kFloat);
    TORCH_CHECK_DTYPE(past_ids, kLong);
    TORCH_CHECK(in_logits.size(0) == 1, "rep. penalties only implemented for bsz 1");  // TODO
    TORCH_CHECK_SHAPES(past_ids, 0, in_logits, 0, 1);
    TORCH_CHECK_SHAPES_FULL(in_logits, out_logits);

    int past_len = past_ids.size(1);
    int vocab_size = in_logits.size(1);
    int num_blocks = CEIL_DIVIDE(vocab_size, BLOCK_VOCAB_SPAN);

    #define kernel_args \
        (const void*) in_logits.data_ptr(), \
        (float*) out_logits.data_ptr(), \
        (const uint64_t*) past_ids.data_ptr(), \
        past_len, \
        vocab_size, \
        pres_p, \
        freq_p, \
        sustain_range, \
        decay_range

    if (in_logits.dtype() == at::kHalf)
        apply_pres_freq_pens_kernel<true><<<num_blocks, NUM_THREADS, 0, stream>>>(kernel_args);
    else
        apply_pres_freq_pens_kernel<false><<<num_blocks, NUM_THREADS, 0, stream>>>(kernel_args);

    #undef kernel_args

    cuda_check(cudaPeekAtLastError());
}