import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.0):
        """
        Initialize the LSTM model.

        :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
        :param output_size: The number of output features
        :param dropout: If non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer, with dropout probability equal to `dropout`
        """
        super(Model, self).__init__()
        # Initialize hidden state with random values
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout, bidirectional=True)
        self.fc = nn.Linear(hidden_size * 2, output_size)
    
    def forward(self, x,h0,c0):
        """
        Forward pass through the LSTM model.

        :param x: The input tensor, shape (batch_size, sequence_length, input_size)
        :return: The output tensor, shape (batch_size, sequence_length, output_size)
        """
        # Forward propagate LSTM
        out, hn = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size)
        
        # Decode the hidden state of the last time step
        out = self.fc(out[:, -1, :])  # out: tensor of shape (batch_size, output_size)
        
        return out

# Test code
batch_size = 10
sequence_length = 512
input_size = 128
hidden_size = 256
num_layers = 6
output_size = 10
dropout = 0.0

def get_inputs():
    return [torch.rand(batch_size, sequence_length, input_size),torch.rand((num_layers*2, batch_size, hidden_size)),torch.rand((num_layers*2, batch_size, hidden_size))]

def get_init_inputs():
    return [input_size, hidden_size, num_layers, output_size, dropout]

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));
}

__global__ void lstm_layer_kernel(
    const float* __restrict__ inp,
    const float* __restrict__ h0,
    const float* __restrict__ c0,
    const float* __restrict__ wihT_f,
    const float* __restrict__ wihT_b,
    const float* __restrict__ whhT_f,
    const float* __restrict__ whhT_b,
    const float* __restrict__ bias_f,
    const float* __restrict__ bias_b,
    float* __restrict__ out,
    int B, int T, int in_dim, int H, int layer)
{
    int b = blockIdx.x;
    int dir = blockIdx.y;
    int i = threadIdx.x;

    extern __shared__ float smem[];
    float* sh_h = smem;
    float* sh_c = smem + H;

    const float* wihT = (dir == 0) ? wihT_f : wihT_b;
    const float* whhT = (dir == 0) ? whhT_f : whhT_b;
    const float* bias = (dir == 0) ? bias_f : bias_b;

    if (i < H) {
        int hc_off = ((layer * 2 + dir) * B + b) * H + i;
        sh_h[i] = h0[hc_off];
        sh_c[i] = c0[hc_off];
    }
    __syncthreads();

    int t0 = (dir == 0) ? 0 : (T - 1);
    int tstep = (dir == 0) ? 1 : -1;
    for (int tt = 0, t = t0; tt < T; ++tt, t += tstep) {
        if (i < H) {
            float gi = bias[i];
            float gf = bias[H + i];
            float gg = bias[2 * H + i];
            float go = bias[3 * H + i];

            const float* xptr = inp + ((b * T + t) * in_dim);
            // W_ih is supplied transposed: [in_dim, 4H]
            for (int k = 0; k < in_dim; ++k) {
                float xv = xptr[k];
                const float* wk = wihT + k * (4 * H);
                gi += xv * wk[i];
                gf += xv * wk[H + i];
                gg += xv * wk[2 * H + i];
                go += xv * wk[3 * H + i];
            }
            // W_hh is supplied transposed: [H, 4H]
            for (int k = 0; k < H; ++k) {
                float hv = sh_h[k];
                const float* wk = whhT + k * (4 * H);
                gi += hv * wk[i];
                gf += hv * wk[H + i];
                gg += hv * wk[2 * H + i];
                go += hv * wk[3 * H + i];
            }

            float in_gate = sigmoidf_stable(gi);
            float forget_gate = sigmoidf_stable(gf);
            float cell_gate = tanhf(gg);
            float out_gate = sigmoidf_stable(go);
            float new_c = forget_gate * sh_c[i] + in_gate * cell_gate;
            float new_h = out_gate * tanhf(new_c);
            // Make sure all threads have finished reading the previous recurrent state.
            __syncthreads();
            sh_c[i] = new_c;
            sh_h[i] = new_h;
            out[((b * T + t) * (2 * H)) + dir * H + i] = new_h;
        } else {
            __syncthreads();
        }
        __syncthreads();
    }
}

