/*
 * 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.
 */

#pragma once

#include "src/fastertransformer/layers/DenseWeight.h"
#include "src/fastertransformer/models/vit/ViTLayerWeight.h"

namespace fastertransformer {

template<typename T>
struct ViTEmbeds {
    const T* class_embed;
    const T* position_embed;
};

#define WEIGHT_N 6
template<typename T>
struct ViTWeight {

    ViTWeight() = delete;
    ViTWeight(const int  embed_dim,
              const int  inter_size,
              const int  num_layer,
              const int  img_size,
              const int  patch_size,
              const int  chn_num,
              const bool with_cls_token,
              const bool hold_buffer = true):
        with_cls_token_(with_cls_token),
        embed_dim_(embed_dim),
        inter_size_(inter_size),
        num_layer_(num_layer),
        img_size_(img_size),
        patch_size_(patch_size),
        chn_num_(chn_num),
        seq_len_(img_size_ * img_size_ / (patch_size_ * patch_size_) + (with_cls_token_ ? 1 : 0))
    {
        weights_size[0] = chn_num_ * patch_size_ * patch_size_ * embed_dim_;
        weights_size[1] = embed_dim_;
        weights_size[2] = with_cls_token_ ? embed_dim_ : 0;
        weights_size[3] = embed_dim_ * seq_len_;
        weights_size[4] = embed_dim_;
        weights_size[5] = embed_dim_;

        if (hold_buffer) {
            for (int i = 0; i < WEIGHT_N; i++) {
                if (weights_size[i] == 0) {
                    continue;
                }

                deviceMalloc(&weights_ptr[i], weights_size[i]);
            }

            setWeightPtr();
        }

        vit_layer_weights.reserve(num_layer_);
        for (int i = 0; i < num_layer_; i++) {
            vit_layer_weights.push_back(ViTLayerWeight<T>(embed_dim_, inter_size_, i, hold_buffer));
        }
    }

    ~ViTWeight()
    {
        if (is_maintain_buffer == true) {
            vit_layer_weights.clear();
            for (int i = 0; i < WEIGHT_N; i++) {
                if (weights_ptr[i] != nullptr) {
                    deviceFree(weights_ptr[i]);
                }
            }

            post_transformer_layernorm_weights.gamma = nullptr;
            post_transformer_layernorm_weights.beta  = nullptr;
            pre_transform_embeds.class_embed         = nullptr;
            pre_transform_embeds.position_embed      = nullptr;
            pre_encoder_conv_weights.kernel          = nullptr;
            pre_encoder_conv_weights.bias            = nullptr;
            is_maintain_buffer                       = false;
        }
    }

    ViTWeight(const ViTWeight& other):
        with_cls_token_(other.with_cls_token_),
        embed_dim_(other.embed_dim_),
        inter_size_(other.inter_size_),
        num_layer_(other.num_layer_),
        img_size_(other.img_size_),
        patch_size_(other.patch_size_),
        chn_num_(other.chn_num_),
        seq_len_(other.seq_len_)
    {
        memcpy(weights_size, other.weights_size, sizeof(size_t) * WEIGHT_N);
        if (other.is_maintain_buffer) {
            for (int i = 0; i < WEIGHT_N; i++) {
                if (!is_maintain_buffer) {
                    deviceMalloc(&weights_ptr[i], weights_size[i]);
                }
                cudaD2Dcpy(weights_ptr[i], other.weights_ptr[i], weights_size[i]);
            }
            setWeightPtr();
        }

        vit_layer_weights.clear();
        vit_layer_weights.reserve(num_layer_);
        for (int i = 0; i < num_layer_; i++) {
            vit_layer_weights.push_back(other.vit_layer_weights[i]);
        }
    }

    ViTWeight& operator=(const ViTWeight& other)
    {
        embed_dim_      = other.embed_dim_;
        inter_size_     = other.inter_size_;
        num_layer_      = other.num_layer_;
        img_size_       = other.img_size_;
        patch_size_     = other.patch_size_;
        chn_num_        = other.chn_num_;
        seq_len_        = other.seq_len_;
        with_cls_token_ = other.with_cls_token_;
        memcpy(weights_size, other.weights_size, sizeof(size_t) * WEIGHT_N);

        if (other.is_maintain_buffer) {
            for (int i = 0; i < WEIGHT_N; i++) {
                if (!is_maintain_buffer) {
                    deviceMalloc(&weights_ptr[i], weights_size[i]);
                }
                cudaD2Dcpy(weights_ptr[i], other.weights_ptr[i], weights_size[i]);
            }
            setWeightPtr();
        }

        vit_layer_weights.clear();
        vit_layer_weights.reserve(num_layer_);
        for (int i = 0; i < num_layer_; i++) {
            vit_layer_weights.push_back(other.vit_layer_weights[i]);
        }

        return *this;
    }

    size_t GetSerializeSize()
    {
        size_t count;
        for (int i = 0; i < WEIGHT_N; i++) {
            count += weights_size[i];
        }
        count *= sizeof(T);

        for (auto& lw : vit_layer_weights) {
            count += lw.GetSerializeSize();
        }

        return count;
    }

    void serialize(void* buffer)
    {
        char* tmp_buf = (char*)buffer;
        for (int i = 0; i < WEIGHT_N; i++) {
            cudaMemcpy(tmp_buf, weights_ptr[i], sizeof(T) * weights_size[i], cudaMemcpyDeviceToHost);
            tmp_buf += sizeof(T) * weights_size[i];
        }

        for (auto& lw : vit_layer_weights) {
            lw.serialize(tmp_buf);
            tmp_buf += lw.GetSerializeSize();
        }
    }

