import os
from functools import lru_cache
from typing import Dict, List, Optional, Tuple, Union

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import ConcatDataset, Dataset

from logger import get_logger
from .polarization import PolarizationProcessor

logger = get_logger(__name__).set_context("DATASET")


class RGBP_Dataset(Dataset):
    """
    Optimized version of RGBP_Dataset dataset with performance improvements:
    1. Reduced tensor/numpy conversions
    2. Optional caching
    3. Simplified processing pipeline
    4. Batch-friendly operations
    5. Consistent image sizing for batching
    6. Support for both single file clockwise and separate polarization files
    7. Flexible scene filtering with include/exclude patterns
    8. RGB-only mode for datasets without polarization data
    9. Per-scene frame subsampling via sample_every_n (load 1/N frames)

    Polarization formats supported:
    - "single_file_clock": Single file with 4 polarization images arranged clockwise
    - "separate_files": Four separate files ending with _000, _045, _090, _135
    - "separate_files_stokes": Three separate .npy files with Stokes parameters _S0, _S1, _S2
    - "single_file_topdown": Single file with 4 polarization images arranged vertically (top to bottom)

    Scene filtering:
    - include: str or List[str] - Include scenes that match (substring or exact)
    - exclude: str or List[str] - Exclude scenes that match (substring or exact)

    RGB-only mode:
    - load_rgb_only: bool - If True, forces loading only RGB data and ignores polarization data
    
    File paths:
    - return_filepaths: bool - If True, includes 'filepaths' dict in output with keys 'raw_path', 'pol_path', 'diffuse_path', 'intrinsics_path'
    """

    def __init__(
        self,
        root_dir: str,
        rho_s: float = 0.6,
        eps: float = 1e-8,
        rgb_ext: str = ".png",
        pol_ext: str = ".npy",
        polarization_format: str = "single_file_clock",  # "single_file_clock", "separate_files" or "mosaic"
        rgb_dir_name: str = "rgb",
        pol_dir_name: str = "pol",
        diffuse_dir_name: str = "diffuse",
        intrinsics_file_name: str = "intrinsics.txt",
        # Polarization data format
        # Image sizing parameters
        transform=None,
        target_size: Optional[Tuple[int, int]] = (874, 1132),  # Default size (H, W)
        resize_mode: str = "crop",  # "crop", "resize", "pad", or "resize+crop"
        # Performance optimizations
        use_cache: bool = True,
        cache_size: int = 100,
        simplify_upsampling: bool = True,  # Use simple bicubic instead of edge-aware
        precompute_stokes: bool = False,  # Precompute Stokes parameters
        # Reduced smoothing options
        smooth_specular: bool = True,
        gaussian_sigma: float = 0.0,
        dolp_min_intensity: float = 0.05,
        dolp_min_value: float = 0.04,
        # Scene filtering
        include: Optional[Union[str, List[str]]] = None,
        exclude: Optional[Union[str, List[str]]] = None,
        # Few images mode for quick testing
        few_images: bool = False,
        overfit_test: bool = False,
        # Per-folder subsampling (load 1/N images per scene folder)
        sample_every_n: int = 1,
        # RGB-only mode
        load_rgb_only: bool = False,  # Force loading only RGB data, ignore polarization
        # Return file paths in output
        return_filepaths: bool = True,  # If True, include file paths in returned dictionary
        # Deprecated parameters (for backward compatibility)
        # Highlight detection (optional)
        highlight_enable: bool = False,
        highlight_brightness_threshold: float = 0.93,
        highlight_return_mask: bool = False,
        highlight_rect_size: Optional[Tuple[int, int]] = None,
        highlight_return_rect: bool = False,
        highlight_return_rect_as_rgb: bool = False,
    ):
        self.root_dir = os.path.expandvars(root_dir)
        self.rho_s = rho_s
        self.eps = eps
        self.rgb_ext = rgb_ext
        self.pol_ext = pol_ext
        self.transform = transform
        self.load_rgb_only = load_rgb_only
        self.return_filepaths = return_filepaths
        self.rgb_dir_name = rgb_dir_name
        self.pol_dir_name = pol_dir_name
        self.diffuse_dir_name = diffuse_dir_name
        self.intrinsics_file_name = intrinsics_file_name

        # Polarization format validation
        self.polarization_format = polarization_format.lower()
        if self.polarization_format not in [
            "single_file_clock",
            "separate_files",
            "separate_files_stokes",
            "single_file_topdown",
            "mosaic",
        ]:
            raise ValueError(
                f"polarization_format must be one of ['single_file_clock', 'separate_files', 'separate_files_stokes', 'single_file_topdown', 'mosaic'], got {polarization_format}"
            )

        # Image sizing parameters
        self.target_size = target_size
        self.resize_mode = resize_mode.lower()
        if self.resize_mode not in ["crop", "resize", "pad", "resize+crop"]:
            raise ValueError(
                f"resize_mode must be one of ['crop', 'resize', 'pad', 'resize+crop'], got {resize_mode}"
            )

        self.use_cache = use_cache
        self.simplify_upsampling = simplify_upsampling
        self.precompute_stokes = precompute_stokes

        # Simplified smoothing config
        self.smooth_specular = smooth_specular
        self.gaussian_sigma = gaussian_sigma
        self.dolp_min_intensity = dolp_min_intensity
        self.dolp_min_value = dolp_min_value

        # Store filtering parameters
        self.include = include
        self.exclude = exclude or []
        self.few_images = few_images
        # Validate and store subsampling factor
        if not isinstance(sample_every_n, int) or sample_every_n < 1:
            raise ValueError(
                f"sample_every_n must be an integer >= 1, got {sample_every_n}"
            )
        self.sample_every_n = sample_every_n

        self.scene_pairs = self._find_scene_pairs()

        # Limit to 32 samples if few_images is True
        if self.few_images and len(self.scene_pairs) > 8:
            original_count = len(self.scene_pairs)
            self.scene_pairs = self.scene_pairs[:8]
            logger.info(
                f"Few images mode: Limited dataset from {original_count} to 8 samples"
            )

        # Setup caching if enabled
        if self.use_cache:
            self._cache_intrinsics = {}
            # Make cached methods
            self._load_intrinsics_cached = lru_cache(maxsize=cache_size)(
                self._load_intrinsics_impl
            )
            if self.precompute_stokes:
                self._load_pol_cached = lru_cache(maxsize=cache_size)(
                    self._load_and_process_polarization_impl
                )
        self.overfit_test = overfit_test
        if self.overfit_test:
            self.scene_pairs = len(self.scene_pairs) * [self.scene_pairs[0]]
            logger.info(
                f"Overfit test mode: Limited dataset from {len(self.scene_pairs)} to 1 sample"
            )

        # Highlight detection configuration
        self.highlight_enabled = highlight_enable
        self.highlight_brightness_threshold = highlight_brightness_threshold
        self.highlight_return_mask = highlight_return_mask
        self.highlight_rect_size = highlight_rect_size
        self.highlight_return_rect = highlight_return_rect
        self.highlight_return_rect_as_rgb = highlight_return_rect_as_rgb

        # Polarization processor (delegates polarization loading logic)
        self._pol_processor = PolarizationProcessor(
            rho_s=self.rho_s,
            eps=self.eps,
            dolp_min_intensity=self.dolp_min_intensity,
            dolp_min_value=self.dolp_min_value,
        )

    def _resize_tensor(
        self, tensor: torch.Tensor, target_size: Tuple[int, int]
    ) -> torch.Tensor:
        """
        Resize tensor to target size using specified mode.

        Args:
            tensor: Input tensor of shape [C, H, W] or [H, W]
            target_size: Target size as (H, W)

        Returns:
            Resized tensor of shape [C, target_H, target_W] or [target_H, target_W]
            
        Resize modes:
            - "crop": Center crop to target size (no resizing)
            - "resize": Resize to target size (may distort aspect ratio)
            - "pad": Pad to target size with zeros (no resizing)
            - "resize+crop": Resize to fit target size while maintaining aspect ratio, then center crop
        """
        if tensor.dim() == 2:
            # Single channel tensor [H, W]
            tensor = tensor.unsqueeze(0)  # [1, H, W]
            was_2d = True
        else:
            was_2d = False

        current_size = tensor.shape[-2:]  # [H, W]

        if self.resize_mode == "crop":
            # Center crop to target size
            if current_size[0] > target_size[0]:
                start_h = (current_size[0] - target_size[0]) // 2
                end_h = start_h + target_size[0]
                tensor = tensor[..., start_h:end_h, :]
            if current_size[1] > target_size[1]:
                start_w = (current_size[1] - target_size[1]) // 2
                end_w = start_w + target_size[1]
                tensor = tensor[..., :, start_w:end_w]

        elif self.resize_mode == "resize":
            # Resize to target size using bilinear interpolation
            tensor = F.interpolate(
                tensor.unsqueeze(0),  # Add batch dimension
                size=target_size,
                mode="bilinear",
                align_corners=False,
            ).squeeze(0)  # Remove batch dimension

        elif self.resize_mode == "resize+crop":
            # Resize to fit target size while maintaining aspect ratio, then center crop
            target_h, target_w = target_size
            current_h, current_w = current_size
            
            # Calculate scale factor to fit the image into target size
            scale_h = target_h / current_h
            scale_w = target_w / current_w
            scale = max(scale_h, scale_w)  # Use max to ensure we can crop to target size
            
            # Calculate intermediate size after scaling
            intermediate_h = int(current_h * scale)
            intermediate_w = int(current_w * scale)
            
            # Resize to intermediate size
            tensor = F.interpolate(
                tensor.unsqueeze(0),  # Add batch dimension
                size=(intermediate_h, intermediate_w),
                mode="bilinear",
                align_corners=False,
            ).squeeze(0)  # Remove batch dimension
            
            # Center crop to target size
            start_h = (intermediate_h - target_h) // 2
            end_h = start_h + target_h
            start_w = (intermediate_w - target_w) // 2
            end_w = start_w + target_w
            
            tensor = tensor[..., start_h:end_h, start_w:end_w]

        elif self.resize_mode == "pad":
            # Pad to target size with zeros
            pad_h = max(0, target_size[0] - current_size[0])
            pad_w = max(0, target_size[1] - current_size[1])

            if pad_h > 0 or pad_w > 0:
                # Pad format: (pad_left, pad_right, pad_top, pad_bottom)
                pad_left = pad_w // 2
                pad_right = pad_w - pad_left
                pad_top = pad_h // 2
                pad_bottom = pad_h - pad_top

                tensor = F.pad(
                    tensor,
                    (pad_left, pad_right, pad_top, pad_bottom),
                    mode="constant",
                    value=0,
                )

        if was_2d:
            tensor = tensor.squeeze(0)  # Remove channel dimension

        return tensor

    def _resize_raw_tensor(self, raw: torch.Tensor) -> torch.Tensor:
        """
        Resize RAW tensor to target size.

        Args:
            raw: RAW tensor of shape [3, H, W] in [0, 1]

        Returns:
            Resized RAW tensor of shape [3, target_H, target_W] in [0, 1]
        """
        if self.target_size is None:
            return raw

        return self._resize_tensor(raw, self.target_size)

    def _resize_polarization_data(
        self, pol_data: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Resize all polarization data tensors to target size.

        Args:
            pol_data: Dictionary containing polarization tensors with shapes [C, H, W]
                     Keys include 'I0', 'I45', 'I90', 'I135', 'S0', 'S1', 'S2', 'DoLP', etc.

        Returns:
            Dictionary with same keys but tensors resized to target size [C, target_H, target_W]
        """
        """
        Resize all polarization data tensors to target size.

        Args:
            pol_data: Dictionary containing polarization data tensors

        Returns:
            Dictionary with all tensors resized to target size
        """
        if self.target_size is None:
            return pol_data

        resized_data = {}
        for key, tensor in pol_data.items():
            resized_data[key] = self._resize_tensor(tensor, self.target_size)
        return resized_data

    def _should_include_scene(self, scene_name: str) -> bool:
        """
        Check if a scene should be included based on include/exclude filters.

        Args:
            scene_name: Name of the scene to check

        Returns:
            True if scene should be included, False otherwise
        """
        """
        Check if a scene should be included based on include/exclude filters.

        Args:
            scene_name: Name of the scene to check

        Returns:
            True if scene should be included, False otherwise
        """
        # Check exclude filter first
        if self.exclude:
            if isinstance(self.exclude, str):
                # Single string - check if scene name contains the exclude string
                if self.exclude in scene_name:
                    return False
            elif isinstance(self.exclude, list):
                # List of strings - check exact matches and substring matches
                for exclude_pattern in self.exclude:
                    if exclude_pattern == scene_name or exclude_pattern in scene_name:
                        return False

        # Check include filter
        if self.include is not None:
            if isinstance(self.include, str):
                # Single string - check if scene name contains the include string
                return self.include in scene_name
            elif isinstance(self.include, list):
                # List of strings - check exact matches and substring matches
                for include_pattern in self.include:
                    if include_pattern == scene_name or include_pattern in scene_name:
                        return True
                # If include list is provided but no match found, exclude the scene
                return False

        # If no include filter specified and not excluded, include the scene
        return True

    def _find_scene_pairs(self) -> List[Tuple[str, Optional[str], Optional[str], Optional[str], bool]]:
        """
        Find matching RAW (RGB), optional diffuse, polarization, and intrinsics entries per scene.

        Scans the root directory structure to find corresponding files:
        - RGB images in rgb/ subdirectory
        - Polarization data in pol/ subdirectory (optional)
        - Camera intrinsics in intrinsics/ subdirectory (optional)

        Returns:
            List of tuples (raw_path, pol_path, diffuse_path, intrinsics_path, has_pol_data)
        """
        """Find valid entries based on polarization format and optional diffuse folder"""
        scene_pairs = []
        for scene_name in os.listdir(self.root_dir):
            # Use new filtering logic
            if not self._should_include_scene(scene_name):
                continue

            scene_path = os.path.join(self.root_dir, scene_name)
            if not os.path.isdir(scene_path):
                continue

            rgb_dir = os.path.join(scene_path, self.rgb_dir_name)
            pol_dir = os.path.join(scene_path, self.pol_dir_name)
            diffuse_dir = (
                os.path.join(scene_path, self.diffuse_dir_name)
                if self.diffuse_dir_name is not None
                else None
            )
            intrinsics_path = os.path.join(scene_path, self.intrinsics_file_name)

            # Check if RGB directory exists (required)
            if not os.path.exists(rgb_dir):
                continue

            # Check if polarization directory exists (optional)
            has_pol_data = os.path.exists(pol_dir)

            # Set intrinsics_path to None if file doesn't exist
            if not os.path.exists(intrinsics_path):
                intrinsics_path = None

            rgb_files = sorted(
                [f for f in os.listdir(rgb_dir) if f.endswith(self.rgb_ext)]
            )

            # Get polarization files if pol directory exists
            pol_files = []
            if has_pol_data:
                pol_files = [f for f in os.listdir(pol_dir) if f.endswith(self.pol_ext)]
            # Get diffuse files if diffuse directory exists
            diffuse_files = []
            if diffuse_dir is not None and os.path.exists(diffuse_dir):
                diffuse_files = [f for f in os.listdir(diffuse_dir) if f.endswith(self.rgb_ext)]

            # Subsample per scene folder if requested (load 1/N frames)
            if self.sample_every_n > 1:
                rgb_files = rgb_files[:: self.sample_every_n]

            for rgb_file in rgb_files:
                raw_path = os.path.join(rgb_dir, rgb_file)
                # Determine diffuse path if available
                diffuse_path = None
                if diffuse_dir is not None and (rgb_file in diffuse_files):
                    diffuse_path = os.path.join(diffuse_dir, rgb_file)

                if not has_pol_data:
                    # No polarization data available - include raw-only sample (diffuse may or may not exist)
                    scene_pairs.append(
                        (
                            raw_path,
                            None,  # No polarization path
                            diffuse_path,
                            intrinsics_path,
                            False,  # No polarization data
                        )
                    )
                elif self.polarization_format == "single_file_clock":
                    # Original behavior: single polarization file
                    pol_file = rgb_file.replace(self.rgb_ext, self.pol_ext)
                    if pol_file in pol_files:
                        scene_pairs.append(
                            (
                                raw_path,
                                os.path.join(pol_dir, pol_file),
                                diffuse_path,
                                intrinsics_path,
                                True,  # Has polarization data
                            )
                        )

                elif self.polarization_format == "separate_files":
                    # New behavior: separate polarization files (_000, _045, _090, _135)
                    base_name = rgb_file.replace(self.rgb_ext, "")
                    pol_files_needed = [
                        f"{base_name}_000{self.pol_ext}",
                        f"{base_name}_045{self.pol_ext}",
                        f"{base_name}_090{self.pol_ext}",
                        f"{base_name}_135{self.pol_ext}",
                    ]

                    # Check if all 4 polarization files exist
                    if all(pol_file in pol_files for pol_file in pol_files_needed):
                        # Store the base path for separate files (we'll construct individual paths later)
                        pol_base_path = os.path.join(pol_dir, base_name)
                        scene_pairs.append(
                            (
                                raw_path,
                                pol_base_path,  # Base path for separate files
                                diffuse_path,
                                intrinsics_path,
                                True,  # Has polarization data
                            )
                        )

                elif self.polarization_format == "separate_files_stokes":
                    # New behavior: separate Stokes parameter files (_S0, _S1, _S2)
                    base_name = rgb_file.replace(self.rgb_ext, "")
                    stokes_files_needed = [
                        f"{base_name}_S0.npy",
                        f"{base_name}_S1.npy",
                        f"{base_name}_S2.npy",
                    ]

                    # Check if all 3 Stokes files exist
                    if all(
                        stokes_file in pol_files for stokes_file in stokes_files_needed
                    ):
                        # Store the base path for separate Stokes files (we'll construct individual paths later)
                        pol_base_path = os.path.join(pol_dir, base_name)
                        scene_pairs.append(
                            (
                                raw_path,
                                pol_base_path,  # Base path for separate Stokes files
                                diffuse_path,
                                intrinsics_path,
                                True,  # Has polarization data
                            )
                        )

                elif self.polarization_format == "single_file_topdown":
                    # New behavior: single file with 4 polarization images arranged vertically
                    pol_file = rgb_file.replace(self.rgb_ext, self.pol_ext)
                    if pol_file in pol_files:
                        scene_pairs.append(
                            (
                                raw_path,
                                os.path.join(pol_dir, pol_file),
                                diffuse_path,
                                intrinsics_path,
                                True,  # Has polarization data
                            )
                        )

        return scene_pairs

    def __len__(self) -> int:
        return len(self.scene_pairs)

    def get_loaded_scenes(self) -> List[str]:
        """
        Get list of scene names that were loaded into the dataset.

        Returns:
            List of scene names (without file extensions) that passed filtering
        """
        """Get list of scene names that are actually loaded in the dataset."""
        loaded_scenes = set()
        for raw_path, _, _, _, _ in self.scene_pairs:
            scene_name = os.path.basename(os.path.dirname(os.path.dirname(raw_path)))
            loaded_scenes.add(scene_name)
        return sorted(list(loaded_scenes))

    @staticmethod
    def _to_luminance_torch(rgb: torch.Tensor) -> torch.Tensor:
        """
        Convert RGB tensor to luminance using standard weights.

        Args:
            rgb: RGB tensor of shape [H, W, 3] in range [0, 1]

        Returns:
            Luminance tensor of shape [H, W] in range [0, 1]
        """
        """Convert RGB to luminance staying in torch. Input: [...,3] in [0,1]."""
        return 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]

    def _load_intrinsics_impl(self, intrinsics_path: str) -> torch.Tensor:
        """
        Load camera intrinsics from file (implementation without caching).

        Args:
            intrinsics_path: Path to intrinsics file (typically .txt or .npy)

        Returns:
            Camera intrinsics matrix of shape [3, 3]
        """
        """Implementation for loading intrinsics (cacheable)"""
        try:
            K = np.loadtxt(intrinsics_path).reshape(3, 3).astype(np.float32)
            return torch.from_numpy(K)
        except Exception:
            # warnings.warn(f"Could not load intrinsics from {intrinsics_path}: {e}")
            return torch.eye(3, dtype=torch.float32)

    def _load_intrinsics(self, intrinsics_path: str) -> torch.Tensor:
        """
        Load camera intrinsics with optional caching.

        Args:
            intrinsics_path: Path to intrinsics file

        Returns:
            Camera intrinsics matrix of shape [3, 3]
        """
        """Load intrinsics with optional caching"""
        if self.use_cache:
            return self._load_intrinsics_cached(intrinsics_path)
        else:
            return self._load_intrinsics_impl(intrinsics_path)

    def _simple_upsample(
        self, f_spec_half: torch.Tensor, target_size: Tuple[int, int]
    ) -> torch.Tensor:
        """
        Simple bilinear upsampling - much faster than edge-aware methods.

        Args:
            f_spec_half: Specular fraction tensor of shape [H, W] at half resolution
            target_size: Target size as (H, W) for upsampling

        Returns:
            Upsampled tensor of shape [target_H, target_W] clamped to [0, 1]
        """
        # Use torch interpolation instead of OpenCV
        f_batch = f_spec_half.unsqueeze(0).unsqueeze(0)  # Add batch and channel dims
        f_up = torch.nn.functional.interpolate(
            f_batch, size=target_size, mode="bilinear", align_corners=False
        )
        return f_up.squeeze(0).squeeze(0).clamp(0, 1)

    def _load_and_process_polarization_impl(
        self, pol_path: str
    ) -> Dict[str, torch.Tensor]:
        """
        Delegates polarization loading to PolarizationProcessor. Kept as a
        method for backward compatibility and caching behavior.
        """
        return self._pol_processor.load(
            pol_path=pol_path,
            polarization_format=self.polarization_format,
            pol_ext=self.pol_ext,
        )

    def _load_single_file_polarization(self, pol_path: str) -> Dict[str, torch.Tensor]:
        """Backward-compatible wrapper to single-file loader."""
        return self._pol_processor.load_single_file_clock(pol_path)

    def _load_and_process_polarization(self, pol_path: str) -> Dict[str, torch.Tensor]:
        """Load polarization with optional caching"""
        if self.use_cache and self.precompute_stokes:
            return self._load_pol_cached(pol_path)
        else:
            return self._load_and_process_polarization_impl(pol_path)

    def _load_raw_and_separate(
        self, raw_path: str, f_spec_half: torch.Tensor
    ) -> Dict[str, torch.Tensor]:
        """Optimized RAW loading and separation at full resolution"""

        # Load RAW directly as torch tensor
        raw_img = Image.open(raw_path).convert("RGB")
        raw = torch.from_numpy(np.asarray(raw_img, dtype=np.float32)) / 255.0
        H, W, _ = raw.shape

        # Upsample specular fraction
        if self.simplify_upsampling:
            f_full = self._simple_upsample(f_spec_half, (H, W))
        else:
            # Fall back to original method if needed
            f_full = self._edge_aware_upsample_original(f_spec_half, raw)

        # Compute specular/diffuse separation
        raw_chw = raw.permute(2, 0, 1)  # 3xHxW
        f_expanded = f_full.unsqueeze(0)  # 1xHxW

        I_spec = (f_expanded * raw_chw).clamp(0, 1)
        I_diff = (raw_chw - I_spec).clamp(0, 1)

        # Return full resolution data - resizing will be done later if needed
        return {
            "raw": raw_chw,
            "specular": I_spec,
            "diffuse": I_diff,
        }

    def _load_raw_only(self, raw_path: str) -> Dict[str, torch.Tensor]:
        """
        Load RAW data only (when polarization data is not available) at full resolution.

        Args:
            raw_path: Path to RAW image file

        Returns:
            Dictionary containing RAW data with dummy specular/diffuse components at full resolution
        """
        # Load RAW directly as torch tensor
        raw_img = Image.open(raw_path).convert("RGB")
        raw = torch.from_numpy(np.asarray(raw_img, dtype=np.float32)) / 255.0

        # Convert to CHW format
        raw_chw = raw.permute(2, 0, 1)  # [3, H, W]

        # Create dummy specular and diffuse components (all zeros for specular, RAW for diffuse)
        I_spec = torch.zeros_like(raw_chw)  # [3, H, W] - no specular component
        I_diff = raw_chw.clone()  # [3, H, W] - all RAW is diffuse

        # Return full resolution data - resizing will be done later if needed
        return {
            "raw": raw_chw,
            "specular": I_spec,
            "diffuse": I_diff,
        }

    def _load_raw_and_diffuse(self, raw_path: str, diffuse_path: str) -> Dict[str, torch.Tensor]:
        """Load RAW and DIFFUSE images and compute SPECULAR as raw - diffuse.

        Returns a dict with keys: 'raw', 'diffuse', 'specular' as CHW tensors in [0,1].
        """
        # Load images
        raw_img = Image.open(raw_path).convert("RGB")
        raw = torch.from_numpy(np.asarray(raw_img, dtype=np.float32)) / 255.0
        diffuse_img = Image.open(diffuse_path).convert("RGB")
        diffuse = torch.from_numpy(np.asarray(diffuse_img, dtype=np.float32)) / 255.0

        # To CHW
        raw_chw = raw.permute(2, 0, 1)
        diffuse_chw = diffuse.permute(2, 0, 1)

        # Specular as residual
        specular = (raw_chw - diffuse_chw).clamp(0.0, 1.0)

        return {
            "raw": raw_chw,
            "diffuse": diffuse_chw,
            "specular": specular,
        }

    def _compute_highlight_mask(self, frame_chw: torch.Tensor) -> torch.Tensor:
        """
        Compute binary highlight mask for a single RGB frame.

        Args:
            frame_chw: Tensor of shape [C, H, W]

        Returns:
            Tensor of shape [1, H, W] with 1 at highlight pixels
        """
        if frame_chw.dim() == 3:
            if frame_chw.shape[0] == 3:
                # Use luminance weights consistent with HighlightDataset
                grayscale = (
                    0.299 * frame_chw[0] + 0.587 * frame_chw[1] + 0.114 * frame_chw[2]
                )
            else:
                grayscale = frame_chw.mean(dim=0)
        elif frame_chw.dim() == 2:
            grayscale = frame_chw
        else:
            raise ValueError(
                f"Unexpected frame dimensions for highlight mask: {frame_chw.shape}"
            )

        mask = (grayscale > self.highlight_brightness_threshold).float().unsqueeze(0)
        return mask

    def _find_rectangle_with_least_highlights(
        self, binary_mask_hw: torch.Tensor, rect_size: Optional[Tuple[int, int]]
    ) -> torch.Tensor:
        """
        Find rect (top, left, bottom, right) with fewest highlights in a binary mask.

        Args:
            binary_mask_hw: Tensor [H, W] of 0/1
            rect_size: (height, width) or None

        Returns:
            Int tensor [4] = (top, left, bottom, right)
        """
        if rect_size is None:
            return torch.tensor([0, 0, 0, 0]).int()

        target_height, target_width = rect_size
        img_height, img_width = binary_mask_hw.shape

        if target_height > img_height or target_width > img_width:
            return torch.tensor([0, 0, 0, 0]).int()

        max_top = img_height - target_height + 1
        max_left = img_width - target_width + 1

        cumsum = torch.cumsum(torch.cumsum(binary_mask_hw, dim=0), dim=1)
        padded_cumsum = F.pad(cumsum, (1, 0, 1, 0), value=0)

        tops = torch.arange(max_top, device=binary_mask_hw.device)[:, None]
        lefts = torch.arange(max_left, device=binary_mask_hw.device)[None, :]
        bottoms = tops + target_height - 1
        rights = lefts + target_width - 1

        highlight_counts = (
            padded_cumsum[bottoms + 1, rights + 1]
            - padded_cumsum[tops, rights + 1]
            - padded_cumsum[bottoms + 1, lefts]
            + padded_cumsum[tops, lefts]
        )

        min_pos = torch.argmin(highlight_counts)
        best_top = (min_pos // max_left).item()
        best_left = (min_pos % max_left).item()
        best_bottom = best_top + target_height - 1
        best_right = best_left + target_width - 1

        return torch.tensor([best_top, best_left, best_bottom, best_right]).int()

    def _crop_rectangle_from_raw(
        self, raw_chw: torch.Tensor, rect_coords: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Crop rectangle from RAW and compute its highlight mask.

        Returns:
            (cropped_raw [C,h,w], cropped_mask [1,h,w])
        """
        top, left, bottom, right = [int(v) for v in rect_coords]

        if top == 0 and left == 0 and bottom == 0 and right == 0:
            if self.highlight_rect_size is not None:
                h, w = self.highlight_rect_size
                empty_frame = torch.zeros(
                    raw_chw.shape[0], h, w, device=raw_chw.device, dtype=raw_chw.dtype
                )
                empty_mask = torch.zeros(
                    1, h, w, device=raw_chw.device, dtype=raw_chw.dtype
                )
                return empty_frame, empty_mask
            else:
                return raw_chw, torch.zeros(
                    1,
                    raw_chw.shape[1],
                    raw_chw.shape[2],
                    device=raw_chw.device,
                    dtype=raw_chw.dtype,
                )

        cropped_raw = raw_chw[:, top : bottom + 1, left : right + 1]
        cropped_mask = self._compute_highlight_mask(cropped_raw)
        return cropped_raw, cropped_mask

    def _edge_aware_upsample_original(
        self, f_half: torch.Tensor, rgb_full: torch.Tensor
    ) -> torch.Tensor:
        """Original edge-aware upsampling (fallback)"""
        # Convert to numpy for OpenCV operations
        f_np = f_half.cpu().numpy()
        rgb_np = rgb_full.cpu().numpy()

        H, W, _ = rgb_np.shape
        f_up = cv2.resize(f_np, (W, H), interpolation=cv2.INTER_CUBIC)

        # Simple bilateral filter instead of guided filter for speed
        f_up = cv2.bilateralFilter(
            f_up.astype(np.float32), d=5, sigmaColor=0.1 * 255, sigmaSpace=5
        )

        return torch.from_numpy(np.clip(f_up, 0.0, 1.0))

    def _load_separate_polarization_files(
        self, pol_base_path: str
    ) -> Dict[str, torch.Tensor]:
        """Backward-compatible wrapper to separate-files loader."""
        return self._pol_processor.load_separate_files(pol_base_path, self.pol_ext)

    def _load_separate_stokes_files(
        self, pol_base_path: str
    ) -> Dict[str, torch.Tensor]:
        """Backward-compatible wrapper to separate Stokes loader."""
        return self._pol_processor.load_separate_stokes(pol_base_path)

    def single_arat_files_topdown(self, pol_path: str) -> Dict[str, torch.Tensor]:
        """Backward-compatible wrapper to top-down single-file loader."""
        return self._pol_processor.load_single_file_topdown(pol_path)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """
        Load and return a single dataset sample with optimized processing.

        Args:
            idx: Index of the sample to load (0 <= idx < len(dataset))

        Returns:
            Dictionary containing:
            - Polarization data: 'I0', 'I45', 'I90', 'I135', 'S0', 'S1', 'S2', 'DoLP', 'AoP', 'f_spec' (if available)
            - Image data: 'raw', 'specular', 'diffuse' (with backward-compatible alias 'rgb' == 'raw')
            - Camera data: 'intrinsics' [3, 3]
            - File paths: 'filepaths' dict with keys 'raw_path', 'pol_path', 'diffuse_path', 'intrinsics_path' (if return_filepaths=True)
            All image tensors have shape [C, H, W] where H, W match target_size if specified
        """
        raw_path, pol_path, diffuse_path, intrinsics_path, has_pol_data = self.scene_pairs[idx]

        # Load data (potentially from cache)
        intrinsics = self._load_intrinsics(intrinsics_path)

        if has_pol_data and not self.load_rgb_only:
            # Load polarization data and use it for RAW separation if diffuse not provided
            pol_data = self._load_and_process_polarization(pol_path)
            # Get specular fraction for RGB processing
            f_spec = pol_data["f_spec"].squeeze(0)  # Remove batch dimension [H, W]
            if diffuse_path is not None:
                raw_data = self._load_raw_and_diffuse(raw_path, diffuse_path)
            else:
                raw_data = self._load_raw_and_separate(raw_path, f_spec)
            # Combine results
            sample = {**pol_data, **raw_data, "intrinsics": intrinsics}
        else:
            # No polarization data available or load_rgb_only is True - load RAW (and diffuse if present)
            if diffuse_path is not None:
                raw_data = self._load_raw_and_diffuse(raw_path, diffuse_path)
            else:
                raw_data = self._load_raw_only(raw_path)
            sample = {**raw_data, "intrinsics": intrinsics}

        # Optional highlight detection and cropping on full resolution
        if self.highlight_enabled and "raw" in sample:
            raw_chw = sample["raw"]  # Full resolution RAW
            mask = self._compute_highlight_mask(raw_chw)
            if self.highlight_return_mask:
                sample["highlight_masks"] = mask
                total_pixels = raw_chw.shape[-2] * raw_chw.shape[-1]
                coverage_percent = (mask.sum() / max(total_pixels, 1)) * 100.0
                sample["highlight_coverage"] = coverage_percent.to(torch.float32)

            if self.highlight_rect_size is not None:
                rect_coords = self._find_rectangle_with_least_highlights(
                    mask.squeeze(0), self.highlight_rect_size
                )
                sample["rect_coords"] = rect_coords

                if self.highlight_return_rect:
                    rect_raw, rect_mask = self._crop_rectangle_from_raw(
                        raw_chw, rect_coords
                    )

                    sample["rect_crop"] = rect_raw
                    sample["rect_mask"] = rect_mask

                    if self.highlight_return_rect_as_rgb:
                        sample["uncropped_raw"] = sample["raw"]
                        sample["raw"] = rect_raw
                        # Also crop specular and diffuse components
                        if "specular" in sample:
                            sample["specular"] = sample["specular"][
                                :,
                                rect_coords[0] : rect_coords[2] + 1,
                                rect_coords[1] : rect_coords[3] + 1,
                            ]
                        if "diffuse" in sample:
                            sample["diffuse"] = sample["diffuse"][
                                :,
                                rect_coords[0] : rect_coords[2] + 1,
                                rect_coords[1] : rect_coords[3] + 1,
                            ]

                        # Also crop polarization data
                        pol_keys = [
                            "I0",
                            "I45",
                            "I90",
                            "I135",
                            "S0",
                            "S1",
                            "S2",
                            "S3",
                            "stokes",
                            "intensity",
                            "DoLP",
                            "AoP",
                            "AoLP",
                            "DoP",
                            "DoCP",
                            "ellipticity_angle",
                            "f_spec",
                        ]
                        for key in pol_keys:
                            if key in sample:
                                sample[key] = sample[key][
                                    :,
                                    rect_coords[0] : rect_coords[2] + 1,
                                    rect_coords[1] : rect_coords[3] + 1,
                                ]

        # Resize all data to target size if specified
        if self.target_size is not None:

            # Resize RAW-related data
            if "raw" in sample:
                sample["raw"] = self._resize_raw_tensor(sample["raw"])
            if "specular" in sample:
                sample["specular"] = self._resize_tensor(
                    sample["specular"], self.target_size
                )
            if "diffuse" in sample:
                sample["diffuse"] = self._resize_tensor(
                    sample["diffuse"], self.target_size
                )
            if "highlight_masks" in sample:
                sample["highlight_masks"] = F.interpolate(
                    sample["highlight_masks"].unsqueeze(0),
                    size=self.target_size,
                    mode="nearest",
                    align_corners=None,
                ).squeeze(0)
            if "rect_crop" in sample and not self.highlight_return_rect_as_rgb:
                # Only resize rect_crop if it's not being used as the main RGB
                sample["rect_crop"] = self._resize_raw_tensor(sample["rect_crop"])
            if "rect_mask" in sample and not self.highlight_return_rect_as_rgb:
                # Only resize rect_mask if it's not being used as the main RGB
                sample["rect_mask"] = F.interpolate(
                    sample["rect_mask"].unsqueeze(0),
                    size=self.target_size,
                    mode="nearest",
                    align_corners=None,
                ).squeeze(0)

            # Resize polarization data
            pol_keys = [
                "I0",
                "I45",
                "I90",
                "I135",
                "S0",
                "S1",
                "S2",
                "S3",
                "stokes",
                "intensity",
                "DoLP",
                "AoP",
                "AoLP",
                "DoP",
                "DoCP",
                "ellipticity_angle",
                "f_spec",
            ]
            for key in pol_keys:
                if key in sample:
                    sample[key] = self._resize_tensor(sample[key], self.target_size)

        if self.transform:
            sample = self.transform(sample)

        # if self.load_rgb_only:
        #     sample_rgbonly = {
        #         "rgb": sample["rgb"],
        #         "specular": sample["specular"],
        #         "diffuse": sample["diffuse"],
        #         "intrinsics": sample["intrinsics"],
        #     }
        #     return sample_rgbonly

        # Backward-compatible alias: provide 'rgb' that mirrors 'raw'
        # if "raw" in sample:
        #     sample["rgb"] = sample["raw"]

        # Add file paths if requested
        if self.return_filepaths:
            filepaths = {
                "raw_path": raw_path,
                # "pol_path": pol_path,
                # "intrinsics_path": intrinsics_path,
            }
            sample["filepaths"] = filepaths

        return sample

def from_config(
    config: Dict, dataset_names: Optional[List[str]] = None
) -> Dict[str, Union[Dataset, None]]:
    """
    Create datasets from configuration file.

    This function reads the configuration file and creates training and validation datasets
    based on the VAL_SCENES parameter. The logic is:
    - VAL_SCENES: defines which scenes to use for validation
    - TRAIN_SCENES: if provided and not None/[], overrides the default training scenes
    - If TRAIN_SCENES is None/[], training uses all scenes except those in VAL_SCENES

    Args:
        config: Configuration dictionary loaded from config file
        dataset_names: Optional list of dataset names to load. If None, loads all available datasets.

    Returns:
        Dictionary with 'training', 'validation', and 'test' keys containing ConcatDataset objects
    """
    if dataset_names is None:
        # Get all available dataset names from config
        dataset_names = []

        datasets_config = config.DATASETS

        if isinstance(datasets_config, dict):
            datasets_value = datasets_config

            if datasets_value is not None:
                dataset_names = [
                    name
                    for name in datasets_value.keys()
                    if isinstance(datasets_value[name], dict)
                ]

    if not dataset_names:
        raise ValueError(
            "No datasets found in configuration. Check DATASETS section in config file."
        )

    # Map dataset names to classes - this is where you add new dataset classes
    dataset_classes = {
        "SCRREAM": SCRREAM_Dataset,
        "HOUSECAT6D": HOUSECAT6D_Dataset,
        "POLARGB": POLARGB_Dataset,
        "SCARED": SCARED_Dataset,
        "STEREOMIS_TRACKING": STEREOMIS_TRACKING_Dataset,
        "CHOLEC80": CHOLEC80_Dataset,
        "CROMO": CROMO_Dataset,
        "SYNTHETIC": SYNTHETIC_Dataset,
        "PSD": PSD_Dataset,
        # Future datasets will be added here by the user
    }

    train_datasets = []
    val_datasets = []

    # Global config parameters
    global_config = config

    # Get global scene configuration
    global_train_scenes = global_config.get("TRAIN_SCENES", {}).get("value")
    global_val_scenes = global_config.get("VAL_SCENES", {}).get("value")

    logger.info(f"Processing {len(dataset_names)} datasets: {dataset_names}")

    for dataset_name in dataset_names:
        if dataset_name not in dataset_classes:
            logger.warning(
                f"Warning: Dataset class for '{dataset_name}' not found. Skipping."
            )
            logger.info(f"Available classes: {list(dataset_classes.keys())}")
            continue

        # Get dataset-specific config
        datasets_value = global_config.DATASETS
        if datasets_value is None:
            raise ValueError("DATASETS['value'] is None in config")
        dataset_config = datasets_value[dataset_name]
        # Get root directory

        # Extract configuration parameters with fallbacks to global config
        def get_config_value(param_name, default_value):
            """Helper to get parameter from dataset config or global config"""
            dataset_value = dataset_config.get(param_name)
            if dataset_value is not None:
                return dataset_value
            global_param = global_config.get(param_name, {})
            if isinstance(global_param, dict) and "value" in global_param:
                return global_param["value"]
            return default_value

        # root_dir = get_config_value("ROOT_", 0.6)
        # os.path.expandvars(dataset_config.ROOT_DIR)
        # if not os.path.exists(root_dir):
        #     logger.warning(
        #         f"Warning: Root directory '{root_dir}' for dataset '{dataset_name}' not found. Skipping."
        #     )
        #     continue
        dataset_params = {
            # "root_dir": get_config_value("ROOT", 0.6),
            "rho_s": get_config_value("RHO_S", 0.6),
            "eps": get_config_value("EPS", 1e-8),
            "target_size": tuple(get_config_value("TARGET_SIZE", [224, 224])),
            "resize_mode": get_config_value("RESIZE_MODE", "crop"),
            "use_cache": get_config_value("USE_CACHE", True),
            "simplify_upsampling": get_config_value("SIMPLIFY_UPSAMPLING", True),
            "few_images": get_config_value("FEW_IMAGES", False),
            "sample_every_n": get_config_value("SAMPLE_EVERY_N", 1),
            # "polarization_format": get_config_value(
            #     "POLARIZATION_FORMAT", "single_file_clock"
            # ),
            "load_rgb_only": get_config_value("LOAD_RGB_ONLY", False),
            # Highlight options
            "highlight_enable": get_config_value("HIGHLIGHT_ENABLE", False),
            "highlight_brightness_threshold": get_config_value(
                "HIGHLIGHT_BRIGHTNESS_THRESHOLD", 0.93
            ),
            "highlight_return_mask": get_config_value("HIGHLIGHT_RETURN_MASK", False),
            "highlight_return_rect": get_config_value("HIGHLIGHT_RETURN_RECT", False),
            "highlight_return_rect_as_rgb": get_config_value(
                "HIGHLIGHT_RETURN_RECT_AS_RGB", False
            ),
        }

        # Handle optional tuple conversion for rect size if provided
        rect_size_val = get_config_value("HIGHLIGHT_RECT_SIZE", None)
        if rect_size_val is not None:
            try:
                dataset_params["highlight_rect_size"] = tuple(rect_size_val)
            except Exception:
                dataset_params["highlight_rect_size"] = None

        # Get scenes configuration with priority: global > dataset-specific
        dataset_train_scenes = dataset_config.get("TRAIN_SCENES", [])
        dataset_val_scenes = dataset_config.get("VAL_SCENES", [])

        # Final scene determination with clear precedence
        val_scenes = (
            global_val_scenes if global_val_scenes is not None else dataset_val_scenes
        )

        # TRAIN_SCENES override logic: if global TRAIN_SCENES is provided and not empty, use it
        if global_train_scenes is not None and len(global_train_scenes) > 0:
            train_scenes = global_train_scenes
            logger.info(f"Using global TRAIN_SCENES for {dataset_name}: {train_scenes}")
        elif dataset_train_scenes and len(dataset_train_scenes) > 0:
            train_scenes = dataset_train_scenes
            logger.info(
                f"Using dataset-specific TRAIN_SCENES for {dataset_name}: {train_scenes}"
            )
        else:
            # Use all scenes except validation scenes
            train_scenes = None
            logger.info(
                f"Using all scenes except VAL_SCENES for {dataset_name} training"
            )

        # Get dataset class
        dataset_class = dataset_classes[dataset_name]

        # Create training dataset
        if train_scenes is not None and len(train_scenes) > 0:
            dataset_params.update({"highlight_enable": True})

            # Use specific training scenes
            train_dataset = dataset_class(
                include=train_scenes,
                # highlight_enable=dataset_config.get("HIGHLIGHT_ENABLE"),
                # highlight_brightness_threshold=dataset_config.get("HIGHLIGHT_BRIGHTNESS_THRESHOLD"),
                # highlight_return_mask=dataset_config.get("HIGHLIGHT_RETURN_MASK"),
                # highlight_rect_size=dataset_config.get("HIGHLIGHT_RECT_SIZE"),
                # highlight_return_rect=dataset_config.get("HIGHLIGHT_RETURN_RECT"),
                # highlight_return_rect_as_rgb=dataset_config.get("HIGHLIGHT_RETURN_RECT_AS_RGB"),
                **dataset_params,
            )
            if len(train_dataset) > 0:
                train_datasets.append(train_dataset)
                logger.info(
                    f"  ✓ Created training dataset for {dataset_name}: {len(train_dataset)} samples from specific scenes"
                )
            else:
                logger.warning(f"  ✗ Training dataset for {dataset_name} is empty")
        else:
            # Use all scenes except validation scenes
            exclude_scenes = val_scenes if val_scenes and len(val_scenes) > 0 else []
            train_dataset = dataset_class(
                exclude=exclude_scenes,
                # highlight_enable=dataset_config.get("HIGHLIGHT_ENABLE"),
                # highlight_brightness_threshold=dataset_config.get("HIGHLIGHT_BRIGHTNESS_THRESHOLD"),
                # highlight_return_mask=dataset_config.get("HIGHLIGHT_RETURN_MASK"),
                # highlight_rect_size=dataset_config.get("HIGHLIGHT_RECT_SIZE"),
                # highlight_return_rect=dataset_config.get("HIGHLIGHT_RETURN_RECT"),
                # highlight_return_rect_as_rgb=dataset_config.get("HIGHLIGHT_RETURN_RECT_AS_RGB"),
                **dataset_params,
            )
            if len(train_dataset) > 0:
                train_datasets.append(train_dataset)
                excluded_text = (
                    f" (excluding {len(exclude_scenes)} val scenes)"
                    if exclude_scenes
                    else ""
                )
                logger.info(
                    f"  ✓ Created training dataset for {dataset_name}: {len(train_dataset)} samples{excluded_text}"
                )
            else:
                logger.warning(f"  ✗ Training dataset for {dataset_name} is empty")

        # Create validation dataset
        if val_scenes and len(val_scenes) > 0:
            # Overrides global highlight_enable. Validation dataset should have original images
            dataset_params.update({"highlight_enable": False})
            val_dataset = dataset_class(
                include=val_scenes,
                **dataset_params,
            )
            if len(val_dataset) > 0:
                val_datasets.append(val_dataset)
                logger.info(
                    f"  ✓ Created validation dataset for {dataset_name}: {len(val_dataset)} samples from {len(val_scenes)} scenes"
                )
            else:
                logger.warning(f"  ✗ Validation dataset for {dataset_name} is empty")
        else:
            logger.warning(f"  ! No validation scenes specified for {dataset_name}")

    # Create ConcatDatasets
    result = {
        "training": ConcatDataset(train_datasets) if train_datasets else None,
        "validation": ConcatDataset(val_datasets) if val_datasets else None,
        "test": ConcatDataset(val_datasets) if val_datasets else None,
    }

    # Print summary
    logger.info("=== Dataset Creation Summary ===")
    logger.info(
        f"Training:   {len(result['training']) if result['training'] else 0} total samples"
    )
    logger.info(
        f"Validation: {len(result['validation']) if result['validation'] else 0} total samples"
    )
    logger.info(
        f"Test:       {len(result['test']) if result['test'] else 0} total samples"
    )

    # Print detailed breakdown if multiple datasets
    # if len(dataset_names) > 1:
    #     logger.info(f"Dataset breakdown:")
    #     for i, dataset_name in enumerate(dataset_names):
    #         if i < len(train_datasets):
    #             logger.info(f"  {dataset_name} - Train: {len(train_datasets[i])}, Val: {len(val_datasets[i]) if i < len(val_datasets) else 0}")

    return result

# Dataset-specific classes inheriting from base RGBP_Dataset class
class SCRREAM_Dataset(RGBP_Dataset):
    """
    SCRREAM dataset implementation for polarization-based reflection removal.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with SCRREAM-specific preprocessing,
    data augmentation, or validation logic as needed.

    The SCRREAM dataset contains RGB images with corresponding polarization
    data for training reflection removal models.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize SCRREAM dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/SCRREAM/",
            rgb_ext=".png",
            pol_ext=".png",
            polarization_format="single_file_clock",  # "single_file_clock", "separate_files" or "mosaic"
            **kwargs,
        )
        # Add any SCRREAM-specific initialization here


class HOUSECAT6D_Dataset(RGBP_Dataset):
    """
    HOUSECAT6D dataset implementation for 6D pose estimation with polarization.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with HOUSECAT6D-specific preprocessing,
    pose annotation loading, or 6D pose-specific data augmentation.

    The HOUSECAT6D dataset provides RGB and polarization data along with
    6D object pose annotations for training pose estimation models.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize HOUSECAT6D dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/HouseCat6D/",
            rgb_ext=".png",
            pol_ext=".png",
            polarization_format="single_file_clock",  # "single_file_clock", "separate_files" or "mosaic"
            **kwargs,
        )
        # Add any HOUSECAT6D-specific initialization here


