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

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

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

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

// One persistent grid computes one layer and one direction.  Blocks are (batch, hidden-chunk).
__global__ void gru_layer_dir_kernel(
    const float* __restrict__ inp,        // [T,B,K], contiguous
    const float* __restrict__ h0,         // [2L,B,H], contiguous, already offset by layer/dir
    const float* __restrict__ wihT,       // [K,3H]
    const float* __restrict__ whhT,       // [H,3H]
    const float* __restrict__ bih,        // [3H] or nullptr-equivalent empty guarded by has_bias
    const float* __restrict__ bhh,        // [3H]
    float* __restrict__ out,              // [T,B,2H]
    float* __restrict__ hn,               // [2L,B,H]
    float* __restrict__ hstate,           // [B,H] scratch for this direction
    int* __restrict__ counter,
    int* __restrict__ sense,
    int T, int B, int K, int H, int outF, int layer_dir_index,
    int reverse, int chunks, int has_bias)
{
    int b = blockIdx.x;
    int ch = blockIdx.y;
    int tid = threadIdx.x;
    int chunk_size = (H + chunks - 1) / chunks;
    int h_start = ch * chunk_size;
    int h_end = min(H, h_start + chunk_size);
    int total_blocks = B * chunks;
    int dir = reverse ? 1 : 0;

    // initialize hidden state from h0
    for (int h = h_start + tid; h < h_end; h += blockDim.x) {
        hstate[b * H + h] = h0[b * H + h];
    }
    grid_barrier(counter, sense, total_blocks);

    for (int step = 0; step < T; ++step) {
        int t = reverse ? (T - 1 - step) : step;
        const float* xrow = inp + ((t * B + b) * K);
        const float* hprev = hstate + b * H;

        for (int h = h_start + tid; h < h_end; h += blockDim.x) {
            int off_r = h;
            int off_z = H + h;
            int off_n = 2 * H + h;
            float ir = has_bias ? bih[off_r] : 0.0f;
            float iz = has_bias ? bih[off_z] : 0.0f;
            float in = has_bias ? bih[off_n] : 0.0f;
            // input contribution, transposed layout gives coalesced loads across h
            #pragma unroll 1
            for (int i = 0; i < K; ++i) {
                float xv = xrow[i];
                const float* wr = wihT + i * (3 * H);
                ir += xv * wr[off_r];
                iz += xv * wr[off_z];
                in += xv * wr[off_n];
            }
            float hr = has_bias ? bhh[off_r] : 0.0f;
            float hz = has_bias ? bhh[off_z] : 0.0f;
            float hnlin = has_bias ? bhh[off_n] : 0.0f;
            #pragma unroll 1
            for (int j = 0; j < H; ++j) {
                float hv = hprev[j];
                const float* wh = whhT + j * (3 * H);
                hr += hv * wh[off_r];
                hz += hv * wh[off_z];
                hnlin += hv * wh[off_n];
            }
            float r = sigmoidf_stable(ir + hr);
            float z = sigmoidf_stable(iz + hz);
            float n = tanhf(in + r * hnlin);
            float oldh = hprev[h];
            float newh = (1.0f - z) * n + z * oldh;
            hstate[b * H + h] = newh;
            out[(t * B + b) * outF + dir * H + h] = newh;
        }
        grid_barrier(counter, sense, total_blocks);
    }

    for (int h = h_start + tid; h < h_end; h += blockDim.x) {
        hn[(layer_dir_index * B + b) * H + h] = hstate[b * H + h];
    }
}

