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 h_n

# 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>

torch::Tensor gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                               std::vector<torch::Tensor> w_ih,
                               std::vector<torch::Tensor> w_hh,
                               std::vector<torch::Tensor> b_ih,
                               std::vector<torch::Tensor> b_hh,
                               bool batch_first);

torch::Tensor gru_forward(torch::Tensor x, torch::Tensor h0,
                          std::vector<torch::Tensor> w_ih,
                          std::vector<torch::Tensor> w_hh,
                          std::vector<torch::Tensor> b_ih,
                          std::vector<torch::Tensor> b_hh,
                          bool batch_first) {
    return gru_forward_cuda(x, h0, w_ih, w_hh, b_ih, b_hh, batch_first);
}
'''

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

namespace cg = cooperative_groups;

__device__ __forceinline__ float sigm(float x) {
    return 1.0f / (1.0f + expf(-x));
}

// cuDNN/PyTorch GRU on H100 uses TF32 GEMMs by default.  Explicitly rounding
// multiplicands to TF32 keeps this custom implementation within fp32 tolerance.
__device__ __forceinline__ float tf32_rn(float x) {
    unsigned u = __float_as_uint(x);
    u += 0x1000u;
    u &= 0xffffe000u;
    return __uint_as_float(u);
}

__device__ __forceinline__ float warp_sum(float v) {
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
    return v;
}

__global__ __launch_bounds__(128, 8)
void gru_layer_kernel(const float* __restrict__ input,
                      const float* __restrict__ h0,
                      const float* __restrict__ w_ih,
                      const float* __restrict__ w_hh,
                      const float* __restrict__ b_ih,
                      const float* __restrict__ b_hh,
                      float* __restrict__ out,
                      float* __restrict__ hs0,
                      float* __restrict__ hs1,
                      float* __restrict__ hn,
                      int T, int B, int in_dim, int H,
                      int layer, int batch_first, int is_first_layer) {
    cg::grid_group grid = cg::this_grid();
    int lane = threadIdx.x & 31;
    int wid  = threadIdx.x >> 5;
    int b = blockIdx.x;
    int h = blockIdx.y * 4 + wid;
    bool ok = h < H;

    if (ok && lane == 0) hs0[b * H + h] = h0[(layer * B + b) * H + h];
    grid.sync();

    for (int t = 0; t < T; ++t) {
        const float* hprev = (t & 1) ? (hs1 + b * H) : (hs0 + b * H);
        float* hcur = (t & 1) ? (hs0 + b * H) : (hs1 + b * H);

        if (ok) {
            const float* xbase;
            if (is_first_layer)
                xbase = batch_first ? input + (b * T + t) * in_dim : input + (t * B + b) * in_dim;
            else
                xbase = input + (t * B + b) * in_dim;

            float gr = 0.0f, gz = 0.0f, gi = 0.0f, gh = 0.0f;
            const float* wr = w_ih + h * in_dim;
            const float* wz = w_ih + (H + h) * in_dim;
            const float* wn = w_ih + (2 * H + h) * in_dim;

            for (int k = lane; k < in_dim; k += 32) {
                float xv = tf32_rn(xbase[k]);
                gr += tf32_rn(wr[k]) * xv;
                gz += tf32_rn(wz[k]) * xv;
                gi += tf32_rn(wn[k]) * xv;
            }

            const float* whr = w_hh + h * H;
            const float* whz = w_hh + (H + h) * H;
            const float* whn = w_hh + (2 * H + h) * H;
            for (int k = lane; k < H; k += 32) {
                float hv = tf32_rn(hprev[k]);
                gr += tf32_rn(whr[k]) * hv;
                gz += tf32_rn(whz[k]) * hv;
                gh += tf32_rn(whn[k]) * hv;
            }

            gr = warp_sum(gr);
            gz = warp_sum(gz);
            gi = warp_sum(gi);
            gh = warp_sum(gh);

            if (lane == 0) {
                gr += b_ih[h] + b_hh[h];
                gz += b_ih[H + h] + b_hh[H + h];
                gi += b_ih[2 * H + h];
                gh += b_hh[2 * H + h];
                float hp = hprev[h];
                float r = sigm(gr);
                float z = sigm(gz);
                float n = tanhf(gi + r * gh);
                float hv = (1.0f - z) * n + z * hp;
                hcur[h] = hv;
                out[(t * B + b) * H + h] = hv;
                if (t == T - 1) hn[(layer * B + b) * H + h] = hv;
            }
        }
        grid.sync();
    }
}

static void launch_layer(const float* input, const float* h0,
                         const float* w_ih, const float* w_hh,
                         const float* b_ih, const float* b_hh,
                         float* out, float* hs0, float* hs1, float* hn,
                         int T, int B, int in_dim, int H, int layer,
                         bool batch_first, bool is_first_layer, cudaStream_t stream) {
    dim3 block(128);
    dim3 grid(B, (H + 3) / 4);
    int bf = batch_first ? 1 : 0;
    int first = is_first_layer ? 1 : 0;
    void* args[] = { (void*)&input, (void*)&h0, (void*)&w_ih, (void*)&w_hh,
                     (void*)&b_ih, (void*)&b_hh, (void*)&out, (void*)&hs0,
                     (void*)&hs1, (void*)&hn, (void*)&T, (void*)&B,
                     (void*)&in_dim, (void*)&H, (void*)&layer, (void*)&bf,
                     (void*)&first };
    cudaLaunchCooperativeKernel((void*)gru_layer_kernel, grid, block, args, 0, stream);
}

torch::Tensor gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                               std::vector<torch::Tensor> w_ih,
                               std::vector<torch::Tensor> w_hh,
                               std::vector<torch::Tensor> b_ih,
                               std::vector<torch::Tensor> b_hh,
                               bool batch_first) {
    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
    int64_t L = h0.size(0);
    int T = batch_first ? (int)x.size(1) : (int)x.size(0);
    int B = batch_first ? (int)x.size(0) : (int)x.size(1);
    int I = (int)x.size(2);
    int H = (int)h0.size(2);

    auto opts = x.options();
    torch::Tensor hn = torch::empty({L, B, H}, opts);
    torch::Tensor buf1 = torch::empty({T, B, H}, opts);
    torch::Tensor buf2 = torch::empty({T, B, H}, opts);
    torch::Tensor hs0 = torch::empty({B, H}, opts);
    torch::Tensor hs1 = torch::empty({B, H}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    for (int l = 0; l < L; ++l) {
        torch::Tensor& out = (l % 2 == 0) ? buf1 : buf2;
        const float* inp = (l == 0) ? x.data_ptr<float>() : ((l % 2 == 1) ? buf1.data_ptr<float>() : buf2.data_ptr<float>());
        int in_dim = (l == 0) ? I : H;
        launch_layer(inp, h0.data_ptr<float>(), w_ih[l].data_ptr<float>(), w_hh[l].data_ptr<float>(),
                     b_ih[l].data_ptr<float>(), b_hh[l].data_ptr<float>(), out.data_ptr<float>(),
                     hs0.data_ptr<float>(), hs1.data_ptr<float>(), hn.data_ptr<float>(),
                     T, B, in_dim, H, l, batch_first, l == 0, stream);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return hn;
}
'''