class POLARGB_Dataset(RGBP_Dataset):
    """
    PolaRGB dataset implementation for polarization-guided RGB processing.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with PolaRGB-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The PolaRGB dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize PolaRGB dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/PolaRGB/",
            rgb_ext=".png",
            pol_ext=".png",
            polarization_format="separate_files",  # "single_file_clock", "separate_files" or "mosaic"
            **kwargs,
        )
        # Add any PolaRGB-specific initialization here

class CROMO_Dataset(RGBP_Dataset):
    """
    CROMO dataset implementation for polarization-guided RGB processing.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with CROMO-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The CROMO dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize CROMO dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/CroMo/",
            rgb_dir_name="rgb/left/data",
            pol_dir_name="polarized/left/data/",
            rgb_ext=".png",
            pol_ext=".png",
            polarization_format="single_file_topdown", 
            **kwargs,
        )
        # Add any SCARED-specific initialization here
        
class SCARED_Dataset(RGBP_Dataset):
    """
    SCARED dataset implementation for polarization-guided RGB processing.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with SCARED-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The SCARED dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize SCARED dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/SCARED/",
            rgb_ext=".png",
            pol_ext=".png",
            **kwargs,
        )
        # Add any SCARED-specific initialization here
        
class STEREOMIS_TRACKING_Dataset(RGBP_Dataset):
    """
    STEREOMIS_TRACKING dataset implementation for polarization-guided RGB processing.
    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with STEREOMIS_TRACKING-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The STEREOMIS_TRACKING dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize SCARED dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/StereoMIS_Tracking/",
            rgb_dir_name="video_frames",
            rgb_ext=".png",
            **kwargs,
        )
        # Add any SCARED-specific initialization here
        
