import numpy as np
import pandas as pd
import joblib
from pathlib import Path
from sklearn.metrics import mean_squared_error, r2_score

# 1. Настройка путей
embeddings_path = Path("./src/emoset_test_embeddings.npy")
csv_path = Path("./NFS/Thesis/Emoset/EmoSet-118K/test/labels.csv")
model_path = Path("./src/music_engine/va_regressor.pkl")

output_dir = Path("./src/scripts/media")
output_file = output_dir / "metrics_output.txt"

# 2. Корректный маппинг 8 классов EmoSet в шкалу DEAM [1.0, 9.0]
# Формула перевода из [-1, 1] в [1, 9]: 5.0 + (X * 4.0)
EMO_TO_VA = {
    "amusement":   [8.2, 6.6],  # Веселье (Высокий позитив, средняя энергия)
    "awe":         [7.0, 7.4],  # Восхищение (Позитив, высокая энергия)
    "contentment": [7.8, 3.4],  # Умиротворение (Позитив, низкая энергия)
    "excitement":  [8.2, 8.2],  # Возбуждение (Макс. позитив, макс. энергия)
    "anger":       [2.2, 7.8],  # Гнев (Глубокий негатив, высокая энергия)
    "disgust":     [2.6, 6.6],  # Отвращение (Негатив, средняя энергия)
    "fear":        [2.6, 8.2],  # Страх (Негатив, максимальная энергия)
    "sadness":     [2.2, 2.6]   # Грусть (Глубокий негатив, низкая энергия)
}

def generate_slide_metrics():
    print("[INFO] Загрузка тестовых артефактов...")
    
    if not all(p.exists() for p in [embeddings_path, csv_path, model_path]):
        print("[ERROR] Проверьте наличие файлов данных или модели регрессора.")
        return

    output_dir.mkdir(parents=True, exist_ok=True)

    # 3. Загрузка эмбеддингов и меток
    X_test = np.load(embeddings_path)
    df = pd.read_csv(csv_path)
    
    if len(X_test) != len(df):
        print(f"[WARN] Корректировка размеров выборки: Эмбеддинги ({len(X_test)}) != Метки ({len(df)})")
        min_len = min(len(X_test), len(df))
        X_test = X_test[:min_len]
        df = df.iloc[:min_len]

    y_test_list = [EMO_TO_VA.get(label.lower().strip(), [5.0, 5.0]) for label in df['label']]
    y_test = np.array(y_test_list)

    # 4. Выполнение инференса
    print("[INFO] Выполнение инференса регрессора на скрытом пространстве признаков...")
    regressor = joblib.load(model_path)
    y_pred = regressor.predict(X_test)
    
    # === БЛОК ДИАГНОСТИКИ ШКАЛЫ ===
    print("\n" + "-"*50)
    print(" ДИАГНОСТИКА ДИАПАЗОНОВ ЗНАЧЕНИЙ ".center(50))
    print("-"*50)
    print(f"Истинные (y_test) -> Мин: {y_test.min():.2f}, Макс: {y_test.max():.2f}, Среднее: {y_test.mean():.2f}")
    print(f"Предсказания (y_pred) -> Мин: {y_pred.min():.2f}, Макс: {y_pred.max():.2f}, Среднее: {y_pred.mean():.2f}")
    print("-"*50 + "\n")
    # ==============================
    
    # 5. Расчет метрик
    mse_v = mean_squared_error(y_test[:, 0], y_pred[:, 0])
    r2_v = r2_score(y_test[:, 0], y_pred[:, 0])
    
    mse_a = mean_squared_error(y_test[:, 1], y_pred[:, 1])
    r2_a = r2_score(y_test[:, 1], y_pred[:, 1])
    
    mse_total = mean_squared_error(y_test, y_pred)
    r2_total = r2_score(y_test, y_pred)
    
    # 6. Вывод и сохранение результатов
    table_content = f"""
==================================================
          ТАБЛИЦА МЕТРИК ДЛЯ СЛАЙДА 10          
==================================================
| Метрика    | Valence (V)  | Arousal (A)  | Общая (Total) |
|------------|--------------|--------------|---------------|
| MSE        | {mse_v:<12.4f} | {mse_a:<12.4f} | {mse_total:<13.4f} |
| R²         | {r2_v:<12.4f} | {r2_a:<12.4f} | {r2_total:<13.4f} |
==================================================

Формула целевой функции для вставки на слайд (LaTeX):
$$Score_{{final}} = D_{{emo}} + 4.0 \cdot Acoustic_{{penalty}}$$
"""

    print(table_content)

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(table_content)
        
    print(f"[SUCCESS] Метрики успешно сохранены в файл: {output_file.absolute()}")

if __name__ == "__main__":
    generate_slide_metrics()