import math
import random
import os
import numpy as np
from neurograd import xp, Tensor
from typing import Optional, List, Tuple, Union, Callable
import glob
from collections import deque



# Try to import DALI first - this determines our capabilities
try:
    import nvidia.dali as dali
    from nvidia.dali import pipeline_def, Pipeline
    import nvidia.dali.fn as fn
    import nvidia.dali.types as types
    from nvidia.dali.plugin.base_iterator import LastBatchPolicy
    DALI_AVAILABLE = True
except ImportError:
    DALI_AVAILABLE = False
    print("INFO: NVIDIA DALI not available. Falling back to OpenCV-based implementation.")
    print("      For maximum performance, install with: pip install --extra-index-url https://developer.download.nvidia.com/compute/redist nvidia-dali-cuda120")

# Always import OpenCV as fallback
try:
    import cv2
    # cv2.setNumThreads(1)  # Uncomment if needed for threading control
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    print("WARNING: OpenCV not available. Image loading will be limited.")

from concurrent.futures import ThreadPoolExecutor

# Image file extensions
IMG_EXTS = (
    '.png', '.jpg', '.jpeg', '.bmp', '.gif',
    '.tif', '.tiff', '.webp', '.jfif', '.avif',
    '.heif', '.heic'
)


class Dataset:
    """Base dataset class for simple tensor data"""
    def __init__(self, X, y, dtype=float32):
        assert len(X) == len(y), "Mismatched input and label lengths"
        self.X = Tensor(X, dtype=dtype)
        self.y = Tensor(y, dtype=dtype)
    
    def __len__(self):
        return len(self.X)
    
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]
    
    def __iter__(self):
        for idx in range(len(self)):
            yield self[idx]
    
    def shuffle(self, seed: Optional[int] = None):
        indices = list(range(len(self)))
        rng = random.Random(seed) if seed is not None else random.Random()
        rng.shuffle(indices)
        self.X = self.X[indices]
        self.y = self.y[indices]
    
    def __repr__(self):
        return f"<Dataset: {len(self)} samples, dtype={self.X.data.dtype}>"
    
    def __str__(self):
        preview_x = self.X[:1]
        preview_y = self.y[:1]
        return (f"Dataset:\n"
                f"  Total samples: {len(self)}\n"
                f"  Input preview: {preview_x}\n"
                f"  Target preview: {preview_y}")