class CHOLEC80_Dataset(RGBP_Dataset):
    """
    CHOLEC80 dataset implementation for polarization-guided RGB processing.
    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with CHOLEC80-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The CHOLEC80 dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize SCARED dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/CHOLEC80/videos/",
            rgb_dir_name="frames",
            rgb_ext=".png",
            **kwargs,
        )
        # Add any SCARED-specific initialization here
        
class PSD_Dataset(RGBP_Dataset):
    """
    PSD dataset implementation for polarization-guided RGB processing.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with PSD-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The PSD dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize PSD dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/PSD_Dataset/",
            rgb_ext=".png",
            rgb_dir_name="specular",
            diffuse_dir_name="diffuse",
            **kwargs,
        )
        # Add any SYNTHETIC-specific initialization here


class SYNTHETIC_Dataset(RGBP_Dataset):
    """
    SYNTHETIC dataset implementation for polarization-guided RGB processing.

    Inherits all functionality from the base RGBP_Dataset class.
    This class can be extended with SYNTHETIC-specific preprocessing,
    polarization analysis, or RGB enhancement techniques.

    The SYNTHETIC dataset combines RGB imagery with polarization measurements
    for improved scene understanding and image enhancement tasks.
    """

    def __init__(self, **kwargs) -> None:
        """
        Initialize SYNTHETIC dataset.

        Args:
            **kwargs: All arguments passed to parent RGBP_Dataset class
        """
        super().__init__(
            root_dir="$DATASET_DIR/SyntheticMitsuba/",
            rgb_ext=".png",
            pol_ext=".npy",
            polarization_format="separate_files_stokes", 
            **kwargs,
        )
        # Add any SYNTHETIC-specific initialization here