_ext = load_inline(name='gru_custom_ext_final',
                   cpp_sources=_cpp_src,
                   cuda_sources=_cuda_src,
                   functions=['gru_forward'],
                   extra_cuda_cflags=['-O3'],
                   verbose=False)

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers=3, bias=True, batch_first=False):
        super().__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers, bias, batch_first, dropout=0, bidirectional=False)
        self.batch_first = batch_first
        self.num_layers = num_layers
        self.bias = bias

    def forward(self, x, h0):
        w_ih = [getattr(self.gru, f'weight_ih_l{i}') for i in range(self.num_layers)]
        w_hh = [getattr(self.gru, f'weight_hh_l{i}') for i in range(self.num_layers)]
        if self.bias:
            b_ih = [getattr(self.gru, f'bias_ih_l{i}') for i in range(self.num_layers)]
            b_hh = [getattr(self.gru, f'bias_hh_l{i}') for i in range(self.num_layers)]
        else:
            b_ih = [torch.zeros(3 * self.gru.hidden_size, device=x.device, dtype=x.dtype) for _ in range(self.num_layers)]
            b_hh = [torch.zeros(3 * self.gru.hidden_size, device=x.device, dtype=x.dtype) for _ in range(self.num_layers)]
        return _ext.gru_forward(x, h0, w_ih, w_hh, b_ih, b_hh, self.batch_first)