import gradio as gr
import joblib
import numpy as np
import spotipy
from spotipy.oauth2 import SpotifyClientCredentials
import random

# Load the saved models and pipelines
anxiety_model = joblib.load('Anxiety_best_model.pkl')
anxiety_pipeline = joblib.load('Anxiety_best_pipeline.pkl')
depression_model = joblib.load('Depression_best_model.pkl')
depression_pipeline = joblib.load('Depression_best_pipeline.pkl')
insomnia_model = joblib.load('Insomnia_best_model.pkl')
insomnia_pipeline = joblib.load('Insomnia_best_pipeline.pkl')
ocd_model = joblib.load('OCD_best_model.pkl')
ocd_pipeline = joblib.load('OCD_best_pipeline.pkl')

# Spotify API credentials
SPOTIPY_CLIENT_ID = '79d5de7b9bec45c4bc2ae857e29d89e4'
SPOTIPY_CLIENT_SECRET = '410907e455a24118810bd89ee99d6f68'

# Initialize Spotify client
sp = spotipy.Spotify(auth_manager=SpotifyClientCredentials(client_id=SPOTIPY_CLIENT_ID,
                                                           client_secret=SPOTIPY_CLIENT_SECRET))

# Define the prediction functions
def predict(model, pipeline, inputs):
    inputs_array = np.array(inputs).reshape(1, -1)
    inputs_scaled = pipeline.transform(inputs_array)
    prediction = model.predict(inputs_scaled)
    return prediction[0]

# Function to recommend songs based on the condition and prediction
def recommend_songs(condition, prediction):
    mood_keywords = {
        "Anxiety": ["relaxing", "calming", "soothing"],
        "Depression": ["uplifting", "motivational", "cheerful"],
        "Insomnia": ["sleep", "ambient", "white noise"],
        "OCD": ["focus", "mindfulness", "meditation"]
    }
    
    # Select a random keyword based on the condition
    keyword = random.choice(mood_keywords[condition])
    
    # Adjust the search query based on the prediction value
    if prediction < 3:
        query = f"mild {keyword} music"
    elif prediction < 6:
        query = f"moderate {keyword} music"
    else:
        query = f"intense {keyword} music"

    results = sp.search(q=query, type='track', limit=5)
    songs = []
    for item in results['tracks']['items']:
        song = {
            'name': item['name'],
            'artist': item['artists'][0]['name'],
            'album': item['album']['name'],
            'image_url': item['album']['images'][0]['url'] if item['album']['images'] else None,
            'preview_url': item['preview_url']
        }
        songs.append(song)
    return songs

# Define the main prediction and recommendation function
def predict_and_recommend(condition, *inputs):
    if condition == "Anxiety":
        prediction = predict(anxiety_model, anxiety_pipeline, inputs)
    elif condition == "Depression":
        prediction = predict(depression_model, depression_pipeline, inputs)
    elif condition == "Insomnia":
        prediction = predict(insomnia_model, insomnia_pipeline, inputs)
    else:  # OCD
        prediction = predict(ocd_model, ocd_pipeline, inputs)
    
    songs = recommend_songs(condition, prediction)
    
    return prediction, f"The predicted {condition} level is {prediction:.2f}", songs

# Function to create HTML for song recommendations
def create_song_html(songs):
    if not songs or not isinstance(songs, list):
        return "No songs to display"
    
    html = "<div style='display: flex; flex-wrap: wrap; justify-content: space-around;'>"
    for song in songs:
        html += f"""
        <div style='width: 200px; margin: 10px; text-align: center;'>
            <img src='{song.get('image_url', '')}' style='width: 150px; height: 150px; object-fit: cover;'>
            <h3>{song.get('name', 'Unknown')}</h3>
            <p>{song.get('artist', 'Unknown Artist')}</p>
            <p>{song.get('album', 'Unknown Album')}</p>
            {f'<audio controls src="{song.get("preview_url", "")}"></audio>' if song.get('preview_url') else 'No preview available'}
        </div>
        """
    html += "</div>"
    return html

# Define the main prediction and recommendation function
def predict_and_recommend(condition, *inputs):
    if condition == "Anxiety":
        prediction = predict(anxiety_model, anxiety_pipeline, inputs)
    elif condition == "Depression":
        prediction = predict(depression_model, depression_pipeline, inputs)
    elif condition == "Insomnia":
        prediction = predict(insomnia_model, insomnia_pipeline, inputs)
    else:  # OCD
        prediction = predict(ocd_model, ocd_pipeline, inputs)
    
    songs = recommend_songs(condition, prediction)
    song_html = create_song_html(songs)
    
    return prediction, f"The predicted {condition} level is {prediction:.2f}", song_html

# Define the Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Music & Mental Health Predictor")
    gr.Markdown("### This innovative application predicts levels of anxiety, depression, insomnia, and OCD on a scale of 0 to 10, providing users with personalized insights into their mental health. Based on these predicted values, the app also recommends tailored music tracks designed to help manage and alleviate specific symptoms. By combining mental health assessment with mood-enhancing music recommendations, this tool offers a unique, holistic approach to supporting emotional well-being and self-care.")
    
    with gr.Row():
        condition = gr.Radio(["Anxiety", "Depression", "Insomnia", "OCD"], label="Select Condition")
    
    with gr.Row():
        with gr.Column():
            age = gr.Number(label="Age")
            hours_per_day = gr.Number(label="Hours per day listening to music")
            depression = gr.Number(label="Depression level (0-10)")
            insomnia = gr.Number(label="Insomnia level (0-10)")
            ocd = gr.Number(label="OCD level (0-10)")
            bpm = gr.Number(label="Preferred BPM")
        
        with gr.Column():
            prediction_value = gr.Number(label="Predicted Value")
            prediction_text = gr.Textbox(label="Prediction Description")
            song_recommendations = gr.HTML(label="Recommended Songs")
    
    predict_btn = gr.Button("Predict and Recommend Songs")
    
    predict_btn.click(
        fn=predict_and_recommend,
        inputs=[condition, age, hours_per_day, depression, insomnia, ocd, bpm],
        outputs=[prediction_value, prediction_text, song_recommendations],
    )

# Launch the interface
demo.launch(debug=True)