"""
PyTorch Dataset Classes for Trading Data
=========================================
Provides Dataset implementations for training ML models with trading data.

This module implements:
- TradingDataset: For attention and transformer models with sequences
- TimeSeriesDataset: For time series models with sliding windows
- Sequential data loading with configurable sequence lengths

Author: ML Pipeline (NEXUS v4.0)
Created: 2026-01-25
"""

from typing import Optional, List, Tuple, Dict, Any, Union
from dataclasses import dataclass, field

import numpy as np
import pandas as pd
from loguru import logger

try:
    import torch
    from torch.utils.data import Dataset, DataLoader
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False
    logger.warning("PyTorch not available. Dataset classes will raise errors if used.")


@dataclass
class DatasetConfig:
    """Configuration for trading datasets."""

    # Sequence parameters
    sequence_length: int = 60
    target_horizon: int = 12
    stride: int = 1

    # Feature generation
    compute_returns: bool = True
    return_horizons: List[int] = field(default_factory=lambda: [1, 5, 12])
    normalize_features: bool = True
    normalization_method: str = 'zscore'  # 'zscore', 'minmax', 'robust'

    # Target configuration
    target_type: str = 'return'  # 'return', 'direction', 'multi_horizon'
    classification_threshold: float = 0.0

    # Data type
    dtype: str = 'float32'


def _check_torch():
    """Check if PyTorch is available."""
    if not TORCH_AVAILABLE:
        raise ImportError(
            "PyTorch is required for Dataset classes. "
            "Install with: pip install torch"
        )


