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

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

torch::Tensor fc_last_cuda(torch::Tensor out, torch::Tensor fc_w, torch::Tensor fc_b);

torch::Tensor lstm_last_fc(torch::Tensor x, torch::Tensor h0, torch::Tensor c0,
                           std::vector<torch::Tensor> params,
                           torch::Tensor fc_w, torch::Tensor fc_b,
                           int64_t num_layers, double dropout) {
    auto hx = std::vector<torch::Tensor>{h0, c0};
    auto res = at::lstm(x, hx, params, true, num_layers, dropout, false, true, true);
    return fc_last_cuda(std::get<0>(res), fc_w, fc_b);
}
'''

_cuda = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>

__global__ void fc_last_kernel(const float* __restrict__ out,
                               const float* __restrict__ w,
                               const float* __restrict__ b,
                               float* __restrict__ y,
                               int B, int T, int F, int O,
                               long s0, long s1, long s2) {
    int bo = blockIdx.x;
    int batch = bo / O;
    int o = bo - batch * O;

    const float* xptr = out + batch * s0 + (T - 1) * s1;
    const float* wptr = w + o * F;
    float sum = 0.0f;

    #pragma unroll
    for (int i = threadIdx.x; i < F; i += 256) {
        sum += xptr[i * s2] * wptr[i];
    }

    __shared__ float smem[256];
    smem[threadIdx.x] = sum;
    __syncthreads();

    #pragma unroll
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) smem[threadIdx.x] += smem[threadIdx.x + stride];
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        y[batch * O + o] = smem[0] + b[o];
    }
}

torch::Tensor fc_last_cuda(torch::Tensor out, torch::Tensor fc_w, torch::Tensor fc_b) {
    c10::cuda::CUDAGuard device_guard(out.device());
    int B = (int)out.size(0);
    int T = (int)out.size(1);
    int F = (int)out.size(2);
    int O = (int)fc_w.size(0);
    auto y = torch::empty({B, O}, out.options());
    fc_last_kernel<<<B * O, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        out.data_ptr<float>(), fc_w.data_ptr<float>(), fc_b.data_ptr<float>(), y.data_ptr<float>(),
        B, T, F, O, out.stride(0), out.stride(1), out.stride(2));
    return y;
}
'''

_ext = load_inline(
    name='lstm_last_fc_ext_final',
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=['lstm_last_fc'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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.num_layers = num_layers
        self.dropout = float(dropout)

    def forward(self, x, h0, c0):
        return _ext.lstm_last_fc(x, h0, c0, self.lstm._flat_weights,
                                 self.fc.weight, self.fc.bias,
                                 self.num_layers, self.dropout)