"""
Performance optimization utilities for FMUS-VID.

This module provides tools for improving processing speed and memory efficiency.
"""

import os
import logging
from typing import Dict, List, Optional, Tuple, Callable, Any
import numpy as np
from pathlib import Path
import threading
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import torch
import cupy
from dataclasses import dataclass
from collections import OrderedDict

logger = logging.getLogger(__name__)

@dataclass
class BatchConfig:
    """Configuration for batch processing."""
    batch_size: int = 32
    num_workers: int = os.cpu_count() or 1
    use_multiprocessing: bool = False
    prefetch_factor: int = 2

class BatchProcessor:
    """
    Handles batch processing of video frames.
    """

    def __init__(self, config: Optional[BatchConfig] = None):
        """
        Initialize batch processor.

        Args:
            config: Batch processing configuration
        """
        self.config = config or BatchConfig()
        self._executor = (
            ProcessPoolExecutor(max_workers=self.config.num_workers)
            if self.config.use_multiprocessing
            else ThreadPoolExecutor(max_workers=self.config.num_workers)
        )

    def process_frames(
        self,
        frames: List[np.ndarray],
        operation: Callable[[np.ndarray], np.ndarray]
    ) -> List[np.ndarray]:
        """
        Process frames in batches.

        Args:
            frames: List of input frames
            operation: Function to apply to each frame

        Returns:
            List of processed frames
        """
        results = []
        for i in range(0, len(frames), self.config.batch_size):
            batch = frames[i:i + self.config.batch_size]
            futures = [
                self._executor.submit(operation, frame)
                for frame in batch
            ]
            results.extend([future.result() for future in futures])
        return results

    def __del__(self):
        """Cleanup executor."""
        self._executor.shutdown()

class FrameCache:
    """
    Cache for frequently accessed frames.
    """

    def __init__(self, max_size_mb: int = 1024):
        """
        Initialize frame cache.

        Args:
            max_size_mb: Maximum cache size in megabytes
        """
        self.max_size = max_size_mb * 1024 * 1024  # Convert to bytes
        self._cache: OrderedDict[str, np.ndarray] = OrderedDict()
        self._size = 0
        self._lock = threading.Lock()

    def get(self, key: str) -> Optional[np.ndarray]:
        """
        Get frame from cache.

        Args:
            key: Cache key

        Returns:
            Cached frame if found, None otherwise
        """
        with self._lock:
            if key in self._cache:
                # Memindahkan item yang diakses ke akhir (MRU)
                frame = self._cache.pop(key)
                self._cache[key] = frame
                return frame
        return None

    def put(self, key: str, frame: np.ndarray) -> None:
        """
        Add frame to cache.

        Args:
            key: Cache key
            frame: Frame to cache
        """
        frame_size = frame.nbytes

        with self._lock:
            # Menghapus frame lama jika cache penuh
            while self._size + frame_size > self.max_size and self._cache:
                old_key, old_frame = self._cache.popitem(last=False)
                self._size -= old_frame.nbytes

            # Menambahkan frame baru
            self._cache[key] = frame
            self._size += frame_size

    def clear(self) -> None:
        """Clear cache."""
        with self._lock:
            self._cache.clear()
            self._size = 0

