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

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

__inline__ __device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_down_sync(0xffffffff, val, offset);
    return val;
}

__inline__ __device__ float warp_reduce_max(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    return val;
}

__global__ void jaro_winkler_softmax_kernel(
    const float* __restrict__ x, 
    const float* __restrict__ target, 
    float* __restrict__ y, 
    int batch_size, 
    int width) 
{
    int row = blockIdx.x;
    int tid = threadIdx.x;

    if (row >= batch_size) return;

    extern __shared__ float s_mem[];
    float* s_target_prob = s_mem;

    float t_max = -3.402823466e+38F;
    for (int i = tid; i < width; i += blockDim.x) {
        t_max = fmaxf(t_max, target[i]);
    }
    t_max = warp_reduce_max(t_max);

    static __shared__ float s_t_max;
    if (tid % 32 == 0) atomicMax((int*)&s_t_max, __float_as_int(t_max)); // atomicMax for float requires casting tricks usually, but here simple reduction is safer via shared mem
    __syncthreads();
    if (tid == 0) s_t_max = -3.402823466e+38F;
    __syncthreads();

    static __shared__ float s_red_max[32];
    int lane = tid % 32;
    int wid = tid / 32;
    if (lane == 0) s_red_max[wid] = t_max;
    __syncthreads();
    if (tid == 0) {
        float val = -3.402823466e+38F;
        for (int i = 0; i < 8; i++) val = fmaxf(val, s_red_max[i]);
        s_t_max = val;
    }
    __syncthreads();
    t_max = s_t_max;

    float t_sum = 0.0f;
    for (int i = tid; i < width; i += blockDim.x) {
        float val = expf(target[i] - t_max);
        s_target_prob[i] = val; // Store exp temporarily
        t_sum += val;
    }
    t_sum = warp_reduce_sum(t_sum);

    static __shared__ float s_red_sum[32];
    if (lane == 0) s_red_sum[wid] = t_sum;
    __syncthreads();
    if (tid == 0) {
        float val = 0.0f;
        for (int i = 0; i < 8; i++) val += s_red_sum[i];
        s_t_max = val; // reuse s_t_max for sum
    }
    __syncthreads();
    t_sum = s_t_max;

    // Normalize target in shared memory
    for (int i = tid; i < width; i += blockDim.x) {
        s_target_prob[i] /= t_sum;
    }
    __syncthreads();


    const float* row_x = x + row * width;

    float x_max = -3.402823466e+38F;
    for (int i = tid; i < width; i += blockDim.x) {
        x_max = fmaxf(x_max, row_x[i]);
    }
    x_max = warp_reduce_max(x_max);

    if (lane == 0) s_red_max[wid] = x_max;
    __syncthreads();
    if (tid == 0) {
        float val = -3.402823466e+38F;
        for (int i = 0; i < 8; i++) val = fmaxf(val, s_red_max[i]);
        s_t_max = val;
    }
    __syncthreads();
    x_max = s_t_max;

    float x_sum = 0.0f;


    for (int i = tid; i < width; i += blockDim.x) {
        x_sum += expf(row_x[i] - x_max);
    }
    x_sum = warp_reduce_sum(x_sum);

    if (lane == 0) s_red_sum[wid] = x_sum;
    __syncthreads();
    if (tid == 0) {
        float val = 0.0f;
        for (int i = 0; i < 8; i++) val += s_red_sum[i];
        s_t_max = val;
    }
    __syncthreads();
    x_sum = s_t_max;

    float m_sum = 0.0f;
    float prefix_sum = 0.0f;

    for (int i = tid; i < width; i += blockDim.x) {
        float val_x = expf(row_x[i] - x_max) / x_sum;
        float val_t = s_target_prob[i];

        float min_val = fminf(val_x, val_t);
        m_sum += min_val;

        if (i < 4) {
            prefix_sum += min_val;
        }
    }

    m_sum = warp_reduce_sum(m_sum);
    prefix_sum = warp_reduce_sum(prefix_sum);

    if (lane == 0) {
        s_red_sum[wid] = m_sum;
        s_red_max[wid] = prefix_sum; // reuse buffer
    }
    __syncthreads();

    if (tid == 0) {
        float total_m = 0.0f;
        float total_prefix = 0.0f;
        for (int i = 0; i < 8; i++) {
            total_m += s_red_sum[i];
            total_prefix += s_red_max[i];
        }

        // Jaro Calculation
        float sim_j = (2.0f * total_m + 1.0f) / 3.0f;
        float jw = sim_j + 0.1f * total_prefix * (1.0f - sim_j);
        y[row] = jw;
    }
}

torch::Tensor launch_jaro_winkler_softmax(torch::Tensor x, torch::Tensor target) {
    auto batch_size = x.size(0);
    auto width = x.size(1);
    auto y = torch::empty({batch_size}, x.options());

    const int threads = 256;
    const int blocks = batch_size;
    int shared_mem = width * sizeof(float);

    jaro_winkler_softmax_kernel<<<blocks, threads, shared_mem>>>(
        x.data_ptr<float>(), 
        target.data_ptr<float>(), 
        y.data_ptr<float>(), 
        batch_size, 
        width
    );
    return y;
}
"""

cpp_source = """
torch::Tensor launch_jaro_winkler_softmax(torch::Tensor x, torch::Tensor target);
"""

jaro_winkler_softmax_module = load_inline(
    name='jaro_winkler_softmax_op',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_jaro_winkler_softmax'],
    verbose=False
)


class ModelNew(nn.Module):
    def __init__(self, target):
        super(ModelNew, self).__init__()
        self.target = nn.Parameter(target)
        self.op = jaro_winkler_softmax_module

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.op.launch_jaro_winkler_softmax(x.contiguous(), self.target.contiguous())