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)
    
    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*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>

void launch_gru_layer(torch::Tensor input, torch::Tensor out, torch::Tensor hn, torch::Tensor h0,
                      torch::Tensor wih_f, torch::Tensor whh_f, torch::Tensor bih_f, torch::Tensor bhh_f,
                      torch::Tensor wih_b, torch::Tensor whh_b, torch::Tensor bih_b, torch::Tensor bhh_b,
                      int S, int B, int input_dim, int H, int layer, bool layer0, bool batch_first,
                      bool has_bias, bool store_output);

torch::Tensor gru_forward(torch::Tensor x, torch::Tensor h0, std::vector<torch::Tensor> params,
                          bool batch_first, int64_t num_layers, int64_t input_size, int64_t hidden_size,
                          bool has_bias) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(h0.is_cuda(), "h0 must be CUDA");
    int64_t H = hidden_size;
    int64_t S = batch_first ? x.size(1) : x.size(0);
    int64_t B = batch_first ? x.size(0) : x.size(1);
    auto hn = torch::empty_like(h0);
    auto opts = x.options();
    auto buf1 = torch::empty({S, B, 2*H}, opts);
    auto buf2 = torch::empty({S, B, 2*H}, opts);
    torch::Tensor dummy = torch::empty({0}, opts);

    for (int l = 0; l < (int)num_layers; ++l) {
        int base = l * 8;
        torch::Tensor inp = (l == 0) ? x : ((l & 1) ? buf1 : buf2);
        torch::Tensor out = (l & 1) ? buf2 : buf1;
        int inpdim = (l == 0) ? (int)input_size : (int)(2*H);
        bool store = (l != (int)num_layers - 1);
        launch_gru_layer(inp, out, hn, h0,
                         params[base+0], params[base+1], params[base+2], params[base+3],
                         params[base+4], params[base+5], params[base+6], params[base+7],
                         (int)S, (int)B, inpdim, (int)H, l, l==0, batch_first, has_bias, store);
    }
    return hn;
}
'''

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

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

__global__ void gru_layer_kernel(const float* __restrict__ input,
                                 float* __restrict__ out,
                                 float* __restrict__ hn,
                                 const float* __restrict__ h0,
                                 const float* __restrict__ wih_f,
                                 const float* __restrict__ whh_f,
                                 const float* __restrict__ bih_f,
                                 const float* __restrict__ bhh_f,
                                 const float* __restrict__ wih_b,
                                 const float* __restrict__ whh_b,
                                 const float* __restrict__ bih_b,
                                 const float* __restrict__ bhh_b,
                                 int S, int B, int I, int H, int layer,
                                 int layer0, int batch_first, int has_bias, int store_output) {
    int bd = blockIdx.x;
    int b = bd >> 1;
    int dir = bd & 1;
    int j = threadIdx.x;
    if (j >= H || b >= B) return;

    extern __shared__ float smem[];
    float* sh = smem;          // H entries: previous hidden
    float* sx = smem + H;      // I entries: current input vector

    const float* wih = dir == 0 ? wih_f : wih_b;
    const float* whh = dir == 0 ? whh_f : whh_b;
    const float* bih = dir == 0 ? bih_f : bih_b;
    const float* bhh = dir == 0 ? bhh_f : bhh_b;

    float hp = h0[((layer * 2 + dir) * B + b) * H + j];

    #pragma unroll 1
    for (int step = 0; step < S; ++step) {
        int t = (dir == 0) ? step : (S - 1 - step);

        for (int k = j; k < I; k += blockDim.x) {
            if (layer0) {
                sx[k] = batch_first ? input[(b * S + t) * I + k]
                                    : input[(t * B + b) * I + k];
            } else {
                sx[k] = input[(t * B + b) * I + k];
            }
        }
        sh[j] = hp;
        __syncthreads();

        const float* wi_r = wih + 0 * H * I + j * I;
        const float* wi_z = wih + 1 * H * I + j * I;
        const float* wi_n = wih + 2 * H * I + j * I;
        const float* wh_r = whh + 0 * H * H + j * H;
        const float* wh_z = whh + 1 * H * H + j * H;
        const float* wh_n = whh + 2 * H * H + j * H;

        float ir = has_bias ? bih[0 * H + j] : 0.0f;
        float iz = has_bias ? bih[1 * H + j] : 0.0f;
        float in = has_bias ? bih[2 * H + j] : 0.0f;
        float hr = has_bias ? bhh[0 * H + j] : 0.0f;
        float hz = has_bias ? bhh[1 * H + j] : 0.0f;
        float hn_lin = has_bias ? bhh[2 * H + j] : 0.0f;

        #pragma unroll 1
        for (int k = 0; k < I; ++k) {
            float xv = sx[k];
            ir = fmaf(wi_r[k], xv, ir);
            iz = fmaf(wi_z[k], xv, iz);
            in = fmaf(wi_n[k], xv, in);
        }
        #pragma unroll 1
        for (int k = 0; k < H; ++k) {
            float hv = sh[k];
            hr = fmaf(wh_r[k], hv, hr);
            hz = fmaf(wh_z[k], hv, hz);
            hn_lin = fmaf(wh_n[k], hv, hn_lin);
        }

        float r = sigmoidf_fast(ir + hr);
        float z = sigmoidf_fast(iz + hz);
        float n = tanhf(in + r * hn_lin);
        hp = fmaf(z, hp - n, n);   // (1-z)*n + z*hp

        if (store_output) {
            out[(t * B + b) * (2 * H) + dir * H + j] = hp;
        }
        __syncthreads();
    }
    hn[((layer * 2 + dir) * B + b) * H + j] = hp;
}

void launch_gru_layer(torch::Tensor input, torch::Tensor out, torch::Tensor hn, torch::Tensor h0,
                      torch::Tensor wih_f, torch::Tensor whh_f, torch::Tensor bih_f, torch::Tensor bhh_f,
                      torch::Tensor wih_b, torch::Tensor whh_b, torch::Tensor bih_b, torch::Tensor bhh_b,
                      int S, int B, int input_dim, int H, int layer, bool layer0, bool batch_first,
                      bool has_bias, bool store_output) {
    dim3 grid(B * 2);
    dim3 block(H);
    size_t shmem = (size_t)(H + input_dim) * sizeof(float);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gru_layer_kernel<<<grid, block, shmem, stream>>>(
        input.data_ptr<float>(), out.data_ptr<float>(), hn.data_ptr<float>(), h0.data_ptr<float>(),
        wih_f.data_ptr<float>(), whh_f.data_ptr<float>(), bih_f.data_ptr<float>(), bhh_f.data_ptr<float>(),
        wih_b.data_ptr<float>(), whh_b.data_ptr<float>(), bih_b.data_ptr<float>(), bhh_b.data_ptr<float>(),
        S, B, input_dim, H, layer, layer0 ? 1 : 0, batch_first ? 1 : 0, has_bias ? 1 : 0, store_output ? 1 : 0);
}
'''

_ext = load_inline(
    name="custom_gru_hn_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["gru_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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=True)

    def _param_list(self, device):
        ps = []
        dummy = None
        if not self.bias:
            dummy = torch.empty((0,), device=device, dtype=self.gru.weight_ih_l0.dtype)
        for l in range(self.num_layers):
            for suffix in ("", "_reverse"):
                ps.append(getattr(self.gru, f"weight_ih_l{l}{suffix}"))
                ps.append(getattr(self.gru, f"weight_hh_l{l}{suffix}"))
                if self.bias:
                    ps.append(getattr(self.gru, f"bias_ih_l{l}{suffix}"))
                    ps.append(getattr(self.gru, f"bias_hh_l{l}{suffix}"))
                else:
                    ps.append(dummy); ps.append(dummy)
        return ps

    def forward(self, x, h0):
        return _ext.gru_forward(x, h0, self._param_list(x.device), self.batch_first,
                                self.num_layers, self.input_size, self.hidden_size, self.bias)