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
        }


df = pd.read_csv("all_genres_clean.csv",low_memory=False)
df_input = df[NUMERICAL_FEATURES]

model = GenrePredictor()
answer = model.predict(df_input.iloc[105148].to_dict())
print(answer)

print("Original :", df[["genre","sub_genres"]].iloc[105148])