import streamlit as st
import streamlit_pianoroll
from datasets import load_dataset
from piano_dataset import PianoTasks
from fortepyan import MidiPiece


def main():
    st.write("# PIANO dataset task review")

    available_tasks = PianoTasks.list_tasks()
    task_name = st.selectbox(
        label="Select PIANO task",
        options=available_tasks,
    )
    piano_task = PianoTasks.get_task(task_name=task_name)

    dataset = load_dataset("epr-labs/maestro-sustain-v2", split="test")
    record_idx = st.number_input(
        label="Record ID",
        min_value=0,
        max_value=len(dataset) - 1,
        value=12,
    )

    record = dataset[record_idx]
    piece = MidiPiece.from_huggingface(record)

    piece_split = piano_task.prompt_target_split(piece.df)

    source_piece = MidiPiece(df=piece_split.source_df)
    target_piece = MidiPiece(df=piece_split.target_df)

    streamlit_pianoroll.from_fortepyan(
        piece=source_piece,
        secondary_piece=target_piece,
    )


if __name__ == "__main__":
    main()