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>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

void causal_conv1d_fwd_launcher(const float* x, const float* w, const float* bias, float* y,
                                int B, int Cin, int L, int Cout, int K, int dilation, bool has_bias,
                                cudaStream_t stream);

torch::Tensor causal_conv_forward(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                                  bool has_bias, int64_t dilation) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(w.is_cuda(), "weight must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 is supported");
    TORCH_CHECK(w.scalar_type() == torch::kFloat32, "only float32 weight is supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(w.is_contiguous(), "weight must be contiguous");
    if (has_bias) {
        TORCH_CHECK(bias.is_cuda(), "bias must be CUDA");
        TORCH_CHECK(bias.scalar_type() == torch::kFloat32, "only float32 bias is supported");
        TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
    }
    const c10::cuda::CUDAGuard device_guard(x.device());
    int B = (int)x.size(0);
    int Cin = (int)x.size(1);
    int L = (int)x.size(2);
    int Cout = (int)w.size(0);
    int K = (int)w.size(2);
    auto y = torch::empty({B, Cout, L}, x.options());
    const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
    causal_conv1d_fwd_launcher(x.data_ptr<float>(), w.data_ptr<float>(), bptr, y.data_ptr<float>(),
                               B, Cin, L, Cout, K, (int)dilation, has_bias,
                               at::cuda::getCurrentCUDAStream());
    return y;
}
'''

_cuda_src = r'''
#include <cuda.h>
#include <cuda_runtime.h>
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#define BM 32
#define BN 128
#define BK 32

__global__ __launch_bounds__(256, 2)
void causal_conv1d_kernel(const float* __restrict__ x,
                          const float* __restrict__ w,
                          const float* __restrict__ bias,
                          float* __restrict__ y,
                          int B, int Cin, int L, int Cout, int K, int dilation, int has_bias) {
    __shared__ float As[BM][BK];
    __shared__ float Bs[BK][BN];

    const int tid = threadIdx.x;
    const int tx = tid & 15;       // 0..15
    const int ty = tid >> 4;       // 0..15

    const int n_block = blockIdx.x * BN;
    const int m_block = blockIdx.y * BM;
    const int row0 = m_block + ty;
    const int row1 = row0 + 16;
    const int totalK = Cin * K;
    const int pad = (K - 1) * dilation;

    float acc00=0.f, acc01=0.f, acc02=0.f, acc03=0.f, acc04=0.f, acc05=0.f, acc06=0.f, acc07=0.f;
    float acc10=0.f, acc11=0.f, acc12=0.f, acc13=0.f, acc14=0.f, acc15=0.f, acc16=0.f, acc17=0.f;

    for (int kb = 0; kb < totalK; kb += BK) {
        // Load A tile: weights [Cout, Cin, K] viewed as [Cout, Cin*K]
        #pragma unroll
        for (int idx = tid; idx < BM * BK; idx += 256) {
            int r = idx / BK;
            int kk = idx - r * BK;
            int p = kb + kk;
            int oc = m_block + r;
            float v = 0.f;
            if (oc < Cout && p < totalK) {
                int ic = p / K;
                int k = p - ic * K;
                v = w[(oc * Cin + ic) * K + k];
            }
            As[r][kk] = v;
        }

        // Load B tile: on-the-fly causal im2col, columns are flattened (batch,time)
        for (int idx = tid; idx < BK * BN; idx += 256) {
            int kk = idx / BN;
            int c = idx - kk * BN;
            int p = kb + kk;
            int n = n_block + c;
            float v = 0.f;
            if (p < totalK && n < B * L) {
                int ic = p / K;
                int k = p - ic * K;
                int b = n / L;
                int t = n - b * L;
                int xt = t + k * dilation - pad;
                if ((unsigned)xt < (unsigned)L) {
                    v = x[(b * Cin + ic) * L + xt];
                }
            }
            Bs[kk][c] = v;
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; ++kk) {
            float a0 = As[ty][kk];
            float a1 = As[ty + 16][kk];
            float b0 = Bs[kk][tx + 0];
            float b1 = Bs[kk][tx + 16];
            float b2 = Bs[kk][tx + 32];
            float b3 = Bs[kk][tx + 48];
            float b4 = Bs[kk][tx + 64];
            float b5 = Bs[kk][tx + 80];
            float b6 = Bs[kk][tx + 96];
            float b7 = Bs[kk][tx + 112];
            acc00 = fmaf(a0, b0, acc00); acc01 = fmaf(a0, b1, acc01);
            acc02 = fmaf(a0, b2, acc02); acc03 = fmaf(a0, b3, acc03);
            acc04 = fmaf(a0, b4, acc04); acc05 = fmaf(a0, b5, acc05);
            acc06 = fmaf(a0, b6, acc06); acc07 = fmaf(a0, b7, acc07);
            acc10 = fmaf(a1, b0, acc10); acc11 = fmaf(a1, b1, acc11);
            acc12 = fmaf(a1, b2, acc12); acc13 = fmaf(a1, b3, acc13);
            acc14 = fmaf(a1, b4, acc14); acc15 = fmaf(a1, b5, acc15);
            acc16 = fmaf(a1, b6, acc16); acc17 = fmaf(a1, b7, acc17);
        }
        __syncthreads();
    }

    if (has_bias) {
        if (row0 < Cout) {
            float bv = bias[row0];
            acc00 += bv; acc01 += bv; acc02 += bv; acc03 += bv; acc04 += bv; acc05 += bv; acc06 += bv; acc07 += bv;
        }
        if (row1 < Cout) {
            float bv = bias[row1];
            acc10 += bv; acc11 += bv; acc12 += bv; acc13 += bv; acc14 += bv; acc15 += bv; acc16 += bv; acc17 += bv;
        }
    }

    const int offs[8] = {0,16,32,48,64,80,96,112};
    float vals0[8] = {acc00,acc01,acc02,acc03,acc04,acc05,acc06,acc07};
    float vals1[8] = {acc10,acc11,acc12,acc13,acc14,acc15,acc16,acc17};

    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        int n = n_block + tx + offs[j];
        if (n < B * L) {
            int b = n / L;
            int t = n - b * L;
            if (row0 < Cout) y[(b * Cout + row0) * L + t] = vals0[j];
            if (row1 < Cout) y[(b * Cout + row1) * L + t] = vals1[j];
        }
    }
}

void causal_conv1d_fwd_launcher(const float* x, const float* w, const float* bias, float* y,
                                int B, int Cin, int L, int Cout, int K, int dilation, bool has_bias,
                                cudaStream_t stream) {
    dim3 block(256);
    dim3 grid((B * L + BN - 1) / BN, (Cout + BM - 1) / BM);
    causal_conv1d_kernel<<<grid, block, 0, stream>>>(x, w, bias, y, B, Cin, L, Cout, K, dilation, has_bias ? 1 : 0);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name="causal_conv1d_k4_fp32_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["causal_conv_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    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(ModelNew, self).__init__()
        self.causal_padding = (kernel_size - 1) * dilation
        self.conv1d = nn.Conv1d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias)
        self.dilation = dilation
        self.register_buffer("_dummy_bias", torch.empty(0), persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv1d.bias
        if b is None:
            return _ext.causal_conv_forward(x, self.conv1d.weight, self._dummy_bias, False, self.dilation)
        else:
            return _ext.causal_conv_forward(x, self.conv1d.weight, b, True, self.dilation)