from typing import Optional

from note_seq.protobuf.music_pb2 import NoteSequence
from note_seq.constants import STANDARD_PPQ


def token_sequence_to_note_sequence(
    token_sequence: str,
    qpm: float = 120.0,
    use_program: bool = True,
    use_drums: bool = True,
    instrument_mapper: Optional[dict] = None,
    only_piano: bool = False,
) -> NoteSequence:
    """
    Converts a sequence of tokens into a sequence of notes.

    Args:
        token_sequence (str): The sequence of tokens to convert.
        qpm (float, optional): The quarter notes per minute. Defaults to 120.0.
        use_program (bool, optional): Whether to use program. Defaults to True.
        use_drums (bool, optional): Whether to use drums. Defaults to True.
        instrument_mapper (Optional[dict], optional): The instrument mapper. Defaults to None.
        only_piano (bool, optional): Whether to only use piano. Defaults to False.

    Returns:
        NoteSequence: The resulting sequence of notes.
    """
    if isinstance(token_sequence, str):
        token_sequence = token_sequence.split()

    note_sequence = empty_note_sequence(qpm)

    # Compute note and bar lengths based on the provided QPM
    note_length_16th = 0.25 * 60 / qpm
    bar_length = 4.0 * 60 / qpm

    # Render all notes.
    current_program = 1
    current_is_drum = False
    current_instrument = 0
    track_count = 0
    for _, token in enumerate(token_sequence):
        if token == "PIECE_START":
            pass
        elif token == "PIECE_END":
            break
        elif token == "TRACK_START":
            current_bar_index = 0
            track_count += 1
            pass
        elif token == "TRACK_END":
            pass
        elif token == "KEYS_START":
            pass
        elif token == "KEYS_END":
            pass
        elif token.startswith("KEY="):
            pass
        elif token.startswith("INST"):
            instrument = token.split("=")[-1]
            if instrument != "DRUMS" and use_program:
                if instrument_mapper is not None:
                    if instrument in instrument_mapper:
                        instrument = instrument_mapper[instrument]
                current_program = int(instrument)
                current_instrument = track_count
                current_is_drum = False
            if instrument == "DRUMS" and use_drums:
                current_instrument = 0
                current_program = 0
                current_is_drum = True
        elif token == "BAR_START":
            current_time = current_bar_index * bar_length
            current_notes = {}
        elif token == "BAR_END":
            current_bar_index += 1
            pass
        elif token.startswith("NOTE_ON"):
            pitch = int(token.split("=")[-1])
            note = note_sequence.notes.add()
            note.start_time = current_time
            note.end_time = current_time + 4 * note_length_16th
            note.pitch = pitch
            note.instrument = current_instrument
            note.program = current_program
            note.velocity = 80
            note.is_drum = current_is_drum
            current_notes[pitch] = note
        elif token.startswith("NOTE_OFF"):
            pitch = int(token.split("=")[-1])
            if pitch in current_notes:
                note = current_notes[pitch]
                note.end_time = current_time
        elif token.startswith("TIME_DELTA"):
            delta = float(token.split("=")[-1]) * note_length_16th
            current_time += delta
        elif token.startswith("DENSITY="):
            pass
        elif token == "[PAD]":
            pass
        else:
            pass

    # Make the instruments right.
    instruments_drums = []
    for note in note_sequence.notes:
        pair = [note.program, note.is_drum]
        if pair not in instruments_drums:
            instruments_drums += [pair]
        note.instrument = instruments_drums.index(pair)

    if only_piano:
        for note in note_sequence.notes:
            if not note.is_drum:
                note.instrument = 0
                note.program = 0

    return note_sequence


def empty_note_sequence(qpm: float = 120.0, total_time: float = 0.0) -> NoteSequence:
    """
    Creates an empty note sequence.

    Args:
        qpm (float, optional): The quarter notes per minute. Defaults to 120.0.
        total_time (float, optional): The total time. Defaults to 0.0.

    Returns:
        NoteSequence: The empty note sequence.
    """
    note_sequence = NoteSequence()
    note_sequence.tempos.add().qpm = qpm
    note_sequence.ticks_per_quarter = STANDARD_PPQ
    note_sequence.total_time = total_time
    return note_sequence