torch::Tensor gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                               std::vector<torch::Tensor> wihTs,
                               std::vector<torch::Tensor> whhTs,
                               std::vector<torch::Tensor> bihs,
                               std::vector<torch::Tensor> bhhs,
                               int64_t num_layers, bool batch_first, bool has_bias) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(h0.is_cuda(), "h0 must be CUDA");
    auto xc = x.contiguous();
    auto h0c = h0.contiguous();
    if (batch_first) {
        xc = xc.transpose(0,1).contiguous();
    }
    int T = (int)xc.size(0);
    int B = (int)xc.size(1);
    int input_size = (int)xc.size(2);
    int H = (int)h0c.size(2);
    auto opts = xc.options();
    auto hn = torch::empty({num_layers * 2, B, H}, opts);
    torch::Tensor bufA, bufB;
    const float* cur_ptr = xc.data_ptr<float>();
    int K = input_size;
    torch::Tensor cur_tensor = xc;

    const int chunks = 4;
    dim3 grid(B, chunks);
    dim3 block(128);
    auto counter = torch::zeros({1}, torch::TensorOptions().device(xc.device()).dtype(torch::kInt32));
    auto sense = torch::zeros({1}, torch::TensorOptions().device(xc.device()).dtype(torch::kInt32));
    auto hscratch = torch::empty({B, H}, opts);

    for (int l = 0; l < (int)num_layers; ++l) {
        int outF = 2 * H;
        torch::Tensor out = torch::empty({T, B, outF}, opts);
        K = (l == 0) ? input_size : (2 * H);
        for (int d = 0; d < 2; ++d) {
            int idx = l * 2 + d;
            // reset barrier state for this persistent kernel
            cudaMemsetAsync(counter.data_ptr<int>(), 0, sizeof(int), at::cuda::getCurrentCUDAStream());
            cudaMemsetAsync(sense.data_ptr<int>(), 0, sizeof(int), at::cuda::getCurrentCUDAStream());
            const float* bihp = has_bias ? bihs[idx].data_ptr<float>() : nullptr;
            const float* bhhp = has_bias ? bhhs[idx].data_ptr<float>() : nullptr;
            gru_layer_dir_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
                cur_tensor.data_ptr<float>(),
                h0c.data_ptr<float>() + ((idx * B) * H),
                wihTs[idx].data_ptr<float>(),
                whhTs[idx].data_ptr<float>(),
                bihp, bhhp,
                out.data_ptr<float>(),
                hn.data_ptr<float>(),
                hscratch.data_ptr<float>(),
                counter.data_ptr<int>(), sense.data_ptr<int>(),
                T, B, K, H, outF, idx, d == 1, chunks, has_bias ? 1 : 0);
        }
        cur_tensor = out;
    }
    return hn;
}
'''

_cpp_src = r'''
#include <torch/extension.h>
#include <vector>

torch::Tensor gru_forward_cuda(torch::Tensor x, torch::Tensor h0,
                               std::vector<torch::Tensor> wihTs,
                               std::vector<torch::Tensor> whhTs,
                               std::vector<torch::Tensor> bihs,
                               std::vector<torch::Tensor> bhhs,
                               int64_t num_layers, bool batch_first, bool has_bias);

torch::Tensor gru_forward(torch::Tensor x, torch::Tensor h0,
                          std::vector<torch::Tensor> wihTs,
                          std::vector<torch::Tensor> whhTs,
                          std::vector<torch::Tensor> bihs,
                          std::vector<torch::Tensor> bhhs,
                          int64_t num_layers, bool batch_first, bool has_bias) {
    return gru_forward_cuda(x, h0, wihTs, whhTs, bihs, bhhs, num_layers, batch_first, has_bias);
}
'''

_ext = load_inline(name='custom_gru_hn_ext_v1',
                   cpp_sources=[_cpp_src],
                   cuda_sources=[_cuda_src],
                   functions=['gru_forward'],
                   extra_cuda_cflags=['-O3', '--use_fast_math'],
                   extra_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=True)
        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_ready = False
        self._cache_ptrs = None
        self._wihTs = None
        self._whhTs = None
        self._bihs = None
        self._bhhs = None

    def _build_cache(self):
        wihTs = []
        whhTs = []
        bihs = []
        bhhs = []
        ptrs = []
        for l in range(self.num_layers):
            for suffix in ('', '_reverse'):
                wih = getattr(self.gru, f'weight_ih_l{l}{suffix}')
                whh = getattr(self.gru, f'weight_hh_l{l}{suffix}')
                ptrs.append(wih.data_ptr()); ptrs.append(whh.data_ptr())
                wihTs.append(wih.detach().t().contiguous())
                whhTs.append(whh.detach().t().contiguous())
                if self.bias:
                    bih = getattr(self.gru, f'bias_ih_l{l}{suffix}')
                    bhh = getattr(self.gru, f'bias_hh_l{l}{suffix}')
                    ptrs.append(bih.data_ptr()); ptrs.append(bhh.data_ptr())
                    bihs.append(bih.detach().contiguous())
                    bhhs.append(bhh.detach().contiguous())
        if not self.bias:
            dev = next(self.gru.parameters()).device
            empty = torch.empty((0,), device=dev, dtype=next(self.gru.parameters()).dtype)
            bihs = [empty for _ in range(self.num_layers * 2)]
            bhhs = [empty for _ in range(self.num_layers * 2)]
        self._wihTs = wihTs
        self._whhTs = whhTs
        self._bihs = bihs
        self._bhhs = bhhs
        self._cache_ptrs = tuple(ptrs)
        self._cache_ready = True

    def forward(self, x, h0):
        # Fast path for the benchmark configuration (CUDA fp32).  Parameters are cached as
        # transposed contiguous matrices for coalesced reads in the CUDA kernel.
        if (not self._cache_ready) or self._wihTs[0].device != x.device:
            self._build_cache()
        return _ext.gru_forward(x, h0, self._wihTs, self._whhTs, self._bihs, self._bhhs,
                                self.num_layers, self.batch_first, self.bias)