import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_source = """
#include <torch/extension.h>
#include <string>

torch::Tensor sphere_loss_cuda_forward(
    const torch::Tensor& cosine, 
    const torch::Tensor& label,
    float s,
    int m,
    std::string reduction);
"""

cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>
#include <cfloat>

#define BLOCK_SIZE 256
#define WARP_SIZE 32
#define PI 3.141592653589793f

// --- Reduction Helpers ---
template<typename T>
__device__ __forceinline__ T warp_reduce_max(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val = max(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

template<typename T>
__device__ __forceinline__ T warp_reduce_sum(T val) {
    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

__device__ __forceinline__ float block_reduce_max(float val) {
    static __shared__ float shared[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;
    val = warp_reduce_max(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : -FLT_MAX;
    if (wid == 0) val = warp_reduce_max(val);
    return val;
}

__device__ __forceinline__ float block_reduce_sum(float val) {
    static __shared__ float shared[32];
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;
    val = warp_reduce_sum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    val = (threadIdx.x < blockDim.x / WARP_SIZE) ? shared[lane] : 0.0f;
    if (wid == 0) val = warp_reduce_sum(val);
    return val;
}

// Helper: Calculate psi(theta) for SphereFace
// psi = (-1)^k * cos(m*theta) - 2k
__device__ __forceinline__ float calculate_psi(float cos_theta, int m) {
    // 1. acos
    // Clamp to avoid NaN due to numerical errors
    float val = fminf(fmaxf(cos_theta, -1.0f + 1e-7f), 1.0f - 1e-7f);
    float theta = acosf(val);
    
    // 2. k
    float k_float = floorf(theta * m / PI);
    int k = (int)k_float;
    
    // 3. cos(m * theta)
    float cos_m_theta = cosf((float)m * theta);
    
    // 4. (-1)^k
    // If k is even, 1. If k is odd, -1.
    float sign = (k % 2 == 0) ? 1.0f : -1.0f;
    
    return sign * cos_m_theta - 2.0f * k_float;
}

struct __align__(16) Float4 {
    float x, y, z, w;
};

// --- Fused SphereFace Kernel ---
__global__ void sphere_loss_kernel(
    float* __restrict__ output,
    const float* __restrict__ cosine,
    const int64_t* __restrict__ label,
    int num_classes,
    float s,
    int m) 
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    const float* row_cosine = cosine + row * num_classes;
    int64_t target_idx = label[row];

    // 1. Pass 1: Find Max
    float local_max = -FLT_MAX;
    
    int i = tid * 4;
    // Vectorized Loop
    while (i < num_classes) {
        // Safe check for boundary
        if (i + 4 <= num_classes) {
            Float4 vec = reinterpret_cast<const Float4*>(&row_cosine[i])[0];
            
            // Unpack and process
            float v0 = vec.x;
            float v1 = vec.y;
            float v2 = vec.z;
            float v3 = vec.w;
            
            // Check if any is target
            int idx0 = i;
            int idx1 = i+1;
            int idx2 = i+2;
            int idx3 = i+3;
            
            if (idx0 == target_idx) v0 = calculate_psi(v0, m);
            if (idx1 == target_idx) v1 = calculate_psi(v1, m);
            if (idx2 == target_idx) v2 = calculate_psi(v2, m);
            if (idx3 == target_idx) v3 = calculate_psi(v3, m);
            
            local_max = fmaxf(local_max, v0 * s);
            local_max = fmaxf(local_max, v1 * s);
            local_max = fmaxf(local_max, v2 * s);
            local_max = fmaxf(local_max, v3 * s);
        } else {
            // Tail handling
            for (int k = 0; k < 4 && (i + k) < num_classes; ++k) {
                float v = row_cosine[i+k];
                if ((i+k) == target_idx) {
                    v = calculate_psi(v, m);
                }
                local_max = fmaxf(local_max, v * s);
            }
        }
        i += gridDim.x * 4; // Not using grid-stride across blocks, block is per row.
        // Actually here blockDim.x is the stride since one block per row.
        i = i - gridDim.x * 4 + blockDim.x * 4; 
    }
    
    // Correct logic for simple stride loop within one row
    // Reset loop
    local_max = -FLT_MAX;
    for (int k = tid; k < num_classes; k += blockDim.x) {
        float val = row_cosine[k];
        if (k == target_idx) {
            val = calculate_psi(val, m);
        }
        val *= s;
        local_max = fmaxf(local_max, val);
    }

    float global_max = block_reduce_max(local_max);
    
    __shared__ float s_max;
    if (tid == 0) s_max = global_max;
    __syncthreads();
    global_max = s_max;

    // 2. Pass 2: Sum Exp
    float local_sum = 0.0f;
    for (int k = tid; k < num_classes; k += blockDim.x) {
        float val = row_cosine[k];
        if (k == target_idx) {
            val = calculate_psi(val, m);
        }
        val *= s;
        local_sum += __expf(val - global_max);
    }

    float global_sum = block_reduce_sum(local_sum);

    // 3. Final Calculation
    if (tid == 0) {
        float t_val = row_cosine[target_idx];
        t_val = calculate_psi(t_val, m) * s;
        
        // Loss = log(sum) - log(exp(target)) 
        //      = log(sum_exp * e^max) - target
        //      = max + log(sum_exp) - target
        float loss = global_max + __logf(global_sum) - t_val;
        output[row] = loss;
    }
}

torch::Tensor sphere_loss_cuda_forward(
    const torch::Tensor& cosine, 
    const torch::Tensor& label,
    float s,
    int m,
    std::string reduction) 
{
    TORCH_CHECK(cosine.is_cuda() && label.is_cuda(), "Inputs must be CUDA");
    TORCH_CHECK(cosine.is_contiguous(), "Cosine must be contiguous");
    
    int batch_size = cosine.size(0);
    int num_classes = cosine.size(1);

    auto output = torch::empty({batch_size}, cosine.options());

    // One Block per Sample
    sphere_loss_kernel<<<batch_size, BLOCK_SIZE>>>(
        output.data_ptr<float>(),
        cosine.data_ptr<float>(),
        label.data_ptr<int64_t>(),
        num_classes,
        s,
        m
    );

    if (reduction == "mean") {
        return output.mean();
    } else if (reduction == "sum") {
        return output.sum();
    }

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, s=64.0, m=4, reduction='none'):
        super(ModelNew, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction
        self.op = load_inline(
            name='sphere_loss_op',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['sphere_loss_cuda_forward'],
            verbose=False,
            extra_cuda_cflags=['-O3']
        )

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        return self.op.sphere_loss_cuda_forward(cosine.contiguous(), label.contiguous(), self.s, self.m, self.reduction)