class TradingDataset:
    """
    PyTorch Dataset for sequence-based trading data.

    This dataset is designed for attention and transformer models
    that require fixed-length sequences of historical data.

    Features:
    - Configurable sequence length
    - Automatic return computation
    - Feature normalization
    - Supports both regression and classification targets

    Usage:
        config = DatasetConfig(sequence_length=60, target_horizon=12)
        dataset = TradingDataset(df, config)
        dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

        for features, targets in dataloader:
            # features shape: (batch_size, sequence_length, num_features)
            # targets shape: (batch_size,) or (batch_size, num_targets)
            predictions = model(features)
    """

    def __init__(
        self,
        data: pd.DataFrame,
        config: Optional[DatasetConfig] = None,
        feature_columns: Optional[List[str]] = None,
        target_column: Optional[str] = None
    ):
        """
        Initialize the trading dataset.

        Args:
            data: DataFrame with OHLCV data (indexed by timestamp)
            config: Dataset configuration
            feature_columns: Columns to use as features (None = auto-detect)
            target_column: Column to use as target (None = compute from close)
        """
        _check_torch()

        self.config = config or DatasetConfig()
        self.data = data.copy()
        self.feature_columns = feature_columns
        self.target_column = target_column

        # Normalization statistics
        self._feature_means: Optional[np.ndarray] = None
        self._feature_stds: Optional[np.ndarray] = None
        self._feature_mins: Optional[np.ndarray] = None
        self._feature_maxs: Optional[np.ndarray] = None

        # Prepare data
        self._prepare_data()

        logger.info(f"TradingDataset initialized")
        logger.info(f"  Samples: {len(self)}")
        logger.info(f"  Sequence length: {self.config.sequence_length}")
        logger.info(f"  Features: {self.features.shape[1]}")

    def _prepare_data(self):
        """Prepare data for training."""
        df = self.data

        # Compute returns if requested
        if self.config.compute_returns:
            df = self._compute_returns(df)

        # Determine feature columns
        if self.feature_columns is None:
            # Auto-detect: exclude OHLCV base columns and target
            exclude_cols = {'open', 'high', 'low', 'close', 'volume', 'vwap', 'timestamp'}
            self.feature_columns = [
                col for col in df.columns
                if col.lower() not in exclude_cols
            ]

            # If no derived features, use OHLCV
            if not self.feature_columns:
                self.feature_columns = ['open', 'high', 'low', 'close', 'volume']

        # Extract features
        self.features = df[self.feature_columns].values.astype(np.float32)

        # Compute or extract target
        if self.target_column and self.target_column in df.columns:
            self.targets = df[self.target_column].values.astype(np.float32)
        else:
            self.targets = self._compute_target(df)

        # Normalize features
        if self.config.normalize_features:
            self._normalize_features()

        # Handle NaN values
        self._handle_nan()

        # Calculate valid indices (sequences that don't include NaN or exceed bounds)
        self._calculate_valid_indices()

    def _compute_returns(self, df: pd.DataFrame) -> pd.DataFrame:
        """Compute return features from price data."""
        df = df.copy()

        if 'close' not in df.columns:
            return df

        # Simple returns
        for horizon in self.config.return_horizons:
            df[f'return_{horizon}'] = df['close'].pct_change(horizon)

        # Log return
        df['log_return'] = np.log(df['close'] / df['close'].shift(1))

        # Volatility (rolling std of returns)
        df['volatility'] = df['log_return'].rolling(window=20).std()

        # Normalized high-low range
        df['range_norm'] = (df['high'] - df['low']) / df['close']

        # Body ratio
        df['body_ratio'] = (df['close'] - df['open']) / (df['high'] - df['low'] + 1e-10)

        return df

    def _compute_target(self, df: pd.DataFrame) -> np.ndarray:
        """Compute target values based on configuration."""
        if 'close' not in df.columns:
            raise ValueError("DataFrame must have 'close' column for target computation")

        horizon = self.config.target_horizon

        if self.config.target_type == 'return':
            # Future return
            target = df['close'].pct_change(horizon).shift(-horizon).values

        elif self.config.target_type == 'direction':
            # Binary direction (1 = up, 0 = down)
            future_return = df['close'].pct_change(horizon).shift(-horizon)
            target = (future_return > self.config.classification_threshold).astype(np.float32).values

        elif self.config.target_type == 'multi_horizon':
            # Multiple horizon targets
            targets = []
            for h in self.config.return_horizons:
                ret = df['close'].pct_change(h).shift(-h).values
                targets.append(ret)
            target = np.stack(targets, axis=1)

        else:
            raise ValueError(f"Unknown target_type: {self.config.target_type}")

        return target.astype(np.float32)

    def _normalize_features(self):
        """Normalize features using the specified method."""
        method = self.config.normalization_method

        if method == 'zscore':
            self._feature_means = np.nanmean(self.features, axis=0)
            self._feature_stds = np.nanstd(self.features, axis=0) + 1e-10
            self.features = (self.features - self._feature_means) / self._feature_stds

        elif method == 'minmax':
            self._feature_mins = np.nanmin(self.features, axis=0)
            self._feature_maxs = np.nanmax(self.features, axis=0)
            range_vals = self._feature_maxs - self._feature_mins + 1e-10
            self.features = (self.features - self._feature_mins) / range_vals

        elif method == 'robust':
            # Use median and IQR for robustness to outliers
            median = np.nanmedian(self.features, axis=0)
            q75 = np.nanpercentile(self.features, 75, axis=0)
            q25 = np.nanpercentile(self.features, 25, axis=0)
            iqr = q75 - q25 + 1e-10
            self.features = (self.features - median) / iqr
            self._feature_means = median
            self._feature_stds = iqr

    def _handle_nan(self):
        """Handle NaN values in features and targets."""
        # Replace NaN with 0 in features (after normalization, 0 is neutral)
        nan_mask = np.isnan(self.features)
        if nan_mask.any():
            logger.debug(f"Replacing {nan_mask.sum()} NaN values in features")
            self.features = np.nan_to_num(self.features, nan=0.0)

    def _calculate_valid_indices(self):
        """Calculate indices of valid sequences."""
        n_samples = len(self.features)
        seq_len = self.config.sequence_length
        horizon = self.config.target_horizon
        stride = self.config.stride

        # Valid range: from seq_len-1 to n_samples - horizon - 1
        min_idx = seq_len - 1
        max_idx = n_samples - horizon - 1

        if max_idx <= min_idx:
            self.valid_indices = np.array([], dtype=np.int64)
            logger.warning("No valid sequences available")
            return

        # Generate indices with stride
        self.valid_indices = np.arange(min_idx, max_idx + 1, stride)

        # Filter out indices with NaN targets
        if self.targets.ndim == 1:
            valid_target_mask = ~np.isnan(self.targets[self.valid_indices])
        else:
            valid_target_mask = ~np.isnan(self.targets[self.valid_indices]).any(axis=1)

        self.valid_indices = self.valid_indices[valid_target_mask]

    def __len__(self) -> int:
        """Return the number of valid samples."""
        return len(self.valid_indices)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get a single sample.

        Args:
            idx: Sample index

        Returns:
            Tuple of (features, target) as torch tensors
            - features: shape (sequence_length, num_features)
            - target: shape () for single target, (num_targets,) for multi
        """
        # Get the actual data index
        data_idx = self.valid_indices[idx]

        # Extract sequence
        seq_start = data_idx - self.config.sequence_length + 1
        seq_end = data_idx + 1
        sequence = self.features[seq_start:seq_end]

        # Extract target
        target = self.targets[data_idx]

        # Convert to tensors
        sequence_tensor = torch.from_numpy(sequence)
        target_tensor = torch.tensor(target)

        return sequence_tensor, target_tensor

    def get_feature_names(self) -> List[str]:
        """Get the list of feature column names."""
        return self.feature_columns.copy()

    def get_normalization_params(self) -> Dict[str, np.ndarray]:
        """Get normalization parameters for inference."""
        return {
            'means': self._feature_means,
            'stds': self._feature_stds,
            'mins': self._feature_mins,
            'maxs': self._feature_maxs,
            'method': self.config.normalization_method
        }

    def create_dataloader(
        self,
        batch_size: int = 32,
        shuffle: bool = True,
        num_workers: int = 0,
        pin_memory: bool = True
    ) -> 'DataLoader':
        """
        Create a PyTorch DataLoader for this dataset.

        Args:
            batch_size: Batch size
            shuffle: Whether to shuffle data
            num_workers: Number of worker processes
            pin_memory: Whether to pin memory (for GPU)

        Returns:
            PyTorch DataLoader
        """
        return DataLoader(
            self,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            pin_memory=pin_memory
        )


class TimeSeriesDataset:
    """
    PyTorch Dataset for sliding window time series data.

    This dataset is optimized for traditional time series models
    that use fixed-size sliding windows.

    Features:
    - Configurable window size and stride
    - Multiple target horizons
    - Efficient memory usage with views

    Usage:
        dataset = TimeSeriesDataset(df, window_size=100, stride=1)
        dataloader = DataLoader(dataset, batch_size=64)

        for window, targets in dataloader:
            # window shape: (batch_size, window_size, num_features)
            predictions = model(window)
    """

    def __init__(
        self,
        data: pd.DataFrame,
        window_size: int = 100,
        stride: int = 1,
        target_horizons: Optional[List[int]] = None,
        feature_columns: Optional[List[str]] = None,
        normalize: bool = True
    ):
        """
        Initialize the time series dataset.

        Args:
            data: DataFrame with OHLCV data
            window_size: Size of each sliding window
            stride: Step size between windows
            target_horizons: List of future horizons for targets
            feature_columns: Columns to use as features
            normalize: Whether to normalize features
        """
        _check_torch()

        self.data = data.copy()
        self.window_size = window_size
        self.stride = stride
        self.target_horizons = target_horizons or [1, 5, 12]
        self.feature_columns = feature_columns
        self.normalize = normalize

        # Statistics for normalization
        self._means: Optional[np.ndarray] = None
        self._stds: Optional[np.ndarray] = None

        self._prepare_data()

        logger.info(f"TimeSeriesDataset initialized")
        logger.info(f"  Windows: {len(self)}")
        logger.info(f"  Window size: {window_size}")
        logger.info(f"  Features: {self.features.shape[1]}")

    def _prepare_data(self):
        """Prepare sliding window data."""
        df = self.data

        # Determine feature columns
        if self.feature_columns is None:
            # Use OHLCV by default
            base_cols = ['open', 'high', 'low', 'close', 'volume']
            self.feature_columns = [c for c in base_cols if c in df.columns]

        # Extract features
        self.features = df[self.feature_columns].values.astype(np.float32)

        # Compute targets (future returns at multiple horizons)
        if 'close' in df.columns:
            targets = []
            for horizon in self.target_horizons:
                ret = df['close'].pct_change(horizon).shift(-horizon).values
                targets.append(ret)
            self.targets = np.stack(targets, axis=1).astype(np.float32)
        else:
            # No targets available
            self.targets = np.zeros((len(df), len(self.target_horizons)), dtype=np.float32)

        # Normalize
        if self.normalize:
            self._means = np.nanmean(self.features, axis=0)
            self._stds = np.nanstd(self.features, axis=0) + 1e-10
            self.features = (self.features - self._means) / self._stds

        # Handle NaN
        self.features = np.nan_to_num(self.features, nan=0.0)

        # Calculate valid window start indices
        max_horizon = max(self.target_horizons)
        n_samples = len(self.features)
        max_start = n_samples - self.window_size - max_horizon

        if max_start < 0:
            self.window_starts = np.array([], dtype=np.int64)
        else:
            self.window_starts = np.arange(0, max_start + 1, self.stride)

    def __len__(self) -> int:
        """Return number of windows."""
        return len(self.window_starts)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get a single window and its targets.

        Args:
            idx: Window index

        Returns:
            Tuple of (window, targets) as torch tensors
            - window: shape (window_size, num_features)
            - targets: shape (num_horizons,)
        """
        start = self.window_starts[idx]
        end = start + self.window_size

        window = self.features[start:end]
        target_idx = end - 1  # Target is computed from last element of window
        targets = self.targets[target_idx]

        return torch.from_numpy(window), torch.from_numpy(targets)

    def get_full_sequence(
        self,
        start_idx: int,
        length: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get a continuous sequence of windows.

        Useful for sequential prediction where windows overlap.

        Args:
            start_idx: Starting window index
            length: Number of consecutive windows

        Returns:
            Tuple of (windows, targets)
            - windows: shape (length, window_size, num_features)
            - targets: shape (length, num_horizons)
        """
        end_idx = min(start_idx + length, len(self))
        actual_length = end_idx - start_idx

        windows = []
        targets = []

        for i in range(start_idx, end_idx):
            w, t = self[i]
            windows.append(w)
            targets.append(t)

        return torch.stack(windows), torch.stack(targets)

    def create_dataloader(
        self,
        batch_size: int = 64,
        shuffle: bool = True,
        num_workers: int = 0
    ) -> 'DataLoader':
        """Create a DataLoader for this dataset."""
        return DataLoader(
            self,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers
        )


def create_train_val_datasets(
    data: pd.DataFrame,
    val_ratio: float = 0.15,
    config: Optional[DatasetConfig] = None
) -> Tuple[TradingDataset, TradingDataset]:
    """
    Create training and validation datasets with temporal split.

    Args:
        data: Full DataFrame with OHLCV data
        val_ratio: Ratio of data to use for validation (from end)
        config: Dataset configuration

    Returns:
        Tuple of (train_dataset, val_dataset)
    """
    n_samples = len(data)
    split_idx = int(n_samples * (1 - val_ratio))

    train_data = data.iloc[:split_idx]
    val_data = data.iloc[split_idx:]

    logger.info(f"Creating train/val split:")
    logger.info(f"  Train: {len(train_data):,} samples")
    logger.info(f"  Val: {len(val_data):,} samples")

    train_dataset = TradingDataset(train_data, config)
    val_dataset = TradingDataset(val_data, config)

    return train_dataset, val_dataset


if __name__ == "__main__":
    # Test the dataset classes
    print("Testing Dataset Classes...")

    # Create sample data
    np.random.seed(42)
    n = 5000

    dates = pd.date_range('2024-01-01', periods=n, freq='5min')
    price = 2650 + np.cumsum(np.random.randn(n) * 0.5)

    df = pd.DataFrame({
        'open': price,
        'high': price + np.abs(np.random.randn(n)) * 2,
        'low': price - np.abs(np.random.randn(n)) * 2,
        'close': price + np.random.randn(n) * 0.5,
        'volume': np.random.randint(100, 1000, n)
    }, index=dates)

    # Test TradingDataset
    print("\n" + "=" * 60)
    print("Testing TradingDataset")
    print("=" * 60)

    config = DatasetConfig(
        sequence_length=60,
        target_horizon=12,
        target_type='return'
    )

    dataset = TradingDataset(df, config)
    print(f"Dataset length: {len(dataset)}")

    # Get a sample
    features, target = dataset[0]
    print(f"Features shape: {features.shape}")
    print(f"Target shape: {target.shape}")
    print(f"Feature names: {dataset.get_feature_names()}")

    # Test DataLoader
    loader = dataset.create_dataloader(batch_size=32)
    for batch_features, batch_targets in loader:
        print(f"Batch features shape: {batch_features.shape}")
        print(f"Batch targets shape: {batch_targets.shape}")
        break

    # Test TimeSeriesDataset
    print("\n" + "=" * 60)
    print("Testing TimeSeriesDataset")
    print("=" * 60)

    ts_dataset = TimeSeriesDataset(
        df,
        window_size=100,
        stride=5,
        target_horizons=[1, 5, 12]
    )
    print(f"Dataset length: {len(ts_dataset)}")

    window, targets = ts_dataset[0]
    print(f"Window shape: {window.shape}")
    print(f"Targets shape: {targets.shape}")

    # Test train/val split
    print("\n" + "=" * 60)
    print("Testing Train/Val Split")
    print("=" * 60)

    train_ds, val_ds = create_train_val_datasets(df, val_ratio=0.15, config=config)
    print(f"Train dataset: {len(train_ds)} samples")
    print(f"Val dataset: {len(val_ds)} samples")

    print("\nTest complete!")