class ImageFolder(Dataset):
    """
    Enhanced ImageFolder that uses DALI when available, OpenCV as fallback.
    
    Automatically detects DALI availability and optimizes accordingly.
    """
    def __init__(
        self,
        root: str,
        img_shape: tuple = None,          # (H, W)
        img_mode: str = "RGB",            # "RGB", "L", etc.
        img_normalize: bool = True,       # /255 -> float
        img_transform: Union[Pipeline, Callable, None] = None,  # DALI pipeline or callable
        target_transform: callable = None,
        img_dtype=xp.float32,
        target_dtype=xp.int64,
        chw: bool = True,                 # return CxHxW if True, else HxWxC
        device: str = "cpu",              # "cpu" or "gpu"/"cuda" (DALI only)
        num_threads: int = 4,             # DALI CPU threads
        prefetch_queue_depth: int = 2,    # DALI prefetch depth
        seed: int = 42
    ):
        self.root = root
        self.img_shape = img_shape
        self.img_mode = img_mode
        self.img_normalize = img_normalize
        self.img_transform = img_transform
        self.target_transform = target_transform
        self.img_dtype = img_dtype
        self.target_dtype = target_dtype
        self.chw = chw
        self.device = "cuda" if device in ["gpu", "cuda"] else "cpu"
        self.num_threads = num_threads
        self.prefetch_queue_depth = prefetch_queue_depth
        self.seed = seed

        # Determine which backend to use
        self.use_dali = DALI_AVAILABLE and isinstance(img_transform, (Pipeline, type(None)))
        if device in ["gpu", "cuda"] and not DALI_AVAILABLE:
            print("WARNING: GPU device requested but DALI not available. Using CPU with OpenCV fallback.")
            self.device = "cpu"

        self.images: List[str] = []
        self.targets: List[str] = []
        self._collect_paths()

        # Stable class mapping
        self.target_names = sorted(set(self.targets))
        self.target_mapping = {name: i for i, name in enumerate(self.target_names)}
        self.num_classes = len(self.target_names)
        
        # Convert targets to numeric labels
        self.numeric_targets = [self.target_mapping[t] for t in self.targets]

        print(f"ImageFolder initialized: {len(self)} samples, {self.num_classes} classes")
        print(f"Backend: {'DALI' if self.use_dali else 'OpenCV'}, Device: {self.device}")

    def _collect_paths(self):
        """Collect image paths and their class labels"""
        if os.path.exists(self.root) and os.path.isdir(self.root):
            # Method 1: Class folders (ImageNet style)
            for class_name in os.listdir(self.root):
                class_path = os.path.join(self.root, class_name)
                if not os.path.isdir(class_path):
                    continue
                
                for ext in IMG_EXTS:
                    pattern = os.path.join(class_path, f"*{ext}")
                    for img_path in glob.glob(pattern):
                        self.images.append(img_path)
                        self.targets.append(class_name)
            
            # Method 2: If no class folders found, walk directory tree
            if not self.images:
                for r, _, files in os.walk(self.root):
                    for f in files:
                        if f.lower().endswith(IMG_EXTS):
                            p = os.path.join(r, f)
                            cls = os.path.basename(os.path.dirname(p))
                            self.images.append(p)
                            self.targets.append(cls)

    def get_dali_pipeline(self, batch_size: int, shuffle: bool = True) -> Optional[Pipeline]:
        """Create DALI pipeline for this dataset (only when DALI is available)"""
        if not self.use_dali:
            return None
            
        @pipeline_def(batch_size=batch_size, num_threads=self.num_threads, device_id=0, seed=self.seed)
        def image_pipeline():
            # Create file reader
            if shuffle:
                images, labels = fn.readers.file(
                    files=self.images,
                    labels=self.numeric_targets,
                    random_shuffle=True,
                    initial_fill=batch_size * 4,  # Buffer size
                    name="Reader"
                )
            else:
                images, labels = fn.readers.file(
                    files=self.images,
                    labels=self.numeric_targets,
                    random_shuffle=False,
                    name="Reader"
                )

            # Decode images
            if self.device == "cuda":
                # Use mixed precision decoding (CPU decode, GPU processing)
                images = fn.decoders.image(images, device="mixed", output_type=types.RGB)
            else:
                images = fn.decoders.image(images, device="cpu", output_type=types.RGB)

            # Resize if needed
            if self.img_shape is not None:
                h, w = self.img_shape
                images = fn.resize(
                    images,
                    device=self.device,
                    size=[h, w],
                    interp_type=types.INTERP_LINEAR
                )

            # Handle grayscale conversion
            if self.img_mode.upper() in ("L", "GRAY", "GREY", "GRAYSCALE"):
                # Convert to grayscale
                images = fn.color_space_conversion(
                    images,
                    device=self.device,
                    image_type=types.RGB,
                    output_type=types.GRAY
                )

            # Apply custom transforms if provided
            if self.img_transform is not None:
                if isinstance(self.img_transform, Pipeline):
                    # If it's already a DALI pipeline, we'd need to integrate it
                    # This is complex, so for now we'll apply it post-DALI
                    pass

            # Normalize to [0, 1] if requested
            if self.img_normalize:
                images = fn.cast(images, device=self.device, dtype=types.FLOAT)
                images = images / 255.0
            else:
                images = fn.cast(images, device=self.device, dtype=types.FLOAT)

            # Convert to CHW format if requested
            if not self.chw:
                # DALI outputs CHW by default after transpose, so we need to handle HWC
                if len(images.shape) == 4:  # Batch dimension
                    images = fn.transpose(images, device=self.device, perm=[0, 2, 3, 1])
                else:
                    images = fn.transpose(images, device=self.device, perm=[1, 2, 0])

            return images, labels.gpu() if self.device == "cuda" else labels

        return image_pipeline()

    def _apply_img_transform(self, arr: np.ndarray) -> np.ndarray:
        """Apply image transforms (OpenCV fallback path)"""
        if self.img_transform is None:
            return arr
        
        # Skip DALI pipelines in fallback mode
        if DALI_AVAILABLE and isinstance(self.img_transform, Pipeline):
            print("WARNING: DALI Pipeline transform provided but using fallback mode. Transform skipped.")
            return arr
            
        # Try Albumentations-style call
        try:
            out = self.img_transform(image=arr)
            if isinstance(out, dict) and "image" in out:
                return out["image"]
            return out
        except TypeError:
            pass
        
        # Fallback: plain callable expecting ndarray
        return self.img_transform(arr)

    def _load_image_opencv(self, path: str) -> np.ndarray:
        """Load image using OpenCV (fallback implementation)"""
        if not CV2_AVAILABLE:
            raise RuntimeError("OpenCV not available for image loading")
            
        # OpenCV-only fast decode/resize
        mode = (self.img_mode or "RGB").upper()
        if mode in ("L", "GRAY", "GREY", "GRAYSCALE"):
            flag = cv2.IMREAD_GRAYSCALE
        elif mode == "RGBA":
            flag = cv2.IMREAD_UNCHANGED  # preserve alpha if present
        else:
            flag = cv2.IMREAD_COLOR  # BGR
        
        # Avoid EXIF orientation work
        try:
            flag |= cv2.IMREAD_IGNORE_ORIENTATION
        except Exception:
            pass

        arr = cv2.imread(path, flag)
        if arr is None:
            raise ValueError(f"Failed to read image: {path}")

        # Convert channel order to match RGB/RGBA expectations
        if mode == "RGB" and arr.ndim == 3 and arr.shape[2] == 3:
            arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
        elif mode == "RGBA" and arr.ndim == 3 and arr.shape[2] == 4:
            arr = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGBA)

        # Resize if requested (cv2 expects (W,H))
        if self.img_shape is not None:
            h, w = self.img_shape
            arr = cv2.resize(arr, (int(w), int(h)), interpolation=cv2.INTER_LINEAR)
        
        if arr.ndim == 2:
            arr = arr[:, :, None]
        
        # Apply transforms
        if self.img_transform:
            arr = self._apply_img_transform(arr)
        
        # Convert to CHW if requested
        if self.chw and arr.ndim == 3:
            arr = np.transpose(arr, (2, 0, 1))  # C,H,W
        
        # Normalize
        if self.img_normalize:
            arr = arr.astype(np.float32) / 255.0
        else:
            arr = arr.astype(np.float32)
        
        return arr

    def __getitem__(self, idx: int):
        """Get single item - uses appropriate backend"""
        img_path = self.images[idx]
        target = self.numeric_targets[idx]
        
        # Load image using appropriate backend
        if self.use_dali:
            # For DALI datasets, single item access uses OpenCV fallback
            # DALI is optimized for batch processing
            image = self._load_image_opencv(img_path)
        else:
            # Standard OpenCV loading
            image = self._load_image_opencv(img_path)
        
        # Apply target transform
        if self.target_transform:
            target = self.target_transform(target)
        
        return Tensor(image, dtype=self.img_dtype), Tensor(target, dtype=self.target_dtype)

    def shuffle(self, seed: Optional[int] = None):
        """Shuffle the dataset"""
        rng = random.Random(seed) if seed is not None else random.Random()
        combined = list(zip(self.images, self.targets, self.numeric_targets))
        rng.shuffle(combined)
        self.images, self.targets, self.numeric_targets = zip(*combined)
        self.images = list(self.images)
        self.targets = list(self.targets)
        self.numeric_targets = list(self.numeric_targets)

    def __len__(self):
        return len(self.images)

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]

    def __repr__(self):
        backend = "DALI" if self.use_dali else "OpenCV"
        shape = None
        if len(self) > 0:
            try:
                image, _ = self[0]
                shape = tuple(image.shape)
            except Exception:
                shape = None
        
        return (f"ImageFolder(root='{self.root}', samples={len(self)}, "
                f"classes={self.num_classes}, backend={backend}, device='{self.device}', "
                f"shape={shape}, img_dtype={self.img_dtype}, target_dtype={self.target_dtype}, "
                f"mode='{self.img_mode}', normalize={self.img_normalize}, chw={self.chw})")


