import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, output_size: int):
        """
        Initialize the Vanilla RNN model.

        :param input_size: The number of input features (int).
        :param hidden_size: The size of the hidden state (int).
        :param output_size: The number of output features (int).
        """
        super(Model, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        # Define the RNN cell components (input to hidden, hidden to hidden, and hidden to output)
        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)  # Input to hidden
        self.h2o = nn.Linear(hidden_size, output_size)  # Hidden to output
        self.tanh = nn.Tanh()  # Activation function for hidden state

    def forward(self, x: torch.Tensor, h0: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the Vanilla RNN.

        :param x: Input tensor of shape (seq_len, batch_size, input_size)
        :param h0: Initial hidden state tensor of shape (batch_size, hidden_size)
        :return: Output tensor of shape (seq_len, batch_size, output_size)
        """
        seq_len, batch_size, _ = x.size()
        hidden = h0.to(x.device)
        outputs = []

        for t in range(seq_len):
            combined = torch.cat((x[t], hidden), dim=1)  # Concatenate input and hidden state
            hidden = self.tanh(self.i2h(combined))  # Update hidden state
            output = self.h2o(hidden)  # Compute output
            outputs.append(output)

        return torch.stack(outputs, dim=0)  # (seq_len, batch_size, output_size)

# === Test configuration ===
batch_size = 8
input_size = 1024
hidden_size = 256
output_size = 128
sequence_length = 256

def get_inputs():
    return [
        torch.rand(sequence_length, batch_size, input_size),
        torch.rand(batch_size, hidden_size)
    ]

def get_init_inputs():
    return [input_size, hidden_size, output_size]

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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor rnn_forward_cuda(torch::Tensor x, torch::Tensor h0, torch::Tensor wi, torch::Tensor bi, torch::Tensor wo, torch::Tensor bo);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;

__global__ void rnn_persist_kernel(const float* __restrict__ x,
                                   const float* __restrict__ h0,
                                   const float* __restrict__ wi,
                                   const float* __restrict__ bi,
                                   const float* __restrict__ wo,
                                   const float* __restrict__ bo,
                                   float* __restrict__ out,
                                   float* __restrict__ buf) {
    extern __shared__ float sw[];
    cg::grid_group grid = cg::this_grid();
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int b = tid >> 5;
    const int hh = blockIdx.x;
    const int gtid = blockIdx.x * blockDim.x + tid;

    float* buf0 = buf;
    float* buf1 = buf + 2048;

    const float* wbase = wi + hh * 1280;
    #pragma unroll
    for (int i = tid; i < 1280; i += 256) sw[i] = wbase[i];
    if (blockIdx.x < 128) {
        const float* wobase = wo + (blockIdx.x << 8);
        sw[1280 + tid] = wobase[tid];
    }
    if (gtid < 2048) buf0[gtid] = h0[gtid];
    grid.sync();

    float* oldh = buf0;
    float* newh = buf1;
    for (int t = 0; t < 256; ++t) {
        const float* xrow = x + ((t * 8 + b) << 10);
        const float* hrow = oldh + (b << 8);
        float sum = (lane == 0) ? bi[hh] : 0.0f;
        #pragma unroll 32
        for (int k = lane; k < 1024; k += 32) sum = fmaf(xrow[k], sw[k], sum);
        #pragma unroll 8
        for (int k = lane; k < 256; k += 32) sum = fmaf(hrow[k], sw[1024 + k], sum);
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) sum += __shfl_down_sync(0xffffffff, sum, off);
        if (lane == 0) newh[(b << 8) + hh] = __tanhf(sum);

        grid.sync();

        if (blockIdx.x < 128) {
            const float* hrow2 = newh + (b << 8);
            float osum = (lane == 0) ? bo[blockIdx.x] : 0.0f;
            #pragma unroll 8
            for (int k = lane; k < 256; k += 32) osum = fmaf(hrow2[k], sw[1280 + k], osum);
            #pragma unroll
            for (int off = 16; off > 0; off >>= 1) osum += __shfl_down_sync(0xffffffff, osum, off);
            if (lane == 0) out[t * 1024 + (b << 7) + blockIdx.x] = osum;
        }

        grid.sync();
        float* tmp = oldh; oldh = newh; newh = tmp;
    }
}

torch::Tensor rnn_forward_cuda(torch::Tensor x, torch::Tensor h0, torch::Tensor wi, torch::Tensor bi, torch::Tensor wo, torch::Tensor bo) {
    auto out = torch::empty({256, 8, 128}, x.options());
    auto buf = torch::empty({4096}, x.options());
    const float *xp = x.data_ptr<float>(), *h0p = h0.data_ptr<float>(), *wip = wi.data_ptr<float>(), *bip = bi.data_ptr<float>(), *wop = wo.data_ptr<float>(), *bop = bo.data_ptr<float>();
    float *outp = out.data_ptr<float>(), *bufp = buf.data_ptr<float>();
    void* args[] = {&xp, &h0p, &wip, &bip, &wop, &bop, &outp, &bufp};
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    C10_CUDA_CHECK(cudaLaunchCooperativeKernel((void*)rnn_persist_kernel, 256, 256, args, 1536 * sizeof(float), stream));
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
'''

_ext = load_inline(name='rnn_custom_persist_final_ext',
                   cpp_sources=[_cpp_src],
                   cuda_sources=[_cuda_src],
                   functions=['rnn_forward_cuda'],
                   extra_cflags=['-O3'],
                   extra_cuda_cflags=['-O3'],
                   with_cuda=True,
                   verbose=False)

class ModelNew(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, output_size: int):
        super(ModelNew, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
        self.h2o = nn.Linear(hidden_size, output_size)
        self.tanh = nn.Tanh()

    def forward(self, x: torch.Tensor, h0: torch.Tensor) -> torch.Tensor:
        return _ext.rnn_forward_cuda(x, h0, self.i2h.weight, self.i2h.bias, self.h2o.weight, self.h2o.bias)