import abc
import atexit
from collections.abc import Sequence
from typing import Dict
import warnings

import numpy as np
import torch
from charm_toolkit.Data import Data
from sklearn.preprocessing import StandardScaler


class BaseTimeSeriesDataset(torch.utils.data.Dataset, metaclass=abc.ABCMeta):
    """
    Abstract base class for time series datasets.
    
    This class provides a foundation for all time series dataset implementations
    with common functionality for data handling and preprocessing.
    
    Parameters:
    -----------
    dataset_name : str
        Name identifier for the dataset.
    descriptions : list
        List of descriptions for each channel in the time series.
    labels : list, optional
        Labels associated with the time series data.
    """
    def __init__(self, dataset_name, descriptions, labels=None):
        self.dataset_name = dataset_name
        self.descriptions = descriptions
        self.labels = labels
        self.scaler = None

    @abc.abstractmethod
    def __len__(self):
        """
        Returns the number of samples in the dataset.
        
        Returns:
        --------
        int
            Number of samples.
        """
        pass

    @abc.abstractmethod
    def __getitem__(self, idx):
        """
        Retrieves a sample from the dataset at the specified index.
        
        Parameters:
        -----------
        idx : int
            Index of the sample to retrieve.
            
        Returns:
        --------
        Data
            The requested data sample.
            
        Raises:
        -------
        StopIteration
            If the index is out of bounds.
        """
        if idx < 0:
            raise IndexError("Negative indices are not allowed in this dataset.")
            
        if idx >= len(self):
            raise StopIteration

    def apply_scaler(self, scaler=None, fit=True):
        """
        Apply a scaler to this dataset.
        If scaler is None, a new StandardScaler will be created and (optionally) fit.
        
        Parameters:
        -----------
        scaler : sklearn.preprocessing.StandardScaler, optional
            Scaler to apply. If None, a new StandardScaler is created.
        fit : bool, default=True
            Whether to fit the scaler on this dataset's data.
            
        Returns:
        --------
        BaseTimeSeriesDataset
            Self, for method chaining.
            
        Raises:
        -------
        ValueError
            If the dataset has no 'data' attribute.
        """
        if not hasattr(self, "data") or self.data is None:
            raise ValueError("Dataset has no 'data' to scale.")

        # Create or reuse scaler
        if scaler is None:
            warnings.warn(
            "No scaler provided. Defaulting to StandardScaler.",
            UserWarning
            )
            scaler = StandardScaler()

        # Fit only if required
        if fit:
            scaler.fit(self.data.numpy())

        # Transform and replace data
        scaled = scaler.transform(self.data.numpy())
        self.data = torch.tensor(scaled, dtype=torch.float32)
        self.scaler = scaler

        return self

    def inverse_transform(self, data):
        """
        Inverse transform a numpy array or torch tensor using the stored scaler.
        
        Parameters:
        -----------
        data : numpy.ndarray or torch.Tensor
            Data to inverse transform.
            
        Returns:
        --------
        numpy.ndarray
            Inverse transformed data.
            
        Raises:
        -------
        ValueError
            If no scaler is available.
        """
        if self.scaler is None:
            raise ValueError("This dataset was not scaled or scaler is unavailable.")
        if isinstance(data, torch.Tensor):
            data = data.numpy()
        return self.scaler.inverse_transform(data)


