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_out_cuda(torch::Tensor score, torch::Tensor out, double gap, double temp);
torch::Tensor sw_forward_out(torch::Tensor score, torch::Tensor out, double gap, double temp) {
    return sw_forward_out_cuda(score, out, gap, temp);
}
'''

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

__global__ void sw_kernel(const float* __restrict__ S, float* __restrict__ out,
                          int M, int N, float gap, float temp) {
    extern __shared__ float H[];
    __shared__ float red[32];
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int W = N + 1;
    int total = (M + 1) * (N + 1);

    for (int idx = tid; idx < total; idx += 32) H[idx] = 0.0f;
    __syncthreads();

    float local_max = 0.0f;
    int base = b * M * N;
    float invt = 1.0f / temp;

    for (int d = 2; d <= M + N; ++d) {
        int ilo = d - N;
        if (ilo < 1) ilo = 1;
        int ihi = d - 1;
        if (ihi > M) ihi = M;
        int len = ihi - ilo + 1;
        if (tid < len) {
            int i = ilo + tid;
            int j = d - i;
            int hidx = i * W + j;
            float match = H[hidx - W - 1] + S[base + (i - 1) * N + (j - 1)];
            float delv  = H[hidx - W] + gap;
            float insv  = H[hidx - 1] + gap;
            float m = fmaxf(0.0f, fmaxf(match, fmaxf(delv, insv)));
            float sum = expf(-m * invt) + expf((match - m) * invt) +
                        expf((delv - m) * invt) + expf((insv - m) * invt);
            float val = m + temp * logf(sum);
            H[hidx] = val;
            local_max = fmaxf(local_max, val);
        }
        __syncthreads();
    }

    red[tid] = local_max;
    __syncthreads();
    for (int s = 16; s > 0; s >>= 1) {
        if (tid < s) red[tid] = fmaxf(red[tid], red[tid + s]);
        __syncthreads();
    }
    if (tid == 0) out[b] = red[0];
}

torch::Tensor sw_forward_out_cuda(torch::Tensor score, torch::Tensor out, double gap, double temp) {
    int B = score.size(0), M = score.size(1), N = score.size(2);
    size_t shmem = (size_t)(M + 1) * (N + 1) * sizeof(float);
    sw_kernel<<<B, 32, shmem>>>(score.data_ptr<float>(), out.data_ptr<float>(), M, N, (float)gap, (float)temp);
    return out;
}
'''

_ext = load_inline(
    name='sw_soft_align_ext_final',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['sw_forward_out'],
    extra_cflags=['-O3'],
    extra_cuda_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
        self._out = None

    def forward(self, score_matrix: torch.Tensor) -> torch.Tensor:
        if not score_matrix.is_contiguous():
            score_matrix = score_matrix.contiguous()
        B = score_matrix.shape[0]
        if (self._out is None or self._out.shape[0] != B or
            self._out.device != score_matrix.device or self._out.dtype != score_matrix.dtype):
            self._out = torch.empty((B,), device=score_matrix.device, dtype=score_matrix.dtype)
        return _ext.sw_forward_out(score_matrix, self._out, float(self.gap_open), float(self.temperature))