import os
import zipfile
import shutil

def setup_dataset(zip_path="test.zip", extract_path="./dataset"):
    """Extract dataset if not already extracted"""
    dataset_path = os.path.join(extract_path, "test")
    
    if not os.path.exists(dataset_path):
        print("Extracting dataset...")
        os.makedirs(extract_path, exist_ok=True)
        
        try:
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(extract_path)
            print("Dataset extracted successfully!")
        except Exception as e:
            print(f"Error extracting dataset: {e}")
            return None
    
    return dataset_path

def cleanup_cache():
    """Clean up any temporary files"""
    cache_dirs = ['__pycache__', '.ipynb_checkpoints']
    for cache_dir in cache_dirs:
        if os.path.exists(cache_dir):
            shutil.rmtree(cache_dir)