import streamlit as st
import pandas as pd
import numpy as np
from catboost import CatBoostRegressor

# Page 
st.set_page_config(
    page_title="BPM Predictor",
    page_icon="🎵",
    layout="centered"
)

# Feature Engineering 
def create_features(df):
    df = df.copy()
    # Energy - Rhythm
    df['Energy_x_Rhythm'] = df['Energy'] * df['RhythmScore']
    df['Energy_x_Loudness'] = df['Energy'] * df['AudioLoudness']
    df['Mood_x_Rhythm'] = df['MoodScore'] * df['RhythmScore']
    # Ratios
    df['Inst_Vocal_Ratio'] = df['InstrumentalScore'] / (df['VocalContent'] + 1)
    df['Loudness_Energy_Ratio'] = df['AudioLoudness'] / (df['Energy'] + 1)
    # milliseconds to minutes 
    df['Duration_Min'] = df['TrackDurationMs'] / 60000
    # Squared for importance
    df['Energy_Squared'] = df['Energy'] ** 2
    df['Rhythm_Squared'] = df['RhythmScore'] ** 2
    
    return df

# Load Model 
@st.cache_resource
def load_model():
    model = CatBoostRegressor()
    model.load_model("src/catboost_model.cbm")
    return model
try:
    model = load_model()
except Exception as e:
    st.error(f"Error loading model. Make sure 'catboost_model.cbm' is uploaded. Error: {e}")
    st.stop()

# UI 
st.title("🎵 Song BPM Predictor")
st.write("Enter the audio features of a song to predict its Beats Per Minute (BPM).")

# User Inputs 
with st.form("prediction_form"):
    st.subheader("Audio Features")
    
    col1, col2 = st.columns(2)
    with col1:
        track_duration = st.number_input("Track Duration (ms)", min_value=10000, max_value=1000000, value=200000, step=1000, help="Length of the song in milliseconds")
        energy = st.slider("Energy", 0.0, 1.0, 0.5)
        rhythm = st.slider("Rhythm Score", 0.0, 1.0, 0.5)
        loudness = st.slider("Audio Loudness (dB)", -60.0, 0.0, -10.0)
        vocal = st.slider("Vocal Content", 0.0, 1.0, 0.5)
    with col2:
        acoustic = st.slider("Acoustic Quality", 0.0, 1.0, 0.5)
        instrumental = st.slider("Instrumental Score", 0.0, 1.0, 0.5)
        live = st.slider("Live Performance Likelihood", 0.0, 1.0, 0.5)
        mood = st.slider("Mood Score", 0.0, 1.0, 0.5)
    submitted = st.form_submit_button("Predict BPM 🚀")

# Prediction 
if submitted:
    input_data = pd.DataFrame({
        'RhythmScore': [rhythm],
        'AudioLoudness': [loudness],
        'VocalContent': [vocal],
        'AcousticQuality': [acoustic],
        'InstrumentalScore': [instrumental],
        'LivePerformanceLikelihood': [live],
        'MoodScore': [mood],
        'TrackDurationMs': [track_duration],
        'Energy': [energy]})
    
    processed_data = create_features(input_data)
    
    prediction = model.predict(processed_data)[0]
    
    st.success(f"Predicted Tempo: **{prediction:.2f} BPM**")
    
    st.progress(min(prediction / 200, 1.0))