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 w_i2h,
                          torch::Tensor b_i2h,
                          torch::Tensor w_h2o,
                          torch::Tensor b_h2o);
"""

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

#define MAX_H 256

__global__ void rnn_recur_kernel(const float* __restrict__ inpart,
                                 const float* __restrict__ h0,
                                 const float* __restrict__ w_hh_base,
                                 const float* __restrict__ b_h,
                                 float* __restrict__ hidden_seq,
                                 int T, int B, int H, int total_in)
{
    int b = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float prev[MAX_H];
    __shared__ float curr[MAX_H];

    if (tid < H) prev[tid] = h0[b * H + tid];
    __syncthreads();

    for (int t = 0; t < T; ++t) {
        if (tid < H) {
            float acc = inpart[(t * B + b) * H + tid] + b_h[tid];
            const float* wrow = w_hh_base + tid * total_in;
            #pragma unroll 4
            for (int k = 0; k < MAX_H; ++k) {
                acc = fmaf(prev[k], wrow[k], acc);
            }
            float hv = tanhf(acc);
            curr[tid] = hv;
            hidden_seq[(t * B + b) * H + tid] = hv;
        }
        __syncthreads();
        if (tid < H) prev[tid] = curr[tid];
        __syncthreads();
    }
}

__global__ void add_bias_kernel(float* __restrict__ y, const float* __restrict__ bias, int n, int O)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) y[idx] += bias[idx % O];
}

static inline void rowmajor_A_times_Wt(cublasHandle_t handle,
                                       const float* A, const float* W,
                                       float* C,
                                       int M, int K, int N, int W_ld,
                                       cudaStream_t stream)
{
    const float alpha = 1.0f;
    const float beta = 0.0f;
    // Row-major C[M,N] = A[M,K] * W[N,K]^T.
    // Interpret row-major buffers as column-major transposes and compute C^T = W * A^T.
    CUBLAS_CHECK(cublasSetStream(handle, stream));
    CUBLAS_CHECK(cublasSgemm(handle,
                             CUBLAS_OP_T, CUBLAS_OP_N,
                             N, M, K,
                             &alpha,
                             W, W_ld,
                             A, K,
                             &beta,
                             C, N));
}

torch::Tensor rnn_forward(torch::Tensor x,
                          torch::Tensor h0,
                          torch::Tensor w_i2h,
                          torch::Tensor b_i2h,
                          torch::Tensor w_h2o,
                          torch::Tensor b_h2o)
{
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(h0.is_cuda(), "h0 must be CUDA");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only fp32 supported");

    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)w_h2o.size(0);
    const int total_in = I + H;
    TORCH_CHECK(H == MAX_H, "this optimized kernel expects hidden_size=256");

    auto opts = x.options();
    torch::Tensor inpart = torch::empty({T * B, H}, opts);
    torch::Tensor hidden_seq = torch::empty({T * B, H}, opts);
    torch::Tensor out = torch::empty({T * B, O}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();

    rowmajor_A_times_Wt(handle,
                        x.data_ptr<float>(), w_i2h.data_ptr<float>(), inpart.data_ptr<float>(),
                        T * B, I, H, total_in, stream);

    rnn_recur_kernel<<<B, 256, 0, stream>>>(inpart.data_ptr<float>(),
                                            h0.data_ptr<float>(),
                                            w_i2h.data_ptr<float>() + I,
                                            b_i2h.data_ptr<float>(),
                                            hidden_seq.data_ptr<float>(),
                                            T, B, H, total_in);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    rowmajor_A_times_Wt(handle,
                        hidden_seq.data_ptr<float>(), w_h2o.data_ptr<float>(), out.data_ptr<float>(),
                        T * B, H, O, H, stream);

    int nout = T * B * O;
    add_bias_kernel<<<(nout + 255) / 256, 256, 0, stream>>>(out.data_ptr<float>(), b_h2o.data_ptr<float>(), nout, O);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out.view({T, B, O});
}
"""

_ext = load_inline(
    name="vanilla_rnn_fused_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["rnn_forward"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas"],
    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:
        if not x.is_contiguous():
            x = x.contiguous()
        h0 = h0.to(device=x.device)
        if not h0.is_contiguous():
            h0 = h0.contiguous()
        return _ext.rnn_forward(x, h0, self.i2h.weight, self.i2h.bias, self.h2o.weight, self.h2o.bias)