import pandas as pd
from pathlib import Path
from tqdm import tqdm

# Конфигурация путей и целевых признаков
BASE_DIR = Path("../../dataset/DEAM")
MUSIC_DB_PATH = BASE_DIR / "music_db.csv"
FEATURES_DIR = BASE_DIR / "features" / "features"
OUTPUT_PATH = BASE_DIR / "music_db_enriched.csv"

# Маппинг низкоуровневых признаков экстрактора (openSMILE/GeMAPS) в дескрипторы системы
TARGET_FEATURES = {
    'pcm_RMSenergy_sma_amean': 'energy',
    'pcm_fftMag_spectralFlux_sma_amean': 'flux',
    'pcm_fftMag_spectralCentroid_sma_amean': 'centroid',
    'F0final_sma_amean': 'pitch',
    'logHNR_sma_amean': 'hnr',
    'pcm_zcr_sma_amean': 'zcr',
    'pcm_fftMag_spectralEntropy_sma_amean': 'entropy',
    'pcm_fftMag_psySharpness_sma_amean': 'sharpness'
}

def aggregate_acoustic_features():
    if not MUSIC_DB_PATH.exists():
        print(f"Базовый файл аннотаций не найден: {MUSIC_DB_PATH}")
        return

    print("Загрузка эмоциональной разметки DEAM...")
    df_main = pd.read_csv(MUSIC_DB_PATH)
    
    print("Агрегация фреймовых акустических признаков...")
    aggregated_data = []

    # Итерация по трекам для сбора покадровых характеристик
    for _, row in tqdm(df_main.iterrows(), total=len(df_main), desc="Обработка аудио-векторов"):
        song_id = int(row['song_id'])
        feature_file = FEATURES_DIR / f"{song_id}.csv"
        
        if feature_file.exists():
            try:
                # Чтение сырых векторов (формат csv с разделителем ';')
                df_feat = pd.read_csv(feature_file, sep=';')
                
                # Усреднение характеристик по временной оси (time frames)
                mean_features = df_feat[list(TARGET_FEATURES.keys())].mean()
                
                # Формирование агрегированной записи
                track_data = {'song_id': song_id}
                for orig_col, new_col in TARGET_FEATURES.items():
                    track_data[new_col] = mean_features[orig_col]
                    
                aggregated_data.append(track_data)
                
            except Exception as e:
                print(f"Ошибка парсинга файла {feature_file.name}: {e}")

    # Слияние акустических дескрипторов с эмоциональными координатами (Inner Join)
    df_features = pd.DataFrame(aggregated_data)
    df_enriched = pd.merge(df_main, df_features, on='song_id', how='inner')
    
    # Очистка возможных артефактов NaN после агрегации
    df_enriched = df_enriched.dropna(subset=list(TARGET_FEATURES.values()))
    
    df_enriched.to_csv(OUTPUT_PATH, index=False)
    print(f"Экспорт завершен. Сформирована обогащенная база: {OUTPUT_PATH.name}")
    print(f"Итоговый размер выборки: {len(df_enriched)} треков.")

if __name__ == "__main__":
    aggregate_acoustic_features()