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 launch_gru_layer(torch::Tensor xg, torch::Tensor h0_l, torch::Tensor whh,
                      c10::optional<torch::Tensor> bhh, torch::Tensor out);

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) {
    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 float32 is supported");

    torch::Tensor xs = batch_first ? x.transpose(0, 1).contiguous() : x.contiguous();
    torch::Tensor h0c = h0.contiguous();

    const int64_t T = xs.size(0);
    const int64_t B = xs.size(1);
    const int64_t H = w_hh[0].size(1);
    const int64_t L = (int64_t)w_ih.size();
    const bool has_bias = (b_ih.size() == w_ih.size()) && (b_hh.size() == w_hh.size());

    torch::Tensor cur;
    for (int64_t l = 0; l < L; ++l) {
        torch::Tensor inp = (l == 0) ? xs : cur;
        torch::Tensor inp2 = inp.reshape({T * B, inp.size(2)});
        torch::Tensor xg;
        if (has_bias) {
            xg = at::addmm(b_ih[l], inp2, w_ih[l].t());
        } else {
            xg = at::mm(inp2, w_ih[l].t());
        }
        xg = xg.contiguous();
        torch::Tensor out = torch::empty({T, B, H}, xs.options());
        torch::Tensor h0_l = h0c.select(0, l);
        if (has_bias) {
            launch_gru_layer(xg, h0_l, w_hh[l].contiguous(), b_hh[l].contiguous(), out);
        } else {
            launch_gru_layer(xg, h0_l, w_hh[l].contiguous(), c10::nullopt, out);
        }
        cur = out;
    }
    return batch_first ? cur.transpose(0, 1) : cur;
}
'''

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

namespace cg = cooperative_groups;

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

// One cooperative grid computes one recurrent layer.  Grid is (ceil(H/16), B).
// Each block owns 16 hidden units for one batch element; 8 threads cooperate on
// each unit's recurrent dot products.  A grid-wide sync separates time steps.
__global__ void gru_layer_kernel(const float* __restrict__ xg,
                                 const float* __restrict__ h0,
                                 const float* __restrict__ whh,
                                 const float* __restrict__ bhh,
                                 float* __restrict__ out,
                                 int T, int B, int H) {
    cg::grid_group grid = cg::this_grid();

    constexpr int ROWS_PER_BLOCK = 16;
    constexpr int LANES_PER_ROW = 8;
    int b = blockIdx.y;
    int tile = blockIdx.x;
    int tid = threadIdx.x;
    int local_row = tid / LANES_PER_ROW;
    int lane = tid & (LANES_PER_ROW - 1);
    int j = tile * ROWS_PER_BLOCK + local_row;

    // The output buffer doubles as the recurrent state area after the sequence;
    // use a compact state stored in the first B*H elements of out temporarily.
    // To avoid clobbering real outputs, place state at the end is impossible
    // without extra storage, so each step reads previous state from either h0 or
    // the preceding output time slice.
    grid.sync();

    for (int t = 0; t < T; ++t) {
        float sr = 0.0f, sz = 0.0f, sn = 0.0f;
        if (j < H) {
            const float* hprev_base = (t == 0) ? (h0 + b * H) : (out + ((t - 1) * B + b) * H);
            const float* wr = whh + j * H;
            const float* wz = whh + (H + j) * H;
            const float* wn = whh + (2 * H + j) * H;
            for (int k = lane; k < H; k += LANES_PER_ROW) {
                float hv = hprev_base[k];
                sr += wr[k] * hv;
                sz += wz[k] * hv;
                sn += wn[k] * hv;
            }
        }
#pragma unroll
        for (int off = LANES_PER_ROW / 2; off > 0; off >>= 1) {
            sr += __shfl_down_sync(0xffffffff, sr, off, LANES_PER_ROW);
            sz += __shfl_down_sync(0xffffffff, sz, off, LANES_PER_ROW);
            sn += __shfl_down_sync(0xffffffff, sn, off, LANES_PER_ROW);
        }
        if (j < H && lane == 0) {
            const int64_t gate_base = ((int64_t)t * B + b) * (3 * H);
            const float bhr = bhh ? bhh[j] : 0.0f;
            const float bhz = bhh ? bhh[H + j] : 0.0f;
            const float bhn = bhh ? bhh[2 * H + j] : 0.0f;
            float r = sigmoid_f(xg[gate_base + j] + sr + bhr);
            float z = sigmoid_f(xg[gate_base + H + j] + sz + bhz);
            float n = tanhf(xg[gate_base + 2 * H + j] + r * (sn + bhn));
            float hp = (t == 0) ? h0[b * H + j] : out[((int64_t)(t - 1) * B + b) * H + j];
            out[((int64_t)t * B + b) * H + j] = (1.0f - z) * n + z * hp;
        }
        grid.sync();
    }
}

void launch_gru_layer(torch::Tensor xg, torch::Tensor h0_l, torch::Tensor whh,
                      c10::optional<torch::Tensor> bhh_opt, torch::Tensor out) {
    const int T = (int)out.size(0);
    const int B = (int)out.size(1);
    const int H = (int)out.size(2);
    dim3 block(128);
    dim3 grid((H + 15) / 16, B);
    const float* bhh = bhh_opt.has_value() ? bhh_opt.value().data_ptr<float>() : nullptr;
    void* args[] = {
        (void*)&xg.data_ptr<float>(),
        (void*)&h0_l.data_ptr<float>(),
        (void*)&whh.data_ptr<float>(),
        (void*)&bhh,
        (void*)&out.data_ptr<float>(),
        (void*)&T,
        (void*)&B,
        (void*)&H
    };
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    C10_CUDA_CHECK(cudaLaunchCooperativeKernel((void*)gru_layer_kernel, grid, block, args, 0, stream));
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}
'''

_ext = load_inline(
    name='gru_custom_coop_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['gru_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--expt-relaxed-constexpr'],
    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.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 = []
        w_hh = []
        b_ih = []
        b_hh = []
        for l in range(self.num_layers):
            w_ih.append(getattr(self.gru, f'weight_ih_l{l}'))
            w_hh.append(getattr(self.gru, f'weight_hh_l{l}'))
            if self.bias:
                b_ih.append(getattr(self.gru, f'bias_ih_l{l}'))
                b_hh.append(getattr(self.gru, f'bias_hh_l{l}'))
        return _ext.gru_forward(x, h0, w_ih, w_hh, b_ih, b_hh, self.batch_first)