import gradio as gr
from similarity import CLAPSimilarity
import matplotlib.pyplot as plt

# Initialize the CLAPSimilarity instance
similarity_calculator = CLAPSimilarity(training_embeddings_prefix='training')

def process_query(text, audio_file, max_tracks):
    if text and not audio_file:
        # Process text input
        similarity_scores = similarity_calculator.compute_similarity(
            input_data=text,
            input_type='text',
            max_tracks=int(max_tracks)  # Use user-defined max_tracks
        )
    elif audio_file and not text:
        # Process audio input
        similarity_scores = similarity_calculator.compute_similarity(
            input_data=audio_file,
            input_type='audio',
            max_tracks=int(max_tracks)  # Use user-defined max_tracks
        )
    else:
        return "Please provide either text or audio input."
    
    # Calculate the total sum of scores
    total_score = sum(similarity_scores.values())

    # Normalize the scores to sum to 100%
    normalized_scores = {
        filename: (score / total_score) * 100
        for filename, score in similarity_scores.items()
    }

    # Prepare the output data with normalized scores
    data = [
        [filename, round(score, 2)] 
        for filename, score in normalized_scores.items()
    ]
    return data

with gr.Blocks() as demo:
    with gr.Row():
        text_input = gr.Textbox(label="Enter text query")
        audio_input = gr.Audio(label="Upload audio", type="filepath")
        max_tracks_input = gr.Number(
            label="Max Tracks", value=10, precision=0
        )  # Add this input field

    with gr.Row():
        submit_btn = gr.Button("Submit")
    
    output_table = gr.Dataframe(
        headers=["Filename", "Score"],
        label="Similarity Results",
        datatype=["str", "number"],
        interactive=False
    )

    submit_btn.click(
        fn=process_query,
        inputs=[text_input, audio_input, max_tracks_input],  # Include max_tracks_input
        outputs=[output_table]
    )

demo.launch()