import abc
from typing import Optional, Union
import torch
import numpy as np
from sklearn.preprocessing import StandardScaler

from charm_toolkit.Data import Data, collate_data
from charm_toolkit.BaseTimesSeries import LazyWindowedDataset, ForecastingDataset, UnwindowedDataset, train_test_split, PreWindowedDataset

# ----------------------------
# Reconstruction dataset
# ----------------------------
class ReconstructionDataset(LazyWindowedDataset):
    """
    Dataset for time series reconstruction tasks.
    
    This class extends LazyWindowedDataset to provide data in a format suitable
    for reconstruction tasks, where the model aims to reconstruct the input.
    
    Parameters:
    -----------
    dataset : UnwindowedDataset
        Base dataset to create windows from.
    window_size : int
        Size of each window.
    stride : int
        Step size between consecutive windows.
    """
    def __getitem__(self, idx):
        """
        Get a sample for reconstruction training.
        
        Parameters:
        -----------
        idx : int
            Index of the sample to retrieve.
            
        Returns:
        --------
        tuple
            (data, target) pair where data is the input and target is the same data
            to be reconstructed.
        """
        data = super().__getitem__(idx)        
        return data, data["ts_array"]


def create_reconstruction_datasets(
    raw_data,
    descriptions=None,
    labels=None,
    dataset_name="dataset",
    window_size=32,
    stride=1,
    train_ratio=0.7,
    val_ratio=0.15,
    sequential=True,
    scale=True,
):
    """
    Returns train/val/test ReconstructionDataset instances with proper scaling.
    
    Parameters:
    -----------
    raw_data : torch.Tensor or numpy.ndarray
        Raw time series data with shape (T, C).
    descriptions : list, optional
        Descriptions for each channel.
    labels : list, optional
        Per-timestep labels (e.g., dates as strings).
    dataset_name : str, default="dataset"
        Name identifier for the dataset (e.g., asset ID).
    window_size : int, default=32
        Size of each window.
    stride : int, default=1
        Step size between consecutive windows.
    train_ratio : float, default=0.7
        Proportion of data to use for training.
    val_ratio : float, default=0.15
        Proportion of data to use for validation.
    sequential : bool, default=True
        If True, split the data sequentially; otherwise, split randomly.
        
    Returns:
    --------
    tuple
        (train_dataset, val_dataset, test_dataset) of ReconstructionDataset instances.
    
    Usage:
    ------
        train_ds, val_ds, test_ds = create_reconstruction_datasets(raw_data)
        
        # For inference (all data in one split):
        _, val_ds, _ = create_reconstruction_datasets(
            raw_data, ..., train_ratio=0, val_ratio=1.0, scale=False
        )
    """
    # Create base UnwindowedDataset
    full_dataset = UnwindowedDataset(raw_data, dataset_name, descriptions, labels=labels)
    
    # Split into train/val/test
    train_ds, temp_ds = train_test_split(full_dataset, train_ratio, sequential)
    relative_val_ratio = val_ratio / (1 - train_ratio)
    val_ds, test_ds = train_test_split(temp_ds, train_ratio=relative_val_ratio, sequential=sequential)
    
    # Fit scaler on train only and apply to all splits
    if scale:
        train_ds.apply_scaler(fit=True)
        scaler = train_ds.scaler              
        val_ds.apply_scaler(scaler, fit=False)
        test_ds.apply_scaler(scaler, fit=False)
    
    # Wrap in ReconstructionDataset with windowing
    train_recon = ReconstructionDataset(train_ds, window_size, stride)
    val_recon = ReconstructionDataset(val_ds, window_size, stride)
    test_recon = ReconstructionDataset(test_ds, window_size, stride)

    return train_recon, val_recon, test_recon

# ----------------------------
# Forecasting dataset
# ----------------------------
class ForecasterDataset(ForecastingDataset):
    """
    Dataset for time series forecasting tasks.
    
    This class extends ForecastingDataset to provide data in a format suitable
    for forecasting tasks, where the model aims to predict future values.
    
    Parameters:
    -----------
    dataset : UnwindowedDataset
        Base dataset to create forecasting pairs from.
    train_horizon : int
        Number of time steps to use as input.
    test_horizon : int
        Number of time steps to predict.
    """
    def __getitem__(self, idx):
        """
        Get a sample for forecasting training.
        
        Parameters:
        -----------
        idx : int
            Index of the sample to retrieve.
            
        Returns:
        --------
        tuple
            (input_data, target_data) pair where input_data is the historical data
            and target_data is the future values to predict.
        """
        train_data, test_data = super().__getitem__(idx)
        return train_data, test_data['ts_array']


