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")
    
    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)