import joblib
import numpy as np
import pandas as pd
from pathlib import Path

class MusicMatcher:
    def __init__(self, db_path: Path | str, model_path: Path | str):
        # Загрузка базы данных музыкальных произведений
        self.music_db = pd.read_csv(db_path)
        self.acoustic_features = ['energy', 'flux', 'centroid', 'pitch', 'hnr', 'zcr']
        
        # Удаление записей с пропущенными целевыми или акустическими признаками
        target_columns = ['valence', 'arousal'] + self.acoustic_features
        self.music_db = self.music_db.dropna(subset=target_columns)
        
        # Масштабирование акустических параметров к диапазону [0, 1]
        self.norm_db = self.music_db.copy()
        for feat in self.acoustic_features:
            f_min = self.norm_db[feat].min()
            f_max = self.norm_db[feat].max()
            if f_max > f_min:
                self.norm_db[f"norm_{feat}"] = (self.norm_db[feat] - f_min) / (f_max - f_min)
            else:
                self.norm_db[f"norm_{feat}"] = 0.0

        # Определение путей к аудиофайлам и загрузка модели регрессии
        self.audio_dir = Path(db_path).parent / "DEAM_audio" / "MEMD_audio"
        
        if Path(model_path).exists():
            self.regressor = joblib.load(model_path)
        else:
            self.regressor = None

    def predict_va(self, embedding: np.ndarray) -> tuple[float, float]:
        # Прогнозирование координат Valence/Arousal по визуальному эмбеддингу
        if not self.regressor:
            return 5.0, 5.0
            
        raw_prediction = self.regressor.predict(embedding.reshape(1, -1))[0]
        valence_pred = np.clip(raw_prediction[0], 1.0, 9.0)
        arousal_pred = np.clip(raw_prediction[1], 1.0, 9.0)
        
        return float(valence_pred), float(arousal_pred)

    def get_audio_path(self, song_id: int | float | str) -> Path | None:
        # Поиск физического пути к аудиофайлу в зависимости от расширения
        if not self.audio_dir.exists():
            return None
            
        clean_id = str(int(float(song_id)))
        for ext in ['.mp3', '.wav']:
            path = self.audio_dir / f"{clean_id}{ext}"
            if path.exists():
                return path
        return None

    def find_nearest_tracks(self, target_v: float, target_a: float, llm_profile: dict = None, top_k: int = 5) -> pd.DataFrame:
        # Расчет евклидова расстояния в эмоциональном пространстве Рассела
        v_dist = (self.norm_db['valence'] - target_v) ** 2
        a_dist = (self.norm_db['arousal'] - target_a) ** 2
        
        # Взвешенное расстояние с приоритетом оси активации (Arousal)
        self.norm_db['emo_distance'] = np.sqrt(1.0 * v_dist + 2.5 * a_dist)
        
        # Ранжирование только по эмоциональному критерию при отсутствии профиля LLM
        if not llm_profile:
            self.norm_db['final_score'] = self.norm_db['emo_distance']
            return self.norm_db.sort_values(by='final_score').head(top_k)

        # Расчет отклонений по вектору акустических параметров LLM
        acoustic_penalty = np.zeros(len(self.norm_db))
        for feat in self.acoustic_features:
            if feat in llm_profile:
                target_val = llm_profile[feat]
                acoustic_penalty += np.abs(self.norm_db[f"norm_{feat}"] - target_val)

        # Нормирование акустической дистанции
        self.norm_db['acoustic_distance'] = acoustic_penalty / len(self.acoustic_features)

        # Вычисление интегральной метрики соответствия (мультимодальный скоринг)
        self.norm_db['final_score'] = self.norm_db['emo_distance'] + (self.norm_db['acoustic_distance'] * 2.0)

        return self.norm_db.sort_values(by='final_score').head(top_k)