import gradio as gr
from sagan_inference import run_sagan, z_score_standardize

AGE_CHOICES = ["kindergarten", "grade_6", "adult"]

def score_audio(audio, age_group):
    # Gradio passes `audio` as (tempfilepath, sr)
    wav_path = audio[0]
    raw = run_sagan(wav_path)
    calibrated = z_score_standardize(raw, age_group)
    # Display-friendly labels
    return {
        "Pitch Accuracy (0–1)": calibrated["pitch"],
        "Rhythm Consistency (0–1)": calibrated["rhythm"],
        "Timbre Score (0–1)": calibrated["timbre"],
    }

# with gr.Blocks() as demo:
#     gr.Markdown("## SAGAN Singing Scorer with Age-Calibrated Z-Scores")
#     audio_input = gr.Audio(source="upload", label="Upload Singing Clip")
#     age_input   = gr.Dropdown(AGE_CHOICES, label="Age Group")
#     output      = gr.Label(label="Calibrated Scores")
#     btn         = gr.Button("Score My Singing")
#     btn.click(fn=score_audio,
#               inputs=[audio_input, age_input],
#               outputs=output)

with gr.Blocks() as demo:
    gr.Markdown("## SAGAN Singing Scorer with Age-Calibrated Z-Scores")
    
    audio_input = gr.Audio(
        sources=["upload"],      # use the plural `sources`
        type="filepath",         # returns a local file path to your fn
        label="Upload Singing Clip"
    )
    
    age_input = gr.Dropdown(
        choices=AGE_CHOICES, 
        label="Age Group"
    )
    
    # output = gr.Label(
    #     label="Calibrated Scores"
    # )
    output = gr.JSON(
        label="Calibrated Scores"
    )
    
    btn = gr.Button("Score My Singing")
    btn.click(
        fn=score_audio,
        inputs=[audio_input, age_input],
        outputs=output
    )


if __name__ == "__main__":
    demo.launch()