    void deserialize(const void* buffer)
    {
        if (!is_maintain_buffer) {
            return;
        }

        char* tmp_buf = (char*)buffer;
        for (int i = 0; i < WEIGHT_N; i++) {
            cudaMemcpy(weights_ptr[i], tmp_buf, sizeof(T) * weights_size[i], cudaMemcpyHostToDevice);
            tmp_buf += sizeof(T) * weights_size[i];
        }

        for (auto& lw : vit_layer_weights) {
            lw.deserialize(tmp_buf);
            tmp_buf += lw.GetSerializeSize();
        }
    }

    void CopyWeightsFromHostBuffers(const T* const*& w)
    {
        cudaMemcpy(
            const_cast<T*>(pre_encoder_conv_weights.kernel), *w++, sizeof(T) * weights_size[0], cudaMemcpyHostToDevice);
        cudaMemcpy(
            const_cast<T*>(pre_encoder_conv_weights.bias), *w++, sizeof(T) * weights_size[1], cudaMemcpyHostToDevice);
        if (with_cls_token_) {
            cudaMemcpy(const_cast<T*>(pre_transform_embeds.class_embed),
                       *w++,
                       sizeof(T) * weights_size[2],
                       cudaMemcpyHostToDevice);
            with_cls_token_ = true;
        }
        cudaMemcpy(const_cast<T*>(pre_transform_embeds.position_embed),
                   *w++,
                   sizeof(T) * weights_size[3],
                   cudaMemcpyHostToDevice);

        for (int i = 0; i < num_layer_; i++) {
            auto& layer_weight = vit_layer_weights[i];
            layer_weight.CopyWeightsFromHostBuffers(w);
        }

        cudaMemcpy(const_cast<T*>(post_transformer_layernorm_weights.gamma),
                   *w++,
                   sizeof(T) * weights_size[4],
                   cudaMemcpyHostToDevice);
        cudaMemcpy(const_cast<T*>(post_transformer_layernorm_weights.beta),
                   *w++,
                   sizeof(T) * weights_size[5],
                   cudaMemcpyHostToDevice);
    }

    inline size_t GetWeightCount()
    {
        size_t weight_count = WEIGHT_N;
        weight_count += num_layer_ * vit_layer_weights[0].GetWeightCount();

        return weight_count;
    }

    void ExportWeights() const
    {
        DataType dtype = DataType::TYPE_INVALID;
        if (std::is_same<T, half>::value) {
            dtype = DataType::TYPE_FP16;
        }
        else if (std::is_same<T, float>::value) {
            dtype = DataType::TYPE_FP32;
        }

        Tensor conv_kernel{MEMORY_GPU,
                           dtype,
                           std::vector<size_t>{embed_dim_, chn_num_, patch_size_, patch_size_},
                           pre_encoder_conv_weights.kernel};  // OIHW
        conv_kernel.saveNpy("./weights/conv_kernel.npy");
        Tensor conv_bias{MEMORY_GPU, dtype, std::vector<size_t>{embed_dim_}, pre_encoder_conv_weights.bias};  // OIHW
        conv_bias.saveNpy("./weights/conv_bias.npy");
        Tensor cls_token{MEMORY_GPU, dtype, std::vector<size_t>{embed_dim_}, pre_transform_embeds.class_embed};
        if (with_cls_token_) {
            cls_token.saveNpy("./weights/cls_token.npy");
        }
        Tensor pos_embed{
            MEMORY_GPU, dtype, std::vector<size_t>{1, seq_len_, embed_dim_}, pre_transform_embeds.position_embed};
        pos_embed.saveNpy("./weights/pos_embed.npy");
        Tensor ln_gamma{MEMORY_GPU, dtype, std::vector<size_t>{embed_dim_}, post_transformer_layernorm_weights.gamma};
        ln_gamma.saveNpy("./weights/enc_ln_scale.npy");
        Tensor ln_beta{MEMORY_GPU, dtype, std::vector<size_t>{embed_dim_}, post_transformer_layernorm_weights.beta};
        ln_beta.saveNpy("./weights/enc_ln_bias.npy");

        for (int i = 0; i < num_layer_; i++) {
            vit_layer_weights[i].ExportWeights(i);
        }
    }

    std::vector<ViTLayerWeight<T>> vit_layer_weights;
    LayerNormWeight<T>             post_transformer_layernorm_weights;
    ViTEmbeds<T>                   pre_transform_embeds;
    DenseWeight<T>                 pre_encoder_conv_weights;
    bool                           with_cls_token_;

private:
    void setWeightPtr()
    {
        pre_encoder_conv_weights.kernel          = weights_ptr[0];
        pre_encoder_conv_weights.bias            = weights_ptr[1];
        pre_transform_embeds.class_embed         = weights_ptr[2];
        pre_transform_embeds.position_embed      = weights_ptr[3];
        post_transformer_layernorm_weights.gamma = weights_ptr[4];
        post_transformer_layernorm_weights.beta  = weights_ptr[5];

        is_maintain_buffer = true;
    }
    int    embed_dim_;
    int    inter_size_;
    int    num_layer_;
    int    img_size_;
    int    patch_size_;
    int    chn_num_;
    int    seq_len_;
    bool   is_maintain_buffer = false;
    T*     weights_ptr[WEIGHT_N]{nullptr};
    size_t weights_size[WEIGHT_N];
};

#undef WEIGHT_N
}  // namespace fastertransformer
