/*
 * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "src/fastertransformer/models/decoding/Decoding.h"
#include "src/fastertransformer/utils/nvtx_utils.h"
#include <cuda_profiler_api.h>
#include <sys/time.h>

using namespace fastertransformer;

template<typename T>
int decodingExample(const size_t batch_size,
                    const size_t beam_width,
                    const size_t head_num,
                    const size_t size_per_head,
                    const size_t inter_size,
                    const size_t vocab_size,
                    const size_t num_layers,
                    const size_t max_seq_len,
                    const size_t memory_max_seq_len,
                    const size_t memory_hidden_units,
                    const int    top_k,
                    const float  top_p);

int main(int argc, char** argv)
{
    if (argc != 14) {
        printf("[ERROR] decoding_example batch_size beam_width head_num size_per_head inter_size vocab_size"
               " num_layers max_seq_len memory_max_seq_len memory_hidden_units top_k top_p data_type\n");
        printf("e.g., ./bin/decoding_example 4 1 8 64 2048 30000 6 32 32 512 0 0.6 1\n");
        return 0;
    }

    int                  batch_size          = atoi(argv[1]);
    int                  beam_width          = atoi(argv[2]);
    int                  head_num            = atoi(argv[3]);
    int                  size_per_head       = atoi(argv[4]);
    int                  inter_size          = atoi(argv[5]);
    int                  vocab_size          = atoi(argv[6]);
    int                  num_layers          = atoi(argv[7]);
    int                  max_seq_len         = atoi(argv[8]);
    int                  memory_max_seq_len  = atoi(argv[9]);
    int                  memory_hidden_units = atoi(argv[10]);
    int                  top_k               = atoi(argv[11]);
    float                top_p               = atof(argv[12]);
    const CublasDataType data_type           = static_cast<CublasDataType>(atoi(argv[13]));  // 0 FP32, 1 FP16, 2 BF 16

    if (data_type == FLOAT_DATATYPE) {
        return decodingExample<float>(batch_size,
                                      beam_width,
                                      head_num,
                                      size_per_head,
                                      inter_size,
                                      vocab_size,
                                      num_layers,
                                      max_seq_len,
                                      memory_max_seq_len,
                                      memory_hidden_units,
                                      top_k,
                                      top_p);
    }
    else if (data_type == HALF_DATATYPE) {
        return decodingExample<half>(batch_size,
                                     beam_width,
                                     head_num,
                                     size_per_head,
                                     inter_size,
                                     vocab_size,
                                     num_layers,
                                     max_seq_len,
                                     memory_max_seq_len,
                                     memory_hidden_units,
                                     top_k,
                                     top_p);
    }
#ifdef ENABLE_BF16
    else if (data_type == BFLOAT16_DATATYPE) {
        return decodingExample<__nv_bfloat16>(batch_size,
                                              beam_width,
                                              head_num,
                                              size_per_head,
                                              inter_size,
                                              vocab_size,
                                              num_layers,
                                              max_seq_len,
                                              memory_max_seq_len,
                                              memory_hidden_units,
                                              top_k,
                                              top_p);
    }
#endif
    else {
        throw std::runtime_error(std::string("[FT][ERROR] is_fp16 should be 0 (use float)"
                                             "or 1 (use half). \n "));
    }
}

template<typename T>
int decodingExample(const size_t batch_size,
                    const size_t beam_width,
                    const size_t head_num,
                    const size_t size_per_head,
                    const size_t inter_size,
                    const size_t vocab_size,
                    const size_t num_layers,
                    const size_t max_seq_len,
                    const size_t memory_max_seq_len,
                    const size_t memory_hidden_units,
                    const int    top_k,
                    const float  top_p)
{
    const size_t hidden_units = head_num * size_per_head;
    const int    start_id     = 0;
    const int    end_id       = 1;

    cudaStream_t     stream;
    cublasHandle_t   cublas_handle;
    cublasLtHandle_t cublaslt_handle;
    cudaStreamCreate(&stream);
    cublasCreate(&cublas_handle);
    cublasLtCreate(&cublaslt_handle);
    cublasSetStream(cublas_handle, stream);
    cublasAlgoMap* cublas_algo_map = new cublasAlgoMap("gemm_config.in");

    Allocator<AllocatorType::CUDA> allocator(getDevice());

    std::mutex*     cublas_wrapper_mutex = new std::mutex();
    cublasMMWrapper cublas_wrapper =
        cublasMMWrapper(cublas_handle, cublaslt_handle, stream, cublas_algo_map, cublas_wrapper_mutex, &allocator);
    if (std::is_same<T, half>::value) {
        cublas_wrapper.setFP16GemmConfig();
    }
#ifdef ENABLE_BF16
    else if (std::is_same<T, __nv_bfloat16>::value) {
        cublas_wrapper.setBF16GemmConfig();
    }
#endif
    else if (std::is_same<T, float>::value) {
        cublas_wrapper.setFP32GemmConfig();
    }
    struct cudaDeviceProp prop;
    check_cuda_error(cudaGetDeviceProperties(&prop, 0));

    fastertransformer::DecodingWeight<T> decoding_weights(
        hidden_units, inter_size, vocab_size, num_layers, max_seq_len, memory_hidden_units);
    Decoding<T> decoding = Decoding<T>(batch_size,
                                       max_seq_len,
                                       memory_max_seq_len,
                                       beam_width,
                                       head_num,
                                       size_per_head,
                                       inter_size,
                                       num_layers,
                                       vocab_size,
                                       start_id,
                                       end_id,
                                       0.0f,
                                       top_k,
                                       top_p,
                                       1.0,   // temperature
                                       1.0f,  // len_penalty,
                                       1.0,   // repetition_penalty
                                       stream,
                                       &cublas_wrapper,
                                       &allocator,
                                       false,
                                       &prop);

    T*   d_memory_tensor;
    int* d_memory_sequence_lengths;
    deviceMalloc(&d_memory_tensor, memory_hidden_units * memory_max_seq_len * batch_size * beam_width);
    deviceMalloc(&d_memory_sequence_lengths, batch_size * beam_width);
    int* h_memory_sequence_lengths = new int[batch_size * beam_width];
    for (int i = 0; i < (int)(batch_size * beam_width); i++) {
        h_memory_sequence_lengths[i] = memory_max_seq_len;
    }
    cudaH2Dcpy(d_memory_sequence_lengths, h_memory_sequence_lengths, batch_size * beam_width);

    int* d_output_ids;
    int* d_parent_ids;
    int* d_sequence_lengths;
    deviceMalloc(&d_output_ids, batch_size * beam_width * max_seq_len, false);
    deviceMalloc(&d_parent_ids, batch_size * beam_width * max_seq_len, false);
    deviceMalloc(&d_sequence_lengths, batch_size * beam_width, false);

    std::vector<Tensor> input_tensors = std::vector<Tensor>{
        Tensor{MEMORY_GPU,
               getTensorType<T>(),
               std::vector<size_t>{batch_size * beam_width, memory_max_seq_len, memory_hidden_units},
               d_memory_tensor},
        Tensor{MEMORY_GPU, TYPE_INT32, std::vector<size_t>{batch_size * beam_width}, d_memory_sequence_lengths}};

    std::vector<Tensor> output_tensors = std::vector<Tensor>{
        Tensor{MEMORY_GPU, TYPE_INT32, std::vector<size_t>{(size_t)max_seq_len, batch_size, beam_width}, d_output_ids},
        Tensor{MEMORY_GPU, TYPE_INT32, std::vector<size_t>{(size_t)max_seq_len, batch_size, beam_width}, d_parent_ids},
        Tensor{MEMORY_GPU, TYPE_INT32, std::vector<size_t>{batch_size, beam_width}, d_sequence_lengths}};
    print_mem_usage();

    cudaProfilerStart();
    const int ite = 10;
    // warm up
    for (int i = 0; i < ite; ++i) {
        decoding.forward(&output_tensors, &input_tensors, &decoding_weights);
    }
    cudaDeviceSynchronize();

    struct timeval start, end;
    cudaDeviceSynchronize();
    gettimeofday(&start, NULL);

    for (int i = 0; i < ite; ++i) {
        decoding.forward(&output_tensors, &input_tensors, &decoding_weights);
    }

    cudaDeviceSynchronize();
    gettimeofday(&end, NULL);

    cudaProfilerStop();

    printf("[INFO] batch_size %ld beam_width %ld head_num %ld size_per_head %ld max_seq_len %ld"
           " num_layers %ld vocab_size %ld, top_k %d, top_p %.3f, FT-CPP-decoding-time %.2f ms\n",
           batch_size,
           beam_width,
           head_num,
           size_per_head,
           max_seq_len,
           num_layers,
           vocab_size,
           top_k,
           top_p,
           ((end.tv_sec - start.tv_sec) * 1000 + (end.tv_usec - start.tv_usec) * 0.001) / ite);

    delete cublas_algo_map;
    delete cublas_wrapper_mutex;
    return 0;
}
