"""PyTorch Dataset classes for cryptocurrency time series data."""

from typing import Literal, Tuple
import numpy as np
import torch
from torch.utils.data import Dataset


class CryptoDataset(Dataset):
    """PyTorch Dataset for cryptocurrency time series prediction."""

    def __init__(
        self,
        X: np.ndarray,
        y: np.ndarray,
        task: Literal['classification', 'regression'] = 'classification'
    ):
        """Initialize dataset.

        Args:
            X: Feature sequences of shape (n_samples, lookback, n_features)
            y: Labels of shape (n_samples,)
            task: Task type ('classification' or 'regression')
        """
        self.X = torch.FloatTensor(X)

        if task == 'classification':
            self.y = torch.LongTensor(y)
        else:  # regression
            self.y = torch.FloatTensor(y)

        self.task = task

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

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

        Args:
            idx: Sample index

        Returns:
            Tuple of (features, label)
        """
        return self.X[idx], self.y[idx]

    def get_feature_dim(self) -> int:
        """Get number of features.

        Returns:
            Number of features
        """
        return self.X.shape[-1]

    def get_sequence_length(self) -> int:
        """Get sequence length (lookback).

        Returns:
            Sequence length
        """
        return self.X.shape[1]


def create_dataloaders(
    train_data: Tuple[np.ndarray, np.ndarray],
    val_data: Tuple[np.ndarray, np.ndarray],
    test_data: Tuple[np.ndarray, np.ndarray],
    task: Literal['classification', 'regression'] = 'classification',
    batch_size: int = 64,
    num_workers: int = 0,  # 0 for M1 compatibility
    pin_memory: bool = False  # False for MPS
) -> Tuple[torch.utils.data.DataLoader, ...]:
    """Create DataLoaders for train/val/test sets.

    Args:
        train_data: Tuple of (X_train, y_train)
        val_data: Tuple of (X_val, y_val)
        test_data: Tuple of (X_test, y_test)
        task: Task type
        batch_size: Batch size
        num_workers: Number of worker processes (0 for M1)
        pin_memory: Whether to pin memory (False for MPS)

    Returns:
        Tuple of (train_loader, val_loader, test_loader)
    """
    # Create datasets
    train_dataset = CryptoDataset(*train_data, task=task)
    val_dataset = CryptoDataset(*val_data, task=task)
    test_dataset = CryptoDataset(*test_data, task=task)

    # Create dataloaders
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=pin_memory
    )

    val_loader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory
    )

    test_loader = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory
    )

    return train_loader, val_loader, test_loader


if __name__ == "__main__":
    # Test the dataset
    from fetcher import BinanceDataFetcher
    from preprocessor import CryptoPreprocessor, split_data

    # Fetch and preprocess data
    fetcher = BinanceDataFetcher()
    df = fetcher.fetch_ohlcv('BTC', '1h')

    preprocessor = CryptoPreprocessor(lookback=168)
    X, y, _ = preprocessor.preprocess_pipeline(df, task='classification')

    # Split data
    train_data, val_data, test_data = split_data(X, y)

    # Create dataloaders
    train_loader, val_loader, test_loader = create_dataloaders(
        train_data, val_data, test_data,
        batch_size=32
    )

    # Test iteration
    for batch_X, batch_y in train_loader:
        print(f"Batch X shape: {batch_X.shape}")
        print(f"Batch y shape: {batch_y.shape}")
        break