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(torch::Tensor x,
                          torch::Tensor h0,
                          torch::Tensor i2h_w,
                          torch::Tensor i2h_b,
                          torch::Tensor h2o_w,
                          torch::Tensor h2o_b);
"""

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

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#endif

__global__ void hidden_step_kernel(const float* __restrict__ inp,
                                   const float* __restrict__ prev,
                                   const float* __restrict__ w,
                                   float* __restrict__ next,
                                   int t, int B, int I, int H) {
    const int h = blockIdx.x;
    const int b = blockIdx.y;
    const int tid = threadIdx.x;
    extern __shared__ float sdata[];

    float sum = 0.0f;
    const float* __restrict__ prow = prev + b * H;
    const float* __restrict__ wrow = w + h * (I + H) + I;
    for (int k = tid; k < H; k += blockDim.x) {
        sum = fmaf(prow[k], wrow[k], sum);
    }
    sdata[tid] = sum;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] += sdata[tid + stride];
        __syncthreads();
    }

    if (tid == 0) {
        float v = inp[(t * B + b) * H + h] + sdata[0];
        next[b * H + h] = tanhf(v);
    }
}

__global__ void output_step_kernel(const float* __restrict__ hidden,
                                   const float* __restrict__ w,
                                   const float* __restrict__ bias,
                                   float* __restrict__ out,
                                   int t, int B, int H, int O) {
    const int o = blockIdx.x;
    const int b = blockIdx.y;
    const int tid = threadIdx.x;
    extern __shared__ float sdata[];

    float sum = 0.0f;
    const float* __restrict__ hrow = hidden + b * H;
    const float* __restrict__ wrow = w + o * H;
    for (int k = tid; k < H; k += blockDim.x) {
        sum = fmaf(hrow[k], wrow[k], sum);
    }
    sdata[tid] = sum;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] += sdata[tid + stride];
        __syncthreads();
    }

    if (tid == 0) {
        out[(t * B + b) * O + o] = sdata[0] + bias[o];
    }
}

torch::Tensor rnn_forward(torch::Tensor x,
                          torch::Tensor h0,
                          torch::Tensor i2h_w,
                          torch::Tensor i2h_b,
                          torch::Tensor h2o_w,
                          torch::Tensor h2o_b) {
    CHECK_CUDA(x); CHECK_CUDA(h0); CHECK_CUDA(i2h_w); CHECK_CUDA(i2h_b); CHECK_CUDA(h2o_w); CHECK_CUDA(h2o_b);
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    c10::cuda::CUDAGuard device_guard(x.device());

    const int T = (int)x.size(0);
    const int B = (int)x.size(1);
    const int I = (int)x.size(2);
    const int H = (int)h0.size(1);
    const int O = (int)h2o_w.size(0);

    auto x2 = x.contiguous().view({T * B, I});
    // Large, time-independent input contribution: x_t @ W_x^T + b_i2h
    auto wx_t = i2h_w.narrow(1, 0, I).transpose(0, 1);
    auto inp = at::matmul(x2, wx_t);
    inp.add_(i2h_b);
    inp = inp.contiguous();

    auto opts = x.options();
    auto out = torch::empty({T, B, O}, opts);
    auto hidden_prev = h0.contiguous().clone();
    auto hidden_next = torch::empty({B, H}, opts);

    const int threads = 256;
    const size_t smem = threads * sizeof(float);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    dim3 grid_h(H, B);
    dim3 grid_o(O, B);
    const float* inp_ptr = inp.data_ptr<float>();
    const float* iw_ptr = i2h_w.data_ptr<float>();
    const float* ow_ptr = h2o_w.data_ptr<float>();
    const float* ob_ptr = h2o_b.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();

    for (int t = 0; t < T; ++t) {
        hidden_step_kernel<<<grid_h, threads, smem, stream>>>(inp_ptr, hidden_prev.data_ptr<float>(), iw_ptr,
                                                              hidden_next.data_ptr<float>(), t, B, I, H);
        output_step_kernel<<<grid_o, threads, smem, stream>>>(hidden_next.data_ptr<float>(), ow_ptr, ob_ptr,
                                                              out_ptr, t, B, H, O);
        std::swap(hidden_prev, hidden_next);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(
    name="rnn_vanilla_fused_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["rnn_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    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)

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