import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class MusicLSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, model='lstm', num_layers=1, dropout_p=0):
        super(MusicLSTM, self).__init__()
        self.model = model
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.num_layers = num_layers

        self.embeddings = nn.Embedding(input_size, hidden_size)
        if self.model == 'lstm':
            self.rnn = nn.LSTM(hidden_size, hidden_size, num_layers)
        elif self.model == 'gru':
            self.rnn = nn.GRU(hidden_size, hidden_size, num_layers)
        else:
            raise NotImplementedError

        self.out = nn.Linear(self.hidden_size, self.output_size)
        self.drop = nn.Dropout(p=dropout_p)

    def init_hidden(self, batch_size=1):
        """Initialize hidden states."""
        if self.model == 'lstm':
            self.hidden = (
                torch.zeros(self.num_layers, batch_size, self.hidden_size),
                torch.zeros(self.num_layers, batch_size, self.hidden_size)
            )
        elif self.model == 'gru':
            self.hidden = torch.zeros(self.num_layers, batch_size, self.hidden_size)

        return self.hidden

    def forward(self, x):
        """Forward pass."""
        # Ensure x is 2D (sequence length, batch size)
        if x.dim() > 2:
            x = x.squeeze()

        batch_size = 1 if x.dim() == 1 else x.size(0)
        x = x.long()

        # Embed the input
        embeds = self.embeddings(x)

        # Initialize hidden state if not already done
        if not hasattr(self, 'hidden'):
            self.init_hidden(batch_size)

        # Ensure embeds is 3D for RNN input (sequence length, batch size, embedding size)
        if embeds.dim() == 2:
            embeds = embeds.unsqueeze(1)

        # RNN processing
        rnn_out, self.hidden = self.rnn(embeds, self.hidden)

        # Dropout and output layer
        rnn_out = self.drop(rnn_out.squeeze(1))
        output = self.out(rnn_out)

        return output