from typing import (
    Union,
    Tuple,
    List,
    Iterator,
    Mapping,
    Sequence,
    TypeAlias,
    cast,
)

from os import path
import math
from enum import Enum

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler, OrdinalEncoder, TargetEncoder

import torch
from torch.utils.data import Dataset, IterableDataset

from tqdm import tqdm
from datasets import Dataset as HFDataset

from litespecformer.utils import (
    convert_df_input_to_list_of_dicts_input,
    left_pad_and_cat_2D,
    validate_and_prepare_single_dict_task,
    convert_list_of_tensors_input_to_list_of_dicts_input,
    convert_tensor_input_to_list_of_dicts_input,
)

TensorOrArray: TypeAlias = torch.Tensor | np.ndarray


class DatasetMode(str, Enum):
    TRAIN = "train"
    VALIDATION = "validation"
    TEST = "test"


class LiteSpecFormerDataset(IterableDataset):
    """
    A dataset wrapper for Chronos-2 models.

    Arguments
    ----------
    inputs
        Time series data. Must be a list of dictionaries where each dictionary may have the following keys.
        - `target` (required): a 1-d or 2-d `torch.Tensor` or `np.ndarray` of shape (history_length,) or (n_variates, history_length).
        Forecasts will be generated for items in `target`.

        - `past_covariates` (optional): a dict of past-only covariates or past values of known future covariates. The keys of the dict
        must be names of the covariates and values must be 1-d `torch.Tensor` or `np.ndarray` with length equal to the `history_length`
        of `target`.

        - `future_covariates` (optional): a dict of future values of known future covariates. The keys of the dict must be names of the
        covariates and values must be 1-d `torch.Tensor` or `np.ndarray` with length equal to the `prediction_length`. All keys in
        `future_covariates` must be a subset of the keys in `past_covariates`.
        Note: when the mode is set to TRAIN, the values inside `future_covariates` are not technically used for training the model;
        however, this key is used to infer which covariates are known into the future. Therefore, if your task contains known future covariates,
        make sure that this key exists in `inputs`. The values of individual future covariates may be set to `None` or an empty array.

    context_length
        The maximum context length used for training or inference TODO: 注意这里是最大的上下文长度
    prediction_length
        The prediction horizon
    batch_size
        The batch size for training the model. Note that the batch size here means the number of time series, including target(s) and
        covariates, that are input into the model. If your data has multiple target and/or covariates, the effective number of time series
        tasks in a batch will be lower than this value.
    output_patch_size
        The output patch size of the model. This is used to compute the number of patches needed to cover `prediction_length`
    min_past
        The minimum number of time steps the context must have during training. All time series shorter than `min_past + prediction_length`
        are filtered out, by default 1
    mode
        `DatasetMode` governing whether to generate training, validation or test samples, by default "train"
    """

    def __init__(
        self,
        inputs: Sequence[
            Mapping[str, TensorOrArray | Mapping[str, TensorOrArray | None]]
        ],
        context_length: int,
        prediction_length: int,
        batch_size: int,
        output_patch_size: int,
        min_past: int = 1,
        mode: str | DatasetMode = DatasetMode.TRAIN,
    ) -> None:
        super().__init__()
        assert mode in {
            DatasetMode.TRAIN,
            DatasetMode.VALIDATION,
            DatasetMode.TEST,
        }, f"Invalid mode: {mode}"

        # 这里输入的数据是一个列表
        # 里面的元素是列表
        self.tasks = LiteSpecFormerDataset._prepare_tasks(
            inputs, prediction_length, min_past, mode
        )
        self.context_length = context_length
        self.prediction_length = prediction_length
        self.batch_size = batch_size
        self.num_output_patches = math.ceil(prediction_length / output_patch_size)
        self.min_past = min_past
        self.mode = mode

    @staticmethod
    def _prepare_tasks(
        inputs: Sequence[
            Mapping[str, TensorOrArray | Mapping[str, TensorOrArray | None]]
        ],
        prediction_length: int,
        min_past: int,
        mode: str | DatasetMode,
    ):
        # 构建承载训练数据的列表
        tasks = []

        for idx, raw_task in enumerate(inputs):
            if mode != DatasetMode.TEST:
                raw_future_covariates = raw_task.get("future_covariates", {})

                raw_future_covariates = cast(
                    dict[str, TensorOrArray | None], raw_future_covariates
                )

                if raw_future_covariates:
                    fixed_future_covariates = {}

                    for key, value in raw_future_covariates.items():
                        fixed_future_covariates[key] = (
                            np.full(prediction_length, np.nan)
                            if value is None or len(value) == 0
                            else value
                        )

                    raw_task = {
                        **raw_task,
                        "future_covariates": fixed_future_covariates,
                    }

            raw_task = cast(
                dict[str, TensorOrArray | Mapping[str, TensorOrArray]], raw_task
            )

            # convert to a format compatible with model's forward
            task = validate_and_prepare_single_dict_task(
                raw_task,
                idx=idx,
                prediction_length=prediction_length,
                min_past=min_past,
            )

            if task is None:
                # filter out time series that are too short to provide the required context and prediction length
                continue

            if (
                mode != DatasetMode.TEST
                and task[0].shape[-1] < min_past + prediction_length
            ):
                # filter tasks based on min_past + prediction_length
                continue
            tasks.append(task)

        # Check if the dataset is empty
        if len(tasks) == 0:
            raise ValueError(
                "The dataset is empty after filtering based on the length of the time series (length >= min_past + prediction_length). "
                "Please provide longer time series or reduce `min_past` or `prediction_length`. "
            )

        return tasks

    def _construct_slice(
        self, task_idx: int
    ) -> tuple[torch.Tensor, torch.Tensor | None, int]:
        # 根据指定的索引从数据集中抽取样本
        # shape:  (task_n_targets + task_n_covariates, history_length)
        (
            task_past_tensor,
            task_future_tensor,
            task_n_targets,
            task_n_covariates,
            task_n_future_covariates,
        ) = self.tasks[task_idx]

        # 对输入的look-back和预测的数值进行备份
        task_past_tensor, task_future_tensor = (
            task_past_tensor.clone(),
            task_future_tensor.clone(),
        )

        # 获取这个样本的长度
        full_length = task_past_tensor.shape[-1]

        if self.mode == DatasetMode.TRAIN:
            # slice a random subsequence from the full series
            # TODO: 这里从我们指定的训练数据中随机抽取一个子序列用于正向传播
            slice_idx = np.random.randint(
                self.min_past, full_length - self.prediction_length + 1
            )

        elif self.mode == DatasetMode.VALIDATION:
            # slice the last window for validation
            slice_idx = full_length - self.prediction_length

        else:
            # slice the full series for prediction
            slice_idx = full_length

        if slice_idx >= self.context_length:
            # slice series, if it is longer than context_length
            # 如果超出了最大的上下文长度
            task_context = task_past_tensor[
                :, slice_idx - self.context_length : slice_idx
            ]
        else:
            task_context = task_past_tensor[:, :slice_idx]

        # In the TEST mode, we have no target available and the task_future_covariates can be directly used
        # In the TRAIN and VALIDATION modes, the target and task_future_covariates need to be constructed from
        # the task_context_tensor by slicing the appropriate indices which we do below
        if self.mode in [DatasetMode.TRAIN, DatasetMode.VALIDATION]:
            # the first task_n_targets elements in task_context_tensor are the targets
            task_future_target = task_past_tensor[
                :, slice_idx : slice_idx + self.prediction_length
            ].clone()
            # mask out all rows corresponding to covariates
            task_future_target[task_n_targets:] = torch.nan

        else:
            task_future_target = None

        # task_context: (task_n_targets + task_n_covariates, min(context_length, history_length))
        # task_future_target: (task_n_targets + task_n_covariates, prediction_length), the future values of known future covariates
        # are ignored during loss computation
        # task_future_covariates: (task_n_targets + task_n_past_only_covariates + task_n_future_covariates, prediction_length),
        # the entries corresponding to targets and past-only covariates are NaNs

        return task_context, task_future_target, task_n_targets

    def _build_batch(
        self, task_indices: list[int]
    ) -> dict[str, torch.Tensor | int | list[tuple[int, int]] | None]:
        """Build a batch from given task indices."""

        batch_context_tensor_list = []
        batch_future_target_tensor_list = []
        target_idx_ranges: list[tuple[int, int]] = []

        target_start_idx = 0
        for task_idx in task_indices:
            # 核心似乎在这里也就是构造训练的切片这里

            # FIXME: 为什么会出现不同长度的输入数据
            task_context, task_future_target, task_n_targets = self._construct_slice(
                task_idx
            )

            batch_context_tensor_list.append(task_context)
            batch_future_target_tensor_list.append(task_future_target)
            target_idx_ranges.append(
                (target_start_idx, target_start_idx + task_n_targets)
            )

        # 返回的内容是一个字典
        return {
            "context": left_pad_and_cat_2D(batch_context_tensor_list),
            "future_target": (
                None
                if self.mode == DatasetMode.TEST
                else torch.cat(
                    cast(list[torch.Tensor], batch_future_target_tensor_list), dim=0
                )
            ),
            "num_output_patches": self.num_output_patches,
            "target_idx_ranges": target_idx_ranges,
        }

    def _generate_train_batches(self):
        """加载用于模型训练的数据"""

        while True:
            current_batch_size = 0
            task_indices = []

            while current_batch_size < self.batch_size:
                # 从中随机抽取索引用于模型的训练

                task_idx = np.random.randint(len(self.tasks))
                task_indices.append(task_idx)
                current_batch_size += self.tasks[task_idx][0].shape[0]

            yield self._build_batch(task_indices)

    def _generate_sequential_batches(self):
        task_idx = 0
        while task_idx < len(self.tasks):
            current_batch_size = 0
            task_indices = []

            while task_idx < len(self.tasks) and current_batch_size < self.batch_size:
                task_indices.append(task_idx)
                current_batch_size += self.tasks[task_idx][0].shape[0]
                task_idx += 1

            yield self._build_batch(task_indices)

    def __iter__(self) -> Iterator:
        """
        Generate batches of data for the Chronos-2 model. In training mode, this iterator is infinite.

        Yields
        ------
        dict
            A dictionary containing:
            - context: torch.Tensor of shape (batch_size, context_length) containing input sequences
            - future_target: torch.Tensor of shape (batch_size, prediction_length) containing future target sequences, None in TEST mode
            - future_covariates: torch.Tensor of shape (batch_size, prediction_length) containing known future covariates
            - group_ids: torch.Tensor of shape (batch_size,) containing the group ID for each sequence
            - num_output_patches: int indicating number of patches the model should output to cover prediction_length
            - target_idx_ranges: (only in TEST mode) list of tuples indicating the start & end indices of targets in context
        """
        if self.mode == DatasetMode.TRAIN:
            for batch in self._generate_train_batches():
                batch.pop("target_idx_ranges")
                yield batch
        elif self.mode == DatasetMode.VALIDATION:
            for batch in self._generate_sequential_batches():
                batch.pop("target_idx_ranges")
                yield batch
        else:
            yield from self._generate_sequential_batches()

    @classmethod
    def convert_inputs(
        cls,
        inputs: (
            TensorOrArray
            | Sequence[TensorOrArray]
            | Sequence[Mapping[str, TensorOrArray | Mapping[str, TensorOrArray | None]]]
        ),
        context_length: int,
        prediction_length: int,
        batch_size: int,
        output_patch_size: int,
        min_past: int = 1,
        mode: str | DatasetMode = DatasetMode.TRAIN,
    ) -> "LiteSpecFormerDataset":
        """获取数据集时直接调用的方法"""
        """Convert from different input formats to a LiteSpecFormerDataset."""
        if isinstance(inputs, (torch.Tensor, np.ndarray)):
            inputs = convert_tensor_input_to_list_of_dicts_input(inputs)
        elif isinstance(inputs, list) and all(
            [isinstance(x, (torch.Tensor, np.ndarray)) for x in inputs]
        ):
            inputs = cast(list[TensorOrArray], inputs)
            inputs = convert_list_of_tensors_input_to_list_of_dicts_input(inputs)
        elif isinstance(inputs, list) and all([isinstance(x, dict) for x in inputs]):
            pass
        else:
            raise ValueError("Unexpected inputs format")

        inputs = cast(list[dict[str, TensorOrArray | dict[str, TensorOrArray]]], inputs)

        return cls(
            inputs,
            context_length=context_length,
            prediction_length=prediction_length,
            batch_size=batch_size,
            output_patch_size=output_patch_size,
            min_past=min_past,
            mode=mode,
        )