class LazyFrameLoader:
    """
    Lazy loading of video frames with preloading support.
    """

    def __init__(
        self,
        video_path: Path,
        preload_seconds: float = 2.0,
        max_preload_frames: int = 60
    ):
        """
        Initialize lazy frame loader.

        Args:
            video_path: Path to video file
            preload_seconds: Seconds of video to preload
            max_preload_frames: Maximum number of frames to preload
        """
        self.video_path = Path(video_path)
        self.preload_seconds = preload_seconds
        self.max_preload_frames = max_preload_frames
        self._preloaded: Dict[int, np.ndarray] = {}
        self._current_frame = 0
        self._preload_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()

    def get_frame(self, frame_number: int) -> np.ndarray:
        """
        Get frame by number, using preloaded frame if available.

        Args:
            frame_number: Frame number to retrieve

        Returns:
            Video frame
        """
        with self._lock:
            if frame_number in self._preloaded:
                frame = self._preloaded.pop(frame_number)
                self._current_frame = frame_number
                self._start_preload()
                return frame

        # Memuat frame jika tidak ada di cache
        frame = self._load_frame(frame_number)
        self._current_frame = frame_number
        self._start_preload()
        return frame

    def _start_preload(self) -> None:
        """Start preloading thread."""
        if (
            self._preload_thread is None
            or not self._preload_thread.is_alive()
        ):
            self._preload_thread = threading.Thread(
                target=self._preload_frames
            )
            self._preload_thread.daemon = True
            self._preload_thread.start()

    def _preload_frames(self) -> None:
        """Preload upcoming frames."""
        start = self._current_frame + 1
        for i in range(start, start + self.max_preload_frames):
            if len(self._preloaded) >= self.max_preload_frames:
                break
            try:
                frame = self._load_frame(i)
                with self._lock:
                    self._preloaded[i] = frame
            except Exception as e:
                logger.debug(f"Preload stopped at frame {i}: {e}")
                break

    def _load_frame(self, frame_number: int) -> np.ndarray:
        """
        Load single frame from video.

        Args:
            frame_number: Frame number to load

        Returns:
            Video frame
        """
        # Implementasi sebenarnya akan menggunakan backend video
        raise NotImplementedError

class GPUAccelerator:
    """
    GPU acceleration utilities.
    """

    def __init__(self):
        """Initialize GPU accelerator."""
        self.cuda_available = torch.cuda.is_available()
        self.cuda_device_count = (
            torch.cuda.device_count() if self.cuda_available else 0
        )

        try:
            import cupy as cp
            self.cupy_available = True
        except ImportError:
            self.cupy_available = False

    def get_optimal_device(self) -> str:
        """
        Get optimal compute device.

        Returns:
            Device name ('cuda', 'cpu')
        """
        if self.cuda_available and self.cuda_device_count > 0:
            return 'cuda'
        return 'cpu'

    def get_device_info(self) -> Dict[str, Any]:
        """
        Get detailed device information.

        Returns:
            Dictionary containing device information
        """
        info = {
            'cuda_available': self.cuda_available,
            'device_count': self.cuda_device_count,
            'cupy_available': self.cupy_available,
        }

        if self.cuda_available:
            info.update({
                'current_device': torch.cuda.current_device(),
                'device_name': torch.cuda.get_device_name(),
                'memory_allocated': torch.cuda.memory_allocated(),
                'memory_cached': torch.cuda.memory_reserved(),
            })

        return info

    def to_gpu(self, data: np.ndarray) -> Any:
        """
        Transfer data to GPU memory.

        Args:
            data: Input data

        Returns:
            Data in GPU memory
        """
        if self.cuda_available:
            if self.cupy_available:
                return cupy.asarray(data)
            return torch.from_numpy(data).cuda()
        return data

    def to_cpu(self, data: Any) -> np.ndarray:
        """
        Transfer data to CPU memory.

        Args:
            data: Input data

        Returns:
            Data in CPU memory
        """
        if isinstance(data, cupy.ndarray):
            return cupy.asnumpy(data)
        if torch.is_tensor(data):
            return data.cpu().numpy()
        return data

class MemoryManager:
    """
    Memory usage optimization.
    """

    def __init__(self, target_usage_mb: int = 1024):
        """
        Initialize memory manager.

        Args:
            target_usage_mb: Target memory usage in megabytes
        """
        self.target_usage = target_usage_mb * 1024 * 1024  # Convert to bytes
        self._cached_objects: Dict[str, Any] = {}

    def register_object(self, key: str, obj: Any) -> None:
        """
        Register object for memory management.

        Args:
            key: Object identifier
            obj: Object to manage
        """
        self._cached_objects[key] = obj

    def clear_cache(self) -> None:
        """Clear all cached objects."""
        self._cached_objects.clear()

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if cupy.cuda.runtime.getDeviceCount() > 0:
            mempool = cupy.get_default_memory_pool()
            mempool.free_all_blocks()
