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 = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>
#include <cmath>

template <typename scalar_t>
__global__ void sw_forward_kernel(
    const scalar_t* __restrict__ score_matrix,
    scalar_t* __restrict__ H,
    scalar_t* __restrict__ output,
    int B, int M, int N, 
    scalar_t gap_open, 
    scalar_t temperature) {
    
    int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (batch_idx >= B) return;
    
    // Each thread handles one batch element
    // H has shape (B, M+1, N+1), first row and column are 0
    
    // Initialize first row to 0
    for (int j = 0; j <= N; j++) {
        H[batch_idx * (M+1) * (N+1) + j] = scalar_t(0);
    }
    // Initialize first column to 0
    for (int i = 1; i <= M; i++) {
        H[batch_idx * (M+1) * (N+1) + i * (N+1)] = scalar_t(0);
    }
    
    scalar_t max_val = scalar_t(0);
    scalar_t inv_temp = scalar_t(1) / temperature;
    
    // Iterate through the DP matrix
    for (int i = 1; i <= M; i++) {
        for (int j = 1; j <= N; j++) {
            int64_t idx_prev_row = batch_idx * (int64_t)(M+1) * (N+1) + (i-1) * (N+1) + (j-1);
            int64_t idx_curr = batch_idx * (int64_t)(M+1) * (N+1) + i * (N+1) + j;
            
            scalar_t h_diag = H[idx_prev_row];  // H[i-1][j-1]
            scalar_t h_up = H[idx_prev_row + 1];  // H[i-1][j]
            scalar_t h_left = H[idx_curr - 1];  // H[i][j-1]
            
            scalar_t s = score_matrix[batch_idx * (int64_t)M * N + (i-1) * N + (j-1)];
            
            scalar_t match = h_diag + s;
            scalar_t del = h_up + gap_open;
            scalar_t ins = h_left + gap_open;
            
            // logsumexp over [0, match, del, ins]
            scalar_t a0 = scalar_t(0);
            scalar_t a1 = match * inv_temp;
            scalar_t a2 = del * inv_temp;
            scalar_t a3 = ins * inv_temp;
            
            scalar_t max_a = fmax(fmax(a0, a1), fmax(a2, a3));
            
            scalar_t sum_exp = exp(a0 - max_a) + exp(a1 - max_a) + 
                               exp(a2 - max_a) + exp(a3 - max_a);
            scalar_t logsumexp = max_a + log(sum_exp);
            
            scalar_t h_val = temperature * logsumexp;
            H[idx_curr] = h_val;
            
            if (h_val > max_val) max_val = h_val;
        }
    }
    
    output[batch_idx] = max_val;
}

torch::Tensor sw_forward_cuda(torch::Tensor score_matrix, double gap_open, double temperature) {
    TORCH_CHECK(score_matrix.is_cuda(), "score_matrix must be a CUDA tensor");
    TORCH_CHECK(score_matrix.dim() == 3, "score_matrix must be 3D");
    
    int B = score_matrix.size(0);
    int M = score_matrix.size(1);
    int N = score_matrix.size(2);
    
    auto options = score_matrix.options();
    auto H = torch::zeros({B, M + 1, N + 1}, options);
    auto output = torch::empty({B}, options);
    
    int threads = 256;
    int blocks = (B + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(score_matrix.scalar_type(), "sw_forward_kernel", ([&] {
        sw_forward_kernel<scalar_t><<<blocks, threads>>>(
            score_matrix.data_ptr<scalar_t>(),
            H.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            B, M, N,
            static_cast<scalar_t>(gap_open),
            static_cast<scalar_t>(temperature)
        );
    }));
    
    C10_CUDA_CHECK(cudaGetLastError());
    return output;
}
"""

cpp_source = """
torch::Tensor sw_forward_cuda(torch::Tensor score_matrix, double gap_open, double temperature);
"""

sw_module = load_inline(
    name='sw_module',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['sw_forward_cuda'],
    verbose=True,
)

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 sw_module.sw_forward_cuda(score_matrix, self.gap_open, self.temperature)