import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    """
    Differentiable Smith-Waterman-style local sequence alignment scoring.
    Uses a soft-max relaxation of the DP recurrence so gradients can flow
    through the alignment matrix.  
    """
    def __init__(self, gap_open=-1.0, temperature=1.0):
        super().__init__()
        self.gap_open = gap_open
        self.temperature = temperature

    def forward(self, score_matrix: torch.Tensor) -> torch.Tensor:
        """
        Args:
            score_matrix: (B, M, N) – per-position similarity scores
        Returns:
            alignment_score: (B,) – soft local-alignment score (log-sum-exp over DP cells)
        """
        B, M, N = score_matrix.shape
        neg_inf = torch.tensor(float("-inf"), device=score_matrix.device, dtype=score_matrix.dtype)

        # H[i][j] = max(0, H[i-1][j-1]+S[i][j], H[i-1][j]+gap, H[i][j-1]+gap)
        H = score_matrix.new_full((B, M + 1, N + 1), 0.0)

        for i in range(1, M + 1):
            for j in range(1, N + 1):
                match = H[:, i - 1, j - 1] + score_matrix[:, i - 1, j - 1]
                delete = H[:, i - 1, j] + self.gap_open
                insert = H[:, i, j - 1] + self.gap_open
                zero = torch.zeros(B, device=score_matrix.device)
                stack = torch.stack([zero, match, delete, insert], dim=-1)
                H[:, i, j] = self.temperature * torch.logsumexp(stack / self.temperature, dim=-1)

        return H.view(B, -1).max(dim=-1).values


seq_m = 32
seq_n = 32
batch_size = 16


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), seq_m, seq_n)]


def get_init_inputs():
    return []

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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void soft_smith_waterman_kernel(
    const float* __restrict__ score_matrix,
    float* __restrict__ output,
    int B, int M, int N,
    float gap_open, float temperature
) {
    int b = blockIdx.x;
    if (b >= B) return;

    extern __shared__ float H[];
    const int rows = M + 1;
    const int cols = N + 1;

    // Initialize H to 0
    for (int idx = threadIdx.x; idx < rows * cols; idx += blockDim.x) {
        H[idx] = 0.0f;
    }
    __syncthreads();

    const float* scores = score_matrix + b * M * N;
    float inv_temp = 1.0f / temperature;

    for (int i = 1; i <= M; ++i) {
        for (int j = 1; j <= N; ++j) {
            float match = H[(i-1)*cols + (j-1)] + scores[(i-1)*N + (j-1)];
            float del   = H[(i-1)*cols + j] + gap_open;
            float ins   = H[i*cols + (j-1)] + gap_open;
            float zero  = 0.0f;

            // Compute logsumexp
            float m = fmaxf(fmaxf(zero, match), fmaxf(del, ins));
            float sum = expf((zero - m) * inv_temp) +
                        expf((match - m) * inv_temp) +
                        expf((del - m) * inv_temp) +
                        expf((ins - m) * inv_temp);
            float val = m + temperature * logf(sum);

            if (threadIdx.x == 0) {
                H[i*cols + j] = val;
            }
            __syncthreads();
        }
    }

    // Find max over all H
    float local_max = 0.0f;
    for (int idx = threadIdx.x; idx < rows * cols; idx += blockDim.x) {
        local_max = fmaxf(local_max, H[idx]);
    }

    // Reduce within block
    __shared__ float smax[32];
    int lane = threadIdx.x % 32;
    int wid = threadIdx.x / 32;
    local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, 16));
    local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, 8));
    local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, 4));
    local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, 2));
    local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, 1));
    if (lane == 0) smax[wid] = local_max;
    __syncthreads();

    if (threadIdx.x == 0) {
        float gmax = 0.0f;
        for (int i = 0; i < (blockDim.x + 31) / 32; ++i) {
            gmax = fmaxf(gmax, smax[i]);
        }
        output[b] = gmax;
    }
}

torch::Tensor soft_smith_waterman_cuda(
    torch::Tensor score_matrix,
    float gap_open,
    float temperature
) {
    int B = score_matrix.size(0);
    int M = score_matrix.size(1);
    int N = score_matrix.size(2);

    auto output = torch::empty({B}, score_matrix.options());

    int threads = 128;
    int shmem = (M + 1) * (N + 1) * sizeof(float);

    soft_smith_waterman_kernel<<<B, threads, shmem>>>(
        score_matrix.data_ptr<float>(),
        output.data_ptr<float>(),
        B, M, N,
        gap_open, temperature
    );

    return output;
}
'''

cpp_source = r'''
torch::Tensor soft_smith_waterman_cuda(torch::Tensor score_matrix, float gap_open, float temperature);
'''

module = load_inline(
    name='soft_sw',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['soft_smith_waterman_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, gap_open=-1.0, temperature=1.0):
        super().__init__()
        self.gap_open = gap_open
        self.temperature = temperature

    def forward(self, score_matrix: torch.Tensor) -> torch.Tensor:
        return module.soft_smith_waterman_cuda(
            score_matrix, self.gap_open, self.temperature
        )
```