# 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 = 0

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
from torch.utils.cpp_extension import load_inline
_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor netvlad_cuda_forward(torch::Tensor x, torch::Tensor clusters, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor clusters2, double eps, bool training);
torch::Tensor netvlad_forward(torch::Tensor x, torch::Tensor clusters, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor clusters2, double eps, bool training) {
    return netvlad_cuda_forward(x, clusters, bn_weight, bn_bias, running_mean, running_var, clusters2, eps, training);
}
'''
_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#define D 512
#define K 32
#define N 100
#define INV_SQRT_K 0.17677669529663687f
__global__ void score_kernel4(const float* __restrict__ x, const float* __restrict__ clusters, float* __restrict__ scores, int rows) {
    int lt = threadIdx.x & 127;
    int sub = threadIdx.x >> 7;
    int row = blockIdx.x * 4 + sub;
    if(row >= rows) return;
    int kk = lt >> 2;
    int lane = lt & 3;
    float sum = 0.0f;
    const float* xr = x + (size_t)row * D;
    for(int d=lane; d<D; d+=4) sum = fmaf(xr[d], clusters[d*K + kk], sum);
    sum += __shfl_xor_sync(0xffffffff, sum, 1, 4);
    sum += __shfl_xor_sync(0xffffffff, sum, 2, 4);
    if(lane == 0) scores[(size_t)row*K + kk] = sum;
}
__global__ void stats_kernel(const float* __restrict__ scores, float* __restrict__ mean, float* __restrict__ invstd, int rows) {
    int k=blockIdx.x, tid=threadIdx.x;
    float s=0.0f, ss=0.0f;
    for(int r=tid; r<rows; r+=256){ float v=scores[(size_t)r*K+k]; s+=v; ss=fmaf(v,v,ss); }
    __shared__ float sm[256], sv[256]; sm[tid]=s; sv[tid]=ss; __syncthreads();
    for(int off=128; off>0; off>>=1){ if(tid<off){ sm[tid]+=sm[tid+off]; sv[tid]+=sv[tid+off]; } __syncthreads(); }
    if(tid==0){ float m=sm[0]/(float)rows; float var=sv[0]/(float)rows - m*m; if(var<0.0f)var=0.0f; mean[k]=m; invstd[k]=rsqrtf(var+1.0e-5f); }
}
__global__ void bn_softmax_kernel8(float* __restrict__ scores, const float* __restrict__ bw, const float* __restrict__ bb, const float* __restrict__ mean, const float* __restrict__ invstd, int rows) {
    int lane = threadIdx.x & 31;
    int wrp = threadIdx.x >> 5;
    int row = blockIdx.x * 8 + wrp;
    if(row >= rows) return;
    float v = (scores[(size_t)row*K + lane] - mean[lane]) * invstd[lane];
    v = v * bw[lane] + bb[lane];
    float mx = v;
    for(int off=16; off>0; off>>=1) mx = fmaxf(mx, __shfl_down_sync(0xffffffff, mx, off));
    mx = __shfl_sync(0xffffffff, mx, 0);
    float e = expf(v - mx), s = e;
    for(int off=16; off>0; off>>=1) s += __shfl_down_sync(0xffffffff, s, off);
    s = __shfl_sync(0xffffffff, s, 0);
    scores[(size_t)row*K + lane] = e / s;
}
__global__ void score_softmax_eval_kernel(const float* __restrict__ x, const float* __restrict__ clusters, const float* __restrict__ bw, const float* __restrict__ bb, const float* __restrict__ rm, const float* __restrict__ rv, float eps, float* __restrict__ assign, int rows) {
    int row=blockIdx.x, tid=threadIdx.x, kk=tid>>2, lane=tid&3; float sum=0.0f; const float* xr=x+(size_t)row*D;
    for(int d=lane; d<D; d+=4) sum=fmaf(xr[d],clusters[d*K+kk],sum);
    sum+=__shfl_xor_sync(0xffffffff,sum,1,4); sum+=__shfl_xor_sync(0xffffffff,sum,2,4);
    __shared__ float logits[K]; if(lane==0) logits[kk]=((sum-rm[kk])*rsqrtf(rv[kk]+eps))*bw[kk]+bb[kk]; __syncthreads();
    if(tid<K){ float v=logits[tid]; float mx=v; for(int off=16;off>0;off>>=1) mx=fmaxf(mx,__shfl_down_sync(0xffffffff,mx,off)); mx=__shfl_sync(0xffffffff,mx,0); float e=expf(v-mx), s=e; for(int off=16;off>0;off>>=1) s+=__shfl_down_sync(0xffffffff,s,off); s=__shfl_sync(0xffffffff,s,0); assign[(size_t)row*K+tid]=e/s; }
}
__global__ void vlad_kernel(const float* __restrict__ x, const float* __restrict__ assign, const float* __restrict__ c2, float* __restrict__ out, int B) {
    size_t idx=(size_t)blockIdx.x*blockDim.x+threadIdx.x, total=(size_t)B*D*K; if(idx>=total) return;
    int k=idx%K, tmp=idx/K, d=tmp%D, b=tmp/D; float acc=0.0f, center=c2[d*K+k]; const float* xb=x+(size_t)b*N*D+d; const float* ab=assign+(size_t)b*N*K+k;
    #pragma unroll 4
    for(int n=0; n<N; ++n) acc=fmaf(ab[n*K], xb[(size_t)n*D]-center, acc);
    out[idx]=acc;
}
__global__ void intra_norm_kernel(float* __restrict__ out, int B) {
    int bid=blockIdx.x, b=bid/K, k=bid-b*K, tid=threadIdx.x; float ss=0.0f;
    for(int d=tid; d<D; d+=256){ float v=out[((size_t)b*D+d)*K+k]; ss=fmaf(v,v,ss); }
    __shared__ float smem[256]; smem[tid]=ss; __syncthreads();
    for(int s=128; s>0; s>>=1){ if(tid<s) smem[tid]+=smem[tid+s]; __syncthreads(); }
    float inv=rsqrtf(smem[0]+1e-12f)*INV_SQRT_K;
    for(int d=tid; d<D; d+=256) out[((size_t)b*D+d)*K+k]*=inv;
}
torch::Tensor netvlad_cuda_forward(torch::Tensor x, torch::Tensor clusters, torch::Tensor bn_weight, torch::Tensor bn_bias, torch::Tensor running_mean, torch::Tensor running_var, torch::Tensor clusters2, double eps, bool training) {
    x=x.contiguous(); int64_t B=x.size(0); int rows=(int)(B*N); auto opts=x.options(); auto assign=torch::empty({rows,K},opts);
    if(training){
        score_kernel4<<<(rows+3)/4,512>>>(x.data_ptr<float>(),clusters.data_ptr<float>(),assign.data_ptr<float>(),rows);
        auto stats=torch::empty({2,K},opts);
        stats_kernel<<<K,256>>>(assign.data_ptr<float>(),stats.data_ptr<float>(),stats.data_ptr<float>()+K,rows);
        bn_softmax_kernel8<<<(rows+7)/8,256>>>(assign.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),stats.data_ptr<float>(),stats.data_ptr<float>()+K,rows);
    } else {
        score_softmax_eval_kernel<<<rows,128>>>(x.data_ptr<float>(),clusters.data_ptr<float>(),bn_weight.data_ptr<float>(),bn_bias.data_ptr<float>(),running_mean.data_ptr<float>(),running_var.data_ptr<float>(),(float)eps,assign.data_ptr<float>(),rows);
    }
    auto out3=torch::empty({B,D,K},opts); size_t total=(size_t)B*D*K;
    vlad_kernel<<<(int)((total+255)/256),256>>>(x.data_ptr<float>(),assign.data_ptr<float>(),clusters2.data_ptr<float>(),out3.data_ptr<float>(),(int)B);
    intra_norm_kernel<<<(int)B*K,256>>>(out3.data_ptr<float>(),(int)B);
    return out3.view({B,D*K});
}
'''
_ext=load_inline(name='netvlad_ext_final',cpp_sources=[_cpp_src],cuda_sources=[_cuda_src],functions=['netvlad_forward'],extra_cuda_cflags=['-O3','--use_fast_math'],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*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):
        return _ext.netvlad_forward(x,self.clusters,self.batch_norm.weight,self.batch_norm.bias,self.batch_norm.running_mean,self.batch_norm.running_var,self.clusters2,self.batch_norm.eps,self.training)