import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=3, bias=True, batch_first=False):
        """
        :param input_size: The number of expected features in the input x
        :param hidden_size: The number of features in the hidden state h
        :param num_layers: Number of recurrent layers (default: 1)
        :param bias: If False, then the layer does not use bias weights b_ih and b_hh (default: True)
        :param batch_first: If True, then the input and output tensors are provided as (batch, seq, feature) (default: False)
        """
        super(Model, self).__init__()
        
        self.gru = nn.GRU(input_size, hidden_size, num_layers, bias, batch_first, dropout=0, bidirectional=False)
    
    def forward(self, x,h0):
        """
        :param x: The input tensor, shape (seq_len, batch_size, input_size) if batch_first=False, otherwise (batch_size, seq_len, input_size)
        :param h_0: The initial hidden state for the input sequence, shape (num_layers * num_directions, batch_size, hidden_size) (default: None)
        :return: output, h_n
            - output: The output features (h_t) from the last layer of the GRU, for each t, shape (seq_len, batch_size, num_directions * hidden_size) if batch_first=False, otherwise (batch_size, seq_len, num_directions * hidden_size)
            - h_n: The hidden state for t = seq_len, shape (num_layers * num_directions, batch_size, hidden_size)
        """
        output, h_n = self.gru(x, h0)
        return output

# Test code
batch_size = 10
seq_len = 512
input_size = 128
hidden_size = 256
num_layers = 6

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

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

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

cpp_src = r"""
#include <torch/extension.h>
#include <vector>

void gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                      torch::Tensor output, torch::Tensor hn,
                      std::vector<torch::Tensor> weights,
                      int64_t num_layers, bool batch_first);

torch::Tensor gru_forward(torch::Tensor x, torch::Tensor h0,
                          std::vector<torch::Tensor> weights,
                          int64_t num_layers, bool batch_first) {
    TORCH_CHECK(x.is_cuda(), "x must be cuda");
    TORCH_CHECK(h0.is_cuda(), "h0 must be cuda");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "float32 only");
    int64_t seq = batch_first ? x.size(1) : x.size(0);
    int64_t batch = batch_first ? x.size(0) : x.size(1);
    int64_t hidden = h0.size(2);
    auto out = torch::empty({batch_first ? batch : seq, batch_first ? seq : batch, hidden}, x.options());
    auto hn = torch::empty_like(h0);
    gru_forward_cuda(x.contiguous(), h0.contiguous(), out, hn, weights, num_layers, batch_first);
    return out;
}
"""

cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <math.h>

#define H 256
#define B 10
#define X0 128

__device__ __forceinline__ float sigmoidf_fast(float x) {
    return 1.0f / (1.0f + __expf(-x));
}

// one block computes one (layer,batch) vector of 256 hidden units for one timestep
__global__ void gru_step_kernel(const float* __restrict__ inp, const float* __restrict__ hprev,
                                float* __restrict__ hnext, float* __restrict__ out,
                                const float* __restrict__ wih, const float* __restrict__ whh,
                                const float* __restrict__ bih, const float* __restrict__ bhh,
                                int in_size, int inp_stride0, int layer, int t, int seq, bool batch_first, bool write_out) {
    int b = blockIdx.x;
    int j = threadIdx.x; // 0..255
    if (j >= H || b >= B) return;
    const float* hp = hprev + (layer*B + b)*H;
    float h_old = hp[j];
    float gi_r = bih[j];
    float gi_z = bih[H+j];
    float gi_n = bih[2*H+j];
    // input contribution: rows are contiguous in PyTorch GRU weight_ih_l[k] shape (3H, input)
    const float* wr = wih + j*in_size;
    const float* wz = wih + (H+j)*in_size;
    const float* wn = wih + (2*H+j)*in_size;
    if (layer == 0) {
        // input is original x, seq-major or batch-major
        const float* xptr;
        if (batch_first) xptr = inp + ((b*seq + t) * in_size);
        else xptr = inp + ((t*B + b) * in_size);
        #pragma unroll 4
        for (int k=0;k<X0;k++) {
            float xv = xptr[k];
            gi_r += wr[k]*xv; gi_z += wz[k]*xv; gi_n += wn[k]*xv;
        }
    } else {
        // input is previous layer state buffer laid out [B,H]
        const float* xptr = inp + b*H;
        #pragma unroll 4
        for (int k=0;k<H;k++) {
            float xv = xptr[k];
            gi_r += wr[k]*xv; gi_z += wz[k]*xv; gi_n += wn[k]*xv;
        }
    }
    float gh_r = bhh[j];
    float gh_z = bhh[H+j];
    float gh_n = bhh[2*H+j];
    const float* hr = whh + j*H;
    const float* hz = whh + (H+j)*H;
    const float* hnw = whh + (2*H+j)*H;
    #pragma unroll 4
    for (int k=0;k<H;k++) {
        float hv = hp[k];
        gh_r += hr[k]*hv; gh_z += hz[k]*hv; gh_n += hnw[k]*hv;
    }
    float r = sigmoidf_fast(gi_r + gh_r);
    float z = sigmoidf_fast(gi_z + gh_z);
    float n = tanhf(gi_n + r*gh_n);
    float hv = n + z*(h_old - n);
    hnext[(layer*B + b)*H + j] = hv;
    if (write_out) {
        if (batch_first) out[(b*seq + t)*H + j] = hv;
        else out[(t*B + b)*H + j] = hv;
    }
}