class LiteSpecFormerTestingDataset(IterableDataset):
    def __init__(
        self,
        tasks: Dataset,
        prediction_length: int,
        test_context_lengths: int = 256,
    ) -> None:
        super().__init__()

        self.data = self.load_data(tasks)

        # 数据的变量数目
        self.num_variables, self.seq_length = self.data.shape

        self.prediction_length = prediction_length

        self.test_context_lengths = test_context_lengths

        # 最长的上下文窗口
        self.max_context_length = max(test_context_lengths)

        # 计算数据集中每次迭代遍历的样本数目
        self.num_samples = sum(
            [
                self.seq_length - 2 * self.prediction_length - context_length + 1
                for context_length in self.test_context_lengths
            ]
        )

    def left_padding(self, tensor: torch.Tensor) -> torch.Tensor:
        """对输入的tensor进行左侧填充，使其长度达到max_context_length"""

        # 获取上下文长度
        context_length = tensor.shape[-1]

        # 如果长度已经超过了最大长度，则直接返回
        if context_length >= self.max_context_length:
            return tensor[:, -self.max_context_length :]

        # 否则，进行填充
        else:
            padding = torch.full(
                (self.num_variables, self.max_context_length - context_length),
                fill_value=torch.nan,
                device=tensor.device,
            )
            return torch.cat([padding, tensor], dim=-1)

    def load_data(self, dataset) -> torch.FloatTensor:
        """
        Load the data from the specified csv file.

        Parameters
        ----------
        dataset_name_or_path (str): The path to the csv file containing the data.

        Returns
        -------
        torch.FloatTensor: A tensor containing the loaded data.
        """

        # Convert the data to a list
        data_list = []

        # Iterate over each sample in the dataset
        for i in range(len(dataset)):
            # Get the target series
            series = dataset[i]["target"]
            data_list.append(((series - series.mean()) / series.std()).unsqueeze(0))

        return torch.concatenate(data_list, dim=0).float()

    def __len__(self) -> int:
        """Return the total number of samples in the dataset."""
        return self.num_samples

    def __iter__(self) -> Iterator:
        for context_length in self.test_context_lengths:
            # 根据context_length对数据进行划分
            data_x, data_y = (
                self.data[:, : -self.prediction_length],
                self.data[:, context_length:],
            )

            # 统计样本数
            num_samples = data_x.shape[-1] - context_length - self.prediction_length + 1

            for idx in range(0, num_samples):
                yield {
                    "context": self.left_padding(data_x[:, idx : idx + context_length]),
                    "future_target": data_y[:, idx : idx + self.prediction_length],
                }


