# Device setup utilities for GPU/CPU detection

import torch
import gc
from config import GPU_MEMORY_FRACTION


def setup_device():
    # GPU or CPU setup, Returns device object and bool indicating if GPU is available

    if torch.cuda.is_available():
        device = torch.device("cuda")
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"Using GPU: {gpu_name} ({gpu_memory:.1f}GB)")
        
        # Seting some fraction of memory to avoid OOM errors
        torch.cuda.set_per_process_memory_fraction(GPU_MEMORY_FRACTION)
        return device, True
    else:
        print("GPU not available, using CPU")
        return torch.device("cpu"), False


def clear_gpu_memory():

    """Clear GPU memory cache to free up VRAM"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()
        print("GPU memory cache cleared")


def get_gpu_memory_info():

    """Get current GPU memory usage information"""
    if torch.cuda.is_available():
        gpu_memory_used = torch.cuda.memory_allocated(0) / 1e9
        gpu_memory_total = torch.cuda.get_device_properties(0).total_memory / 1e9
        return gpu_memory_used, gpu_memory_total
    return 0, 0