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=True)
        self.h0 = torch.randn((num_layers * 2, batch_size, hidden_size))
    
    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*2, 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> weights,
                               int64_t num_layers, int64_t input_size,
                               int64_t hidden_size, bool has_bias,
                               bool batch_first);
'''

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

namespace cg = cooperative_groups;

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

__device__ __forceinline__ float load_inp(const float* __restrict__ inp,
                                          int t, int b, int k,
                                          int S, int B, int K,
                                          bool input_is_original,
                                          bool batch_first) {
    if (input_is_original && batch_first) {
        return inp[(b * S + t) * K + k];
    } else {
        return inp[(t * B + b) * K + k];
    }
}

extern "C" __global__ void gru_dir_kernel(
    const float* __restrict__ inp,
    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__ hidden_tmp,
    int S, int B, int K, int H,
    int layer, int dir,
    bool input_is_original,
    bool batch_first,
    bool has_bias,
    bool out_batch_first)
{
    cg::grid_group grid = cg::this_grid();
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int stride = blockDim.x * gridDim.x;
    const int total = B * H;

    // ping buffer 0 receives the initial hidden state for this layer/direction
    for (int idx = tid; idx < total; idx += stride) {
        int b = idx / H;
        int h = idx - b * H;
        hidden_tmp[idx] = h0[((layer * 2 + dir) * B + b) * H + h];
    }
    grid.sync();

    for (int it = 0; it < S; ++it) {
        int t = (dir == 0) ? it : (S - 1 - it);
        const int old_base = (it & 1) ? total : 0;
        const int new_base = (it & 1) ? 0 : total;

        for (int idx = tid; idx < total; idx += stride) {
            int b = idx / H;
            int h = idx - b * H;

            float ir = has_bias ? b_ih[h] : 0.0f;
            float iz = has_bias ? b_ih[H + h] : 0.0f;
            float in = has_bias ? b_ih[2 * H + h] : 0.0f;

            // input affine, rows are laid out as [reset, update, new]
            const float* wr = w_ih + h * K;
            const float* wz = w_ih + (H + h) * K;
            const float* wn = w_ih + (2 * H + h) * K;
            for (int k = 0; k < K; ++k) {
                float xv = load_inp(inp, t, b, k, S, B, K, input_is_original, batch_first);
                ir += xv * wr[k];
                iz += xv * wz[k];
                in += xv * wn[k];
            }

            float hr = has_bias ? b_hh[h] : 0.0f;
            float hz = has_bias ? b_hh[H + h] : 0.0f;
            float hn = has_bias ? b_hh[2 * H + h] : 0.0f;
            const float* rr = w_hh + h * H;
            const float* rz = w_hh + (H + h) * H;
            const float* rn = w_hh + (2 * H + h) * H;
            const float* hprev_vec = hidden_tmp + old_base + b * H;
            for (int k = 0; k < H; ++k) {
                float hv = hprev_vec[k];
                hr += hv * rr[k];
                hz += hv * rz[k];
                hn += hv * rn[k];
            }

            float r = sigmoid_f(ir + hr);
            float z = sigmoid_f(iz + hz);
            float n = tanhf(in + r * hn);
            float hprev = hprev_vec[h];
            float hnew = (1.0f - z) * n + z * hprev;
            hidden_tmp[new_base + b * H + h] = hnew;

            if (out_batch_first) {
                out[(b * S + t) * (2 * H) + dir * H + h] = hnew;
            } else {
                out[(t * B + b) * (2 * H) + dir * H + h] = hnew;
            }
        }
        grid.sync();
    }
}

torch::Tensor gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                               std::vector<torch::Tensor> weights,
                               int64_t num_layers64, int64_t input_size64,
                               int64_t hidden_size64, bool has_bias,
                               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 supported");
    TORCH_CHECK(h0.scalar_type() == at::kFloat, "only float32 supported");
    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
    TORCH_CHECK(h0.is_contiguous(), "h0 must be contiguous");
    for (auto &w : weights) TORCH_CHECK(w.is_cuda() && w.is_contiguous() && w.scalar_type() == at::kFloat, "weights must be contiguous CUDA float tensors");

    const c10::cuda::CUDAGuard device_guard(x.device());
    int S = (int)(batch_first ? x.size(1) : x.size(0));
    int B = (int)(batch_first ? x.size(0) : x.size(1));
    int input_size = (int)input_size64;
    int H = (int)hidden_size64;
    int L = (int)num_layers64;

    auto opts = x.options();
    torch::Tensor buf_a = torch::empty({S, B, 2 * H}, opts);
    torch::Tensor buf_b = (L > 2) ? torch::empty({S, B, 2 * H}, opts) : torch::Tensor();
    torch::Tensor hidden_tmp = torch::empty({2, B, H}, opts);
    torch::Tensor final_out;

    const int threads = 128;
    int blocks_n = (B * H + threads - 1) / threads;
    dim3 grid(blocks_n);
    dim3 block(threads);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const int per_layer = has_bias ? 8 : 4;
    const int per_dir = has_bias ? 4 : 2;

    torch::Tensor current;
    for (int layer = 0; layer < L; ++layer) {
        int K = (layer == 0) ? input_size : (2 * H);
        bool input_orig = (layer == 0);
        const float* inp_ptr = input_orig ? x.data_ptr<float>() : current.data_ptr<float>();

        bool is_last = (layer == L - 1);
        torch::Tensor out;
        if (is_last) {
            if (batch_first) out = torch::empty({B, S, 2 * H}, opts);
            else out = torch::empty({S, B, 2 * H}, opts);
        } else {
            out = (layer & 1) ? buf_b : buf_a;
        }

        for (int dir = 0; dir < 2; ++dir) {
            int wi = layer * per_layer + dir * per_dir;
            const float* wih = weights[wi + 0].data_ptr<float>();
            const float* whh = weights[wi + 1].data_ptr<float>();
            const float* bih = has_bias ? weights[wi + 2].data_ptr<float>() : nullptr;
            const float* bhh = has_bias ? weights[wi + 3].data_ptr<float>() : nullptr;
            bool out_bf = is_last && batch_first;

            void* args[] = {
                (void*)&inp_ptr, (void*)&(h0.data_ptr<float>()), (void*)&wih, (void*)&whh,
                (void*)&bih, (void*)&bhh, (void*)&(out.data_ptr<float>()), (void*)&(hidden_tmp.data_ptr<float>()),
                (void*)&S, (void*)&B, (void*)&K, (void*)&H,
                (void*)&layer, (void*)&dir, (void*)&input_orig, (void*)&batch_first,
                (void*)&has_bias, (void*)&out_bf
            };
            C10_CUDA_CHECK(cudaLaunchCooperativeKernel((void*)gru_dir_kernel, grid, block, args, 0, stream));
        }
        current = out;
        if (is_last) final_out = out;
    }
    return final_out;
}
'''

_ext = load_inline(
    name='custom_gru_bidir_fwd_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['gru_forward_cuda'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    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=bias,
                          batch_first=batch_first, dropout=0, bidirectional=True)

    def _weight_list(self):
        ws = []
        for l in range(self.num_layers):
            for suf in ('', '_reverse'):
                ws.append(getattr(self.gru, 'weight_ih_l%d%s' % (l, suf)))
                ws.append(getattr(self.gru, 'weight_hh_l%d%s' % (l, suf)))
                if self.bias:
                    ws.append(getattr(self.gru, 'bias_ih_l%d%s' % (l, suf)))
                    ws.append(getattr(self.gru, 'bias_hh_l%d%s' % (l, suf)))
        return ws

    def forward(self, x, h0):
        return _ext.gru_forward_cuda(x, h0, self._weight_list(), self.num_layers,
                                     self.input_size, self.hidden_size,
                                     self.bias, self.batch_first)