# 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(torch::Tensor logits,
                           torch::Tensor x,
                           torch::Tensor clusters2,
                           torch::Tensor bn_weight,
                           torch::Tensor bn_bias,
                           torch::Tensor running_mean,
                           torch::Tensor running_var,
                           bool training,
                           double momentum,
                           double eps,
                           int64_t cluster_size);

torch::Tensor netvlad_forward(torch::Tensor logits,
                              torch::Tensor x,
                              torch::Tensor clusters2,
                              torch::Tensor bn_weight,
                              torch::Tensor bn_bias,
                              torch::Tensor running_mean,
                              torch::Tensor running_var,
                              bool training,
                              double momentum,
                              double eps,
                              int64_t cluster_size) {
    return netvlad_cuda(logits, x, clusters2, bn_weight, bn_bias, running_mean, running_var,
                        training, momentum, eps, cluster_size);
}
'''

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

__global__ void bn_stats_kernel(const float* __restrict__ logits,
                                float* __restrict__ mean,
                                float* __restrict__ invstd,
                                float* __restrict__ running_mean,
                                float* __restrict__ running_var,
                                int M, int C, float momentum, float eps) {
    int c = blockIdx.x;
    __shared__ double ssum[256];
    __shared__ double ssq[256];
    double sum = 0.0;
    double sq = 0.0;
    for (int m = threadIdx.x; m < M; m += blockDim.x) {
        float v = logits[m * C + c];
        sum += (double)v;
        sq += (double)v * (double)v;
    }
    ssum[threadIdx.x] = sum;
    ssq[threadIdx.x] = sq;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            ssum[threadIdx.x] += ssum[threadIdx.x + stride];
            ssq[threadIdx.x] += ssq[threadIdx.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        double mu = ssum[0] / (double)M;
        double var = ssq[0] / (double)M - mu * mu;
        if (var < 0.0) var = 0.0;
        mean[c] = (float)mu;
        invstd[c] = rsqrtf((float)var + eps);
        if (running_mean != nullptr) {
            running_mean[c] = running_mean[c] * (1.0f - momentum) + ((float)mu) * momentum;
            float unbiased = (M > 1) ? ((float)var * ((float)M / (float)(M - 1))) : (float)var;
            running_var[c] = running_var[c] * (1.0f - momentum) + unbiased * momentum;
        }
    }
}

__global__ void invstd_eval_kernel(const float* __restrict__ running_var,
                                   float* __restrict__ invstd,
                                   int C, float eps) {
    int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c < C) invstd[c] = rsqrtf(running_var[c] + eps);
}

__global__ void bn_softmax_kernel(const float* __restrict__ logits,
                                  const float* __restrict__ mean,
                                  const float* __restrict__ invstd,
                                  const float* __restrict__ weight,
                                  const float* __restrict__ bias,
                                  float* __restrict__ probs,
                                  int M, int C, int K) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M) return;
    float vals[128];
    float mx = -3.402823466e38f;
    int base = row * C;
    #pragma unroll
    for (int c = 0; c < 128; ++c) {
        if (c < C) {
            float y = (logits[base + c] - mean[c]) * invstd[c];
            y = y * weight[c] + bias[c];
            vals[c] = y;
            mx = fmaxf(mx, y);
        }
    }
    float denom = 0.0f;
    #pragma unroll
    for (int c = 0; c < 128; ++c) {
        if (c < C) {
            float e = expf(vals[c] - mx);
            vals[c] = e;
            denom += e;
        }
    }
    float rden = 1.0f / denom;
    int obase = row * K;
    #pragma unroll
    for (int c = 0; c < 128; ++c) {
        if (c < K) probs[obase + c] = vals[c] * rden;
    }
}

__global__ void vlad_norm_kernel(const float* __restrict__ probs,
                                 const float* __restrict__ x,
                                 const float* __restrict__ clusters2,
                                 float* __restrict__ out,
                                 int B, int N, int D, int K) {
    int b = blockIdx.x;
    int k = blockIdx.y;
    int tid = threadIdx.x;
    __shared__ float ps[128];
    __shared__ float red[256];
    if (tid < N) ps[tid] = probs[(b * N + tid) * K + k];
    __syncthreads();

    float v0 = 0.0f, v1 = 0.0f;
    int d0 = tid;
    int d1 = tid + 256;
    if (d0 < D) {
        float c2 = clusters2[d0 * K + k];
        float s = 0.0f;
        #pragma unroll 4
        for (int n = 0; n < 100; ++n) {
            if (n < N) s += ps[n] * (x[(b * N + n) * D + d0] - c2);
        }
        v0 = s;
    }
    if (d1 < D) {
        float c2 = clusters2[d1 * K + k];
        float s = 0.0f;
        #pragma unroll 4
        for (int n = 0; n < 100; ++n) {
            if (n < N) s += ps[n] * (x[(b * N + n) * D + d1] - c2);
        }
        v1 = s;
    }
    float ss = v0 * v0 + v1 * v1;
    red[tid] = ss;
    __syncthreads();
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) red[tid] += red[tid + stride];
        __syncthreads();
    }
    float scale = rsqrtf(red[0] + 1.0e-12f) * rsqrtf((float)K);
    if (d0 < D) out[b * (D * K) + d0 * K + k] = v0 * scale;
    if (d1 < D) out[b * (D * K) + d1 * K + k] = v1 * scale;
}

torch::Tensor netvlad_cuda(torch::Tensor logits,
                           torch::Tensor x,
                           torch::Tensor clusters2,
                           torch::Tensor bn_weight,
                           torch::Tensor bn_bias,
                           torch::Tensor running_mean,
                           torch::Tensor running_var,
                           bool training,
                           double momentum,
                           double eps,
                           int64_t cluster_size) {
    const int B = (int)x.size(0);
    const int N = (int)x.size(1);
    const int D = (int)x.size(2);
    const int M = B * N;
    const int C = (int)logits.size(1);
    const int K = (int)cluster_size;
    auto opts = logits.options();
    auto mean = torch::empty({C}, opts);
    auto invstd = torch::empty({C}, opts);
    auto probs = torch::empty({M, K}, opts);
    auto out = torch::empty({B, D * K}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        bn_stats_kernel<<<C, 256, 0, stream>>>(logits.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
                                               running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
                                               M, C, (float)momentum, (float)eps);
    } else {
        // In eval mode BatchNorm uses the stored running statistics directly.
        mean.copy_(running_mean);
        invstd_eval_kernel<<<(C + 255) / 256, 256, 0, stream>>>(running_var.data_ptr<float>(), invstd.data_ptr<float>(), C, (float)eps);
    }
    bn_softmax_kernel<<<(M + 255) / 256, 256, 0, stream>>>(logits.data_ptr<float>(), mean.data_ptr<float>(), invstd.data_ptr<float>(),
                                                           bn_weight.data_ptr<float>(), bn_bias.data_ptr<float>(),
                                                           probs.data_ptr<float>(), M, C, K);
    dim3 grid(B, K);
    vlad_norm_kernel<<<grid, 256, 0, stream>>>(probs.data_ptr<float>(), x.data_ptr<float>(),
                                               clusters2.data_ptr<float>(), out.data_ptr<float>(),
                                               B, N, D, K);
    return out;
}
'''

_ext = load_inline(
    name='netvlad_custom_ext_v1',
    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):
        # Reference ignores mask.
        x_c = x.contiguous()
        x_flat = x_c.view(-1, self.feature_size)
        logits = torch.matmul(x_flat, self.clusters)
        return _ext.netvlad_forward(
            logits, x_c, self.clusters2, self.batch_norm.weight, self.batch_norm.bias,
            self.batch_norm.running_mean, self.batch_norm.running_var,
            bool(self.batch_norm.training), float(self.batch_norm.momentum), float(self.batch_norm.eps),
            int(self.cluster_size)
        )