def create_forecasting_datasets(
    raw_data,
    descriptions=None,
    labels=None,
    dataset_name="dataset",
    train_horizon=50,
    test_horizon=10,
    stride=1,
    train_ratio=0.7,
    val_ratio=0.15,
    sequential=True,
    scale=True,
):
    """
    Returns train/val/test ForecasterDataset instances with proper scaling.
    
    Parameters:
    -----------
    raw_data : torch.Tensor or numpy.ndarray
        Raw time series data with shape (T, C).
    descriptions : list, optional
        Descriptions for each channel.
    labels : list, optional
        Per-timestep labels (e.g., dates as strings).
    dataset_name : str, default="dataset"
        Name identifier for the dataset (e.g., asset ID).
    train_horizon : int, default=50
        Number of time steps to use as input (context size).
    test_horizon : int, default=10
        Number of time steps to predict (forecast horizon).
    stride : int, default=1
        Step size between consecutive samples.
    train_ratio : float, default=0.7
        Proportion of data to use for training.
    val_ratio : float, default=0.15
        Proportion of data to use for validation.
    sequential : bool, default=True
        If True, split the data sequentially; otherwise, split randomly.
        
    Returns:
    --------
    tuple
        (train_dataset, val_dataset, test_dataset) of ForecasterDataset instances.
        
    Usage:
    ------
        # For inference (all data in one split):
        _, val_ds, _ = create_forecasting_datasets(
            raw_data, ..., train_ratio=0, val_ratio=1.0, scale=False
        )
    """
    full_dataset = UnwindowedDataset(raw_data, dataset_name, descriptions, labels)

    # Split into train/val/test
    train_ds, temp_ds = train_test_split(full_dataset, train_ratio, sequential)
    relative_val_ratio = val_ratio / (1 - train_ratio)
    val_ds, test_ds = train_test_split(temp_ds, train_ratio=relative_val_ratio, sequential=sequential)
    
    # Fit scaler on train only and apply to all splits
    if scale:
        train_ds.apply_scaler(fit=True)
        scaler = train_ds.scaler              
        val_ds.apply_scaler(scaler, fit=False)
        test_ds.apply_scaler(scaler, fit=False)

    # Wrap in ForecasterDataset
    train_forecast = ForecasterDataset(train_ds, train_horizon, test_horizon, stride)
    val_forecast = ForecasterDataset(val_ds, train_horizon, test_horizon, stride)
    test_forecast = ForecasterDataset(test_ds, train_horizon, test_horizon, stride)

    return train_forecast, val_forecast, test_forecast

# ----------------------------
# Classification dataset
# ----------------------------
class ClassificationDataset(PreWindowedDataset):
    """
    Dataset for time series classification tasks.
    
    This class extends PreWindowedDataset to provide data in a format suitable
    for classification tasks, where the model aims to classify time series.
    
    Parameters:
    -----------
    base_dataset : PreWindowedDataset
        Base dataset containing the time series data and labels.
    """
    def __init__(self, base_dataset):
        """
        Initialize a ClassificationDataset.
        
        Parameters:
        -----------
        base_dataset : PreWindowedDataset
            Base dataset containing the time series data and labels.
        """
        data = base_dataset.data
        dataset_name = base_dataset.dataset_name
        descriptions = base_dataset.descriptions
        labels = base_dataset.labels
    
        # Expand each label to a list of length T (per-timestep label)
        if labels is not None:
            T = data.shape[1]
            labels = [[str(label)] * T for label in labels]
    
        super().__init__(data, dataset_name, descriptions, labels)

    def __getitem__(self, idx):
        """
        Get a sample for classification training.
        
        Parameters:
        -----------
        idx : int
            Index of the sample to retrieve.
            
        Returns:
        --------
        tuple
            (data, label) pair where data is the input time series and label is the
            classification target.
        """
        data_obj = super().__getitem__(idx)
        label_val = torch.tensor(int(data_obj["labels"][0][0]), dtype=torch.long)
        label_val = label_val.unsqueeze(0)

        return data_obj, label_val


def create_classification_datasets(
    raw_data,
    labels,  # one label per time series
    descriptions,
    train_ratio=0.7,
    val_ratio=0.15,
    sequential=True,
):
    """
    Returns train/val/test ClassificationDataset instances.
    
    Parameters:
    -----------
    raw_data : torch.Tensor or numpy.ndarray
        Raw time series data with shape (N, T, C) where N is number of samples,
        T is time steps, and C is channels.
    labels : list
        List of labels, one per time series.
    descriptions : list
        Descriptions for each channel.
    train_ratio : float, default=0.7
        Proportion of data to use for training.
    val_ratio : float, default=0.15
        Proportion of data to use for validation.
    sequential : bool, default=True
        If True, split the data sequentially; otherwise, split randomly.
        
    Returns:
    --------
    tuple
        (train_dataset, val_dataset, test_dataset) of ClassificationDataset instances.
    """
    
    full_dataset = PreWindowedDataset(raw_data, "dataset", descriptions, labels)

    # Split train/val/test
    train_ds, temp_ds = train_test_split(full_dataset, train_ratio, sequential)
    relative_val_ratio = val_ratio / (1 - train_ratio)
    val_ds, test_ds = train_test_split(temp_ds, train_ratio=relative_val_ratio, sequential=sequential)
    
    # Wrap in ClassificationDataset
    
    train_cls = ClassificationDataset(train_ds)
    val_cls   = ClassificationDataset(val_ds)
    test_cls  = ClassificationDataset(test_ds)

    return train_cls, val_cls, test_cls


def collator(batch):
    """
    Collate function for DataLoader with classification datasets.
    
    Parameters:
    -----------
    batch : list
        List of (data, label) tuples.
        
    Returns:
    --------
    tuple
        (collated_data, collated_labels) where collated_data is a Data object
        and collated_labels is a tensor of labels.
    """
    xs, ys = zip(*batch)  
    x = collate_data(list(xs))
    y = torch.stack(list(ys), dim=0).squeeze(1)
    return x, y

    