class SpectrumLibraryDataset(Dataset):
    """
    The universal dataset class for spectrum prediction tasks.
    """

    def __init__(
        self,
        seq_length: int,
        prediction_length: int,
        dataset_name_or_path: str,
        split_ratio: float = 0.6,
        standard_scale: bool = True,
        flag: str = "train",
    ) -> None:
        super().__init__()

        # The length of the input sequence and the prediction sequence
        self.seq_length, self.prediction_length = seq_length, prediction_length

        # Assert the flag
        assert flag in ["train", "test"]
        self.flag = flag

        # Create the scaler for time series normalization
        self.standard_scale = standard_scale
        self.scaler = StandardScaler()

        # The path to load the csv data
        self.dataset_name_or_path = dataset_name_or_path

        # The ratio to split the data into train and test sets, only used when flag is "train"
        assert 0 < split_ratio < 1, "split_ratio must be between 0 and 1"
        self.split_ratio = split_ratio

        # Load the data from the csv file
        self.data, self.data_x, self.data_y = self.preprocessing(
            self.dataset_name_or_path
        )

        # Transform the data to torch tensors
        self.data_x, self.data_y = (
            torch.from_numpy(self.data_x).float(),
            torch.from_numpy(self.data_y).float(),
        )

        # The number of the channels
        self._num_variables = self.data_x.shape[1]

    def load_data(self, dataset_name_or_path: str) -> np.ndarray:
        """
        Load the data from the specified path.
        The data can be in the form of a .npy file, a .csv file, or a Hugging Face Dataset.

        Arguments
        ----------
        dataset_name_or_path
            The path to the data file or directory.
            If the path is a file, it can be either a .npy file or a .csv file.
            If the path is a directory, it is assumed to be a Hugging Face Dataset.

        Returns
        -------
        np.ndarray
            A numpy array containing the loaded data. The shape of the array should be (seq_length, num_variables).
        """
        if not path.exists(dataset_name_or_path):
            raise FileNotFoundError(f"Dataset not found at {dataset_name_or_path}")

        # 判断数据文件的类型
        if dataset_name_or_path.endswith(".npy"):
            # The shape of data should be (seq_length, num_variables)
            data = np.load(dataset_name_or_path)

        elif dataset_name_or_path.endswith(".csv"):
            # Remove the first column which is the index column
            data = pd.read_csv(dataset_name_or_path).iloc[:, 1:].values

        elif path.isdir(dataset_name_or_path):
            # If the data path is a directory,
            # we assume it is a Hugging Face Dataset and try to load it using the datasets library
            hf_dataset = HFDataset.load_from_disk(
                dataset_path=dataset_name_or_path, keep_in_memory=True
            )  # FIXME: 这里需要进行修改 可以直接从网络上下载数据
            hf_dataset.set_format("numpy")

            # Get the data from the Hugging Face Dataset, we assume that the target series is stored under the key "target"
            data = np.stack(
                [hf_dataset[i]["target"] for i in range(len(hf_dataset))], axis=0
            ).transpose(
                1, 0
            )  # shape: (seq_length, num_variables)

        else:
            raise ValueError(
                f"Unsupported data format for file: {dataset_name_or_path}"
            )

        return data

    def preprocessing(
        self, dataset_name_or_path: str
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Preprocess the data for training or testing.

        Arguments
        ----------
        dataset_name_or_path
            The path to the data file or directory.

        Returns
        -------
        Tuple[np.ndarray, np.ndarray, np.ndarray]
            A tuple containing the preprocessed data, the input sequences, and the target sequences.
        """

        # Load the data from the dataset_name_or_path
        data = self.load_data(dataset_name_or_path=dataset_name_or_path)

        # Scale the data if required
        if self.standard_scale:
            self.scaler.fit(data)
            data = self.scaler.transform(data)

        # Get the length of the data
        length = len(data)

        # Split the data into train and test sets
        split_index = int(length * self.split_ratio)

        # get the data for the given flag
        if self.flag == "train":
            data = data[: split_index + self.prediction_length]
        else:
            data = data[split_index - self.seq_length :]

        # split the data into input and target
        data_x = data[0 : -self.prediction_length]
        data_y = data[self.seq_length :]

        return data, data_x, data_y

    @property
    def num_variables(self) -> int:
        """Get the number of variables in the dataset."""
        return self._num_variables

    def __len__(self) -> int:
        """Get the number of samples in the dataset."""
        return len(self.data_x) - self.seq_length - self.prediction_length + 1

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """Get the input and target sequences for a given index."""
        # Get the start and end indices for the input sequences
        s_begin = index
        s_end = s_begin + self.seq_length

        # Get the start and end indices for the target sequences
        r_begin = s_begin
        r_end = r_begin + self.prediction_length

        # Return the input and target sequences
        return self.data_x[s_begin:s_end], self.data_y[r_begin:r_end]


if __name__ == "__main__":
    from torch.utils.data import DataLoader

    from matplotlib import pyplot as plt
    import numpy as np

    from accelerate import Accelerator

    # data = torch.arange(256) / 256

    # tasks = [{"target": data} for _ in range(16)]

    # dataset = LiteSpecFormerDataset(
    #     inputs=tasks,
    #     context_length=300,
    #     prediction_length=16,
    #     batch_size=4,
    #     output_patch_size=16,
    #     min_past=1,
    # )

    # # 在训练的时候batch_size无法指定为None
    # data_loader = DataLoader(dataset=dataset, batch_size=1, shuffle=False)

    # shape_size_dict = {}

    # accelerator = Accelerator(
    #     device_placement=True,
    #     # split_batches=True,
    #     mixed_precision="fp16",
    #     gradient_accumulation_steps=1,
    # )

    # data_loader = accelerator.prepare(data_loader)

    # data_list = []

    # for batch in data_loader:
    #     data_list.append(batch)

    #     if len(data_list) == 6:
    #         break

    # for data in data_list:
    #     print("training context:", data["context"].shape)

    # for i in range(10):
    #     print()

    # eval_dataset = LiteSpecFormerDataset(
    #     inputs=tasks,
    #     context_length=300,
    #     prediction_length=16,
    #     batch_size=4,
    #     output_patch_size=16,
    #     min_past=1,
    #     mode=DatasetMode.TEST,
    # )

    # # 在测试的时候batch_size可以指定为None，因为我们在测试的时候是顺序生成样本的，不需要随机抽取
    # data_loader = DataLoader(dataset=eval_dataset, batch_size=None)

    # for batch in data_loader:
    #     print(batch["context"].shape)

    # tasks = []
    # for i in range(120):
    #     tasks.append({"target": torch.randn(1024)})

    from datasets import Dataset

    tasks = Dataset.load_from_disk("/home/flyvideo/Sata/wwx/LSPD/test")
    tasks.set_format("torch")
    print(len(tasks))

    dataset = LiteSpecFormerTestingDataset(
        tasks=tasks,
        prediction_length=16,
        test_context_lengths=[128],
    )

    print(len(dataset))

    index = 0

    data_loader = DataLoader(
        dataset=dataset, batch_size=256, num_workers=0, pin_memory=True
    )

    print(len(data_loader))

    for batch in data_loader:
        print(batch["context"].shape, batch["future_target"].shape)

        index += 1

    print(index)

    index = 0
    data_loader = DataLoader(
        dataset=dataset, batch_size=256, num_workers=16, pin_memory=True
    )

    print(len(data_loader))

    for batch in data_loader:
        print(batch["context"].shape, batch["future_target"].shape)
        # from time import sleep
        # sleep(1)
        index += 1

    print(index)