__global__ void fc_kernel(
    const float* __restrict__ last,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ y,
    int B, int T, int H2, int O)
{
    int b = blockIdx.x;
    int o = blockIdx.y;
    int tid = threadIdx.x;
    __shared__ float s[256];
    float acc = 0.0f;
    const float* x = last + ((b * T + (T - 1)) * H2);
    for (int k = tid; k < H2; k += blockDim.x) {
        acc += x[k] * w[o * H2 + k];
    }
    s[tid] = acc;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) s[tid] += s[tid + stride];
        __syncthreads();
    }
    if (tid == 0) y[b * O + o] = s[0] + bias[o];
}

torch::Tensor lstm_fc_forward(
    torch::Tensor x,
    torch::Tensor h0,
    torch::Tensor c0,
    std::vector<torch::Tensor> wihT,
    std::vector<torch::Tensor> whhT,
    std::vector<torch::Tensor> bsum,
    torch::Tensor fc_w,
    torch::Tensor fc_b)
{
    const int B = (int)x.size(0);
    const int T = (int)x.size(1);
    const int input_size = (int)x.size(2);
    const int H = (int)h0.size(2);
    const int L = (int)wihT.size() / 2;
    const int O = (int)fc_w.size(0);

    auto opts = x.options();
    torch::Tensor cur;
    const float* input_ptr = x.data_ptr<float>();
    int in_dim = input_size;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    size_t shmem = (size_t)2 * H * sizeof(float);

    for (int l = 0; l < L; ++l) {
        auto out = torch::empty({B, T, 2 * H}, opts);
        dim3 grid(B, 2);
        int threads = H;
        lstm_layer_kernel<<<grid, threads, shmem, stream>>>(
            input_ptr,
            h0.data_ptr<float>(), c0.data_ptr<float>(),
            wihT[2*l].data_ptr<float>(), wihT[2*l+1].data_ptr<float>(),
            whhT[2*l].data_ptr<float>(), whhT[2*l+1].data_ptr<float>(),
            bsum[2*l].data_ptr<float>(), bsum[2*l+1].data_ptr<float>(),
            out.data_ptr<float>(),
            B, T, in_dim, H, l);
        cur = out;
        input_ptr = cur.data_ptr<float>();
        in_dim = 2 * H;
    }

    auto y = torch::empty({B, O}, opts);
    dim3 fgrid(B, O);
    fc_kernel<<<fgrid, 256, 0, stream>>>(cur.data_ptr<float>(), fc_w.data_ptr<float>(), fc_b.data_ptr<float>(), y.data_ptr<float>(), B, T, 2 * H, O);
    return y;
}
'''

_ext = load_inline(
    name='custom_lstm_last_fc_ext_v1',
    cpp_sources='',
    cuda_sources=_cuda_src,
    functions=['lstm_fc_forward'],
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3'],
    verbose=False,
)

class ModelNew(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.0):
        super(ModelNew, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout, bidirectional=True)
        self.fc = nn.Linear(hidden_size * 2, output_size)
        self._cache_ready = False
        self._wihT = None
        self._whhT = None
        self._bsum = None

    def _build_cache(self):
        wihT = []
        whhT = []
        bsum = []
        L = self.lstm.num_layers
        with torch.no_grad():
            for l in range(L):
                for suffix in ('', '_reverse'):
                    wih = getattr(self.lstm, f'weight_ih_l{l}{suffix}')
                    whh = getattr(self.lstm, f'weight_hh_l{l}{suffix}')
                    bih = getattr(self.lstm, f'bias_ih_l{l}{suffix}')
                    bhh = getattr(self.lstm, f'bias_hh_l{l}{suffix}')
                    wihT.append(wih.detach().t().contiguous())
                    whhT.append(whh.detach().t().contiguous())
                    bsum.append((bih.detach() + bhh.detach()).contiguous())
        self._wihT = wihT
        self._whhT = whhT
        self._bsum = bsum
        self._cache_ready = True

    def forward(self, x, h0, c0):
        if not self._cache_ready:
            self._build_cache()
        return _ext.lstm_fc_forward(x.contiguous(), h0.contiguous(), c0.contiguous(), self._wihT, self._whhT, self._bsum, self.fc.weight.contiguous(), self.fc.bias.contiguous())