import tensorflow as tf
from tensorflow.keras.layers import LSTM, Dense, Embedding
from tensorflow.keras.models import Sequential
import numpy as np

class MusicGenerator:
    def __init__(self, vocab_size):
        self.model = Sequential([
            Embedding(vocab_size, 128, input_length=1),
            LSTM(256, return_sequences=True),
            LSTM(256),
            Dense(vocab_size, activation='softmax')
        ])
        self.model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
        self.vocab_size = vocab_size

    def train(self, sequences, epochs=10):
        X = np.array([s[:-1] for s in sequences])
        y = np.array([s[1:] for s in sequences])
        self.model.fit(X, y, epochs=epochs)

    def generate(self, start_sequence, length=10):
        result = start_sequence
        current_sequence = np.array(start_sequence)
        for _ in range(length):
            predicted_probs = self.model.predict(current_sequence, verbose=0)[0]
            predicted_id = np.random.choice(len(predicted_probs), p=predicted_probs)
            result.append(predicted_id)
            current_sequence = np.array([predicted_id])
        return result

if __name__ == '__main__':
    # Example usage
    vocab_size = 12  # C, C#, D, ...
    generator = MusicGenerator(vocab_size)

    # Dummy training data (replace with real data)
    sequences = [
        [0, 4, 7, 0], # Cmaj -> C
        [5, 9, 0, 5]  # Fmaj -> F
    ]
    generator.train(sequences)

    # Generate a new sequence
    start_sequence = [0] # Start with C
    generated_sequence = generator.generate(start_sequence)
    print("Generated sequence:", generated_sequence)