class UnwindowedDataset(BaseTimeSeriesDataset):
    """
    Dataset for raw, unwindowed time series data.
    
    This class represents a time series dataset without any windowing applied.
    It serves as a base for creating windowed or forecasting datasets.
    
    Parameters:
    -----------
    data : numpy.ndarray or torch.Tensor
        Time series data with shape (T, C) where T is time steps and C is channels.
    dataset_name : str
        Name identifier for the dataset.
    descriptions : list
        List of descriptions for each channel.
    labels : list, optional
        Labels associated with the time series (e.g., dates).
    scaler : sklearn.preprocessing.StandardScaler, optional
        Pre-fitted scaler to apply to the data.
    """
    def __init__(self, data, dataset_name, descriptions, labels=None, scaler=None):
        super().__init__(dataset_name, descriptions, labels=labels)

        if isinstance(data, np.ndarray):
            data = torch.as_tensor(data)

        assert isinstance(data, torch.Tensor), "data must be a torch.Tensor"
        assert data.ndim == 2, "data must be a 2D tensor (T, C)"

        # Apply scaler if provided
        if scaler is not None:
            scaled_data = scaler.transform(data.numpy())
            self.data = torch.tensor(scaled_data, dtype=torch.float32)
            self.scaler = scaler
        else:
            self.data = data.to(dtype=torch.float32, copy=False)

    def __getitem__(self, idx):
        """
        Not supported for UnwindowedDataset.
        
        Raises:
        -------
        ValueError
            Always raises this error as direct indexing is not supported.
        """
        raise ValueError(
            """
            UnwindowedDataset does not support __getitem__.
            Either use LazyWindowedDataset or PreWindowedDataset.
            Do so by calling:
            - dataset.window(window_size, stride)
            - dataset.forecasting(train_horizon, test_horizon)
            """
        )

    def __len__(self):
        """
        Returns the number of time steps in the dataset.
        
        Returns:
        --------
        int
            Number of time steps.
        """
        return len(self.data)

    def window(self, window_size, stride):
        """
        Create a windowed version of this dataset.
        
        Parameters:
        -----------
        window_size : int
            Size of each window.
        stride : int
            Step size between consecutive windows.
            
        Returns:
        --------
        LazyWindowedDataset
            A windowed version of this dataset.
        """
        if window_size is None:
            return self
        return LazyWindowedDataset(self, window_size, stride)

    def forecasting(self, train_horizon, test_horizon):
        """
        Create a forecasting dataset from this dataset.
        
        Parameters:
        -----------
        train_horizon : int
            Number of time steps to use as input.
        test_horizon : int
            Number of time steps to predict.
            
        Returns:
        --------
        ForecastingDataset
            A forecasting dataset based on this dataset.
        """
        return ForecastingDataset(self, train_horizon, test_horizon)



class PreWindowedDataset(BaseTimeSeriesDataset):
    """
    Dataset for pre-windowed time series data.
    
    This class represents a time series dataset where the data is already
    organized into windows/segments.
    
    Parameters:
    -----------
    data : torch.Tensor
        Pre-windowed time series data with shape (N, T, C) where N is number of windows,
        T is time steps per window, and C is channels.
    dataset_name : str
        Name identifier for the dataset.
    descriptions : list
        List of descriptions for each channel.
    labels : list, optional
        Labels associated with each window.
    """
    def __init__(self, data, dataset_name, descriptions, labels=None):
        assert data.ndim == 3, "data must be a 3D tensor, N, T, C"
        # data: (N, T, C)
        super().__init__(dataset_name, descriptions, labels)
        self.data = data
        self._length = data.shape[0]

    def __len__(self):
        """
        Returns the number of windows in the dataset.
        
        Returns:
        --------
        int
            Number of windows.
        """
        return self._length

    def __getitem__(self, idx):
        """
        Retrieves a window from the dataset at the specified index.
        
        Parameters:
        -----------
        idx : int
            Index of the window to retrieve.
            
        Returns:
        --------
        Data
            Data object containing the requested window.
        """
        super().__getitem__(idx)

        return Data(
            ts_array=self.data[idx],
            descriptions=self.descriptions,
            labels=self.labels[idx],
            dataset_name=self.dataset_name,
        )


