import streamlit as st
import numpy as np
import pandas as pd
import joblib
from predict import extract_features
import os
import tempfile
from huggingface_hub import hf_hub_download, list_repo_files
import logging
import traceback
import sklearn
import pkg_resources

# 版本检查
required_versions = {
    'numpy': '1.23.5',
    'scipy': '1.10.1',
    'scikit-learn': '1.2.2'
}

def check_versions():
    """检查包版本是否符合要求"""
    for package, required_version in required_versions.items():
        try:
            installed_version = pkg_resources.get_distribution(package).version
            if installed_version != required_version:
                logger.warning(f"Warning: {package} version mismatch. Required: {required_version}, Installed: {installed_version}")
                return False
        except pkg_resources.DistributionNotFound:
            logger.error(f"Error: {package} not found!")
            return False
    return True

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Page configuration
st.set_page_config(
    page_title="Healing Music Classifier",
    page_icon="🎵",
    layout="centered"
)

@st.cache_resource
def load_model():
    """Load model from Hugging Face Hub"""
    try:
        # 检查版本
        if not check_versions():
            logger.error("Package version requirements not met")
            return None, None

        # 创建临时目录
        os.makedirs("temp_models", exist_ok=True)

        # 下载模型文件
        try:
            model_path = hf_hub_download(
                repo_id="404Brain-Not-Found-yeah/healing-music-classifier",
                filename="models/model.joblib",
                local_dir="temp_models"
            )
            scaler_path = hf_hub_download(
                repo_id="404Brain-Not-Found-yeah/healing-music-classifier",
                filename="models/scaler.joblib",
                local_dir="temp_models"
            )
        except Exception as e:
            logger.error(f"Error downloading model files: {str(e)}")
            return None, None
        
        # 加载模型文件
        try:
            if not os.path.exists(model_path) or not os.path.exists(scaler_path):
                logger.error("Model files not found")
                return None, None

            # 尝试使用不同的pickle协议加载
            try:
                model = joblib.load(model_path)
                scaler = joblib.load(scaler_path)
            except Exception as load_error:
                logger.warning(f"Standard loading failed: {str(load_error)}")
                # 尝试使用兼容模式加载
                import pickle
                with open(model_path, 'rb') as f:
                    model = pickle.load(f, encoding='latin1')
                with open(scaler_path, 'rb') as f:
                    scaler = pickle.load(f, encoding='latin1')

            return model, scaler
        except Exception as e:
            logger.error(f"Error loading model/scaler files: {str(e)}")
            return None, None
            
    except Exception as e:
        logger.error(f"Unexpected error in load_model: {str(e)}")
        return None, None

def main():
    st.title("🎵 Healing Music Classifier")
    st.write("""
    Upload your music file, and AI will analyze its healing potential!
    Supports mp3, wav formats.
    """)

    # Add file upload component
    uploaded_file = st.file_uploader("Choose an audio file...", type=['mp3', 'wav'])

    if uploaded_file is not None:
        # Create progress bar
        progress_bar = st.progress(0)
        status_text = st.empty()

        try:
            # Create temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
                tmp_file.write(uploaded_file.getvalue())
                tmp_file_path = tmp_file.name

            # Update status
            status_text.text("Analyzing music...")
            progress_bar.progress(30)

            # Load model
            model, scaler = load_model()
            if model is None or scaler is None:
                st.error("Model loading failed. Please try again later.")
                return

            progress_bar.progress(50)

            # Extract features
            features = extract_features(tmp_file_path)
            if features is None:
                st.error("Failed to extract audio features. Please ensure the file is a valid audio file.")
                return

            progress_bar.progress(70)

            # Predict
            try:
                scaled_features = scaler.transform([features])
                healing_probability = model.predict_proba(scaled_features)[0][1]
                progress_bar.progress(90)
            except Exception as e:
                logger.error(f"Error during prediction: {str(e)}")
                st.error("Error during prediction. Please try again.")
                return

            # Display results
            st.subheader("Analysis Results")

            # Create visualization progress bar
            healing_percentage = healing_probability * 100
            st.progress(healing_probability)

            # Display percentage
            st.write(f"Healing Index: {healing_percentage:.1f}%")

            # Provide explanation
            if healing_percentage >= 75:
                st.success("This music has strong healing properties! 🌟")
            elif healing_percentage >= 50:
                st.info("This music has moderate healing effects. ✨")
            else:
                st.warning("This music has limited healing potential. 🎵")

        except Exception as e:
            st.error("An unexpected error occurred. Please try again.")
            logger.exception("Unexpected error")

        finally:
            # Clean up temporary file
            try:
                if 'tmp_file_path' in locals() and os.path.exists(tmp_file_path):
                    os.unlink(tmp_file_path)
            except Exception as e:
                logger.error(f"Failed to clean up temporary file: {str(e)}")

            # Complete progress bar
            progress_bar.progress(100)
            status_text.text("Analysis complete!")

if __name__ == "__main__":
    main()