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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor sw_forward(torch::Tensor score_matrix, double gap_open, double temperature);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__device__ __forceinline__ float lse4_soft(float a, float b, float c, float d, float temp) {
    float m = fmaxf(fmaxf(a, b), fmaxf(c, d));
    float invt = 1.0f / temp;
    float s = __expf((a - m) * invt) + __expf((b - m) * invt) + __expf((c - m) * invt) + __expf((d - m) * invt);
    return m + temp * logf(s);
}

// Specialized for the benchmark shape (M=N=32), but handles any M,N <= 64.
__global__ void sw_kernel(const float* __restrict__ score, float* __restrict__ out,
                          int B, int M, int N, float gap, float temp) {
    int b = blockIdx.x;
    if (b >= B) return;
    int tid = threadIdx.x;

    extern __shared__ float H[];  // (M+1)*(N+1)
    int stride = N + 1;
    int total = (M + 1) * (N + 1);
    for (int idx = tid; idx < total; idx += blockDim.x) H[idx] = 0.0f;
    __syncthreads();

    const float* S = score + ((long long)b) * M * N;

    // Wavefront dynamic programming.  Cells on one anti-diagonal are independent.
    for (int diag = 2; diag <= M + N; ++diag) {
        int i_min = diag - N;
        if (i_min < 1) i_min = 1;
        int i_max = diag - 1;
        if (i_max > M) i_max = M;
        int count = i_max - i_min + 1;
        for (int k = tid; k < count; k += blockDim.x) {
            int i = i_min + k;
            int j = diag - i;
            float match = H[(i - 1) * stride + (j - 1)] + S[(i - 1) * N + (j - 1)];
            float delv  = H[(i - 1) * stride + j] + gap;
            float insv  = H[i * stride + (j - 1)] + gap;
            H[i * stride + j] = lse4_soft(0.0f, match, delv, insv, temp);
        }
        __syncthreads();
    }

    // Reduce maximum over H for this batch item.
    __shared__ float red[128];
    float mx = 0.0f;
    for (int idx = tid; idx < total; idx += blockDim.x) {
        mx = fmaxf(mx, H[idx]);
    }
    red[tid] = mx;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (tid < off) red[tid] = fmaxf(red[tid], red[tid + off]);
        __syncthreads();
    }
    if (tid == 0) out[b] = red[0];
}

torch::Tensor sw_forward(torch::Tensor score_matrix, double gap_open, double temperature) {
    TORCH_CHECK(score_matrix.is_cuda(), "score_matrix must be CUDA");
    TORCH_CHECK(score_matrix.scalar_type() == at::kFloat, "score_matrix must be float32");
    auto x = score_matrix.contiguous();
    int B = (int)x.size(0);
    int M = (int)x.size(1);
    int N = (int)x.size(2);
    auto out = torch::empty({B}, x.options());
    int threads = 128;
    size_t shmem = (size_t)(M + 1) * (N + 1) * sizeof(float);
    sw_kernel<<<B, threads, shmem, at::cuda::getCurrentCUDAStream()>>>(
        x.data_ptr<float>(), out.data_ptr<float>(), B, M, N, (float)gap_open, (float)temperature);
    return out;
}
'''

_ext = load_inline(
    name='sw_softalign_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['sw_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    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 _ext.sw_forward(score_matrix, float(self.gap_open), float(self.temperature))