import json
import gradio as gr
import numpy as np

from taiko import (
    preprocess,
    generate_taiko_wav,
    COURSE
)

def handle(chart_path, music_path):
    data = json.loads(open(chart_path, "r").read())

    if len(data["data"]) < 0 and len(data["data"]) > 5:
        raise("Issue occur: Json data")

    for d in data["data"]:
        chart = preprocess(d["chart"])
        audio = generate_taiko_wav(chart, music_path)
        COURSE[d["course"]]["audio"] = audio
    
    c = 2
    if music_path is None: c = 1

    return \
        (COURSE[0]["audio"].frame_rate * c, np.array(COURSE[0]["audio"].get_array_of_samples())), \
        (COURSE[1]["audio"].frame_rate * c, np.array(COURSE[1]["audio"].get_array_of_samples())), \
        (COURSE[2]["audio"].frame_rate * c, np.array(COURSE[2]["audio"].get_array_of_samples())), \
        (COURSE[3]["audio"].frame_rate * c, np.array(COURSE[3]["audio"].get_array_of_samples())), \
        (COURSE[4]["audio"].frame_rate * c, np.array(COURSE[4]["audio"].get_array_of_samples()))

if __name__ == "__main__":
    inputs = [
        gr.File(label="太鼓達人譜面Json/Taiko Chart Json Data"),
        gr.File(label="譜面音樂/Chart Music (Optional)")
    ]
    outputs = [gr.Audio(label=course["label"]) for course in COURSE]

    demo = gr.Interface(fn=handle, inputs=inputs, outputs=outputs, title="`程設二作業HW0105 / Taiko Music Generator")
    demo.launch(share=True)