import os
from pathlib import Path
import pandas as pd
import numpy as np
import streamlit as st

from music_engine.matcher import MusicMatcher
from music_engine.image_processor import ImageProcessor

BASE_DIR = Path(__file__).resolve().parent

@st.cache_resource
def load_music_engine():
    # Инициализация базы данных и регрессора для музыкального мэтчинга
    db_path = BASE_DIR.parent / "dataset" / "DEAM" / "music_db.csv"
    model_path = BASE_DIR / "music_engine" / "va_regressor.pkl"
    
    if not db_path.exists():
        print(f"Музыкальная БД не найдена: {db_path}")
        return None
        
    return MusicMatcher(db_path=db_path, model_path=model_path)

@st.cache_resource
def load_image_processor():
    # Модуль обработки визуальных признаков
    model_path = BASE_DIR / "emoset_resnet50_best.pth"
    
    # Обработка пути при вызове из корневой директории
    if not model_path.exists():
        model_path = BASE_DIR.parent / "emoset_resnet50_best.pth"
        
    return ImageProcessor(model_path=model_path)

@st.cache_data
def load_emoset_data():
    # Выборка данных датасета для вкладки отладки
    dataset_root = BASE_DIR.parent / "dataset" / "EmoSet-118K" / "test"
    
    csv_path = dataset_root / "labels.csv"
    img_dir = dataset_root / "images"
    emb_path = BASE_DIR / "emoset_test_embeddings.npy"
    lbl_path = BASE_DIR / "emoset_test_labels.npy"

    if not all([csv_path.exists(), emb_path.exists(), lbl_path.exists()]):
        print("Тестовые файлы датасета не найдены, вкладка отладки может работать некорректно")
        return None, None, None, None

    labels_df = pd.read_csv(csv_path)
    
    test_filenames = labels_df['filename'].tolist()
    test_embeddings = np.load(emb_path)
    test_labels = np.load(lbl_path)
    
    return test_filenames, test_embeddings, test_labels, img_dir