# Copyright 2018 Antoine Miech 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.

"""
Code modified from here
https://github.com/albanie/collaborative-experts/blob/master/model/net_vlad.py
"""


import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch as th


class Model(nn.Module):
    def __init__(self, cluster_size, feature_size, ghost_clusters):
        super(Model, self).__init__()

        self.feature_size = feature_size
        self.cluster_size = cluster_size
        self.ghost_clusters = ghost_clusters

        init_sc = (1 / math.sqrt(feature_size))
        clusters = cluster_size + ghost_clusters

        # The `clusters` weights are the `(w,b)` in the paper
        self.clusters = nn.Parameter(init_sc * th.randn(feature_size, clusters))
        self.batch_norm = nn.BatchNorm1d(clusters)
        # The `clusters2` weights are the visual words `c_k` in the paper
        self.clusters2 = nn.Parameter(init_sc * th.randn(1, feature_size, cluster_size))
        self.out_dim = self.cluster_size * feature_size

    def forward(self, x, mask=None):
        """Aggregates feature maps into a fixed size representation.  In the following
        notation, B = batch_size, N = num_features, K = num_clusters, D = feature_size.

        Args:
            x (th.Tensor): B x N x D

        Returns:
            (th.Tensor): B x DK
        """
        max_sample = x.size()[1]
        x = x.view(-1, self.feature_size)  # B x N x D -> BN x D

        if x.device != self.clusters.device:
            msg = f"x.device {x.device} != cluster.device {self.clusters.device}"
            raise ValueError(msg)

        assignment = th.matmul(x, self.clusters)  # (BN x D) x (D x (K+G)) -> BN x (K+G)
        assignment = self.batch_norm(assignment)

        assignment = F.softmax(assignment, dim=1)  # BN x (K+G) -> BN x (K+G)
        # remove ghost assigments
        assignment = assignment[:, :self.cluster_size]
        assignment = assignment.view(-1, max_sample, self.cluster_size)  # -> B x N x K
        a_sum = th.sum(assignment, dim=1, keepdim=True)  # B x N x K -> B x 1 x K
        a = a_sum * self.clusters2

        assignment = assignment.transpose(1, 2)  # B x N x K -> B x K x N

        x = x.view(-1, max_sample, self.feature_size)  # BN x D -> B x N x D
        vlad = th.matmul(assignment, x)  # (B x K x N) x (B x N x D) -> B x K x D
        vlad = vlad.transpose(1, 2)  # -> B x D x K
        vlad = vlad - a

        # L2 intra norm
        vlad = F.normalize(vlad)

        # flattening + L2 norm
        vlad = vlad.reshape(-1, self.cluster_size * self.feature_size)  # -> B x DK
        vlad = F.normalize(vlad)
        return vlad  # B x DK

batch_size = 2048
num_features = 100
num_clusters = 32
feature_size = 512
ghost_clusters = 16

def get_inputs():
  return [torch.rand(batch_size, num_features, feature_size)]

def get_init_inputs():
  return [num_clusters, feature_size, ghost_clusters]

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor post_forward(torch::Tensor x, torch::Tensor assignment, torch::Tensor clusters2, int64_t cluster_size);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void softmax_store_kernel(const float* __restrict__ a,
                                     float* __restrict__ p,
                                     float* __restrict__ asum,
                                     int rows, int C, int K, int N) {
    int row = blockIdx.x;
    if (row >= rows) return;
    int tid = threadIdx.x;
    extern __shared__ float smem[];

    float mx = -CUDART_INF_F;
    const float* ar = a + (long long)row * C;
    for (int c = tid; c < C; c += blockDim.x) {
        float v = ar[c];
        mx = fmaxf(mx, v);
    }
    smem[tid] = mx;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    mx = smem[0];

    float sum = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) {
        sum += expf(ar[c] - mx);
    }
    smem[tid] = sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    float inv = 1.0f / smem[0];

    int b = row / N;
    int n = row - b * N;
    float* pr = p + ((long long)b * N + n) * K;
    float* as = asum + (long long)b * K;
    for (int k = tid; k < K; k += blockDim.x) {
        float pv = expf(ar[k] - mx) * inv;
        pr[k] = pv;
        atomicAdd(as + k, pv);
    }
}

