import gradio as gr
import joblib
import json
import numpy as np
import pandas as pd

from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer

# טוען את המודל והקבצים
gmm = joblib.load("gmm_model.pkl")

with open("cluster_to_emotion.json", "r") as f:
    cluster_to_emotion = json.load(f)

# טוען את מאגר השירים
song_db = pd.read_parquet("hf://datasets/johanf/taylor-swift/data/train-00000-of-00001.parquet")
song_db = song_db[["lyrics", "title"]].dropna().drop_duplicates()
song_db["lyrics"] = song_db["lyrics"].str.strip()
song_db["title"] = song_db["title"].str.strip()
song_db = song_db.reset_index(drop=True)

# מחשב embedding לכל השירים
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
lyrics_list = song_db["lyrics"].tolist()
lyrics_embeddings = embedding_model.encode(lyrics_list, show_progress_bar=True)

# מודל להמרת טקסט לרגש
emotion_model = SentenceTransformer("j-hartmann/emotion-english-distilroberta-base")

def predict_emotion(text):
    embedding = emotion_model.encode([text])
    cluster = gmm.predict(embedding)[0]
    return cluster_to_emotion[str(cluster)]

def find_matching_song_by_emotion(user_input):
    emotion = predict_emotion(user_input)

    # מוצא שירים שמתאימים לרגש הזה
    candidates = song_db[song_db["lyrics"].str.lower().str.contains(emotion.lower())]

    if candidates.empty:
        candidates = song_db

    user_embedding = embedding_model.encode([user_input])
    candidate_lyrics = candidates["lyrics"].tolist()
    candidate_embeddings = embedding_model.encode(candidate_lyrics)

    similarities = cosine_similarity(user_embedding, candidate_embeddings)[0]
    top_idx = np.argmax(similarities)
    
    title = candidates.iloc[top_idx]["title"]
    lyrics_snippet = candidates.iloc[top_idx]["lyrics"][:200].replace("\n", " ")
    score = similarities[top_idx]

    return f"**{title}**  (match: {score:.2f})\n\n`{lyrics_snippet}...`\n\n_Emotion: {emotion}_"

demo = gr.Interface(
    fn=find_matching_song_by_emotion,
    inputs=gr.Textbox(placeholder="Tell me something that happened today"),
    outputs="markdown",
    title="Taylor Swift Mood Matcher",
    description="Tell me what you're feeling and I’ll match you with a Taylor Swift song that fits your mood."
)

demo.launch()