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

cuda_source = '''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void causal_conv1d_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch_size,
    int in_channels,
    int out_channels,
    int seq_len,
    int kernel_size,
    int dilation,
    bool has_bias
) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long total_outputs = (long long)batch_size * out_channels * seq_len;
    if (idx >= total_outputs) return;

    int b = idx / ((long long)out_channels * seq_len);
    long long rem = idx % ((long long)out_channels * seq_len);
    int oc = rem / seq_len;
    int t = rem % seq_len;

    float sum = 0.0f;
    long long in_stride = (long long)in_channels * seq_len;
    for (int ic = 0; ic < in_channels; ++ic) {
        for (int k = 0; k < kernel_size; ++k) {
            int input_t = t - (kernel_size - 1 - k) * dilation;
            if (input_t >= 0) {
                float val = input[b * in_stride + (long long)ic * seq_len + input_t];
                float w = weight[oc * (long long)in_channels * kernel_size + (long long)ic * kernel_size + k];
                sum += val * w;
            }
        }
    }
    if (has_bias) {
        sum += bias[oc];
    }
    output[b * (long long)out_channels * seq_len + (long long)oc * seq_len + t] = sum;
}

torch::Tensor causal_conv1d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int kernel_size,
    int dilation
) {
    int batch_size = input.size(0);
    int in_channels = input.size(1);
    int seq_len = input.size(2);
    int out_channels = weight.size(0);
    bool has_bias = bias.defined() && bias.numel() > 0;

    auto output = torch::empty({batch_size, out_channels, seq_len}, input.options());

    long long total = (long long)batch_size * out_channels * seq_len;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    causal_conv1d_kernel<<<blocks, threads>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        has_bias ? bias.data_ptr<float>() : nullptr,
        output.data_ptr<float>(),
        batch_size, in_channels, out_channels, seq_len, kernel_size, dilation, has_bias
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor causal_conv1d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias,
    int kernel_size,
    int dilation
);
'''

module = load_inline(
    name='causal_conv1d',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['causal_conv1d_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, dilation: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias)
        self.kernel_size = kernel_size
        self.dilation = dilation

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bias = self.conv1d.bias if self.conv1d.bias is not None else torch.empty(0, device=x.device)
        return module.causal_conv1d_cuda(
            x, self.conv1d.weight, bias, self.kernel_size, self.dilation
        )