app.py

import gradio as gr
import random
import os
from pathlib import Path
from typing import Dict
from storage import load_data, log_vote, get_leaderboard, get_vote_counts

from typing import Tuple, List, Optional, Any

# Paths
SCRIPT_DIR = Path(__file__).parent
# In HF Space, data/ is in the same dir as app.py. Locally, it's two levels up.
DATA_DIR = SCRIPT_DIR / "data" if (SCRIPT_DIR / "data").exists() else SCRIPT_DIR.parent.parent / "data"

MODELS_FILE = DATA_DIR / "models.json"
SENTENCES_FILE = DATA_DIR / "sentences.json"  # Transformed
ORIGINAL_SENTENCES_FILE = DATA_DIR / "original_sentences.json"  # Original

# Load Initial Data
models_data = load_data(str(MODELS_FILE))
sentences_data = load_data(str(SENTENCES_FILE))
original_sentences_data = load_data(str(ORIGINAL_SENTENCES_FILE))

# Exclude specific sentences that should not be picked in the arena
EXCLUDED_SENTENCE_IDS = {"sent-02", "sent-10", "sent-17"}
sentences_data = [s for s in sentences_data if s["id"] not in EXCLUDED_SENTENCE_IDS]
original_sentences_data = [s for s in original_sentences_data if s["id"] not in EXCLUDED_SENTENCE_IDS]

# Create a mapping for easy lookup
original_text_map = {s["id"]: s["text"] for s in original_sentences_data}
transformed_text_map = {s["id"]: s["text"] for s in sentences_data}

def pick_random_pair() -> Tuple[str, str, Optional[str], Optional[str], Optional[str], Optional[str], Optional[str]]:
    """
    Selects two random models and a sentence for comparison.
    """

    if len(models_data) < 2:
        return "Error", "Error", "Error", None, None, None, None
    
    sentence = random.choice(sentences_data)
    s_id = sentence["id"]
    
    # Find models that actually have audio files and are active
    available_models = []
    for m in models_data:
        if not m.get("active", True):  # Skip inactive models (default to active if not specified)
            continue
        audio_path = DATA_DIR / "audio" / m["id"] / f"{s_id}.wav"
        if audio_path.exists():
            available_models.append(m)

    if len(available_models) < 2:
         return "Error: Not enough models audio", "Error", "Error", None, None, None, None

    # Calculate weights: 1 / (votes + 1)
    # Higher weight for models with fewer votes
    vote_counts = get_vote_counts(models_data)
    
    def get_weight(m_id: str) -> float:
        # Inverse frequency weighting: 1 / (v + 1)
        return 1.0 / (vote_counts.get(m_id, 0) + 1)

    # Weighted sampling for the first model
    weights = [get_weight(m["id"]) for m in available_models]
    model_a = random.choices(available_models, weights=weights, k=1)[0]
    
    # Pick the second model from the remaining available models
    other_models = [m for m in available_models if m["id"] != model_a["id"]]
    # check if other_models is empty
    if not other_models:
         return "Error: Only one model available", "Error", "Error", None, None, None, None
         
    other_weights = [get_weight(m["id"]) for m in other_models]
    model_b = random.choices(other_models, weights=other_weights, k=1)[0]
    
    return (
        original_text_map.get(s_id, "Unknown"),
        transformed_text_map.get(s_id, "Unknown"),
        f"Model A", 
        f"Model B",
        model_a["id"],
        model_b["id"],
        s_id
    )

