import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a causal 1D convolution: left-padded so each output only depends on current and past inputs.
    Used in autoregressive sequence models (Mamba, Hyena, WaveNet, TCN) for causal sequence processing.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        kernel_size (int): Size of the convolution kernel.
        dilation (int, optional): Dilation rate. Defaults to 1.
        bias (bool, optional): If True, adds a learnable bias. Defaults to False.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.causal_padding = (kernel_size - 1) * dilation
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.nn.functional.pad(x, (self.causal_padding, 0))
        return self.conv1d(x)

batch_size = 32
in_channels = 256
out_channels = 256
kernel_size = 4
dilation = 1
seq_len = 8192

def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), in_channels, popcorn_pri.jitter_int(seq_len))]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, dilation]

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

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

torch::Tensor causal_conv1d_forward(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> bias_opt, int64_t dilation);
'''

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

#define TILE_M 16
#define TILE_N 16
#define TILE_K 32

// Custom CUDA implementation of causal Conv1d for float32 NCL tensors.
__global__ void causal_conv1d_kernel(const float* __restrict__ x,
                                     const float* __restrict__ w,
                                     const float* __restrict__ bias,
                                     float* __restrict__ y,
                                     int B, int IC, int L, int OC, int K,
                                     int dilation, int pad) {
    __shared__ float As[TILE_M][TILE_K + 1];
    __shared__ float Bs[TILE_K][TILE_N + 1];

    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int tid = ty * TILE_N + tx;
    int row = blockIdx.x * TILE_M + ty;      // flattened (batch,time)
    int oc = blockIdx.y * TILE_N + tx;
    int total_rows = B * L;
    float acc = 0.0f;

    for (int p0 = 0; p0 < IC * K; p0 += TILE_K) {
        for (int idx = tid; idx < TILE_M * TILE_K; idx += TILE_M * TILE_N) {
            int r = idx / TILE_K;
            int pk = idx - r * TILE_K;
            int p = p0 + pk;
            float val = 0.0f;
            int grow = blockIdx.x * TILE_M + r;
            if (grow < total_rows && p < IC * K) {
                int b = grow / L;
                int t = grow - b * L;
                int ic = p / K;
                int kk = p - ic * K;
                int xt = t + kk * dilation - pad;
                if ((unsigned)xt < (unsigned)L) {
                    val = x[(b * IC + ic) * L + xt];
                }
            }
            As[r][pk] = val;
        }
        for (int idx = tid; idx < TILE_K * TILE_N; idx += TILE_M * TILE_N) {
            int pk = idx / TILE_N;
            int cn = idx - pk * TILE_N;
            int p = p0 + pk;
            int goc = blockIdx.y * TILE_N + cn;
            float val = 0.0f;
            if (p < IC * K && goc < OC) {
                int ic = p / K;
                int kk = p - ic * K;
                val = w[(goc * IC + ic) * K + kk];
            }
            Bs[pk][cn] = val;
        }
        __syncthreads();
        #pragma unroll
        for (int pk = 0; pk < TILE_K; ++pk) {
            acc = fmaf(As[ty][pk], Bs[pk][tx], acc);
        }
        __syncthreads();
    }
    if (row < total_rows && oc < OC) {
        if (bias != nullptr) acc += bias[oc];
        int b = row / L;
        int t = row - b * L;
        y[(b * OC + oc) * L + t] = acc;
    }
}

torch::Tensor causal_conv1d_forward(torch::Tensor x, torch::Tensor w, c10::optional<torch::Tensor> bias_opt, int64_t dilation64) {
    x = x.contiguous();
    w = w.contiguous();
    const int B = x.size(0);
    const int IC = x.size(1);
    const int L = x.size(2);
    const int OC = w.size(0);
    const int K = w.size(2);
    const int dilation = (int)dilation64;
    const int pad = (K - 1) * dilation;
    auto y = torch::empty({B, OC, L}, x.options());
    const float* bias_ptr = nullptr;
    torch::Tensor bias_tensor;
    if (bias_opt.has_value() && bias_opt.value().defined()) {
        bias_tensor = bias_opt.value().contiguous();
        bias_ptr = bias_tensor.data_ptr<float>();
    }
    dim3 block(TILE_N, TILE_M);
    dim3 grid((B * L + TILE_M - 1) / TILE_M, (OC + TILE_N - 1) / TILE_N);
    causal_conv1d_kernel<<<grid, block>>>(x.data_ptr<float>(), w.data_ptr<float>(), bias_ptr, y.data_ptr<float>(), B, IC, L, OC, K, dilation, pad);
    return y;
}
'''

_ext = load_inline(name='causal_conv1d_custom_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['causal_conv1d_forward'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int = 1, bias: bool = False):
        super().__init__()
        self.causal_padding = (kernel_size - 1) * dilation
        self.dilation = dilation
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _ext.causal_conv1d_forward(x, self.conv1d.weight, self.conv1d.bias, self.dilation)