import tensorflow.keras as keras
import json
import numpy as np


class MelodyGenerator:
    """
    This class represents a melody generator. It uses a pre-trained model to generate new melodies based on a given seed.
    """

    def __init__(self, model_path="model.h5", mapping_path="mapping.json", sequence_length=64):
        """
        Initializes the MelodyGenerator object.

        :param model_path: Path to the trained model (default: "model.h5").
        :param mapping_path: Path to the mapping file for symbols to integers (default: "mapping.json").
        :param sequence_length: The length of input sequences for the model (default: 64).
        """
        self.model = keras.models.load_model(model_path)  # Load the pre-trained model
        self.sequence_length = sequence_length  # Store the sequence length

        # Load the mappings from symbols (e.g., "60", "r", "_") to integers
        with open(mapping_path, "r") as fp:
            self._mappings = json.load(fp)

        # Initialize the seed with the start symbol "/" repeated for the sequence length
        self._start_symbols = ["/"] * sequence_length

    def generate_melody(self, seed, num_steps, max_sequence_length, temperature):
        """
        Generates a melody based on the given seed.

        :param seed: Initial sequence of musical symbols (e.g., "60 _ _ r").
        :param num_steps: Number of steps (time units) to generate.
        :param max_sequence_length: Maximum length of the input sequence for the model.
        :param temperature: Controls the randomness of the generated melody. Higher temperature -> more random.
        :return: The generated melody as a list of symbols.
        """
        seed = seed.split()  # Split the seed into individual symbols
        melody = seed  # Initialize the melody with the seed
        seed = self._start_symbols + seed  # Prepend start symbols to the seed

        # Convert seed symbols to their corresponding integer representation
        seed = [self._mappings[symbol] for symbol in seed]

        # Generate melody step by step
        for _ in range(num_steps):
            seed = seed[-max_sequence_length:]  # Keep only the last max_sequence_length elements
            onehot_seed = keras.utils.to_categorical(seed, num_classes=len(self._mappings))  # One-hot encode the seed
            onehot_seed = onehot_seed[np.newaxis, ...]  # Add a batch dimension

            # Predict probabilities for the next symbol
            probabilities = self.model.predict(onehot_seed)[0]

            # Sample the next symbol based on temperature
            output_int = self._sample_with_temperature(probabilities, temperature)
            seed.append(output_int)  # Add the new symbol to the seed

            # Convert the integer back to its symbol representation
            output_symbol = [k for k, v in self._mappings.items() if v == output_int][0]

            # Check for end of sequence symbol
            if output_symbol == "/":
                break

            melody.append(output_symbol)

        return melody  # Return the generated melody

    def _sample_with_temperature(self, probabilities, temperature):
        """
        Samples an index from the given probabilities with temperature adjustment.

        :param probabilities: List of probabilities for each symbol.
        :param temperature: The temperature for sampling.
        :return: The sampled index.
        """
        # Adjust probabilities with temperature
        predictions = np.log(probabilities) / temperature
        probabilities = np.exp(predictions) / np.sum(np.exp(predictions))

        # Sample an index from the adjusted probabilities
        choices = range(len(probabilities))
        index = np.random.choice(choices, p=probabilities)

        return index  # Return the sampled index


# Helper function to load a MelodyGenerator instance
def load_model(model_path="model.h5", mapping_path="mapping.json"):
    return MelodyGenerator(model_path, mapping_path)