import warnings
warnings.filterwarnings("ignore")
import pandas as pd
import numpy as np
import pickle

MODEL_PATH = "genre_model.pkl"
SCALER_PATH = "genre_scaler.pkl"
ENCODER_PATH = "genre_encoder.pkl"

NUMERICAL_FEATURES = [
    "melody_complexity (vocals)",
    "melody_range (vocals)",
    "melody_variability (vocals)",
    "tempo_bpm_original (mix)",
    "danceability custom (mix)",
    "loudness_integrated_lufs custom (mix)",
    "loudness_range_lu custom (mix)",
    "energy_librosa (mix)",
    "energy_librosa_std (mix)",
    "energy_essentia (mix)",
    "energy_essentia_std (mix)",
    "energy_combined (mix)",
    "spectral_centroid_mean custom (mix)",
    "mfcc_mean_1 (mix)",
    "mfcc_mean_2 (mix)",
    "chroma_mean (mix)",
    "spectral_contrast_mean (mix)",
    "repetition_score custom (mix)",
    "pitch_mean (mix)",
    "pitch_std (mix)",
    "rms_energy_mean (mix)",
    "rms_energy_std (mix)",
    "zero_crossing_rate (mix)",
]


def engineer_features(df, feature_cols):
    df = df.copy()
    df['energy_per_tempo'] = df['energy_combined (mix)'] / (df['tempo_bpm_original (mix)'] + 1)
    df['dance_energy_ratio'] = df['danceability custom (mix)'] * df['energy_combined (mix)']
    df['loudness_range_ratio'] = df['loudness_range_lu custom (mix)'] / (abs(df['loudness_integrated_lufs custom (mix)']) + 1)
    df['melody_energy'] = df['melody_variability (vocals)'] * df['energy_combined (mix)']
    df['spectral_complexity'] = df['spectral_centroid_mean custom (mix)'] * df['spectral_contrast_mean (mix)']
    df['mfcc_ratio'] = df['mfcc_mean_1 (mix)'] / (abs(df['mfcc_mean_2 (mix)']) + 1)
    df['rhythm_strength'] = df['tempo_bpm_original (mix)'] * df['danceability custom (mix)']
    df['pitch_variation'] = df['pitch_std (mix)'] / (df['pitch_mean (mix)'] + 1)
    df['rms_energy_ratio'] = df['rms_energy_mean (mix)'] / (df['rms_energy_std (mix)'] + 1)
    df['chroma_energy'] = df['chroma_mean (mix)'] * df['energy_combined (mix)']
    df['zero_tempo'] = df['zero_crossing_rate (mix)'] * df['tempo_bpm_original (mix)']
    df['tempo_category'] = np.where(df['tempo_bpm_original (mix)'] < 100, 0, np.where(df['tempo_bpm_original (mix)'] < 130, 1, 2))
    df['energy_category'] = np.where(df['energy_combined (mix)'] < 0.3, 0, np.where(df['energy_combined (mix)'] < 0.6, 1, 2))
    df['dance_category'] = np.where(df['danceability custom (mix)'] < 0.5, 0, np.where(df['danceability custom (mix)'] < 0.75, 1, 2))
    engineered = ['energy_per_tempo', 'dance_energy_ratio', 'loudness_range_ratio', 'melody_energy', 'spectral_complexity', 
                  'mfcc_ratio', 'rhythm_strength', 'pitch_variation', 'rms_energy_ratio', 'chroma_energy', 'zero_tempo',
                  'tempo_category', 'energy_category', 'dance_category']
    return df, feature_cols + engineered


class GenrePredictor:
    def __init__(self):
        with open(MODEL_PATH, "rb") as f:
            self.model = pickle.load(f)
        with open(SCALER_PATH, "rb") as f:
            self.scaler = pickle.load(f)
        with open(ENCODER_PATH, "rb") as f:
            self.genre_encoder, self.all_subgenres, self.all_features = pickle.load(f)
    
    def predict(self, feature_dict):
        df = pd.DataFrame([feature_dict])
        df, _ = engineer_features(df, NUMERICAL_FEATURES)
        
        for col in self.all_features:
            if col not in df.columns:
                df[col] = 0
        
        values = df[self.all_features].values[0]
        values = np.nan_to_num(values, nan=0, posinf=0, neginf=0)
        
        input_scaled = self.scaler.transform(values.reshape(1, -1))
        genre_idx = self.model.predict(input_scaled)[0]
        genre = self.genre_encoder.inverse_transform([genre_idx])[0]
        
        genre_probs = self.model.predict_proba(input_scaled)[0]
        top_indices = np.argsort(genre_probs)[::-1][:5]
        similar = [(self.genre_encoder.classes_[i], genre_probs[i]) for i in top_indices]
        
        related_subs = [s for s in self.all_subgenres if genre.lower() in s.lower()]
        if not related_subs:
            related_subs = self.all_subgenres[:10]
        
        return {
            "genre": genre,
            "similar_genres": similar,
            "subgenres": related_subs
        }


if __name__ == "__main__":
    predictor = GenrePredictor()
    
    # Example: Pass a row dict to predict
    song_features = {
        "melody_complexity (vocals)": 2.5,
        "melody_range (vocals)": 30.0,
        "melody_variability (vocals)": 0.55,
        "tempo_bpm_original (mix)": 140.0,
        "danceability custom (mix)": 0.70,
        "loudness_integrated_lufs custom (mix)": -12.0,
        "loudness_range_lu custom (mix)": 5.0,
        "energy_librosa (mix)": 0.5,
        "energy_librosa_std (mix)": 0.15,
        "energy_essentia (mix)": 0.3,
        "energy_essentia_std (mix)": 0.15,
        "energy_combined (mix)": 0.45,
        "spectral_centroid_mean custom (mix)": 0.13,
        "mfcc_mean_1 (mix)": 150.0,
        "mfcc_mean_2 (mix)": -20.0,
        "chroma_mean (mix)": 0.45,
        "spectral_contrast_mean (mix)": 20.0,
        "repetition_score custom (mix)": 0.007,
        "pitch_mean (mix)": 200.0,
        "pitch_std (mix)": 80.0,
        "rms_energy_mean (mix)": 0.5,
        "rms_energy_std (mix)": 0.15,
        "zero_crossing_rate (mix)": 0.05,
    }
    
    print("\n" + "=" * 60)
    print("GENRE PREDICTION WITH DICT INPUT")
    print("=" * 60)
    print("\nInput Dictionary:")
    for k, v in song_features.items():
        print(f"  {k}: {v}")
    
    result = predictor.predict(song_features)
    
    print("\n" + "=" * 60)
    print("PREDICTION RESULT")
    print("=" * 60)
    print(f"\n  GENRE: {result['genre']}")
    print("\n  Similar Genres:")
    for g, prob in result['similar_genres']:
        print(f"    - {g}: {prob:.1%}")
    print("\n  Sub-genres:")
    for sub in result['subgenres'][:8]:
        print(f"    - {sub}")