class DataLoader:
    """
    Enhanced DataLoader with DALI acceleration when available, OpenCV fallback otherwise.
    
    Automatically detects capabilities and optimizes accordingly.
    """
    def __init__(
        self,
        dataset: Union[ImageFolder, Dataset],
        batch_size: int = 32,
        shuffle: bool = True,
        num_threads: int = 4,
        device: str = "cpu",
        prefetch_queue_depth: int = 2,
        drop_last: bool = False,
        seed: Optional[int] = None,
        num_workers: Optional[int] = None,  # For fallback threading
        prefetch_batches: int = 2           # For fallback prefetch
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.num_threads = num_threads
        self.device = "cuda" if device in ["gpu", "cuda"] else "cpu"
        self.prefetch_queue_depth = prefetch_queue_depth
        self.drop_last = drop_last
        self.seed = seed or 42
        self.prefetch_batches = max(0, int(prefetch_batches))
        
        # Determine threading for fallback mode
        if num_workers is None:
            cores = os.cpu_count() or 2
            self.num_workers = max(1, min(8, cores - 1))
        else:
            self.num_workers = int(num_workers)
        
        self._pipeline = None
        self._dali_iter = None
        self._executor: Optional[ThreadPoolExecutor] = None
        
        # Initialize DALI pipeline if available and applicable
        self.use_dali = (
            DALI_AVAILABLE and 
            isinstance(dataset, ImageFolder) and 
            dataset.use_dali
        )
        
        if self.use_dali:
            self._init_dali_pipeline()
        
        # Warn if GPU requested but not available
        if device in ["gpu", "cuda"] and not self.use_dali:
            print("WARNING: GPU device requested but DALI not available. Using CPU fallback.")
            self.device = "cpu"

    def _init_dali_pipeline(self):
        """Initialize the DALI pipeline"""
        if not self.use_dali or not isinstance(self.dataset, ImageFolder):
            return
        
        # Update dataset's DALI settings
        self.dataset.device = self.device
        self.dataset.num_threads = self.num_threads
        self.dataset.prefetch_queue_depth = self.prefetch_queue_depth
        self.dataset.seed = self.seed
        
        # Create pipeline
        self._pipeline = self.dataset.get_dali_pipeline(
            batch_size=self.batch_size,
            shuffle=self.shuffle
        )
        
        if self._pipeline is not None:
            # Build the pipeline
            self._pipeline.build()
            print(f"DALI pipeline initialized: batch_size={self.batch_size}, device={self.device}")

    def __len__(self):
        n = len(self.dataset)
        if self.drop_last:
            return n // self.batch_size
        return (n + self.batch_size - 1) // self.batch_size

    def __iter__(self):
        """Main iteration - uses DALI when available, falls back otherwise"""
        if self.use_dali and self._pipeline is not None:
            return self._dali_iterator()
        else:
            return self._regular_iterator()

    def _dali_iterator(self):
        """DALI-based iteration - maximum performance!"""
        # Reset pipeline for new epoch
        self._pipeline.reset()
        
        # Create DALI iterator
        last_batch_policy = LastBatchPolicy.DROP if self.drop_last else LastBatchPolicy.FILL
        
        class DALIGenericIterator:
            def __init__(self, pipeline, output_map, size, last_batch_policy):
                self.pipeline = pipeline
                self.output_map = output_map
                self.size = size
                self.last_batch_policy = last_batch_policy
                self._counter = 0
            
            def __iter__(self):
                return self
            
            def __next__(self):
                if self._counter >= self.size:
                    raise StopIteration
                
                try:
                    outputs = self.pipeline.run()
                    self._counter += 1
                    
                    # Extract images and labels
                    images = outputs[0].as_array() if hasattr(outputs[0], 'as_array') else outputs[0]
                    labels = outputs[1].as_array() if hasattr(outputs[1], 'as_array') else outputs[1]
                    
                    # Convert to CPU if on GPU
                    if self.pipeline.device == "cuda":
                        if hasattr(images, 'as_cpu'):
                            images = images.as_cpu().as_array()
                        if hasattr(labels, 'as_cpu'):
                            labels = labels.as_cpu().as_array()
                    
                    # Convert to Tensors
                    X = Tensor(images, dtype=self.dataset.img_dtype)
                    y = Tensor(labels, dtype=self.dataset.target_dtype)
                    
                    return X, y
                    
                except StopIteration:
                    raise StopIteration
        
        iterator = DALIGenericIterator(
            pipeline=self._pipeline,
            output_map=["images", "labels"],
            size=len(self),
            last_batch_policy=last_batch_policy
        )
        
        for batch in iterator:
            yield batch

    def _regular_iterator(self):
        """Regular iteration for non-DALI datasets with threading and prefetch"""
        batches = list(self._batch_indices())
        window = deque()
        next_to_submit = 0
        total = len(batches)

        # Prime the prefetch window
        pre = self.prefetch_batches if self.prefetch_batches > 0 else 0
        for _ in range(min(pre, total)):
            futs = self._schedule_batch(batches[next_to_submit])
            window.append(futs)
            next_to_submit += 1

        # Iterate in order; keep the window full
        for b in range(total):
            # If window is empty (prefetch=0) or drained, schedule current batch now
            if not window:
                futs = self._schedule_batch(batches[next_to_submit])
                window.append(futs)
                next_to_submit += 1

            futs = window.popleft()

            # Immediately schedule the next batch to keep the window full
            if next_to_submit < total and len(window) < self.prefetch_batches:
                next_futs = self._schedule_batch(batches[next_to_submit])
                window.append(next_futs)
                next_to_submit += 1

            # This blocks only if this batch isn't finished yet
            X, y = self._gather_batch(futs)
            yield X, y

    def _batch_indices(self):
        """Generate batch indices with optional shuffling"""
        n = len(self.dataset)
        order = list(range(n))
        if self.shuffle:
            rng = random.Random(self.seed) if self.seed is not None else random.Random()
            rng.shuffle(order)
        if self.drop_last:
            limit = (n // self.batch_size) * self.batch_size
        else:
            limit = n
        for start in range(0, limit, self.batch_size):
            end = min(start + self.batch_size, limit)
            yield order[start:end]

    def _ensure_executor(self):
        """Ensure thread executor is ready for fallback mode"""
        if self.num_workers > 0 and self._executor is None:
            self._executor = ThreadPoolExecutor(max_workers=self.num_workers)

    def _schedule_batch(self, idxs):
        """Schedule all sample loads in this batch and return the list of futures"""
        if self.num_workers > 0:
            self._ensure_executor()
            return [self._executor.submit(self.dataset.__getitem__, i) for i in idxs]
        else:
            # synchronous path for num_workers=0
            return [(self.dataset[i], None) for i in idxs]  # (result, None) to unify interface

    def _gather_batch(self, futures_or_results):
        """Block until the batch is ready, then stack into (X, y) Tensors"""
        if self.num_workers > 0:
            batch = [f.result() for f in futures_or_results]
        else:
            batch = [r for (r, _) in futures_or_results]

        Xs, ys = zip(*batch)
        X = Tensor(xp.stack([x.data for x in Xs], axis=0), dtype=Xs[0].dtype)
        y = Tensor(xp.stack([y.data for y in ys], axis=0), dtype=ys[0].dtype)
        return X, y

    def __getitem__(self, idx):
        """Get specific batch by index"""
        if idx < 0 or idx >= len(self):
            raise IndexError(f"Batch index {idx} out of range [0, {len(self)})")
        
        # Get all batch indices for consistent ordering
        batches = list(self._batch_indices())
        batch_idxs = batches[idx]
        
        # Load the batch synchronously
        batch_data = [self.dataset[i] for i in batch_idxs]
        Xs, ys = zip(*batch_data)
        
        # Stack into tensors
        X = Tensor(xp.stack([x.data for x in Xs], axis=0), dtype=Xs[0].dtype)
        y = Tensor(xp.stack([y.data for y in ys], axis=0), dtype=ys[0].dtype)
        return X, y

    def reset(self):
        """Reset the iterator - useful for multi-epoch training"""
        if self._pipeline is not None:
            self._pipeline.reset()

    def close(self):
        """Clean up resources"""
        if self._executor is not None:
            self._executor.shutdown(wait=True)
            self._executor = None

    def __repr__(self):
        backend = "DALI" if self.use_dali else "OpenCV+Threading"
        device_str = f", device='{self.device}'" if self.use_dali else ""
        return (f"<DataLoader: {len(self)} batches, "
                f"batch_size={self.batch_size}, shuffle={self.shuffle}, "
                f"backend={backend}{device_str}, "
                f"prefetch_depth={self.prefetch_queue_depth if self.use_dali else self.prefetch_batches}>")


# DALI transform creation helper (only available when DALI is installed)
def create_dali_transforms(
    device: str = "cpu",
    brightness: float = 0.0,
    contrast: float = 0.0,
    saturation: float = 0.0,
    hue: float = 0.0,
    horizontal_flip_prob: float = 0.0,
    vertical_flip_prob: float = 0.0,
    rotation_angle: float = 0.0
):
    """
    Create common DALI augmentation transforms.
    Returns None if DALI is not available.
    """
    if not DALI_AVAILABLE:
        print("WARNING: DALI not available. Transform creation skipped.")
        return None
        
    def apply_transforms(images):
        # Color jitter
        if brightness != 0.0 or contrast != 0.0 or saturation != 0.0 or hue != 0.0:
            images = fn.color_twist(
                images,
                device=device,
                brightness=fn.random.uniform(range=[-brightness, brightness]) if brightness > 0 else None,
                contrast=fn.random.uniform(range=[1-contrast, 1+contrast]) if contrast > 0 else None,
                saturation=fn.random.uniform(range=[1-saturation, 1+saturation]) if saturation > 0 else None,
                hue=fn.random.uniform(range=[-hue, hue]) if hue > 0 else None,
            )
        
        # Flips
        if horizontal_flip_prob > 0.0:
            images = fn.flip(images, device=device, horizontal=fn.random.coin_flip(probability=horizontal_flip_prob))
        if vertical_flip_prob > 0.0:
            images = fn.flip(images, device=device, vertical=fn.random.coin_flip(probability=vertical_flip_prob))
        
        # Rotation
        if rotation_angle != 0.0:
            images = fn.rotate(
                images,
                device=device,
                angle=fn.random.uniform(range=[-rotation_angle, rotation_angle]),
                fill_value=0
            )
        
        return images
    
    return apply_transforms


# Example usage demonstrating graceful fallback:
"""
# Basic usage - works with or without DALI
dataset = ImageFolder(
    root="/path/to/imagenet",
    img_shape=(224, 224),
    img_mode="RGB", 
    img_normalize=True,
    device="cuda" if DALI_AVAILABLE else "cpu",  # Automatically adapts
    num_threads=8,
    prefetch_queue_depth=3
)

dataloader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,
    device="cuda" if DALI_AVAILABLE else "cpu",  # Automatically adapts
    prefetch_queue_depth=3,
    num_threads=8
)

print(f"Using backend: {'DALI' if dataloader.use_dali else 'OpenCV+Threading'}")

# Training loop - works at maximum speed regardless of backend
for epoch in range(num_epochs):
    for batch_idx, (images, targets) in enumerate(dataloader):
        # Your training code here
        # Performance automatically optimized based on available libraries
        pass
    dataloader.reset()  # Reset for next epoch
    
# Clean up when done
dataloader.close()
"""