"""
Data Loading and Preprocessing Utilities
=========================================

Provides memory-efficient data loading for large-scale datasets with:
- Memory-mapped file support for datasets exceeding RAM
- DDP-safe data preparation with proper synchronization
- Thread-safe DataLoader worker initialization
- Multi-format support (NPZ, HDF5, MAT)

Author: Ductho Le (ductho.le@outlook.com)
Version: 1.0.0
"""

import gc
import logging
import os
import pickle
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any

import h5py
import numpy as np
import torch
from accelerate import Accelerator
from scipy.sparse import issparse
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, Dataset
from tqdm.auto import tqdm


# Optional scipy.io for MATLAB files
try:
    import scipy.io

    HAS_SCIPY_IO = True
except ImportError:
    HAS_SCIPY_IO = False


# ==============================================================================
# DATA SOURCE ABSTRACTION
# ==============================================================================

# Supported key names for input/output arrays (priority order, pairwise aligned)
INPUT_KEYS = ["input_train", "input_test", "X", "data", "inputs", "features", "x"]
OUTPUT_KEYS = ["output_train", "output_test", "Y", "labels", "outputs", "targets", "y"]


class LazyDataHandle:
    """
    Context manager wrapper for memory-mapped data handles.

    Provides proper cleanup of file handles returned by load_mmap() methods.
    Can be used either as a context manager (recommended) or with explicit close().

    Usage:
        # Context manager (recommended)
        with source.load_mmap(path) as (inputs, outputs):
            # Use inputs and outputs
            pass  # File automatically closed

        # Manual cleanup
        handle = source.load_mmap(path)
        inputs, outputs = handle.inputs, handle.outputs
        # ... use data ...
        handle.close()

    Attributes:
        inputs: Input data array/dataset
        outputs: Output data array/dataset
    """

    def __init__(self, inputs, outputs, file_handle=None):
        """
        Initialize the handle.

        Args:
            inputs: Input array or lazy dataset
            outputs: Output array or lazy dataset
            file_handle: Optional file handle to close on cleanup
        """
        self.inputs = inputs
        self.outputs = outputs
        self._file = file_handle
        self._closed = False

    def __enter__(self):
        """Return (inputs, outputs) tuple for unpacking."""
        return self.inputs, self.outputs

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Close file handle on context exit."""
        self.close()
        return False  # Don't suppress exceptions

    def close(self):
        """
        Close the underlying file handle.

        Safe to call multiple times.
        """
        if self._closed:
            return
        self._closed = True

        # Close the file handle if we have one
        if self._file is not None:
            try:
                self._file.close()
            except Exception:
                pass
            self._file = None

        # Also close any _TransposedH5Dataset wrappers
        for data in (self.inputs, self.outputs):
            if hasattr(data, "close"):
                try:
                    data.close()
                except Exception:
                    pass

    def __repr__(self) -> str:
        status = "closed" if self._closed else "open"
        return f"LazyDataHandle(status={status})"


class DataSource(ABC):
    """
    Abstract base class for data loaders supporting multiple file formats.

    Subclasses must implement the `load()` method to return input/output arrays,
    and optionally `load_outputs_only()` for memory-efficient target loading.
    """

    @abstractmethod
    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
        """
        Load input and output arrays from a file.

        Args:
            path: Path to the data file

        Returns:
            Tuple of (inputs, outputs) as numpy arrays
        """
        pass

    @abstractmethod
    def load_outputs_only(self, path: str) -> np.ndarray:
        """
        Load only output/target arrays from a file (memory-efficient).

        This avoids loading large input arrays when only targets are needed,
        which is critical for HPC environments with memory constraints.

        Args:
            path: Path to the data file

        Returns:
            Output/target array
        """
        pass

    @staticmethod
    def detect_format(path: str) -> str:
        """
        Auto-detect file format from extension.

        Args:
            path: Path to data file

        Returns:
            Format string: 'npz', 'hdf5', or 'mat'
        """
        ext = Path(path).suffix.lower()
        format_map = {
            ".npz": "npz",
            ".h5": "hdf5",
            ".hdf5": "hdf5",
            ".mat": "mat",
        }
        if ext not in format_map:
            raise ValueError(
                f"Unsupported file extension: '{ext}'. "
                f"Supported formats: .npz, .h5, .hdf5, .mat"
            )
        return format_map[ext]

    @staticmethod
    def _find_key(available_keys: list[str], candidates: list[str]) -> str | None:
        """Find first matching key from candidates in available keys."""
        for key in candidates:
            if key in available_keys:
                return key
        return None


class NPZSource(DataSource):
    """Load data from NumPy .npz archives."""

    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
        data = np.load(path, allow_pickle=True)
        keys = list(data.keys())

        input_key = self._find_key(keys, INPUT_KEYS)
        output_key = self._find_key(keys, OUTPUT_KEYS)

        if input_key is None or output_key is None:
            raise KeyError(
                f"NPZ must contain input and output arrays. "
                f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
                f"Found: {keys}"
            )

        inp = data[input_key]
        outp = data[output_key]

        # Handle object arrays (e.g., sparse matrices stored as objects)
        if inp.dtype == object:
            inp = np.array([x.toarray() if hasattr(x, "toarray") else x for x in inp])

        return inp, outp

    def load_mmap(self, path: str) -> tuple[np.ndarray, np.ndarray]:
        """
        Load data using memory-mapped mode for zero-copy access.

        This allows processing large datasets without loading them entirely
        into RAM. Critical for HPC environments with memory constraints.

        Note: Returns memory-mapped arrays - do NOT modify them.
        """
        data = np.load(path, allow_pickle=True, mmap_mode="r")
        keys = list(data.keys())

        input_key = self._find_key(keys, INPUT_KEYS)
        output_key = self._find_key(keys, OUTPUT_KEYS)

        if input_key is None or output_key is None:
            raise KeyError(
                f"NPZ must contain input and output arrays. "
                f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
                f"Found: {keys}"
            )

        inp = data[input_key]
        outp = data[output_key]

        return inp, outp

    def load_outputs_only(self, path: str) -> np.ndarray:
        """Load only targets from NPZ (avoids loading large input arrays)."""
        data = np.load(path, allow_pickle=True)
        keys = list(data.keys())

        output_key = self._find_key(keys, OUTPUT_KEYS)
        if output_key is None:
            raise KeyError(
                f"NPZ must contain output array. "
                f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
            )

        return data[output_key]


class HDF5Source(DataSource):
    """Load data from HDF5 (.h5, .hdf5) files."""

    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
        with h5py.File(path, "r") as f:
            keys = list(f.keys())

            input_key = self._find_key(keys, INPUT_KEYS)
            output_key = self._find_key(keys, OUTPUT_KEYS)

            if input_key is None or output_key is None:
                raise KeyError(
                    f"HDF5 must contain input and output datasets. "
                    f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
                    f"Found: {keys}"
                )

            # Load into memory (HDF5 datasets are lazy by default)
            inp = f[input_key][:]
            outp = f[output_key][:]

        return inp, outp

    def load_mmap(self, path: str) -> LazyDataHandle:
        """
        Load HDF5 file with lazy/memory-mapped access.

        Returns a LazyDataHandle that reads from disk on-demand,
        avoiding loading the entire file into RAM.

        Usage:
            with source.load_mmap(path) as (inputs, outputs):
                # Use inputs and outputs
                pass  # File automatically closed
        """
        f = h5py.File(path, "r")  # Keep file open for lazy access
        keys = list(f.keys())

        input_key = self._find_key(keys, INPUT_KEYS)
        output_key = self._find_key(keys, OUTPUT_KEYS)

        if input_key is None or output_key is None:
            f.close()
            raise KeyError(
                f"HDF5 must contain input and output datasets. "
                f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
                f"Found: {keys}"
            )

        # Return wrapped handle for proper cleanup
        return LazyDataHandle(f[input_key], f[output_key], file_handle=f)

    def load_outputs_only(self, path: str) -> np.ndarray:
        """Load only targets from HDF5 (avoids loading large input arrays)."""
        with h5py.File(path, "r") as f:
            keys = list(f.keys())

            output_key = self._find_key(keys, OUTPUT_KEYS)
            if output_key is None:
                raise KeyError(
                    f"HDF5 must contain output dataset. "
                    f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
                )

            outp = f[output_key][:]

        return outp


class _TransposedH5Dataset:
    """
    Lazy transpose wrapper for h5py datasets.

    MATLAB stores arrays in column-major (Fortran) order, while Python/NumPy
    expects row-major (C) order. This wrapper provides a transposed view
    without loading the entire dataset into memory.

    Supports:
        - len(): Returns the transposed first dimension
        - []: Returns slices with automatic transpose
        - shape: Returns the transposed shape
        - dtype: Returns the underlying dtype

    This is critical for MATSource.load_mmap() to return consistent axis
    ordering with the eager loader (MATSource.load()).

    IMPORTANT: Holds a strong reference to the h5py.File to prevent
    premature garbage collection while datasets are in use.
    """

    def __init__(self, h5_dataset, file_handle=None):
        """
        Args:
            h5_dataset: The h5py dataset to wrap
            file_handle: Optional h5py.File reference to keep alive
        """
        self._dataset = h5_dataset
        self._file = file_handle  # Keep file alive to prevent GC
        # Transpose shape: MATLAB (cols, rows, ...) -> Python (rows, cols, ...)
        self.shape = tuple(reversed(h5_dataset.shape))
        self.dtype = h5_dataset.dtype

        # Precompute transpose axis order for efficiency
        # For shape (A, B, C) -> reversed (C, B, A), transpose axes are (2, 1, 0)
        self._transpose_axes = tuple(range(len(h5_dataset.shape) - 1, -1, -1))

    def __len__(self) -> int:
        return self.shape[0]

    def __getitem__(self, idx):
        """
        Fetch data with automatic full transpose.

        Handles integer indexing, slices, and fancy indexing.
        All operations return data with fully reversed axes to match .T behavior.
        """
        if isinstance(idx, int | np.integer):
            # Single sample: index into last axis of h5py dataset (column-major)
            # Result needs full transpose of remaining dimensions
            data = self._dataset[..., idx]
            if data.ndim == 0:
                return data
            elif data.ndim == 1:
                return data  # 1D doesn't need transpose
            else:
                # Full transpose: reverse all axes
                return np.transpose(data)

        elif isinstance(idx, slice):
            # Slice indexing: fetch from last axis, then fully transpose
            start, stop, step = idx.indices(self.shape[0])
            data = self._dataset[..., start:stop:step]

            # Handle special case: 1D result (e.g., row vector)
            if data.ndim == 1:
                return data

            # Full transpose: reverse ALL axes (not just moveaxis)
            # This matches the behavior of .T on a numpy array
            return np.transpose(data, axes=self._transpose_axes)

        elif isinstance(idx, list | np.ndarray):
            # Fancy indexing: load samples one at a time (h5py limitation)
            # This is slower but necessary for compatibility
            samples = [self[i] for i in idx]
            return np.stack(samples, axis=0)

        else:
            raise TypeError(f"Unsupported index type: {type(idx)}")

    def close(self):
        """Close the underlying file handle if we own it."""
        if self._file is not None:
            try:
                self._file.close()
            except Exception:
                pass


class MATSource(DataSource):
    """
    Load data from MATLAB .mat files (v7.3+ only, which uses HDF5 format).

    Note: MAT v7.3 files are HDF5 files under the hood, so we use h5py for
    memory-efficient lazy loading. Save with: save('file.mat', '-v7.3')

    Supports MATLAB sparse matrices (automatically converted to dense).

    For older MAT files (v5/v7), convert to NPZ or save with -v7.3 flag.
    """

    @staticmethod
    def _is_sparse_dataset(dataset) -> bool:
        """Check if an HDF5 dataset/group represents a MATLAB sparse matrix."""
        # MATLAB v7.3 stores sparse matrices as groups with 'data', 'ir', 'jc' keys
        if hasattr(dataset, "keys"):
            keys = set(dataset.keys())
            return {"data", "ir", "jc"}.issubset(keys)
        return False

    @staticmethod
    def _load_sparse_to_dense(group) -> np.ndarray:
        """Convert MATLAB sparse matrix (CSC format in HDF5) to dense numpy array."""
        from scipy.sparse import csc_matrix

        data = np.array(group["data"])
        ir = np.array(group["ir"])  # row indices
        jc = np.array(group["jc"])  # column pointers

        # Get shape from MATLAB attributes or infer
        if "MATLAB_sparse" in group.attrs:
            nrows = group.attrs["MATLAB_sparse"]
        else:
            nrows = ir.max() + 1 if len(ir) > 0 else 0
        ncols = len(jc) - 1

        sparse_mat = csc_matrix((data, ir, jc), shape=(nrows, ncols))
        return sparse_mat.toarray()

    def _load_dataset(self, f, key: str) -> np.ndarray:
        """Load a dataset, handling sparse matrices automatically."""
        dataset = f[key]

        if self._is_sparse_dataset(dataset):
            # Sparse matrix: convert to dense
            arr = self._load_sparse_to_dense(dataset)
        else:
            # Regular dense array
            arr = np.array(dataset)

        # Transpose for MATLAB column-major -> Python row-major
        return arr.T

    def load(self, path: str) -> tuple[np.ndarray, np.ndarray]:
        """Load MAT v7.3 file using h5py."""
        try:
            with h5py.File(path, "r") as f:
                keys = list(f.keys())

                input_key = self._find_key(keys, INPUT_KEYS)
                output_key = self._find_key(keys, OUTPUT_KEYS)

                if input_key is None or output_key is None:
                    raise KeyError(
                        f"MAT file must contain input and output arrays. "
                        f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
                        f"Found: {keys}"
                    )

                # Load with sparse matrix support
                inp = self._load_dataset(f, input_key)
                outp = self._load_dataset(f, output_key)

                # Handle 1D outputs that become (1, N) after transpose
                if outp.ndim == 2 and outp.shape[0] == 1:
                    outp = outp.T

        except OSError as e:
            raise ValueError(
                f"Failed to load MAT file: {path}. "
                f"Ensure it's saved as v7.3: save('file.mat', '-v7.3'). "
                f"Original error: {e}"
            )

        return inp, outp

    def load_mmap(self, path: str) -> LazyDataHandle:
        """
        Load MAT v7.3 file with lazy/memory-mapped access.

        Returns a LazyDataHandle that reads from disk on-demand,
        avoiding loading the entire file into RAM.

        Note: For sparse matrices, this will load and convert them.
        For dense arrays, returns a transposed view wrapper for consistent axis ordering.

        Usage:
            with source.load_mmap(path) as (inputs, outputs):
                # Use inputs and outputs
                pass  # File automatically closed
        """
        try:
            f = h5py.File(path, "r")  # Keep file open for lazy access
            keys = list(f.keys())

            input_key = self._find_key(keys, INPUT_KEYS)
            output_key = self._find_key(keys, OUTPUT_KEYS)

            if input_key is None or output_key is None:
                f.close()
                raise KeyError(
                    f"MAT file must contain input and output arrays. "
                    f"Supported keys: {INPUT_KEYS} / {OUTPUT_KEYS}. "
                    f"Found: {keys}"
                )

            # Check for sparse matrices - must load them eagerly
            inp_dataset = f[input_key]
            outp_dataset = f[output_key]

            if self._is_sparse_dataset(inp_dataset):
                inp = self._load_sparse_to_dense(inp_dataset).T
            else:
                # Wrap h5py dataset with transpose view for consistent axis order
                # MATLAB stores column-major, Python expects row-major
                # Pass file handle to keep it alive
                inp = _TransposedH5Dataset(inp_dataset, file_handle=f)

            if self._is_sparse_dataset(outp_dataset):
                outp = self._load_sparse_to_dense(outp_dataset).T
            else:
                # Wrap h5py dataset with transpose view (shares same file handle)
                outp = _TransposedH5Dataset(outp_dataset, file_handle=f)

            # Return wrapped handle for proper cleanup
            return LazyDataHandle(inp, outp, file_handle=f)

        except OSError as e:
            raise ValueError(
                f"Failed to load MAT file: {path}. "
                f"Ensure it's saved as v7.3: save('file.mat', '-v7.3'). "
                f"Original error: {e}"
            )

    def load_outputs_only(self, path: str) -> np.ndarray:
        """Load only targets from MAT v7.3 file (avoids loading large input arrays)."""
        try:
            with h5py.File(path, "r") as f:
                keys = list(f.keys())

                output_key = self._find_key(keys, OUTPUT_KEYS)
                if output_key is None:
                    raise KeyError(
                        f"MAT file must contain output array. "
                        f"Supported keys: {OUTPUT_KEYS}. Found: {keys}"
                    )

                # Load with sparse matrix support
                outp = self._load_dataset(f, output_key)

                # Handle 1D outputs
                if outp.ndim == 2 and outp.shape[0] == 1:
                    outp = outp.T

        except OSError as e:
            raise ValueError(
                f"Failed to load MAT file: {path}. "
                f"Ensure it's saved as v7.3: save('file.mat', '-v7.3'). "
                f"Original error: {e}"
            )

        return outp


def get_data_source(format: str) -> DataSource:
    """
    Factory function to get the appropriate DataSource for a format.

    Args:
        format: One of 'npz', 'hdf5', 'mat'

    Returns:
        DataSource instance
    """
    sources = {
        "npz": NPZSource,
        "hdf5": HDF5Source,
        "mat": MATSource,
    }

    if format not in sources:
        raise ValueError(
            f"Unsupported format: {format}. Supported: {list(sources.keys())}"
        )

    return sources[format]()


def load_training_data(
    path: str, format: str = "auto"
) -> tuple[np.ndarray, np.ndarray]:
    """
    Load training data from file with automatic format detection.

    Supports:
        - NPZ: NumPy compressed archives (.npz)
        - HDF5: Hierarchical Data Format (.h5, .hdf5)
        - MAT: MATLAB files (.mat)

    Flexible key detection supports: input_train/X/data and output_train/y/labels.

    Args:
        path: Path to data file
        format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)

    Returns:
        Tuple of (inputs, outputs) arrays
    """
    if format == "auto":
        format = DataSource.detect_format(path)

    source = get_data_source(format)
    return source.load(path)


def load_outputs_only(path: str, format: str = "auto") -> np.ndarray:
    """
    Load only output/target arrays from file (memory-efficient).

    This function avoids loading large input arrays when only targets are needed,
    which is critical for HPC environments with memory constraints during DDP.

    Args:
        path: Path to data file
        format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)

    Returns:
        Output/target array
    """
    if format == "auto":
        format = DataSource.detect_format(path)

    source = get_data_source(format)
    return source.load_outputs_only(path)


def load_test_data(
    path: str,
    format: str = "auto",
    input_key: str | None = None,
    output_key: str | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """
    Load test/inference data and return PyTorch tensors ready for model input.

    This is the unified data loading function for inference. It:
    - Auto-detects file format from extension
    - Handles custom key names for non-standard datasets
    - Adds channel dimension if missing (dimension-agnostic)
    - Returns None for targets if not present in file

    Supports any input dimensionality:
        - 1D: (N, L) → (N, 1, L)
        - 2D: (N, H, W) → (N, 1, H, W)
        - 3D: (N, D, H, W) → (N, 1, D, H, W)
        - Already has channel: (N, C, ...) → unchanged

    Args:
        path: Path to data file (NPZ, HDF5, or MAT v7.3)
        format: Format hint ('npz', 'hdf5', 'mat', or 'auto' for detection)
        input_key: Custom key for input data (overrides auto-detection)
        output_key: Custom key for output data (overrides auto-detection)

    Returns:
        Tuple of:
            - X: Input tensor with channel dimension (N, 1, *spatial_dims)
            - y: Target tensor (N, T) or None if targets not present

    Example:
        >>> X, y = load_test_data("test_data.npz")
        >>> X, y = load_test_data(
        ...     "data.mat", input_key="waveforms", output_key="params"
        ... )
    """
    if format == "auto":
        format = DataSource.detect_format(path)

    source = get_data_source(format)

    # Build custom key lists if provided
    if input_key:
        custom_input_keys = [input_key] + INPUT_KEYS
    else:
        # Prioritize test keys for inference
        custom_input_keys = ["input_test"] + [
            k for k in INPUT_KEYS if k != "input_test"
        ]

    if output_key:
        custom_output_keys = [output_key] + OUTPUT_KEYS
    else:
        custom_output_keys = ["output_test"] + [
            k for k in OUTPUT_KEYS if k != "output_test"
        ]

    # Load data using appropriate source
    try:
        inp, outp = source.load(path)
    except KeyError:
        # Try with just inputs if outputs not found
        if format == "npz":
            data = np.load(path, allow_pickle=True)
            keys = list(data.keys())
            inp_key = DataSource._find_key(keys, custom_input_keys)
            if inp_key is None:
                raise KeyError(
                    f"Input key not found. Tried: {custom_input_keys}. Found: {keys}"
                )
            inp = data[inp_key]
            if inp.dtype == object:
                inp = np.array(
                    [x.toarray() if hasattr(x, "toarray") else x for x in inp]
                )
            out_key = DataSource._find_key(keys, custom_output_keys)
            outp = data[out_key] if out_key else None
        else:
            raise

    # Handle sparse matrices
    if issparse(inp):
        inp = inp.toarray()
    if outp is not None and issparse(outp):
        outp = outp.toarray()

    # Convert to tensors
    X = torch.tensor(np.asarray(inp), dtype=torch.float32)

    if outp is not None:
        y = torch.tensor(np.asarray(outp), dtype=torch.float32)
        # Normalize target shape: (N,) → (N, 1)
        if y.ndim == 1:
            y = y.unsqueeze(1)
    else:
        y = None

    # Add channel dimension if needed (dimension-agnostic)
    # X.ndim == 2: 1D data (N, L) → (N, 1, L)
    # X.ndim == 3: 2D data (N, H, W) → (N, 1, H, W)
    # X.ndim == 4: Check if already has channel dim (C <= 16 heuristic)
    if X.ndim == 2:
        X = X.unsqueeze(1)  # 1D signal: (N, L) → (N, 1, L)
    elif X.ndim == 3:
        X = X.unsqueeze(1)  # 2D image: (N, H, W) → (N, 1, H, W)
    elif X.ndim == 4:
        # Could be 3D volume (N, D, H, W) or 2D with channel (N, C, H, W)
        # Heuristic: if dim 1 is small (<=16), assume it's already a channel dim
        if X.shape[1] > 16:
            X = X.unsqueeze(1)  # 3D volume: (N, D, H, W) → (N, 1, D, H, W)
    # X.ndim >= 5: assume channel dimension already exists

    return X, y


# ==============================================================================
# WORKER INITIALIZATION
# ==============================================================================
def memmap_worker_init_fn(worker_id: int):
    """
    Worker initialization function for proper memmap handling in multi-worker DataLoader.

    Each DataLoader worker process runs this function after forking. It:
    1. Resets the memmap file handle to None, forcing each worker to open its own
       read-only handle (prevents file descriptor sharing issues and race conditions)
    2. Seeds numpy's random state per worker to ensure statistical diversity in
       random augmentations (prevents all workers from applying identical "random"
       transformations to their batches)

    Args:
        worker_id: Worker index (0 to num_workers-1), provided by DataLoader

    Usage:
        DataLoader(dataset, num_workers=8, worker_init_fn=memmap_worker_init_fn)
    """
    worker_info = torch.utils.data.get_worker_info()
    if worker_info is not None:
        dataset = worker_info.dataset
        # Force re-initialization of memmap in each worker
        dataset.data = None

        # Seed numpy RNG per worker using PyTorch's worker seed for reproducibility
        # This ensures random augmentations (noise, shifts, etc.) are unique per worker
        np.random.seed(worker_info.seed % (2**32 - 1))


# ==============================================================================
# MEMORY-MAPPED DATASET
# ==============================================================================
class MemmapDataset(Dataset):
    """
    Zero-copy memory-mapped dataset for large-scale training.

    Uses numpy memory mapping to load data directly from disk, allowing training
    on datasets that exceed available RAM. The memmap is only opened when first
    accessed (lazy initialization), and each DataLoader worker maintains its own
    file handle for thread safety.

    Args:
        memmap_path: Path to the memory-mapped data file
        targets: Pre-loaded target tensor (small enough to fit in memory)
        shape: Full shape of the memmap array (N, C, H, W)
        indices: Indices into the memmap for this split (train/val)

    Thread Safety:
        When using with DataLoader num_workers > 0, must use memmap_worker_init_fn
        as the worker_init_fn to ensure each worker gets its own file handle.

    Example:
        dataset = MemmapDataset("cache.dat", y_tensor, (10000, 1, 500, 500), train_indices)
        loader = DataLoader(dataset, num_workers=8, worker_init_fn=memmap_worker_init_fn)
    """

    def __init__(
        self,
        memmap_path: str,
        targets: torch.Tensor,
        shape: tuple[int, ...],
        indices: np.ndarray,
    ):
        self.memmap_path = memmap_path
        self.targets = targets
        self.shape = shape
        self.indices = indices
        self.data: np.memmap | None = None  # Lazy initialization

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

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        if self.data is None:
            # Mode 'r' = read-only, prevents accidental data modification
            self.data = np.memmap(
                self.memmap_path, dtype="float32", mode="r", shape=self.shape
            )

        real_idx = self.indices[idx]

        # .copy() detaches from mmap buffer - essential for PyTorch pinned memory
        x = torch.from_numpy(self.data[real_idx].copy()).contiguous()
        y = self.targets[real_idx]

        return x, y

    def __repr__(self) -> str:
        return (
            f"MemmapDataset(path='{self.memmap_path}', "
            f"samples={len(self)}, shape={self.shape})"
        )


# ==============================================================================
# DATA PREPARATION
# ==============================================================================
def prepare_data(
    args: Any,
    logger: logging.Logger,
    accelerator: Accelerator,
    cache_dir: str = ".",
    val_split: float = 0.2,
) -> tuple[DataLoader, DataLoader, StandardScaler, tuple[int, ...], int]:
    """
    Prepare DataLoaders with DDP synchronization guarantees.

    This function handles:
    1. Loading raw data and creating memmap cache (Rank 0 only)
    2. Fitting StandardScaler on training set only (no data leakage)
    3. Synchronizing all ranks before proceeding
    4. Creating thread-safe DataLoaders for DDP training

    Supports any input dimensionality:
        - 1D: (N, L) → returns in_shape = (L,)
        - 2D: (N, H, W) → returns in_shape = (H, W)
        - 3D: (N, D, H, W) → returns in_shape = (D, H, W)

    Args:
        args: Argument namespace with data_path, seed, batch_size, workers
        logger: Logger instance for status messages
        accelerator: Accelerator instance for DDP coordination
        cache_dir: Directory for cache files (default: current directory)
        val_split: Validation set fraction (default: 0.2)

    Returns:
        Tuple of:
            - train_dl: Training DataLoader
            - val_dl: Validation DataLoader
            - scaler: Fitted StandardScaler (for inverse transforms)
            - in_shape: Input spatial dimensions - (L,), (H, W), or (D, H, W)
            - out_dim: Number of output targets

    Cache Files Created:
        - train_data_cache.dat: Memory-mapped input data
        - scaler.pkl: Fitted StandardScaler
        - data_metadata.pkl: Shape and dimension metadata
    """
    CACHE_FILE = os.path.join(cache_dir, "train_data_cache.dat")
    SCALER_FILE = os.path.join(cache_dir, "scaler.pkl")
    META_FILE = os.path.join(cache_dir, "data_metadata.pkl")

    # ==========================================================================
    # PHASE 1: DATA GENERATION (Rank 0 Only)
    # ==========================================================================
    # Check cache existence and validity (data path must match)
    cache_exists = (
        os.path.exists(CACHE_FILE)
        and os.path.exists(SCALER_FILE)
        and os.path.exists(META_FILE)
    )

    # Validate cache matches current data_path (prevents stale cache corruption)
    if cache_exists:
        try:
            with open(META_FILE, "rb") as f:
                meta = pickle.load(f)
            cached_data_path = meta.get("data_path", None)
            if cached_data_path != os.path.abspath(args.data_path):
                if accelerator.is_main_process:
                    logger.warning(
                        f"⚠️  Cache was created from different data file!\n"
                        f"   Cached: {cached_data_path}\n"
                        f"   Current: {os.path.abspath(args.data_path)}\n"
                        f"   Invalidating cache and regenerating..."
                    )
                cache_exists = False
        except Exception:
            cache_exists = False

    if not cache_exists:
        if accelerator.is_main_process:
            # RANK 0: Create cache (can take a long time for large datasets)
            # Other ranks will wait at the barrier below

            # Detect format from extension
            data_format = DataSource.detect_format(args.data_path)
            logger.info(
                f"⚡ [Rank 0] Initializing Data Processing from: {args.data_path} (format: {data_format})"
            )

            # Validate data file exists
            if not os.path.exists(args.data_path):
                raise FileNotFoundError(
                    f"CRITICAL: Data file not found: {args.data_path}"
                )

            # Load raw data using memory-mapped mode for all formats
            # This avoids loading the entire dataset into RAM at once
            try:
                if data_format == "npz":
                    source = NPZSource()
                    inp, outp = source.load_mmap(args.data_path)
                elif data_format == "hdf5":
                    source = HDF5Source()
                    _lazy_handle = source.load_mmap(args.data_path)
                    inp, outp = _lazy_handle.inputs, _lazy_handle.outputs
                elif data_format == "mat":
                    source = MATSource()
                    _lazy_handle = source.load_mmap(args.data_path)
                    inp, outp = _lazy_handle.inputs, _lazy_handle.outputs
                else:
                    inp, outp = load_training_data(args.data_path, format=data_format)
                logger.info("   Using memory-mapped loading (low memory mode)")
            except Exception as e:
                logger.error(f"Failed to load data file: {e}")
                raise

            # Detect shape (handle sparse matrices) - DIMENSION AGNOSTIC
            num_samples = len(inp)

            # Handle 1D targets: (N,) -> treat as single output
            if outp.ndim == 1:
                out_dim = 1
            else:
                out_dim = outp.shape[1]

            sample_0 = inp[0]
            if issparse(sample_0) or hasattr(sample_0, "toarray"):
                sample_0 = sample_0.toarray()

            # Detect dimensionality and validate single-channel assumption
            raw_shape = sample_0.shape

            # Heuristic for ambiguous multi-channel detection:
            # Trigger when shape is EXACTLY 3D (could be C,H,W or D,H,W) with:
            #   - First dim small (<=16) - looks like channels
            #   - BOTH remaining dims large (>16) - confirming it's an image, not a tiny patch
            # This distinguishes (3, 256, 256) from (8, 1024) or (128, 128)
            # Use --single_channel to confirm shallow 3D volumes like (8, 128, 128)
            is_ambiguous_shape = (
                len(raw_shape) == 3  # Exactly 3D: could be (C, H, W) or (D, H, W)
                and raw_shape[0] <= 16  # First dim looks like channels
                and raw_shape[1] > 16
                and raw_shape[2] > 16  # Both spatial dims are large
            )

            # Check for user confirmation via --single_channel flag
            user_confirmed_single_channel = getattr(args, "single_channel", False)

            if is_ambiguous_shape and not user_confirmed_single_channel:
                raise ValueError(
                    f"Ambiguous input shape detected: sample shape {raw_shape}. "
                    f"This could be either:\n"
                    f"  - Multi-channel 2D data (C={raw_shape[0]}, H={raw_shape[1]}, W={raw_shape[2]})\n"
                    f"  - Single-channel 3D volume (D={raw_shape[0]}, H={raw_shape[1]}, W={raw_shape[2]})\n\n"
                    f"If this is single-channel 3D/shallow volume data, use --single_channel flag.\n"
                    f"If this is multi-channel 2D data, reshape to (N*C, H, W) with adjusted targets."
                )

            spatial_shape = raw_shape
            full_shape = (
                num_samples,
                1,
            ) + spatial_shape  # Add channel dim: (N, 1, ...)

            dim_names = {1: "1D (L)", 2: "2D (H, W)", 3: "3D (D, H, W)"}
            dim_type = dim_names.get(len(spatial_shape), f"{len(spatial_shape)}D")
            logger.info(
                f"   Shape Detected: {full_shape} [{dim_type}] | Output Dim: {out_dim}"
            )

            # Save metadata (including data path for cache validation)
            with open(META_FILE, "wb") as f:
                pickle.dump(
                    {
                        "shape": full_shape,
                        "out_dim": out_dim,
                        "data_path": os.path.abspath(args.data_path),
                    },
                    f,
                )

            # Create memmap cache
            if not os.path.exists(CACHE_FILE):
                logger.info("   Writing Memmap Cache (one-time operation)...")
                fp = np.memmap(CACHE_FILE, dtype="float32", mode="w+", shape=full_shape)

                chunk_size = 2000
                pbar = tqdm(
                    range(0, num_samples, chunk_size),
                    desc="Caching",
                    disable=not accelerator.is_main_process,
                )

                for i in pbar:
                    end = min(i + chunk_size, num_samples)
                    batch = inp[i:end]

                    # Handle sparse/dense conversion
                    if issparse(batch[0]) or hasattr(batch[0], "toarray"):
                        data_chunk = np.stack(
                            [x.toarray().astype(np.float32) for x in batch]
                        )
                    else:
                        data_chunk = np.array(batch).astype(np.float32)

                    # Add channel dimension if needed (handles 1D, 2D, and 3D spatial data)
                    # data_chunk shape: (batch, *spatial) -> need (batch, 1, *spatial)
                    # full_shape is (N, 1, *spatial), so expected ndim = len(full_shape)
                    if data_chunk.ndim == len(full_shape) - 1:
                        # Missing channel dim: (batch, *spatial) -> (batch, 1, *spatial)
                        data_chunk = np.expand_dims(data_chunk, 1)

                    fp[i:end] = data_chunk

                    # Periodic flush to disk
                    if i % 10000 == 0:
                        fp.flush()

                fp.flush()
                del fp
                gc.collect()

            # Train/Val split and scaler fitting
            indices = np.arange(num_samples)
            tr_idx, val_idx = train_test_split(
                indices, test_size=val_split, random_state=args.seed
            )

            if not os.path.exists(SCALER_FILE):
                logger.info("   Fitting StandardScaler (training set only)...")

                # Convert lazy datasets to numpy for reliable indexing
                # (h5py and _TransposedH5Dataset may not support fancy indexing)
                if hasattr(outp, "_dataset") or hasattr(outp, "file"):
                    # Lazy h5py or _TransposedH5Dataset - load training subset
                    outp_train = np.array([outp[i] for i in tr_idx])
                else:
                    # Already numpy array
                    outp_train = outp[tr_idx]

                # Ensure 2D for StandardScaler: (N,) -> (N, 1)
                if outp_train.ndim == 1:
                    outp_train = outp_train.reshape(-1, 1)

                scaler = StandardScaler()
                scaler.fit(outp_train)
                with open(SCALER_FILE, "wb") as f:
                    pickle.dump(scaler, f)

            # Cleanup: close file handles BEFORE deleting references
            if "_lazy_handle" in dir() and _lazy_handle is not None:
                try:
                    _lazy_handle.close()
                except Exception:
                    pass
            del inp, outp
            gc.collect()

            logger.info("   ✔ Cache creation complete, synchronizing ranks...")
        else:
            # NON-MAIN RANKS: Wait for cache creation
            # Log that we're waiting (helps with debugging)
            import time

            wait_start = time.time()
            while not (
                os.path.exists(CACHE_FILE)
                and os.path.exists(SCALER_FILE)
                and os.path.exists(META_FILE)
            ):
                time.sleep(5)  # Check every 5 seconds
                elapsed = time.time() - wait_start
                if elapsed > 60 and int(elapsed) % 60 < 5:  # Log every ~minute
                    logger.info(
                        f"   [Rank {accelerator.process_index}] Waiting for cache creation... ({int(elapsed)}s)"
                    )
            # Small delay to ensure files are fully written
            time.sleep(2)

    # ==========================================================================
    # PHASE 2: SYNCHRONIZED LOADING (All Ranks)
    # ==========================================================================
    accelerator.wait_for_everyone()

    # Load metadata
    with open(META_FILE, "rb") as f:
        meta = pickle.load(f)
        full_shape = meta["shape"]
        out_dim = meta["out_dim"]

    # Load and validate scaler
    with open(SCALER_FILE, "rb") as f:
        scaler = pickle.load(f)

    if not hasattr(scaler, "scale_") or scaler.scale_ is None:
        raise RuntimeError("CRITICAL: Scaler is not properly fitted (scale_ is None)")

    # Load targets only (memory-efficient - avoids loading large input arrays)
    # This is critical for HPC environments with memory constraints during DDP
    outp = load_outputs_only(args.data_path)

    # Ensure 2D for StandardScaler: (N,) -> (N, 1)
    if outp.ndim == 1:
        outp = outp.reshape(-1, 1)

    y_scaled = scaler.transform(outp).astype(np.float32)
    y_tensor = torch.tensor(y_scaled)

    # Regenerate indices (deterministic with same seed)
    indices = np.arange(full_shape[0])
    tr_idx, val_idx = train_test_split(
        indices, test_size=val_split, random_state=args.seed
    )

    # Create datasets
    tr_ds = MemmapDataset(CACHE_FILE, y_tensor, full_shape, tr_idx)
    val_ds = MemmapDataset(CACHE_FILE, y_tensor, full_shape, val_idx)

    # Create DataLoaders with thread-safe configuration
    loader_kwargs = {
        "batch_size": args.batch_size,
        "num_workers": args.workers,
        "pin_memory": True,
        "persistent_workers": (args.workers > 0),
        "prefetch_factor": 2 if args.workers > 0 else None,
        "worker_init_fn": memmap_worker_init_fn if args.workers > 0 else None,
    }

    train_dl = DataLoader(tr_ds, shuffle=True, **loader_kwargs)
    val_dl = DataLoader(val_ds, shuffle=False, **loader_kwargs)

    # Return spatial shape (H, W) for model initialization
    in_shape = full_shape[2:]

    return train_dl, val_dl, scaler, in_shape, out_dim