with gr.Blocks(title="Lithuanian TTS Arena") as demo:
    gr.Markdown("# 🇱🇹 Lithuanian TTS Arena")
    
    with gr.Tabs():
        with gr.Tab("Voting Arena"):
            gr.Markdown("Listen to two models and vote for the better one!")
            
            current_model_a = gr.State()
            current_model_b = gr.State()
            current_sentence_id = gr.State()
            
            with gr.Row():
                sentence_text = gr.Textbox(label="Original Lithuanian Sentence", interactive=False)
                
            with gr.Accordion("Phonetic Hint (Transformed Text)", open=False):
                phonetic_hint = gr.Textbox(label="Transformed", interactive=False)
                
            with gr.Row():
                out1 = gr.Audio(label="Model A", interactive=False)
                out2 = gr.Audio(label="Model B", interactive=False)
                
            with gr.Row():
                btn_a = gr.Button("👈 A is better")
                btn_tie = gr.Button("🤝 Tie")
                btn_b = gr.Button("B is better 👉")
                
            btn_pick = gr.Button("Next Pair / Start", variant="primary")

        with gr.Tab("Leaderboard"):
            gr.Markdown("## Current ELO Rankings")
            
            with gr.Row():
                filter_radio = gr.Radio(
                    choices=["Show All", "Original Only", "Transformed Only"],
                    value="Show All",
                    label="Filter Models"
                )
                
            leaderboard_df = gr.Dataframe(
                headers=["Model", "Version", "Glicko-2 Rating", "Votes", "Status"],
                datatype=["str", "str", "str", "number", "str"],
                interactive=False
            )
            btn_refresh = gr.Button("Refresh Leaderboard")


    def next_round() -> Tuple[str, str, Optional[str], Optional[str], Optional[str], Optional[str], Optional[str]]:
        """ Picks the next pair of models and audio to display. """
        results = pick_random_pair()
        if results[0].startswith("Error"):
             return results[0], None, None, None, None, None, None
        
        orig_text, trans_text, m_a_label, m_b_label, m_a_id, m_b_id, s_id = results
        
        # Construct absolute paths for Gradio
        path_a = DATA_DIR / "audio" / m_a_id / f"{s_id}.wav"
        path_b = DATA_DIR / "audio" / m_b_id / f"{s_id}.wav"
        
        audio_a = str(path_a) if path_a.exists() else None
        audio_b = str(path_b) if path_b.exists() else None
        
        if not audio_a and not audio_b:
             print(f"Warning: No audio files found for sentence {s_id}")

        return orig_text, trans_text, audio_a, audio_b, m_a_id, m_b_id, s_id

    def submit_vote(winner_type: str, model_a: str, model_b: str, sentence_id: str) -> Tuple[str, str, Optional[str], Optional[str], Optional[str], Optional[str], Optional[str]]:
        """ Logs the vote and proceeds to the next round. """
        if not model_a or not model_b or not sentence_id:
            return next_round()
        log_vote(model_a, model_b, winner_type, sentence_id)
        return next_round()

    def update_leaderboard(filter_val: str) -> List[List[Any]]:
        """ Refreshes and filters the leaderboard data. """
        full_leaderboard = get_leaderboard(models_data)
        if filter_val == "Original Only":
            return [row for row in full_leaderboard if row[1] == "Original"]
        elif filter_val == "Transformed Only":
            return [row for row in full_leaderboard if row[1] == "Transformed"]
        return full_leaderboard

    btn_pick.click(
        fn=next_round, 
        outputs=[sentence_text, phonetic_hint, out1, out2, current_model_a, current_model_b, current_sentence_id]
    )
    
    btn_a.click(
        fn=lambda a, b, s: submit_vote("model_a", a, b, s),
        inputs=[current_model_a, current_model_b, current_sentence_id],
        outputs=[sentence_text, phonetic_hint, out1, out2, current_model_a, current_model_b, current_sentence_id]
    )

    btn_b.click(
        fn=lambda a, b, s: submit_vote("model_b", a, b, s),
        inputs=[current_model_a, current_model_b, current_sentence_id],
        outputs=[sentence_text, phonetic_hint, out1, out2, current_model_a, current_model_b, current_sentence_id]
    )
    
    btn_tie.click(
        fn=lambda a, b, s: submit_vote("tie", a, b, s),
        inputs=[current_model_a, current_model_b, current_sentence_id],
        outputs=[sentence_text, phonetic_hint, out1, out2, current_model_a, current_model_b, current_sentence_id]
    )
    
    btn_refresh.click(
        fn=update_leaderboard,
        inputs=[filter_radio],
        outputs=leaderboard_df
    )
    
    filter_radio.change(
        fn=update_leaderboard,
        inputs=[filter_radio],
        outputs=leaderboard_df
    )
    
    # Load leaderboard on start
    demo.load(fn=update_leaderboard, inputs=[filter_radio], outputs=leaderboard_df)

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