__global__ void copy_kernel(const float* __restrict__ src, float* __restrict__ dst, int n) {
    int i = blockIdx.x*blockDim.x + threadIdx.x;
    if (i<n) dst[i]=src[i];
}

void gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                      torch::Tensor output, torch::Tensor hn,
                      std::vector<torch::Tensor> weights,
                      int64_t num_layers, bool batch_first) {
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int seq = batch_first ? x.size(1) : x.size(0);
    int batch = batch_first ? x.size(0) : x.size(1);
    int hidden = h0.size(2);
    TORCH_CHECK(batch == B && hidden == H, "only fixed batch/hidden supported");
    float* bufA; float* bufB;
    cudaMalloc(&bufA, num_layers*batch*hidden*sizeof(float));
    cudaMalloc(&bufB, num_layers*batch*hidden*sizeof(float));
    int total = num_layers*batch*hidden;
    copy_kernel<<<(total+255)/256,256,0,stream>>>(h0.data_ptr<float>(), bufA, total);
    // temp buffers for inter-layer activations at current timestep
    float* layA; float* layB;
    cudaMalloc(&layA, batch*hidden*sizeof(float));
    cudaMalloc(&layB, batch*hidden*sizeof(float));
    dim3 grid(batch); dim3 block(hidden);
    for (int t=0; t<seq; ++t) {
        float* prev_states = (t & 1) ? bufB : bufA;
        float* next_states = (t & 1) ? bufA : bufB;
        const float* layer_input_x = x.data_ptr<float>();
        for (int l=0; l<num_layers; ++l) {
            const float* wih = weights[l*4+0].data_ptr<float>();
            const float* whh = weights[l*4+1].data_ptr<float>();
            const float* bih = weights[l*4+2].data_ptr<float>();
            const float* bhh = weights[l*4+3].data_ptr<float>();
            int in_size = (l==0) ? x.size(2) : hidden;
            const float* inp_ptr;
            if (l==0) inp_ptr = layer_input_x;
            else inp_ptr = (l & 1) ? layA : layB;
            float* out_layer = (l & 1) ? layB : layA;
            bool write_out = (l == num_layers-1);
            gru_step_kernel<<<grid,block,0,stream>>>(inp_ptr, prev_states, next_states, out_layer, wih, whh, bih, bhh,
                                                     in_size, 0, l, t, seq, batch_first, write_out);
        }
    }
    float* final_states = (seq & 1) ? bufB : bufA;
    copy_kernel<<<(total+255)/256,256,0,stream>>>(final_states, hn.data_ptr<float>(), total);
    cudaFree(bufA); cudaFree(bufB); cudaFree(layA); cudaFree(layB);
}
"""

_ext = load_inline(name="gru_custom_ext_v1", cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=["gru_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"],
                   extra_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=3, bias=True, batch_first=False):
        super(ModelNew, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bias = bias
        self.batch_first = batch_first
        self.gru = nn.GRU(input_size, hidden_size, num_layers, bias, batch_first, dropout=0, bidirectional=False)
    def forward(self, x, h0):
        if (x.is_cuda and h0.is_cuda and x.dtype is torch.float32 and self.bias and
            self.hidden_size == 256 and x.size(1 if not self.batch_first else 0) == 10):
            weights=[]
            for l in range(self.num_layers):
                weights.append(getattr(self.gru, f'weight_ih_l{l}').contiguous())
                weights.append(getattr(self.gru, f'weight_hh_l{l}').contiguous())
                weights.append(getattr(self.gru, f'bias_ih_l{l}').contiguous())
                weights.append(getattr(self.gru, f'bias_hh_l{l}').contiguous())
            return _ext.gru_forward(x, h0, weights, self.num_layers, self.batch_first)
        else:
            return self.gru(x,h0)[0]