class LazyWindowedDataset(BaseTimeSeriesDataset):
    """
    Dataset that creates windows on-the-fly from an UnwindowedDataset.
    
    This class efficiently creates sliding windows from a base dataset
    without pre-computing all windows, saving memory.
    
    Parameters:
    -----------
    dataset : UnwindowedDataset
        Base dataset to create windows from.
    window_size : int
        Size of each window.
    stride : int
        Step size between consecutive windows.
    """
    def __init__(self, dataset, window_size, stride):
        assert isinstance(dataset, UnwindowedDataset), f"dataset must be an UnwindowedDataset, got {type(dataset)} instead"

        # data: (T, C)
        super().__init__(dataset.dataset_name, dataset.descriptions, dataset.labels)
        self.data = dataset.data
        self.dataset = dataset

        self.window_size = window_size
        self.stride = stride
        T = len(self.data)
        self._length = max(0, (T - window_size) // stride + 1)

    def __len__(self):
        """
        Returns the number of windows that can be created.
        
        Returns:
        --------
        int
            Number of windows.
        """
        return self._length

    def __getitem__(self, idx):
        """
        Creates and returns a window at the specified index.
        
        Parameters:
        -----------
        idx : int
            Index of the window to retrieve.
            
        Returns:
        --------
        Data
            Data object containing the requested window.
        """
        super().__getitem__(idx)

        start = idx * self.stride
        end = start + self.window_size
        window_data = self.data[start:end]  # (window_size, C)

        return Data(
            ts_array=window_data,
            descriptions=self.descriptions,
            labels=self.labels[start:end] if self.labels is not None else None,
            dataset_name=self.dataset_name,
        )


class ForecastingDataset(BaseTimeSeriesDataset):
    """
    Dataset for time series forecasting tasks.
    
    This class creates pairs of input and target sequences for forecasting,
    where the target sequence follows the input sequence in time.
    
    Parameters:
    -----------
    dataset : UnwindowedDataset
        Base dataset to create forecasting pairs from.
    train_horizon : int
        Number of time steps to use as input (context size).
    test_horizon : int
        Number of time steps to predict (forecast horizon).
    stride : int, default=1
        Step size between consecutive samples.
    """
    def __init__(self, dataset, train_horizon, test_horizon, stride=1):
        assert isinstance(dataset, UnwindowedDataset), "dataset must be a UnwindowedDataset"
        super().__init__(dataset.dataset_name, dataset.descriptions, dataset.labels)
        self.dataset = dataset
        self.train_horizon = train_horizon
        self.test_horizon = test_horizon
        self.stride = stride
        
        window = train_horizon + test_horizon
        T = dataset.data.shape[0]
        self._length = max(0, (T - window) // stride + 1)

    def __len__(self):
        return self._length

    def __getitem__(self, idx):
        super().__getitem__(idx)
        
        start = idx * self.stride
        input_end = start + self.train_horizon
        target_end = input_end + self.test_horizon
        
        train_data = self.dataset.data[start:input_end]
        test_data = self.dataset.data[input_end:target_end]

        return (
            Data(
                ts_array=train_data,
                descriptions=self.descriptions,
                labels=self.labels[start:input_end] if self.labels is not None else None,
                dataset_name=self.dataset_name,
            ),
            Data(
                ts_array=test_data,
                descriptions=self.descriptions,
                labels=self.labels[input_end:target_end] if self.labels is not None else None,
                dataset_name=self.dataset_name,
            )
        )


def train_test_split(dataset, train_ratio, sequential=True):
    """
    Split a dataset into training and testing subsets.
    
    Parameters:
    -----------
    dataset : BaseTimeSeriesDataset or dict or list
        Dataset(s) to split. Can be a single dataset, a dictionary of datasets,
        or a list of datasets.
    train_ratio : float
        Proportion of the dataset to include in the training split (0.0 to 1.0).
    sequential : bool, default=True
        If True, split the data sequentially (first part for training, second for testing).
        If False, split the data randomly.
        
    Returns:
    --------
    tuple
        (train_dataset, test_dataset) if input is a single dataset
        (train_dict, test_dict) if input is a dictionary
        (train_list, test_list) if input is a list
    """
    
    # check if dict or list
    if isinstance(dataset, dict):
        train_out = {}
        test_out = {}
        for ds_name, ds in dataset.items():
            train_ds, test_ds = split_helper(ds, train_ratio, sequential)
            train_out[ds_name] = train_ds
            test_out[ds_name] = test_ds
        return train_out, test_out
    elif isinstance(dataset, list):
        train_out = []
        test_out = []
        for ds in dataset:
            train_ds, test_ds = split_helper(ds, train_ratio, sequential)
            train_out.append(train_ds)
            test_out.append(test_ds)
        return train_out, test_out
    else:
        return split_helper(dataset, train_ratio, sequential)
        
def split_helper(dataset, train_ratio, sequential):
    """
    Helper function to split a single dataset into training and testing subsets.
    
    Parameters:
    -----------
    dataset : BaseTimeSeriesDataset
        Dataset to split.
    train_ratio : float
        Proportion of the dataset to include in the training split (0.0 to 1.0).
    sequential : bool
        If True, split the data sequentially (first part for training, second for testing).
        If False, split the data randomly.
        
    Returns:
    --------
    tuple
        (train_dataset, test_dataset)
    """
    assert hasattr(dataset, "data"), "Dataset must have a .data attribute"

    n = len(dataset)
    if sequential:
        split_idx = int(n * train_ratio)
        idx_train = np.arange(0, split_idx)
        idx_test = np.arange(split_idx, n)
    else:
        indices = np.random.permutation(n)
        split_idx = int(n * train_ratio)
        idx_train = indices[:split_idx]
        idx_test = indices[split_idx:]

    data_train = dataset.data[idx_train].clone()
    data_test = dataset.data[idx_test].clone()

    labels_train = labels_test = None
    if dataset.labels is not None:
        labels_train = [dataset.labels[i] for i in idx_train]
        labels_test  = [dataset.labels[i] for i in idx_test]


    def make_subset(data, labels):
        kwargs = dict(
            data=data,
            dataset_name=dataset.dataset_name,
            descriptions=dataset.descriptions,
        )
        
        if labels is not None:
            kwargs["labels"] = labels

        if getattr(dataset, "scaler", None) is not None:
            kwargs["scaler"] = dataset.scaler
        return dataset.__class__(**kwargs)

    return make_subset(data_train, labels_train), make_subset(data_test, labels_test)
