import tensorflow as tf
from typing import List, Dict, Optional
from dataclasses import dataclass

from tqdm.auto import tqdm

@dataclass
class GPUMemoryStats:
    total: int
    used: int
    free: int

class GPUMemoryMonitor:
    """Monitor GPU memory usage with safe CPU fallback."""
    def __init__(self):
        self.has_gpu = False
        try:
            gpus = tf.config.list_physical_devices('GPU')
            self.has_gpu = len(gpus) > 0
        except:
            pass
            
    def get_memory_stats(self) -> Optional[GPUMemoryStats]:
        """Get current GPU memory statistics."""
        if not self.has_gpu:
            return None
            
        try:
            memory_info = tf.config.experimental.get_memory_info('GPU:0')
            return GPUMemoryStats(
                total=memory_info['peak'],
                used=memory_info['current'],
                free=memory_info['peak'] - memory_info['current']
            )
        except:
            return None
    
    def get_memory_usage(self) -> float:
        """Get current GPU memory usage as a percentage."""
        if not self.has_gpu:
            return 0.0
        stats = self.get_memory_stats()
        if stats is None or stats.total == 0:
            return 0.0
        return stats.used / stats.total

    def should_reduce_batch_size(self) -> bool:
        """Check if batch size should be reduced based on memory usage."""
        if not self.has_gpu:
            return False
        usage = self.get_memory_usage()
        return usage > 0.90

    def can_increase_batch_size(self) -> bool:
        """Check if batch size can be increased based on memory usage."""
        if not self.has_gpu:
            return True
        usage = self.get_memory_usage()
        return usage < 0.70