from flask import Flask, render_template, request, jsonify
import joblib
import pandas as pd
import numpy as np
import librosa
from werkzeug.utils import secure_filename
import os

app = Flask(__name__)

app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # Max 16_mb file
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

model = joblib.load("music_genre_classifier.pkl")
scaler = joblib.load("scaler.pkl")
le = joblib.load("label_encoder.pkl")

# Helper function
def extract_features(file_path):
    y, sr = librosa.load(file_path, sr = 22050, duration = 30) # Max 30 sec audio
    
    # Calculate Features
    chroma_stft = librosa.feature.chroma_stft(y = y, sr = sr)
    rms = librosa.feature.rms(y = y)
    spec_cent = librosa.feature.spectral_centroid(y = y, sr = sr)
    spec_bw = librosa.feature.spectral_bandwidth(y = y, sr = sr)
    rolloff = librosa.feature.spectral_rolloff(y = y, sr = sr)
    zcr = librosa.feature.zero_crossing_rate(y)
    harmony, perceptr = librosa.effects.hpss(y)
    tempo, _ = librosa.beat.beat_track(y = y, sr = sr)
    mfcc = librosa.feature.mfcc(y = y, sr = sr, n_mfcc = 20)
    
    features = {
        'length': len(y),
        'chroma_stft_mean': np.mean(chroma_stft),
        'chroma_stft_var': np.var(chroma_stft),
        'rms_mean': np.mean(rms),
        'rms_var': np.var(rms),
        'spectral_centroid_mean': np.mean(spec_cent),
        'spectral_centroid_var': np.var(spec_cent),
        'spectral_bandwidth_mean': np.mean(spec_bw),
        'spectral_bandwidth_var': np.var(spec_bw),
        'rolloff_mean': np.mean(rolloff),
        'rolloff_var': np.var(rolloff),
        'zero_crossing_rate_mean': np.mean(zcr),
        'zero_crossing_rate_var': np.var(zcr),
        'harmony_mean': np.mean(harmony),
        'harmony_var': np.var(harmony),
        'perceptr_mean': np.mean(perceptr),
        'perceptr_var': np.var(perceptr),
        'tempo': tempo[0] if isinstance(tempo, np.ndarray) else tempo,
    }
    
    for i in range(1, 21):
        features[f'mfcc{i}_mean'] = np.mean(mfcc[i - 1])
        features[f'mfcc{i}_var'] = np.var(mfcc[i - 1])
        
    return features

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/predict', methods = ['POST'])
def predict():
    # Check if a file was uploaded
    if 'audio_file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400
        
    file = request.files['audio_file']
    
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400
        
    if file:
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        try:
            features_dict = extract_features(filepath)
            
            input_df = pd.DataFrame([features_dict])
            
            expected_features = scaler.feature_names_in_
            input_df = input_df[expected_features]
            
            scaled_data = scaler.transform(input_df)
            prediction_idx = model.predict(scaled_data)[0]
            
            probs = model.predict_proba(scaled_data)[0]
            confidence = np.max(probs) * 100
            
            genre = le.inverse_transform([prediction_idx])[0]
            
            # Clean up the uploaded file
            os.remove(filepath)
            
            return jsonify({
                'prediction': genre,
                'confidence': confidence
            })
            
        except Exception as e:
            if os.path.exists(filepath):
                os.remove(filepath)
            return jsonify({'error': str(e)}), 500

if __name__ == "__main__":
    app.run(host = "0.0.0.0", port = 7860)