import gin
import numpy as np
import torch
from tensordict import NonTensorStack, TensorDict


class Data(TensorDict):
    """
    A specialized TensorDict for time series data.
    
    This class extends TensorDict to provide a structured container for time series data
    with associated metadata like descriptions, labels, and dataset names.
    
    Parameters:
    -----------
    ts_array : torch.Tensor or array-like
        Time series data. Can be 2D (T, C) or 3D (B, T, C) where:
        - B is batch size
        - T is number of time steps
        - C is number of channels/features
    descriptions : list or numpy.ndarray
        Descriptions for each channel. Should be a list of strings for 2D data
        or a list of lists of strings for 3D data.
    labels : list, optional
        Labels for the time series. Should match the time dimension.
    dataset_name : str or list, optional
        Name(s) of the dataset(s).
    positions : torch.Tensor, optional
        Position information for the time series data.
    """
    def __init__(self, *, ts_array, descriptions, labels=None, dataset_name=None, positions=None, embeds=None):

        ts_array = torch.as_tensor(ts_array)

        if ts_array.ndim == 2:

            assert isinstance(descriptions, np.ndarray) or is_list_of_strings(descriptions), f"Descriptions must be a list of strings, got {type(descriptions)}, {descriptions}"
            assert labels is None or is_list_of_strings(labels), f"Labels must be a list of strings, got {type(labels)}, {labels}"
            assert labels is None or len(labels)==ts_array.shape[-2], "Number of labels should be the same as number of data points"
            assert dataset_name is None or is_string(dataset_name), f"Dataset name must be a list of strings, got {type(dataset_name)}, {dataset_name}"

            ts_array = ts_array.unsqueeze(0)
            if positions is not None: positions = positions.unsqueeze(0)

            descriptions = [descriptions]
            labels = [[None] * ts_array.shape[-2]] if labels is None else [labels]
            dataset_name = [dataset_name]

        elif ts_array.ndim == 3:
            # ensure that string inputs are correct format
            assert isinstance(descriptions, np.ndarray) or is_list_of_lists(descriptions), f"Descriptions must be a list of lists, got {type(descriptions)}, {descriptions}"
            assert labels is None or is_list_of_lists(labels), f"Labels must be a list of lists, got {type(labels)}, {labels}"
            assert labels is None or len(labels)==ts_array.shape[-3], "Labels should be list of length number of batches"
            assert labels is None or len(labels[0])==ts_array.shape[-2], "Number of labels should be the same as number of data points for each labels list in the batch"
            assert labels is None or all(len(sublist) == len(labels[0]) for sublist in labels), "Labels for every batch should have the same length"
            assert dataset_name is None or is_list_of_strings(dataset_name), f"Dataset name must be a list of strings, got {type(dataset_name)}, {dataset_name}"

            # descriptions must be B, C
            assert len(descriptions) == ts_array.shape[0], f"Number of descriptions must be equal to batch size, got {len(descriptions)} and {ts_array.shape[0]}"
            assert len(descriptions[0]) == ts_array.shape[-1], f"Number of descriptions must be equal to number of channels, got {len(descriptions[0])} and {ts_array.shape[-1]}"

            # labels and dataset_name must be B
            # if labels is not None: assert len(labels) == ts_array.shape[0], f"Number of labels must be equal to batch size, got {len(labels)} and {ts_array.shape[0]}"
            if dataset_name is not None: assert len(dataset_name) == ts_array.shape[0], f"Number of dataset names must be equal to batch size, got {len(dataset_name)} and {ts_array.shape[0]}"

            ts_array = ts_array
            labels = [[None] * ts_array.shape[1] for _ in range(ts_array.shape[0])] if labels is None else labels
            dataset_name =  [None] * ts_array.shape[0] if dataset_name is None else dataset_name

            

        else:
            raise ValueError(f"Data must be 2 or 3 dimensional, got {ts_array.ndim} dimensional data")
        
       
        # Validate embeds if provided
        if embeds is not None:
            embeds = torch.as_tensor(embeds)
            assert embeds.shape[0] == ts_array.shape[0], \
                f"Embeddings batch size must match ts_array: {embeds.shape[0]} vs {ts_array.shape[0]}"

        descriptions = NonTensorStack(*descriptions)
        labels = NonTensorStack(*labels)
        dataset_name = NonTensorStack(*dataset_name)

        td = {
            'ts_array': ts_array,
            'positions': positions,
            'descriptions': descriptions,
            'labels': labels,
            'dataset_name': dataset_name,
        }

        if embeds is not None:
            td['embeds'] = embeds 

        super().__init__(td, batch_size=[ts_array.shape[0]])

def is_list(x):
    """
    Check if an object is a list.
    
    Parameters:
    -----------
    x : object
        Object to check.
        
    Returns:
    --------
    bool
        True if the object is a list, False otherwise.
    """
    return isinstance(x, list)

def is_list_of_lists(x):
    """
    Check if an object is a list of lists.
    
    Parameters:
    -----------
    x : object
        Object to check.
        
    Returns:
    --------
    bool
        True if the object is a list of lists, False otherwise.
    """
    return isinstance(x, list) and all(isinstance(i, list) for i in x)

def is_list_of_strings(x):
    """
    Check if an object is a list of strings.
    
    Parameters:
    -----------
    x : object
        Object to check.
        
    Returns:
    --------
    bool
        True if the object is a list of strings, False otherwise.
    """
    return isinstance(x, list) and all(isinstance(i, str) for i in x)

def is_string(x):
    """
    Check if an object is a string.
    
    Parameters:
    -----------
    x : object
        Object to check.
        
    Returns:
    --------
    bool
        True if the object is a string, False otherwise.
    """
    return isinstance(x, str)

def collate_data(data_list):
    """
    Collate a list of Data objects into a single Data object.
    
    Parameters:
    -----------
    data_list : list or Data
        List of Data objects to collate, or a single Data object.
        
    Returns:
    --------
    Data
        Collated Data object.
        
    Raises:
    -------
    ValueError
        If data_list is neither a list of Data objects nor a single Data object.
    """
    try:
        if isinstance(data_list, list):
            return Data.cat(data_list, dim=0)
        elif isinstance(data_list, Data):
            return data_list
        else:
            raise ValueError(f"Data must be a list of Data objects or a single Data object, got {type(data_list)}")
    except Exception as e:
        print("couldn't concat", e)
        print(data_list)
        return data_list