import gradio as gr
import joblib
import librosa
import numpy as np
import json

from librosa import feature as lf
model = joblib.load("genre_model.pkl")

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

# Full 55-feature extraction :(
def extract_features(file_path):
    y, sr = librosa.load(file_path, duration=30)
    
    # Chroma
    stft = np.abs(librosa.stft(y))
    chroma = librosa.feature.chroma_stft(S=stft, sr=sr)
    chroma_mean = np.mean(chroma)        
    chroma_var = np.var(chroma)         

    # Loudness (RMS)
    rms = librosa.feature.rms(y=y)
    loudness = np.mean(rms)
    loudness_variance = np.var(rms)

    # Spectral
    spectral_centroid = librosa.feature.spectral_centroid(y=y, sr=sr)
    spectral_centroid_mean = np.mean(spectral_centroid)
    spectral_centroid_var = np.var(spectral_centroid)

    spectral_bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr)
    spectral_bandwidth_var = np.var(spectral_bandwidth)

    rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)
    rolloff_var = np.var(rolloff)

    # Zero-crossing rate
    zcr = librosa.feature.zero_crossing_rate(y)
    zcr_mean = np.mean(zcr)
    zcr_var = np.var(zcr)

    # Harmony
    y_harm = librosa.effects.harmonic(y)
    harmony = librosa.feature.chroma_cqt(y=y_harm, sr=sr)
    harmony_mean = np.mean(harmony)
    harmony_var = np.var(harmony)

    perceptr_mean = 0
    perceptr_var = 0

    # Tempo
    onset_env = librosa.onset.onset_strength(y=y, sr=sr)
    
    tempo = lf.tempo(onset_envelope=onset_env, sr=sr)[0]

    # MFCCs
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=20)
    mfcc_mean = np.mean(mfcc, axis=1)    # length 20
    mfcc_var = np.var(mfcc, axis=1)      # length 20

    # Combine in correct order: total 55 features
    features = np.hstack([
        chroma_mean, chroma_var, loudness, loudness_variance,
        spectral_centroid_mean, spectral_centroid_var,
        spectral_bandwidth_var, rolloff_var,
        zcr_mean, zcr_var, harmony_mean, harmony_var,
        perceptr_mean, perceptr_var, tempo,
        mfcc_mean, mfcc_var
    ])
    
    return features.reshape(1, -1)  # shape (1,55)




def predict_genre(audio):
    features = extract_features(audio)
    prediction = model.predict(features)[0]
    return f"Predicted Genre is: {prediction}"  # no indexing needed



demo = gr.Interface(
    fn=predict_genre,
    inputs=gr.Audio(type="filepath", label="Upload a song clip"),
    outputs="text",
    title="GTZAN Music Genre Classifier",
    description="Upload a short audio clip to get the predicted genre."
)

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