# utils.py
"""
Utility functions for the music separation project
"""

import os
import torch
import logging
from pathlib import Path
from config import LOG_DIR, MODEL_DIR
from datetime import datetime

def setup_logging():
    """Setup logging configuration"""
    log_file = LOG_DIR / f"music_separator_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )
    
    return logging.getLogger(__name__)

def save_model(model, model_name):
    """Save model weights"""
    try:
        model_path = MODEL_DIR / f"{model_name}.pth"
        torch.save(model.state_dict(), model_path)
        print(f"✅ Model saved to {model_path}")
        return True
    except Exception as e:
        print(f"❌ Error saving model: {str(e)}")
        return False

def load_model(model, model_name):
    """Load model weights"""
    try:
        model_path = MODEL_DIR / f"{model_name}.pth"
        if model_path.exists():
            model.load_state_dict(torch.load(model_path, map_location='cpu'))
            print(f"✅ Model loaded from {model_path}")
            return True
        else:
            print(f"⚠️ Model file {model_path} not found")
            return False
    except Exception as e:
        print(f"❌ Error loading model: {str(e)}")
        return False

def get_system_info():
    """Get comprehensive system information"""
    info = {
        'pytorch_version': torch.__version__,
        'cuda_available': torch.cuda.is_available(),
        'cuda_version': torch.version.cuda if torch.cuda.is_available() else 'N/A',
        'device_count': torch.cuda.device_count() if torch.cuda.is_available() else 0,
    }
    
    if torch.cuda.is_available():
        for i in range(torch.cuda.device_count()):
            info[f'cuda_device_{i}'] = torch.cuda.get_device_name(i)
            info[f'cuda_memory_{i}'] = f"{torch.cuda.get_device_properties(i).total_memory / 1024**3:.1f}GB"
    
    return info

def format_time(seconds):
    """Format seconds into human readable time"""
    if seconds < 60:
        return f"{seconds:.1f}s"
    elif seconds < 3600:
        return f"{seconds/60:.1f}m"
    else:
        return f"{seconds/3600:.1f}h"

def get_audio_duration(file_path):
    """Get audio file duration in seconds"""
    try:
        import soundfile as sf
        info = sf.info(file_path)
        return info.duration
    except:
        return 0