# 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 as th
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor vlad_forward_cuda(torch::Tensor x, torch::Tensor logits, torch::Tensor clusters2, int64_t B, int64_t N, int64_t D, int64_t K, int64_t C);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define D_TILE 32
#define K_TILE 8
#define WARPS_PER_BLOCK 4

__inline__ __device__ float warp_sum0(float v) {
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__inline__ __device__ float warp_max0(float v) {
    for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, o));
    return v;
}

__global__ void softmax_keep_kernel(const float* __restrict__ logits,
                                    float* __restrict__ probs,
                                    int rows, int C, int K) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int wid = tid >> 5;
    int row = blockIdx.x * WARPS_PER_BLOCK + wid;
    if (row >= rows) return;

    const float* in = logits + (long long)row * C;
    float m = -INFINITY;
    for (int c = lane; c < C; c += 32) m = fmaxf(m, in[c]);
    m = warp_max0(m);
    m = __shfl_sync(0xffffffff, m, 0);

    float s = 0.0f;
    for (int c = lane; c < C; c += 32) s += expf(in[c] - m);
    s = warp_sum0(s);
    s = __shfl_sync(0xffffffff, s, 0);
    float inv = 1.0f / s;

    if (lane < K) probs[(long long)row * K + lane] = expf(in[lane] - m) * inv;
}

__global__ void raw_vlad_kernel(const float* __restrict__ x,
                                const float* __restrict__ probs,
                                const float* __restrict__ c2,
                                float* __restrict__ raw,
                                float* __restrict__ partial,
                                int B, int N, int D, int K, int num_dtiles) {
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int d = blockIdx.x * D_TILE + tx;
    int k = blockIdx.y * K_TILE + ty;
    int b = blockIdx.z;
    float val = 0.0f;

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

    float ss = warp_sum0(val * val);
    if (tx == 0 && k < K) {
        partial[((long long)b * K + k) * num_dtiles + blockIdx.x] = ss;
    }
}

__global__ void invnorm_kernel(const float* __restrict__ partial,
                               float* __restrict__ invs,
                               int B, int K, int num_dtiles) {
    int tid = threadIdx.x;
    int k = blockIdx.x;
    int b = blockIdx.y;
    float s = 0.0f;
    if (tid < num_dtiles) s = partial[((long long)b * K + k) * num_dtiles + tid];
    s = warp_sum0(s);
    if (tid == 0) {
        invs[(long long)b * K + k] = (s > 1.0e-24f) ? (rsqrtf(s) * rsqrtf((float)K)) : 0.0f;
    }
}

__global__ void apply_kernel(const float* __restrict__ raw,
                             const float* __restrict__ invs,
                             float* __restrict__ out,
                             long long total, int D, int K) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int k = idx % K;
    long long bd = idx / K;
    int b = bd / D;
    out[idx] = raw[idx] * invs[(long long)b * K + k];
}

torch::Tensor vlad_forward_cuda(torch::Tensor x, torch::Tensor logits, torch::Tensor clusters2,
                                int64_t B, int64_t N, int64_t D, int64_t K, int64_t C) {
    auto opts = x.options();
    auto probs = torch::empty({B * N, K}, opts);
    auto raw = torch::empty({B, D, K}, opts);
    int64_t nd = (D + D_TILE - 1) / D_TILE;
    auto partial = torch::empty({B, K, nd}, opts);
    auto invs = torch::empty({B, K}, opts);
    auto out = torch::empty({B, D * K}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    softmax_keep_kernel<<<(int)((B * N + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK),
                          WARPS_PER_BLOCK * 32, 0, stream>>>(
        logits.data_ptr<float>(), probs.data_ptr<float>(), (int)(B * N), (int)C, (int)K);

    dim3 block(D_TILE, K_TILE);
    dim3 grid((unsigned)nd, (unsigned)((K + K_TILE - 1) / K_TILE), (unsigned)B);
    raw_vlad_kernel<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), probs.data_ptr<float>(), clusters2.data_ptr<float>(),
        raw.data_ptr<float>(), partial.data_ptr<float>(), (int)B, (int)N, (int)D, (int)K, (int)nd);

    dim3 gridn((unsigned)K, (unsigned)B);
    invnorm_kernel<<<gridn, 32, 0, stream>>>(partial.data_ptr<float>(), invs.data_ptr<float>(), (int)B, (int)K, (int)nd);

    long long total = B * D * K;
    apply_kernel<<<(unsigned)((total + 255) / 256), 256, 0, stream>>>(
        raw.data_ptr<float>(), invs.data_ptr<float>(), out.data_ptr<float>(), total, (int)D, (int)K);
    return out;
}
'''

_ext = load_inline(
    name='netvlad_ext_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['vlad_forward_cuda'],
    extra_cuda_cflags=['-O3'],
    extra_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 * th.randn(feature_size, clusters))
        self.batch_norm = nn.BatchNorm1d(clusters)
        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):
        B = x.size(0)
        N = x.size(1)
        D = self.feature_size
        K = self.cluster_size
        C = K + self.ghost_clusters
        xc = x.contiguous()
        assignment = th.matmul(xc.view(-1, D), self.clusters)
        assignment = self.batch_norm(assignment)
        return _ext.vlad_forward_cuda(xc, assignment.contiguous(), self.clusters2.contiguous(), B, N, D, K, C)