__global__ void vlad_bmm_kernel(const float* __restrict__ x,
                                const float* __restrict__ p,
                                const float* __restrict__ asum,
                                const float* __restrict__ c2,
                                float* __restrict__ tmp,
                                int B, int N, int D, int K) {
    int d = blockIdx.x * 16 + threadIdx.x;
    int k = blockIdx.y * 16 + threadIdx.y;
    int b = blockIdx.z;
    if (d >= D || k >= K || b >= B) return;

    float acc = 0.0f;
    const float* xb = x + (long long)b * N * D + d;
    const float* pb = p + (long long)b * N * K + k;
    #pragma unroll 4
    for (int n = 0; n < N; ++n) {
        acc += pb[(long long)n * K] * xb[(long long)n * D];
    }
    acc -= asum[(long long)b * K + k] * c2[(long long)d * K + k];
    tmp[((long long)b * K + k) * D + d] = acc;
}

__global__ void norm_output_kernel(const float* __restrict__ tmp,
                                   float* __restrict__ out,
                                   int B, int D, int K) {
    int b = blockIdx.x;
    int k = blockIdx.y;
    int tid = threadIdx.x;
    extern __shared__ float smem[];

    float ss = 0.0f;
    const float* base = tmp + ((long long)b * K + k) * D;
    for (int d = tid; d < D; d += blockDim.x) {
        float v = base[d];
        ss += v * v;
    }
    smem[tid] = ss;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    float norm = sqrtf(smem[0]);
    float scale = (norm > 1.0e-12f ? 1.0f / norm : 1.0e12f) * rsqrtf((float)K);
    float* ob = out + (long long)b * D * K + k;
    for (int d = tid; d < D; d += blockDim.x) {
        ob[(long long)d * K] = base[d] * scale;
    }
}

torch::Tensor post_forward(torch::Tensor x, torch::Tensor assignment, torch::Tensor clusters2, int64_t cluster_size) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(assignment.is_cuda(), "assignment must be CUDA");
    TORCH_CHECK(clusters2.is_cuda(), "clusters2 must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(assignment.scalar_type() == torch::kFloat32, "only float32 supported");
    c10::cuda::CUDAGuard guard(x.device());

    int B = (int)x.size(0);
    int N = (int)x.size(1);
    int D = (int)x.size(2);
    int C = (int)assignment.size(1);
    int K = (int)cluster_size;
    int rows = B * N;

    auto opts = x.options();
    auto probs = torch::empty({B, N, K}, opts);
    auto asum = torch::zeros({B, K}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int sm_threads = 64;
    softmax_store_kernel<<<rows, sm_threads, sm_threads * sizeof(float), stream>>>(
        assignment.data_ptr<float>(), probs.data_ptr<float>(), asum.data_ptr<float>(), rows, C, K, N);

    auto tmp = torch::empty({B, K, D}, opts);
    dim3 block(16, 16, 1);
    dim3 grid((D + 15) / 16, (K + 15) / 16, B);
    vlad_bmm_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), probs.data_ptr<float>(), asum.data_ptr<float>(),
        clusters2.data_ptr<float>(), tmp.data_ptr<float>(), B, N, D, K);

    auto out = torch::empty({B, D * K}, opts);
    dim3 grid2(B, K, 1);
    int nt = 256;
    norm_output_kernel<<<grid2, nt, nt * sizeof(float), stream>>>(
        tmp.data_ptr<float>(), out.data_ptr<float>(), B, D, K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="netvlad_post_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["post_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, cluster_size, feature_size, ghost_clusters):
        super(ModelNew, self).__init__()
        self.feature_size = feature_size
        self.cluster_size = cluster_size
        self.ghost_clusters = ghost_clusters

        init_sc = (1 / math.sqrt(feature_size))
        clusters = cluster_size + ghost_clusters
        self.clusters = nn.Parameter(init_sc * torch.randn(feature_size, clusters))
        self.batch_norm = nn.BatchNorm1d(clusters)
        self.clusters2 = nn.Parameter(init_sc * torch.randn(1, feature_size, cluster_size))
        self.out_dim = self.cluster_size * feature_size

    def forward(self, x, mask=None):
        max_sample = x.size(1)
        x_contig = x.contiguous()
        x_flat = x_contig.view(-1, self.feature_size)
        if x_flat.device != self.clusters.device:
            msg = f"x.device {x_flat.device} != cluster.device {self.clusters.device}"
            raise ValueError(msg)
        assignment = torch.matmul(x_flat, self.clusters)
        assignment = self.batch_norm(assignment)
        assignment = assignment.contiguous()
        c2 = self.clusters2.contiguous()
        return _ext.post_forward(x_contig, assignment, c2, self.cluster_size)