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
        """
        super(Model, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers,
                            batch_first=True, dropout=dropout, bidirectional=False)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, h0=None, c0=None):
        """
        Forward pass through the LSTM model.

        :param x: The input tensor, shape (batch_size, sequence_length, input_size)
        :param h0: Optional initial hidden state (num_layers, batch_size, hidden_size)
        :param c0: Optional initial cell state (num_layers, batch_size, hidden_size)
        :return: The output tensor, shape (batch_size, output_size)
        """
        batch_size = x.size(0)

        if h0 is None:
            h0 = torch.randn(self.num_layers, batch_size, self.hidden_size, device=x.device)
        if c0 is None:
            c0 = torch.randn(self.num_layers, batch_size, self.hidden_size, device=x.device)

        out, _ = self.lstm(x, (h0, c0))  # out: (batch_size, seq_length, hidden_size)
        out = self.fc(out[:, -1, :])     # out: (batch_size, output_size)

        return out

# === Test configuration ===
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)]

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_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <vector>

torch::Tensor fc_last_cuda(torch::Tensor hy, torch::Tensor weight, torch::Tensor bias);

torch::Tensor lstm_fc(torch::Tensor x,
                      std::vector<torch::Tensor> hx,
                      std::vector<torch::Tensor> params,
                      torch::Tensor fc_w,
                      torch::Tensor fc_b,
                      int64_t num_layers,
                      double dropout,
                      bool train) {
    auto r = at::lstm(x, hx, params, true, num_layers, dropout, train, false, true);
    torch::Tensor hy = std::get<1>(r);
    return fc_last_cuda(hy, fc_w, fc_b);
}
'''

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

__global__ void fc_last_kernel(const float* __restrict__ hy,
                               const float* __restrict__ w,
                               const float* __restrict__ bias,
                               float* __restrict__ out,
                               int B, int H, int O, int L) {
    int b = blockIdx.x;
    int o = blockIdx.y;
    int tid = threadIdx.x;
    const float* h = hy + ((L - 1) * B + b) * H;
    const float* ww = w + o * H;
    float sum = 0.0f;
    for (int k = tid; k < H; k += blockDim.x) {
        sum = fmaf(h[k], ww[k], sum);
    }
    __shared__ float sh[256];
    sh[tid] = sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) sh[tid] += sh[tid + s];
        __syncthreads();
    }
    if (tid == 0) out[b * O + o] = sh[0] + bias[o];
}

torch::Tensor fc_last_cuda(torch::Tensor hy, torch::Tensor weight, torch::Tensor bias) {
    const int L = (int)hy.size(0);
    const int B = (int)hy.size(1);
    const int H = (int)hy.size(2);
    const int O = (int)weight.size(0);
    auto out = torch::empty({B, O}, hy.options());
    dim3 grid(B, O);
    fc_last_kernel<<<grid, 256, 0, at::cuda::getCurrentCUDAStream()>>>(
        hy.data_ptr<float>(), weight.data_ptr<float>(), bias.data_ptr<float>(), out.data_ptr<float>(), B, H, O, L);
    return out;
}
'''

_ext = load_inline(
    name='lstm_fc_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['lstm_fc'],
    extra_cflags=['-O3'],
    extra_cuda_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.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers,
                            batch_first=True, dropout=dropout, bidirectional=False)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, h0=None, c0=None):
        batch_size = x.size(0)
        if h0 is None:
            h0 = torch.randn(self.num_layers, batch_size, self.hidden_size, device=x.device)
        if c0 is None:
            c0 = torch.randn(self.num_layers, batch_size, self.hidden_size, device=x.device)
        return _ext.lstm_fc(x, [h0, c0], list(self.lstm._flat_weights), self.fc.weight, self.fc.bias,
                            self.num_layers, float(self.dropout), bool(self.training))