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_sources = r'''
#include <torch/extension.h>
torch::Tensor sw_forward_cuda(torch::Tensor score, float gap_open, float temperature);
torch::Tensor sw_forward(torch::Tensor score, double gap_open, double temperature) {
    return sw_forward_cuda(score.contiguous(), (float)gap_open, (float)temperature);
}
'''

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

__device__ __forceinline__ float lse4_one(float b, float c, float d) {
    float m = fmaxf(0.0f, fmaxf(b, fmaxf(c, d)));
    float s = expf(-m) + expf(b - m) + expf(c - m) + expf(d - m);
    return m + logf(s);
}

__device__ __forceinline__ float lse4_temp(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);
}

__device__ __forceinline__ float warp_max(float v) {
    unsigned mask = 0xffffffffu;
    v = fmaxf(v, __shfl_down_sync(mask, v, 16));
    v = fmaxf(v, __shfl_down_sync(mask, v, 8));
    v = fmaxf(v, __shfl_down_sync(mask, v, 4));
    v = fmaxf(v, __shfl_down_sync(mask, v, 2));
    v = fmaxf(v, __shfl_down_sync(mask, v, 1));
    return v;
}

__global__ void sw_kernel_32_t1(const float* __restrict__ score, float* __restrict__ out, float gap_open) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float H[];
    if (tid == 0) H[32] = 0.0f;
    H[tid] = 0.0f;
    H[(tid + 1) * 33] = 0.0f;
    __syncthreads();
    float local_max = 0.0f;
    const float* S = score + ((long long)b) * 1024;
    #pragma unroll 1
    for (int diag = 2; diag <= 64; ++diag) {
        int i_min = max(1, diag - 32);
        int i_max = min(32, diag - 1);
        int count = i_max - i_min + 1;
        if (tid < count) {
            int i = i_min + tid;
            int j = diag - i;
            float match = H[(i - 1) * 33 + (j - 1)] + S[(i - 1) * 32 + (j - 1)];
            float del = H[(i - 1) * 33 + j] + gap_open;
            float ins = H[i * 33 + (j - 1)] + gap_open;
            float v = lse4_one(match, del, ins);
            H[i * 33 + j] = v;
            local_max = fmaxf(local_max, v);
        }
        __syncthreads();
    }
    local_max = warp_max(local_max);
    if (tid == 0) out[b] = local_max;
}

__global__ void sw_kernel_generic(const float* __restrict__ score, float* __restrict__ out,
                              int M, int N, float gap_open, float temperature) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int W = N + 1;
    extern __shared__ float H[];
    for (int j = tid; j <= N; j += blockDim.x) H[j] = 0.0f;
    for (int i = tid + 1; i <= M; i += blockDim.x) H[i * W] = 0.0f;
    __syncthreads();
    float local_max = 0.0f;
    const float* S = score + ((long long)b) * M * N;
    for (int diag = 2; diag <= M + N; ++diag) {
        int i_min = max(1, diag - N);
        int i_max = min(M, diag - 1);
        int count = i_max - i_min + 1;
        for (int t = tid; t < count; t += blockDim.x) {
            int i = i_min + t;
            int j = diag - i;
            float match = H[(i - 1) * W + (j - 1)] + S[(i - 1) * N + (j - 1)];
            float del = H[(i - 1) * W + j] + gap_open;
            float ins = H[i * W + (j - 1)] + gap_open;
            float v = lse4_temp(0.0f, match, del, ins, temperature);
            H[i * W + j] = v;
            local_max = fmaxf(local_max, v);
        }
        __syncthreads();
    }
    local_max = warp_max(local_max);
    if (tid == 0) out[b] = local_max;
}

torch::Tensor sw_forward_cuda(torch::Tensor score, float gap_open, float temperature) {
    TORCH_CHECK(score.is_cuda(), "score must be CUDA");
    TORCH_CHECK(score.scalar_type() == torch::kFloat32, "score must be float32");
    int B = (int)score.size(0), M = (int)score.size(1), N = (int)score.size(2);
    auto out = torch::empty({B}, score.options());
    if (M == 32 && N == 32 && temperature == 1.0f) {
        sw_kernel_32_t1<<<B, 32, 1089 * sizeof(float)>>>(score.data_ptr<float>(), out.data_ptr<float>(), gap_open);
    } else {
        size_t shmem = (size_t)(M + 1) * (N + 1) * sizeof(float);
        sw_kernel_generic<<<B, 32, shmem>>>(score.data_ptr<float>(), out.data_ptr<float>(), M, N, gap_open, temperature);
    }
    return out;
}
'''

_ext = load_inline(name='sw_softalign_ext_final', cpp_sources=cpp_sources, cuda_sources=cuda_sources,
                   functions=['sw_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], 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))