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 <ATen/cuda/CUDAContext.h>
#include <vector>

void launch_gru_layer(const at::Tensor& inp,
                      at::Tensor& out,
                      const at::Tensor& hinit,
                      at::Tensor& hfinal,
                      const at::Tensor& wihT,
                      const at::Tensor& whhT,
                      const at::Tensor& bih,
                      const at::Tensor& bhh,
                      at::Tensor& hbuf,
                      at::Tensor& syncbuf,
                      int T, int B, int I, int H,
                      bool input_batch_first,
                      cudaStream_t stream);

torch::Tensor gru_forward(torch::Tensor x,
                          torch::Tensor h0,
                          std::vector<torch::Tensor> wihT,
                          std::vector<torch::Tensor> whhT,
                          std::vector<torch::Tensor> bih,
                          std::vector<torch::Tensor> bhh,
                          bool batch_first) {
    TORCH_CHECK(x.is_cuda() && h0.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only fp32 is supported");
    const int64_t L64 = h0.size(0);
    const int L = (int)L64;
    const int B = (int)h0.size(1);
    const int H = (int)h0.size(2);
    const int T = batch_first ? (int)x.size(1) : (int)x.size(0);

    auto hn = at::empty_like(h0);
    at::Tensor cur = x;
    std::vector<at::Tensor> keep;
    keep.reserve(L);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    for (int l = 0; l < L; ++l) {
        int I = (int)wihT[l].size(0);
        auto out = at::empty({T, B, H}, x.options());
        auto hbuf = at::empty({2, B, H}, x.options());
        auto syncbuf = at::zeros({2}, x.options().dtype(at::kInt));
        at::Tensor hinit_l = h0.select(0, l);
        at::Tensor hfinal_l = hn.select(0, l);
        launch_gru_layer(cur, out, hinit_l, hfinal_l,
                         wihT[l], whhT[l], bih[l], bhh[l],
                         hbuf, syncbuf, T, B, I, H,
                         (l == 0) && batch_first, stream);
        keep.push_back(out);
        cur = keep.back();
    }
    return hn;
}
'''

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

#define CHUNK_H 32

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

__device__ __forceinline__ void grid_barrier(int total_blocks, int* counter, int* sense, int target) {
    __syncthreads();
    if (threadIdx.x == 0) {
        __threadfence();
        int old = atomicAdd(counter, 1);
        if (old == total_blocks - 1) {
            *counter = 0;
            __threadfence();
            atomicExch(sense, target);
        } else {
            while (atomicAdd(sense, 0) != target) { }
        }
    }
    __syncthreads();
}

__global__ void gru_layer_kernel(const float* __restrict__ inp,
                                 float* __restrict__ out,
                                 const float* __restrict__ hinit,
                                 float* __restrict__ hfinal,
                                 const float* __restrict__ wihT,
                                 const float* __restrict__ whhT,
                                 const float* __restrict__ bih,
                                 const float* __restrict__ bhh,
                                 float* __restrict__ hbuf,
                                 int* __restrict__ syncbuf,
                                 int T, int B, int I, int H,
                                 int chunks, int input_batch_first) {
    __shared__ float sr[CHUNK_H];
    __shared__ float sz[CHUNK_H];
    __shared__ float sni[CHUNK_H];
    __shared__ float snh[CHUNK_H];

    const int b = blockIdx.x / chunks;
    const int c = blockIdx.x - b * chunks;
    const int hstart = c * CHUNK_H;
    const int tid = threadIdx.x;
    const int total_blocks = gridDim.x;
    int sense_val = 0;

    // initialize recurrent state buffer 0 from h0
    if (tid < CHUNK_H) {
        int h = hstart + tid;
        if (h < H) {
            hbuf[b * H + h] = hinit[b * H + h];
        }
    }
    sense_val = 1 - sense_val;
    grid_barrier(total_blocks, syncbuf, syncbuf + 1, sense_val);

    int cur = 0;
    for (int t = 0; t < T; ++t) {
        if (tid < 3 * CHUNK_H) {
            const int gate = tid / CHUNK_H;
            const int off  = tid - gate * CHUNK_H;
            const int h = hstart + off;
            if (h < H) {
                float si = 0.0f;
                float sh = 0.0f;
                const int row = gate * H + h;

                // input contribution, weights are cached as transposed: [I, 3H]
                for (int k = 0; k < I; ++k) {
                    float xv;
                    if (input_batch_first) {
                        xv = inp[(b * T + t) * I + k];
                    } else {
                        xv = inp[(t * B + b) * I + k];
                    }
                    si = fmaf(xv, wihT[k * (3 * H) + row], si);
                }
                // recurrent contribution from previous hidden buffer
                const float* hp = hbuf + cur * (B * H) + b * H;
                for (int k = 0; k < H; ++k) {
                    sh = fmaf(hp[k], whhT[k * (3 * H) + row], sh);
                }

                if (gate == 0) {
                    sr[off] = sigmoid_f(si + sh + bih[row] + bhh[row]);
                } else if (gate == 1) {
                    sz[off] = sigmoid_f(si + sh + bih[row] + bhh[row]);
                } else {
                    sni[off] = si + bih[row];
                    snh[off] = sh + bhh[row];
                }
            }
        }
        __syncthreads();

        if (tid < CHUNK_H) {
            int h = hstart + tid;
            if (h < H) {
                float oldh = hbuf[cur * (B * H) + b * H + h];
                float n = tanhf(sni[tid] + sr[tid] * snh[tid]);
                float z = sz[tid];
                float newh = fmaf(z, oldh, (1.0f - z) * n);
                hbuf[(1 - cur) * (B * H) + b * H + h] = newh;
                out[(t * B + b) * H + h] = newh;
            }
        }

        sense_val = 1 - sense_val;
        grid_barrier(total_blocks, syncbuf, syncbuf + 1, sense_val);
        cur = 1 - cur;
    }

    if (tid < CHUNK_H) {
        int h = hstart + tid;
        if (h < H) {
            hfinal[b * H + h] = hbuf[cur * (B * H) + b * H + h];
        }
    }
}

void launch_gru_layer(const at::Tensor& inp,
                      at::Tensor& out,
                      const at::Tensor& hinit,
                      at::Tensor& hfinal,
                      const at::Tensor& wihT,
                      const at::Tensor& whhT,
                      const at::Tensor& bih,
                      const at::Tensor& bhh,
                      at::Tensor& hbuf,
                      at::Tensor& syncbuf,
                      int T, int B, int I, int H,
                      bool input_batch_first,
                      cudaStream_t stream) {
    int chunks = (H + CHUNK_H - 1) / CHUNK_H;
    dim3 grid(B * chunks);
    dim3 block(128);
    gru_layer_kernel<<<grid, block, 0, stream>>>(
        inp.data_ptr<float>(), out.data_ptr<float>(),
        hinit.data_ptr<float>(), hfinal.data_ptr<float>(),
        wihT.data_ptr<float>(), whhT.data_ptr<float>(),
        bih.data_ptr<float>(), bhh.data_ptr<float>(),
        hbuf.data_ptr<float>(), syncbuf.data_ptr<int>(),
        T, B, I, H, chunks, input_batch_first ? 1 : 0);
}
'''

_ext = load_inline(
    name='gru_return_hn_fused_ext',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['gru_forward'],
    extra_cflags=['-O3'],
    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(ModelNew, self).__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers, bias, batch_first, dropout=0, bidirectional=False)
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bias = bias
        self.batch_first = batch_first
        self._cache_versions = None
        self._cache_device = None
        self._wihT = None
        self._whhT = None
        self._bih = None
        self._bhh = None

    def _refresh_cache(self, device):
        params = []
        for l in range(self.num_layers):
            params.append(getattr(self.gru, f'weight_ih_l{l}'))
            params.append(getattr(self.gru, f'weight_hh_l{l}'))
            if self.bias:
                params.append(getattr(self.gru, f'bias_ih_l{l}'))
                params.append(getattr(self.gru, f'bias_hh_l{l}'))
        versions = tuple(p._version for p in params)
        if (self._cache_versions == versions and self._cache_device == device and
            self._wihT is not None):
            return
        self._wihT = []
        self._whhT = []
        self._bih = []
        self._bhh = []
        for l in range(self.num_layers):
            wih = getattr(self.gru, f'weight_ih_l{l}').detach()
            whh = getattr(self.gru, f'weight_hh_l{l}').detach()
            self._wihT.append(wih.t().contiguous())
            self._whhT.append(whh.t().contiguous())
            if self.bias:
                self._bih.append(getattr(self.gru, f'bias_ih_l{l}').detach().contiguous())
                self._bhh.append(getattr(self.gru, f'bias_hh_l{l}').detach().contiguous())
            else:
                self._bih.append(torch.zeros(3 * self.hidden_size, device=device, dtype=torch.float32))
                self._bhh.append(torch.zeros(3 * self.hidden_size, device=device, dtype=torch.float32))
        self._cache_versions = versions
        self._cache_device = device

    def forward(self, x, h0):
        self._refresh_cache(x.device)
        return _ext.gru_forward(x.contiguous(), h0.contiguous(), self._wihT, self._whhT, self._bih, self._bhh, self.batch_first)