import logging
from copy import deepcopy
from datetime import datetime
import time
import math
import os
from os import path

from enum import Enum
import warnings
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    Mapping,
    Sequence,
    Dict,
    List,
    Optional,
    Tuple,
    Union,
    Callable,
    Literal,
    Mapping,
    Sequence,
)

from accelerate import Accelerator
from colorama import Fore, Style

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from einops import rearrange, repeat
from transformers import AutoConfig, PreTrainedModel
from transformers.utils.import_utils import is_peft_available
from transformers.utils.peft_utils import find_adapter_config_file

if TYPE_CHECKING:
    import pandas as pd
    from peft import LoraConfig
    from transformers.trainer_callback import TrainerCallback

logger = logging.getLogger(__name__)

from litespecformer.config import LiteSpecFormerConfig
from litespecformer.model import LiteSpecFormerModel
from litespecformer.dataset import (
    LiteSpecFormerDataset,
    SpectrumLibraryDataset,
    DatasetMode,
    TensorOrArray,
)
from litespecformer.utils import (
    left_pad_and_stack_1D,
    convert_df_input_to_list_of_dicts_input,
    convert_tensor_input_to_list_of_dicts_input,
    get_num_output_patches,
    interpolate_quantiles,
    weighted_quantile,
    logging_results,
)
from litespecformer.metrics import calculate_metrics


class ForecastType(Enum):
    SAMPLES = "samples"
    QUANTILES = "quantiles"


class PipelineRegistry(type):
    REGISTRY: Dict[str, "PipelineRegistry"] = {}

    def __new__(cls, name, bases, attrs):
        """See, https://github.com/faif/python-patterns."""
        new_cls = type.__new__(cls, name, bases, attrs)
        if name is not None:
            cls.REGISTRY[name] = new_cls

        return new_cls


class BaseForecastPipeline(metaclass=PipelineRegistry):
    """
    The BaseForecastPipeline class provides a common interface for
    forecasting tasks. It defines the basic methods that all forecasting
    pipelines must implement.

    The code is adapted from the Chronos2Pipeline class in the Chronos library,
    with modifications to fit the LiteSpecFormer model and use case.
    """

    forecast_type: ForecastType
    dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32}

    def __init__(self, inner_model: "PreTrainedModel"):
        """
        Parameters
        ----------
        inner_model : PreTrainedModel
            A hugging-face transformers PreTrainedModel, e.g., T5ForConditionalGeneration
        """
        # for easy access to the inner HF-style model
        self.inner_model = inner_model

    def _prepare_and_validate_context(
        self, context: Union[torch.Tensor, List[torch.Tensor]]
    ) -> Union[torch.Tensor, List[torch.Tensor]]:
        """
        Prepare and validate the context tensor.

        Parameters
        ----------
        context
            Input series. This is either a 1D tensor, or a list
            of 1D tensors, or a 2D tensor whose first dimension
            is batch. In the latter case, use left-padding with
            ``torch.nan`` to align series of different lengths.

        Returns
        -------
        context
            The prepared and validated context tensor.
            Is a 2D time series tensor.
        """
        if isinstance(context, list):
            # 遍历这个列表中所有的张量并根据最长的张量进行填充
            context: torch.FloatTensor = left_pad_and_stack_1D(context)

        # 处理输入数据是一个样本的情况
        assert isinstance(context, torch.Tensor)
        if context.ndim == 1:
            # duplicate the series to make it a batch of size 1
            context = context.unsqueeze(0)

        # 确保输入数据是二维的
        assert context.ndim == 2

        return context

    def predict(
        self,
        inputs: Union[torch.Tensor, List[torch.Tensor]],
        prediction_length: Optional[int] = None,
    ):
        """
        Get forecasts for the given time series. Predictions will be
        returned in fp32 on the cpu.

        Parameters
        ----------
        inputs
            Input series. This is either a 1D tensor, or a list
            of 1D tensors, or a 2D tensor whose first dimension
            is batch. In the latter case, use left-padding with
            ``torch.nan`` to align series of different lengths.
        prediction_length
            Time steps to predict. Defaults to a model-dependent
            value if not given.

        Returns
        -------
        forecasts
            Tensor containing forecasts. The layout and meaning
            of the forecasts values depends on ``self.forecast_type``.
        """
        raise NotImplementedError()

    def predict_quantiles(
        self,
        inputs: Union[torch.Tensor, List[torch.Tensor]],
        prediction_length: Optional[int] = None,
        quantile_levels: List[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Get quantile and mean forecasts for given time series.
        Predictions will be returned in fp32 on the cpu.

        Parameters
        ----------
        inputs : Union[torch.Tensor, List[torch.Tensor]]
            Input series. This is either a 1D tensor, or a list
            of 1D tensors, or a 2D tensor whose first dimension
            is batch. In the latter case, use left-padding with
            ``torch.nan`` to align series of different lengths.
        prediction_length : Optional[int], optional
            Time steps to predict. Defaults to a model-dependent
            value if not given.
        quantile_levels : List[float], optional
            Quantile levels to compute, by default [0.1, 0.2, ..., 0.9]

        Returns
        -------
        quantiles
            Tensor containing quantile forecasts. Shape
            (batch_size, prediction_length, num_quantiles)
        mean
            Tensor containing mean (point) forecasts. Shape
            (batch_size, prediction_length)
        """
        raise NotImplementedError()

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Union[str, Path],
        *model_args,
        force_s3_download=False,
        **kwargs,
    ):
        """
        Load the model, either from a local path, S3 prefix, or from the HuggingFace Hub.
        Supports the same arguments as ``AutoConfig`` and ``AutoModel`` from ``transformers``.
        """
        from transformers import AutoConfig

        torch_dtype = kwargs.get("torch_dtype", "auto")
        if torch_dtype != "auto" and isinstance(torch_dtype, str):
            kwargs["torch_dtype"] = cls.dtypes[torch_dtype]

        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        is_valid_config = hasattr(config, "litespecformer_pipeline_class") or hasattr(
            config, "forecasting_config"
        )

        if not is_valid_config:
            raise ValueError("Not a LiteSpecFormer config file")

        pipeline_class_name = getattr(
            config, "litespecformer_pipeline_class", "LiteSpecFormerPipeline"
        )
        class_ = PipelineRegistry.REGISTRY.get(pipeline_class_name)
        if class_ is None:
            raise ValueError(
                f"Trying to load unknown pipeline class: {pipeline_class_name}"
            )

        return class_.from_pretrained(  # type: ignore[attr-defined]
            pretrained_model_name_or_path, *model_args, **kwargs
        )


class LiteSpecFormerPipeline(BaseForecastPipeline):
    forecast_type: ForecastType = ForecastType.QUANTILES
    default_context_length: int = 768

    def __init__(self, model: LiteSpecFormerModel):
        super().__init__(inner_model=model)
        self.model = model

    @staticmethod
    def _get_prob_mass_per_quantile_level(
        quantile_levels: torch.Tensor,
    ) -> torch.Tensor:
        """
        Computes normalized probability masses for quantile levels using trapezoidal rule approximation.

        Each quantile receives probability mass proportional to the width of its surrounding interval,
        creating a piecewise uniform distribution. The mass for quantile q_i is computed as
        (q_{i+1} - q_{i-1}) / 2, where q_0 = 0 and q_{n+1} = 1.

        Parameters
        ----------
        quantile_levels : torch.Tensor
            The quantile levels, must be strictly in (0, 1)

        Returns
        -------
        torch.Tensor
            The normalized probability mass per quantile
        """
        assert quantile_levels.ndim == 1
        assert quantile_levels.min() > 0.0 and quantile_levels.max() < 1.0

        device = quantile_levels.device
        boundaries = torch.cat(
            [
                torch.tensor([0.0], device=device),
                quantile_levels,
                torch.tensor([1.0], device=device),
            ]
        )
        prob_mass = (boundaries[2:] - boundaries[:-2]) / 2
        return prob_mass / prob_mass.sum()

    @property
    def model_context_length(self) -> int:
        """Get the context length of the model in the forecasting pipline."""
        return self.model.context_length

    @property
    def model_output_patch_size(self) -> int:
        """Get the output patch size of the model in the forecasting pipeline."""
        return self.model.forecasting_config.output_patch_size

    @property
    def model_prediction_length(self) -> int:
        """Get the prediction length of the model in the forecasting pipeline."""
        return (
            self.model.forecasting_config.max_output_patches
            * self.model.forecasting_config.output_patch_size
        )

    @property
    def quantiles(self) -> list[float]:
        """Get the the list of quantile levels for the model in the forecasting pipeline."""
        return self.model.forecasting_config.quantiles

    @property
    def max_output_patches(self) -> int:
        """Get the maximum number of output patches for the model in the forecasting pipeline."""
        return self.model.forecasting_config.max_output_patches

    def fit(
        self,
        inputs: (
            TensorOrArray
            | Sequence[TensorOrArray]
            | Sequence[Mapping[str, TensorOrArray | Mapping[str, TensorOrArray | None]]]
        ),
        prediction_length: int,
        validation_inputs: (
            TensorOrArray
            | Sequence[TensorOrArray]
            | Sequence[Mapping[str, TensorOrArray | Mapping[str, TensorOrArray | None]]]
            | None
        ) = None,
        finetune_mode: Literal["full", "lora"] = "lora",
        lora_config: "LoraConfig | dict | None" = None,
        context_length: int | None = None,
        learning_rate: float = 1e-6,
        num_steps: int = 2560,
        batch_size: int = 256,
        output_dir: Path | str | None = None,
        min_past: int | None = None,
        finetuned_ckpt_name: str = "finetuned-ckpt",
        callbacks: list["TrainerCallback"] | None = None,
        remove_printer_callback: bool = False,
        disable_data_parallel: bool = True,
        **extra_trainer_kwargs,
    ) -> "LiteSpecFormerPipeline":
        """
        Fine-tune a copy of the current `LiteSpecFormer` model on the given inputs and return a new pipeline.

        Parameters
        ----------
        inputs
            The time series on which the model will be fine-tuned. The allowed formats of inputs are the same as `LiteSpecFormerPipeline.predict()`.
            Note: when `inputs` is a list of dicts, 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.
        prediction_length
            The prediction horizon for which the model will be fine-tuned
        validation_inputs
            The time series used for validation and model selection. The format of `validation_inputs` is exactly the same as `inputs`, by default None which
            means that no validation is performed. Note that enabling validation may slow down fine-tuning for large datasets.
        finetune_mode
            One of "full" (performs full fine-tuning) or "lora" (performs Low Rank Adaptation (LoRA) fine-tuning), by default "full"
        lora_config
            The configuration to use for LoRA fine-tuning when finetune_mode="lora". Can be a `LoraConfig` object or a dict which is used to initialize `LoraConfig`.
            When unspecified and finetune_mode="lora", a default configuration is used
        context_length
            The maximum context length used during fine-tuning, by default set to the model's default context length
        learning_rate
            The learning rate for the optimizer, by default 1e-6
            When finetune_mode="lora", we recommend using a higher value of the learning rate, such as 1e-5
        num_steps
            The number of steps to fine-tune for, by default 2560
        batch_size
            The batch size used for fine-tuning. Note that the batch size here means the number of time series, including target(s) and covariates,
            which 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, by default 256
        output_dir
            The directory in which outputs from the `Trainer` will be saved, by default set to `chronos-2-finetuned/{%Y-%m-%d_%H-%M-%S}`
        min_past
            The minimum number of time steps the context must have during fine-tuning. All time series shorter than `min_past + prediction_length`
            are filtered out, by default set equal to prediction_length
        finetuned_ckpt_name
            The name of the directory inside `output_dir` in which the final fine-tuned checkpoint will be saved, by default "finetuned-ckpt"
        callbacks
            A list of `TrainerCallback`s which will be forwarded to the HuggingFace `Trainer`
        remove_printer_callback
            If True, all instances of `PrinterCallback` are removed from callbacks
        disable_data_parallel
            If True, ensures that DataParallel is disabled and training happens on a single GPU
        **extra_trainer_kwargs
            Extra kwargs are directly forwarded to `TrainingArguments`

        Returns
        -------
        A new `Chronos2Pipeline` with the fine-tuned modelchronos_config
        """

        import torch.cuda
        from transformers.trainer_callback import PrinterCallback
        from transformers.training_args import TrainingArguments

        if finetune_mode == "lora":
            if is_peft_available():
                from peft import LoraConfig, get_peft_model
            else:
                warnings.warn(
                    "`peft` is required for `finetune_mode='lora'`. Please install it with `pip install peft`. Falling back to `finetune_mode='full'`."
                )
                finetune_mode = "full"
                lora_config = None

        # Use the trainer from chronos2 to fine-tune the model
        from chronos.chronos2.trainer import (
            Chronos2Trainer,
            EvaluateAndSaveFinalStepCallback,
        )

        # Check that finetune_mode is valid, and set the default lora_config if necessary
        assert finetune_mode in [
            "full",
            "lora",
        ], f"finetune_mode must be one of ['full', 'lora'], got {finetune_mode}"

        # Check that lora_config is not specified when finetune_mode="full"
        if finetune_mode == "full" and lora_config is not None:
            raise ValueError(
                "lora_config should not be specified when `finetune_mode='full'`. To enable LoRA, set `finetune_mode='lora'`."
            )

        # Create a copy of the model to avoid modifying the original
        config: LiteSpecFormerConfig = deepcopy(self.model.config)
        model = LiteSpecFormerModel(config=config).to(self.model.device)  # type: ignore
        # Load the pre-trained model weights into the new model
        model.load_state_dict(self.model.state_dict())

        if finetune_mode == "lora":
            if lora_config is None:
                # If lora_config is not specified, use a default configuration
                lora_config = LoraConfig(
                    r=8,
                    lora_alpha=16,
                    target_modules=[
                        "self_attention.q",
                        "self_attention.v",
                        "self_attention.k",
                        "self_attention.o",
                        "output_patch_embedding.output_layer",
                    ],
                )
            elif isinstance(lora_config, dict):
                # If lora_config is a dict, use it to initialize a LoraConfig
                lora_config = LoraConfig(**lora_config)
            else:
                assert isinstance(
                    lora_config, LoraConfig
                ), f"lora_config must be an instance of LoraConfig or a dict, got {type(lora_config)}"

            # 一个预训练的基础模型（如 LLaMA、BERT、GPT），与你定义的 PEFT 配置（如 LoRA、IA3、Prefix Tuning）结合，生成一个 “仅微调部分参数” 的 PEFT 模型
            model = get_peft_model(model, lora_config)
            n_trainable_params, n_params = model.get_nb_trainable_parameters()
            logger.info(
                f"Using LoRA. Number of trainable parameters: {n_trainable_params}, total parameters: {n_params}."
            )

        if context_length is None:
            # Get the default context length from the model's config if not specified
            context_length = self.model_context_length

        if min_past is None:
            # Get the default min_past from the model's config if not specified
            min_past = prediction_length

        # Create the training dataset for fine-tuning of LiteSpecFormer model.
        # Note that the dataset will filter out time series shorter than `min_past + prediction_length`.
        train_dataset = LiteSpecFormerDataset.convert_inputs(
            inputs=inputs,
            context_length=context_length,
            prediction_length=prediction_length,
            batch_size=batch_size,
            output_patch_size=self.model_output_patch_size,
            min_past=min_past,
            mode=DatasetMode.TRAIN,
        )

        # Create the output directory for the fine-tuned model
        if output_dir is None:
            output_dir = Path("./results/LiteSpecFormer-Finetuned") / time.strftime(
                "%Y-%m-%d_%H-%M-%S"
            )
        elif isinstance(output_dir, str):
            output_dir = Path(output_dir)

        assert isinstance(output_dir, Path)

        # check if the model is on CPU or GPU, and if GPU is available
        use_cpu = str(self.model.device) == "cpu"
        has_sm80 = (
            torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
        )

        # warn user if a cuda device is available and CPU fine-tuning is used
        if use_cpu and torch.cuda.is_available():
            warnings.warn(
                "The model is being fine-tuned on the CPU, but a CUDA device is available. "
                "We recommend using the GPU for faster fine-tuning.",
                category=UserWarning,
                stacklevel=2,
            )

        # Create the training arguments for the fine-tuning
        training_kwargs: dict = dict(
            output_dir=str(output_dir),
            per_device_train_batch_size=batch_size,
            per_device_eval_batch_size=batch_size,
            learning_rate=learning_rate,
            lr_scheduler_type="linear",
            warmup_ratio=0.0,
            optim="adamw_torch_fused",
            logging_strategy="steps",
            logging_steps=100,
            disable_tqdm=False,
            report_to="none",
            max_steps=num_steps,
            gradient_accumulation_steps=1,
            dataloader_num_workers=0,
            tf32=has_sm80 and not use_cpu,
            bf16=has_sm80 and not use_cpu,
            save_only_model=True,
            prediction_loss_only=True,
            save_total_limit=1,
            save_strategy="no",
            save_steps=None,
            eval_strategy="no",
            eval_steps=None,
            load_best_model_at_end=False,
            metric_for_best_model=None,
            use_cpu=use_cpu,
        )

        # Check if validation inputs are provided
        eval_dataset = None
        callbacks = callbacks or []
        if validation_inputs is not None:
            # construct validation dataset
            eval_dataset = LiteSpecFormerDataset.convert_inputs(
                inputs=validation_inputs,
                context_length=context_length,
                prediction_length=prediction_length,
                batch_size=batch_size,
                output_patch_size=self.model_output_patch_size,
                mode=DatasetMode.VALIDATION,
            )

            # set validation parameters
            training_kwargs["save_strategy"] = "steps"
            training_kwargs["save_steps"] = 100
            training_kwargs["eval_strategy"] = "steps"
            training_kwargs["eval_steps"] = 100
            training_kwargs["load_best_model_at_end"] = True
            training_kwargs["metric_for_best_model"] = "eval_loss"
            training_kwargs["label_names"] = ["future_target"]

            # add callback to ensure that the final model is evaluated
            callbacks.append(EvaluateAndSaveFinalStepCallback())

        # Update training arguments with extra_trainer_kwargs for model evaluation
        training_kwargs.update(extra_trainer_kwargs)

        if training_kwargs["tf32"]:
            # setting tf32=True changes these global properties, we copy them here so that
            # we can restore them after fine-tuning
            matmul_tf32 = torch.backends.cuda.matmul.allow_tf32
            cudnn_tf32 = torch.backends.cudnn.allow_tf32

        # Create the trainer for fine-tuning
        training_args = TrainingArguments(**training_kwargs)

        if disable_data_parallel and not use_cpu:
            # This is a hack to disable the default `transformers` behavior of using DataParallel
            training_args._n_gpu = 1
            assert training_args.n_gpu == 1  # Ensure that the hack worked

        # Create the trainer for fine-tuning from the chronos2
        trainer = Chronos2Trainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            callbacks=callbacks,
        )

        # Remove the PrinterCallback if it was added
        if remove_printer_callback:
            trainer.pop_callback(PrinterCallback)

        # ====== Start fine-tuning ======
        trainer.train()

        # update context_length and max_output_patches, if the model was fine-tuned with larger values
        model.chronos_config.context_length = max(
            model.chronos_config.context_length, context_length
        )
        model.chronos_config.max_output_patches = max(
            model.chronos_config.max_output_patches,
            math.ceil(prediction_length / self.model_output_patch_size),
        )
        # update chronos_config in model's config, so it is saved correctly
        model.config.chronos_config = model.chronos_config.__dict__

        # Create a new pipeline with the fine-tuned model
        finetuned_pipeline = LiteSpecFormerPipeline(model=model)

        # Save fine-tuned model
        finetuned_path = output_dir / finetuned_ckpt_name
        finetuned_pipeline.save_pretrained(finetuned_path)
        logger.info(f"Finetuned model saved to {finetuned_path}")

        if training_kwargs["tf32"]:
            # restore tf32 settings
            torch.backends.cuda.matmul.allow_tf32 = matmul_tf32
            torch.backends.cudnn.allow_tf32 = cudnn_tf32

        return finetuned_pipeline

    def _prepare_inputs_for_long_horizon_unrolling(
        self,
        context: torch.Tensor,
        # group_ids: torch.Tensor,
        # future_covariates: torch.Tensor,
        unrolled_quantiles: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        在进行单次长度预测后如何仍未满足预测的需求
        则进行进一步的自回归滑动预测

        TODO: 判断这里主要调整了什么?
        """

        # Expand the context, future_covariates and group_ids along a new "quantile" axis
        # if future_covariates is not None:
        #     future_covariates = repeat(
        #         future_covariates, "b t -> b q t", q=len(unrolled_quantiles)
        #     )
        context = repeat(context, "b t -> b q t", q=len(unrolled_quantiles))
        # group_ids = repeat(group_ids, "b -> b q", q=len(unrolled_quantiles))

        # Shift the group_ids so that mixing is enabled only for time series with the same group_id and
        # at the same quantile level, e.g., if the group_ids were [0, 0, 1, 1, 1] initially, after expansion
        # and shifting they will be:
        # [[0,  1,  2,  3,  4,  5,  6,  7,  8],
        #  [0,  1,  2,  3,  4,  5,  6,  7,  8],
        #  [9, 10, 11, 12, 13, 14, 15, 16, 17],
        #  [9, 10, 11, 12, 13, 14, 15, 16, 17],
        #  [9, 10, 11, 12, 13, 14, 15, 16, 17]]
        # group_ids = group_ids * len(unrolled_quantiles) + torch.arange(
        #     len(unrolled_quantiles), device=self.model.device
        # ).unsqueeze(0)

        # We unroll the quantiles in unrolled_quantiles to the future and each unrolled quantile gives
        # len(self.quantiles) predictions, so we end up with len(unrolled_quantiles) * len(self.quantiles)
        # "samples". unrolled_sample_weights specifies the amount of probability mass covered by each sample.
        # Note that this effectively leads to shrinking of the probability space but it is better heuristic
        # than just using the median to unroll, which leads to uncertainty collapse.
        unrolled_sample_weights = torch.outer(
            self._get_prob_mass_per_quantile_level(unrolled_quantiles),
            self._get_prob_mass_per_quantile_level(torch.tensor(self.quantiles)),
        )

        # return context, group_ids, future_covariates, unrolled_sample_weights
        return context, unrolled_sample_weights

    def _autoregressive_unroll_for_long_horizon(
        self,
        context: torch.Tensor,
        # group_ids: torch.Tensor,
        # future_covariates: torch.Tensor,
        prediction: torch.Tensor,
        unrolled_quantiles: torch.Tensor,
        unrolled_sample_weights: torch.Tensor,
        num_output_patches: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """进行更长时间序列滚动预测的函数"""

        # Get unrolled_quantiles from prediction and append it to the expanded context
        prediction_unrolled = interpolate_quantiles(
            query_quantile_levels=unrolled_quantiles,
            original_quantile_levels=self.quantiles,
            original_values=rearrange(prediction, "b q h -> b h q"),
        )

        prediction_unrolled = rearrange(prediction_unrolled, "b h q -> b q h")
        context = torch.cat([context, prediction_unrolled], dim=-1)[
            ..., -self.model_context_length :
        ]
        n_paths = len(unrolled_quantiles)

        # Shift future_covariates by prediction.shape[-1] while replacing the predicted values
        # of future covariates in the context with their actual values, if known
        # TODO: 注意需要调整这里的函数
        # if future_covariates is not None:
        #     context, future_covariates = self._slide_context_and_future_covariates(
        #         context=context,
        #         future_covariates=future_covariates,
        #         slide_by=prediction.shape[-1],
        #     )

        # Reshape (batch, n_paths, context_length) -> (batch * n_paths, context_length)

        # TODO: 注意这里仍然是需要调用单步预测的函数
        prediction: torch.Tensor = self._predict_step(
            context=rearrange(context, "b n t -> (b n) t"),
            num_output_patches=num_output_patches,
        )

        # Reshape predictions from (batch * n_paths, n_quantiles, length) to (batch, n_paths * n_quantiles, length)
        prediction = rearrange(prediction, "(b n) q h -> b (n q) h", n=n_paths)

        # Reduce `n_paths * n_quantiles` to n_quantiles and transpose back
        prediction = weighted_quantile(
            query_quantile_levels=self.quantiles,
            sample_weights=rearrange(unrolled_sample_weights, "n q -> (n q)"),
            samples=rearrange(prediction, "b (n q) h -> b h (n q)", n=n_paths),
        )
        prediction = rearrange(prediction, "b h q -> b q h")

        return prediction, context

    @torch.no_grad()
    def predict(
        self,
        inputs: (
            TensorOrArray
            | Sequence[TensorOrArray]
            | Sequence[Mapping[str, TensorOrArray | Mapping[str, TensorOrArray]]]
        ),
        prediction_length: int | None = None,
        batch_size: int = 256,
        context_length: int | None = None,
        # cross_learning: bool = False,
        limit_prediction_length: bool = False,
        **kwargs,
    ) -> list[torch.Tensor]:
        """
        Generate forecasts for the given time series.

        Parameters
        ----------
        inputs
            The time series to generate forecasts for, can be one of:
            - A 3-dimensional `torch.Tensor` or `np.ndarray` of shape (batch, n_variates, history_length). When `n_variates > 1`, information
            will be shared among the different variates of each time series in the batch and the model will perform multivariate forecasting.
            - A list of `torch.Tensor` or `np.ndarray` where each element can either be 1-dimensional of shape (history_length,)
            or 2-dimensional of shape (n_variates, history_length). The history_lengths may be different across elements; left-padding
            will be applied, if needed. The model will perform univariate and multivariate inference for 1-d and 2-d elements, respectively.
            A mixture of 1-d and 2-d elements can be provided in the same list.
            - A list of dictionaries where each dictionary may have the following keys.
                1. `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`.
                2. `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`.
                3. `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`.

            Examples:
            ```python

            # Batch of univariate time series
            inputs = torch.randn(32, 1, 100)

            # Batch of multivariate time series
            inputs = torch.randn(32, 3, 100)

            # List of time series with different lengths and n_variates
            inputs = [
                torch.randn(100),  # univariate series of length 100
                torch.randn(2, 150),  # bivariate series of length 150
                torch.randn(120),  # univariate series of length 120
            ]

            # List of dictionaries with covariates
            prediction_length = 24
            inputs = [
                {
                    # task with 1-d target, one past-only covariate and one known future covariate
                    "target": torch.randn(100),
                    "past_covariates": {"temperature": torch.randn(100), "precipitation": torch.randn(100)},
                    "future_covariates": {"temperature": torch.randn(prediction_length)},
                },
                {
                    # task with 2-d target and one past-only covariate
                    "target": torch.randn(2, 150),
                    "past_covariates": {"wind_speed": torch.randn(150)},
                },
                {
                    # task with 1-d target, two numeric covariates one of which is known into the future
                    # and one categorical covariate known into the future
                    # Note: categorical covariates are only supported as numpy arrays as torch does not support str dtype
                    "target": np.random.randn(150),
                    "past_covariates": {
                        "numeric_covariate_1": np.random.rand(150),
                        "numeric_covariate_2": np.random.rand(150),
                        "cat_covariate": np.random.choice(["A", "B", "C", "D", "E"], size=150),
                    },
                    "future_covariates": {
                        "numeric_covariate_1": np.random.rand(prediction_length),
                        "cat_covariate": np.random.choice(["A", "B", "C", "D", "E"], size=prediction_length),
                    },
                },
                {
                    # task with only a 1-d target
                    "target": torch.randn(1, 150)
                },
            ]
            ```
        prediction_length
            The number of time steps to predict for, defaults to the model's default prediction length
        batch_size
            The batch size used for prediction. Note that the batch size here means the number of time series, including target(s) and covariates,
            which 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, by default 256
        context_length
            The maximum context length used during for inference, by default set to the model's default context length
        # cross_learning: bool = False,
        #     If True, cross-learning is enabled, i.e., all the tasks in `inputs` will be predicted jointly and the model will share information across all inputs, by default False
        #     The following must be noted when using cross-learning:
        #     - Cross-learning doesn't always improve forecast accuracy and must be tested for individual use cases.
        #     - Results become dependent on batch size. Very large batch sizes may not provide benefits as they deviate from the maximum group size used during pretraining.
        #     For optimal results, consider using a batch size around 100 (as used in the Chronos-2 technical report).
        #     - Cross-learning is most helpful when individual time series have limited historical context, as the model can leverage patterns from related series in the batch.
        limit_prediction_length
            If True, an error is raised when prediction_length is greater than model's default prediction length, by default False

        Returns
        -------
        The model's predictions, a list of `torch.Tensor` where each element has shape (n_variates, n_quantiles, prediction_length) and the number of
        elements are equal to the number of target time series (univariate or multivariate) in the `inputs`.

        """
        # 获取模型的单次自回归的最大预测长度
        model_prediction_length = self.model_prediction_length
        if prediction_length is None:
            prediction_length = model_prediction_length

        if kwargs.get("predict_batches_jointly") is not None:
            warnings.warn(
                "The `predict_batches_jointly` argument is deprecated and will be removed in a future version. "
                "Please use `cross_learning=True` to enable the cross-learning mode.",
                category=FutureWarning,
                stacklevel=2,
            )
            cross_learning = kwargs.pop("predict_batches_jointly")

        # The maximum number of output patches to generate in a single forward pass before the long-horizon heuristic kicks in. Note: A value larger
        # than the model's default max_output_patches may lead to degradation in forecast accuracy, defaults to a model-specific value
        max_output_patches = kwargs.pop("max_output_patches", self.max_output_patches)
        # The set of quantiles to use when making long-horizon predictions; must be a subset of the model's default quantiles. These quantiles
        # are appended to the historical context and input into the model autoregressively to generate long-horizon predictions. Note that the
        # effective batch size increases by a factor of `len(unrolled_quantiles)` when making long-horizon predictions,

        # by default [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        unrolled_quantiles = kwargs.pop(
            "unrolled_quantiles",
            [
                0.05,
                0.1,
                0.2,
                0.3,
                0.4,
                0.5,
                0.6,
                0.7,
                0.8,
                0.9,
                0.95,
            ],
        )

        # A callback which is called after each batch has been processed
        after_batch_callback: Callable = kwargs.pop("after_batch", lambda: None)

        if len(kwargs) > 0:
            raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}.")

        if not set(unrolled_quantiles).issubset(self.quantiles):
            raise ValueError(
                f"Unrolled quantiles must be a subset of the model's quantiles. "
                f"Found: {unrolled_quantiles=}, model_quantiles={self.quantiles}"
            )
        unrolled_quantiles_tensor = torch.tensor(unrolled_quantiles)

        if prediction_length > model_prediction_length:
            msg = (
                f"We recommend keeping prediction length <= {model_prediction_length}. "
                "The quality of longer predictions may degrade since the model is not optimized for it. "
            )
            if limit_prediction_length:
                msg += "You can turn off this check by setting `limit_prediction_length=False`."
                raise ValueError(msg)
            warnings.warn(msg)

        # 设置模型预测的上下文长度
        if context_length is None:
            context_length = self.model_context_length

        # If the specified context length is greater than the model's default context length,
        # it may lead to suboptimal forecasts since the model is not optimized for such long contexts.
        # Then a warning will be issued and the context length will be reset to the model's default value.
        if context_length > self.model_context_length:
            warnings.warn(
                f"The specified context_length {context_length} is greater than the model's default context length {self.model_context_length}. "
                f"Resetting context_length to {self.model_context_length}."
            )
            context_length = self.model_context_length

        # Convert the input to a LiteSpecFormerDataset and create a DataLoader.
        test_dataset = LiteSpecFormerDataset.convert_inputs(
            inputs=inputs,
            context_length=context_length,
            prediction_length=prediction_length,
            batch_size=batch_size,
            output_patch_size=self.model_output_patch_size,
            mode=DatasetMode.TEST,
        )

        # 构建用于测试的DataLoader
        test_loader = DataLoader(
            test_dataset,
            batch_size=None,
            # pin_memory=self.model.device.type
            # == "cuda",  # 如果使用GPU，则启用pin_memory以加速数据传输
            shuffle=False,
            drop_last=False,
        )

        # 记录预测的结果
        all_predictions: list[torch.Tensor] = []

        for batch in test_loader:
            assert batch["future_target"] is None
            batch_context = batch["context"]
            batch_target_idx_ranges = batch["target_idx_ranges"]

            # FIXME: 这个里面的一些变量需要进行修改
            batch_prediction = self._predict_batch(
                context=batch_context,
                # group_ids=batch_group_ids,
                # future_covariates=batch_future_covariates,
                unrolled_quantiles_tensor=unrolled_quantiles_tensor,
                prediction_length=prediction_length,
                max_output_patches=max_output_patches,
                target_idx_ranges=batch_target_idx_ranges,
            )
            all_predictions.extend(batch_prediction)
            after_batch_callback()

        return all_predictions

    def _predict_batch(
        self,
        context: torch.Tensor,
        unrolled_quantiles_tensor: torch.Tensor,
        prediction_length: int,
        max_output_patches: int,
        target_idx_ranges: list[Tuple[int, int]],
    ) -> list[torch.Tensor]:
        """
        Predict a batch of time series and return the predictions in the original order of the input time series.

        当通过测试数据构建好一个
        """
        context = context.to(device=self.model.device, dtype=torch.float32)
        # group_ids = group_ids.to(device=self.model.device)
        # future_covariates = future_covariates.to(
        #     device=self.model.device, dtype=torch.float32
        # )

        # 存储预测的结果和待预测的长度
        predictions = []
        remaining = prediction_length

        # predict first set of patches up to max_output_patches
        prediction: torch.Tensor = self._predict_step(
            context=context,
            num_output_patches=get_num_output_patches(
                remaining_horizon=remaining,
                model_output_patch_size=self.model_output_patch_size,
                max_output_patches=max_output_patches,
            ),
        )

        # Record the prediction and update the remaining horizon
        predictions.append(prediction)
        # 计算剩余的需要预测的长度
        remaining -= prediction.shape[-1]

        # prepare inputs for long horizon prediction
        if remaining > 0:
            (
                context,
                unrolled_sample_weights,
            ) = self._prepare_inputs_for_long_horizon_unrolling(
                context=context,
                # group_ids=group_ids,
                # future_covariates=future_covariates,
                unrolled_quantiles=unrolled_quantiles_tensor,
            )

        # long horizon heuristic
        while remaining > 0:
            prediction, context = self._autoregressive_unroll_for_long_horizon(
                context=context,
                # group_ids=group_ids,
                # future_covariates=future_covariates,
                prediction=prediction,
                unrolled_quantiles=unrolled_quantiles_tensor,
                unrolled_sample_weights=unrolled_sample_weights,
                num_output_patches=get_num_output_patches(
                    remaining_horizon=remaining,
                    model_output_patch_size=self.model_output_patch_size,
                    max_output_patches=max_output_patches,
                ),
            )
            predictions.append(prediction)
            remaining -= prediction.shape[-1]

        batch_prediction = torch.cat(predictions, dim=-1)[..., :prediction_length].to(
            dtype=torch.float32, device="cpu"
        )

        return [batch_prediction[start:end] for (start, end) in target_idx_ranges]

    def _predict_step(
        self,
        context: torch.Tensor,
        # group_ids: torch.Tensor,
        # future_covariates: torch.Tensor | None,
        num_output_patches: int,
    ) -> torch.Tensor:
        """Predict the next `num_output_patches` patches given the context time series inputs."""

        # TODO: 这里移除了未来协变量的输入
        # kwargs = {}
        # if future_covariates is not None:
        #     output_size = num_output_patches * self.model_output_patch_size

        #     if output_size > future_covariates.shape[1]:
        #         batch_size = len(future_covariates)
        #         padding_size = output_size - future_covariates.shape[1]
        #         padding_tensor = torch.full(
        #             (batch_size, padding_size),
        #             fill_value=torch.nan,
        #             device=future_covariates.device,
        #         )
        #         future_covariates = torch.cat(
        #             [future_covariates, padding_tensor], dim=1
        #         )

        #     else:
        #         future_covariates = future_covariates[..., :output_size]
        #     kwargs["future_covariates"] = future_covariates
        with torch.no_grad():
            prediction: torch.Tensor = self.model(
                context=context,
                # group_ids=group_ids,
                num_output_patches=num_output_patches,
                # **kwargs,
            ).quantile_preds.to(context)

        return prediction

    @staticmethod
    def _slide_context_and_future_covariates(
        context: torch.Tensor, future_covariates: torch.Tensor, slide_by: int
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # replace context with future_covariates, where the values of future covariates are known (not NaN)
        future_covariates_slice = future_covariates[..., :slide_by]
        context[..., -slide_by:] = torch.where(
            torch.isnan(future_covariates_slice),
            context[..., -slide_by:],
            future_covariates_slice,
        )
        # shift future_covariates
        future_covariates = future_covariates[..., slide_by:]

        return context, future_covariates

    def predict_quantiles(  # type: ignore[override]
        self,
        inputs: (
            TensorOrArray
            | Sequence[TensorOrArray]
            | Sequence[Mapping[str, TensorOrArray | Mapping[str, TensorOrArray]]]
        ),
        prediction_length: int | None = None,
        quantile_levels: list[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
        **predict_kwargs,
    ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
        """
        Refer to ``Chronos2Pipeline.predict`` for shared parameters.

        Additional parameters
        ---------------------
        quantile_levels
            Quantile levels to compute, by default [0.1, 0.2, ..., 0.9]

        Returns
        -------
        quantiles
            A list of torch tensors containing quantile forecasts. Each element of the list has shape (n_variates, prediction_length, len(quantile_levels))
            and the number of elements are equal to the number of target time series (univariate or multivariate) in the `inputs`.
        meanPipelineRegistry
            A list of torch tensors containing containing mean (point) forecasts. Each element of the list has shape (n_variates, prediction_length)
            and the number of elements are equal to the number of target time series (univariate or multivariate) in the `inputs`.
        """
        training_quantile_levels = self.quantiles

        # Invoke predict to get quantile predictions
        # with shape (n_variates, len(training_quantile_levels), prediction_length)
        # for each time series in the batch
        predictions: list[torch.Tensor] = self.predict(
            inputs, prediction_length=prediction_length, **predict_kwargs
        )

        # Swap quantile and time axes for each prediction
        predictions = [rearrange(pred, "... q h -> ... h q") for pred in predictions]

        if set(quantile_levels).issubset(training_quantile_levels):
            # no need to perform intra/extrapolation
            quantile_indices = [
                training_quantile_levels.index(q) for q in quantile_levels
            ]
            quantiles = [pred[..., quantile_indices] for pred in predictions]
        else:
            # we interpolate quantiles if quantiles that Chronos-2 was trained on were not provided
            if min(quantile_levels) < min(training_quantile_levels) or max(
                quantile_levels
            ) > max(training_quantile_levels):
                logger.warning(
                    f"\tQuantiles to be predicted ({quantile_levels}) are not within the range of "
                    f"quantiles that Chronos-2 was trained on ({training_quantile_levels}). "
                    "Quantile predictions will be set to the minimum/maximum levels at which Chronos-2 "
                    "was trained on. This may significantly affect the quality of the predictions."
                )

            quantiles = [
                interpolate_quantiles(quantile_levels, training_quantile_levels, pred)
                for pred in predictions
            ]

        # NOTE: the median is returned as the mean here
        mean = [pred[..., training_quantile_levels.index(0.5)] for pred in predictions]

        return quantiles, mean

    def predict_df(
        self,
        df: "pd.DataFrame",
        future_df: "pd.DataFrame | None" = None,
        id_column: str = "item_id",
        timestamp_column: str = "timestamp",
        target: str | list[str] = "target",
        prediction_length: int | None = None,
        quantile_levels: list[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
        batch_size: int = 256,
        context_length: int | None = None,
        # cross_learning: bool = False,
        validate_inputs: bool = True,
        freq: str | None = None,
        **predict_kwargs,
    ) -> "pd.DataFrame":
        """
        Perform forecasting on time series data in a long-format pandas DataFrame.

        Parameters
        ----------
        df
            Time series data in long format with an id column, a timestamp, and at least one target column.
            The remaining columns in df will be treated as past-only covariates unless they are also
            present in future_df
        future_df
            Future covariates data with an id column, a timestamp, and any number of covariate columns,
            all of these columns will be treated as known future covariates
        id_column
            The name of the column which contains the unique time series identifiers, by default "item_id"
        timestamp_column
            The name of the column which contains timestamps, by default "timestamp"
            All time series in the dataframe must have regular timestamps with the same frequency (no gaps)
        target
            The name of the column(s) which contain the target variables to be forecasted, by default "target"
        prediction_length
            Number of steps to predict for each time series
        quantile_levels
            Quantile levels to compute
        batch_size
            The batch size used for prediction. Note that the batch size here means the number of time series, including target(s) and covariates,
            which 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, by default 256
        context_length
            The maximum context length used during for inference, by default set to the model's default context length
        cross_learning
            If True, cross-learning is enabled, i.e., all the tasks in `inputs` will be predicted jointly and the model will share information across all inputs, by default False
            The following must be noted when using cross-learning:
            - Cross-learning doesn't always improve forecast accuracy and must be tested for individual use cases.
            - Results become dependent on batch size. Very large batch sizes may not provide benefits as they deviate from the maximum group size used during pretraining.
            For optimal results, consider using a batch size around 100 (as used in the Chronos-2 technical report).
            - Cross-learning is most helpful when individual time series have limited historical context, as the model can leverage patterns from related series in the batch.
        validate_inputs
            [ADVANCED] When True (default), validates dataframes before prediction. Setting to False removes the
            validation overhead, but may silently lead to wrong predictions if data is misformatted. When False, you
            must ensure: (1) all dataframes are sorted by (id_column, timestamp_column); (2) future_df (if provided)
            has the same item IDs as df with exactly prediction_length rows of future timestamps per item; (3) all
            timestamps are regularly spaced (e.g., with hourly frequency).
        freq
            Frequency string for timestamp generation (e.g., "h", "D", "W"). Can only be used when
            validate_inputs=False. When provided, skips frequency inference from the data.
        **predict_kwargs
            Additional arguments passed to predict_quantiles

        Returns
        -------
        The forecasts dataframe generated by the model with the following columns
        - `id_column`: The time series ID
        - `timestamp_column`: Future timestamps
        - "target_name": The name of the target column
        - "predictions": The point predictions generated by the model
        - One column for predictions at each quantile level in `quantile_levels`
        """
        try:
            import pandas as pd
        except ImportError:
            raise ImportError(
                "pandas is required for predict_df. Please install it with `pip install pandas`."
            )

        if prediction_length is None:
            prediction_length = self.model_prediction_length

        if not isinstance(target, list):
            target = [target]

        (
            inputs,
            original_order,
            prediction_timestamps,
        ) = convert_df_input_to_list_of_dicts_input(
            df=df,
            future_df=future_df,
            id_column=id_column,
            timestamp_column=timestamp_column,
            target_columns=target,
            prediction_length=prediction_length,
            freq=freq,
            validate_inputs=validate_inputs,
        )

        # Generate forecasts
        quantiles, mean = self.predict_quantiles(
            inputs=inputs,
            prediction_length=prediction_length,
            quantile_levels=quantile_levels,
            limit_prediction_length=False,
            batch_size=batch_size,
            context_length=context_length,
            # cross_learning=cross_learning,
            **predict_kwargs,
        )

        # since predict_df tasks are homogenous by input design, we can safely stack the list of tensors into a single tensor
        quantiles_np = torch.stack(
            quantiles
        ).numpy()  # [n_tasks, n_variates, horizon, num_quantiles]
        mean_np = torch.stack(mean).numpy()  # [n_tasks, n_variates, horizon]

        n_tasks = len(prediction_timestamps)
        n_variates = len(target)

        series_ids = list(prediction_timestamps.keys())
        future_ts = list(prediction_timestamps.values())

        data = {
            id_column: np.repeat(series_ids, n_variates * prediction_length),
            timestamp_column: np.concatenate(
                [np.tile(ts, n_variates) for ts in future_ts]
            ),
            "target_name": np.tile(np.repeat(target, prediction_length), n_tasks),
            "predictions": mean_np.ravel(),
        }

        quantiles_flat = quantiles_np.reshape(-1, len(quantile_levels))
        for q_idx, q_level in enumerate(quantile_levels):
            data[str(q_level)] = quantiles_flat[:, q_idx]

        predictions_df = pd.DataFrame(data)
        # If validate_inputs=False, the df is used as-is without sorting by item_id, no reordering required
        if validate_inputs:
            predictions_df.set_index(id_column, inplace=True)
            predictions_df = predictions_df.loc[original_order]
            predictions_df.reset_index(inplace=True)

        return predictions_df

    @torch.no_grad()
    def embed(
        self,
        inputs: TensorOrArray | Sequence[TensorOrArray],
        batch_size: int = 256,
        context_length: int | None = None,
    ) -> tuple[list[torch.Tensor], list[tuple[torch.Tensor, torch.Tensor]]]:
        """
        Get encoder embeddings for the given time series.

        Parameters
        ----------
        inputs
            The time series to get embeddings for, can be one of:
            - A 3-dimensional `torch.Tensor` or `np.ndarray` of shape (batch, n_variates, history_length). When `n_variates > 1`, information
            will be shared among the different variates of each time series in the batch.
            - A list of `torch.Tensor` or `np.ndarray` where each element can either be 1-dimensional of shape (history_length,)
            or 2-dimensional of shape (n_variates, history_length). The history_lengths may be different across elements; left-padding
            will be applied, if needed.
        batch_size
            The batch size used for generating embeddings. Note that the batch size here means the total number of time series which are input into the model.
            If your data has multiple variates, the effective number of time series tasks in a batch will be lower than this value, by default 256
        context_length
            The maximum context length used during for inference, by default set to the model's default context length

        Returns
        -------
        embeddings
            a list of `torch.Tensor` where each element has shape (n_variates, num_patches + 2, d_model) and the number of elements are equal to the number
            of target time series (univariate or multivariate) in the `inputs`. The extra +2 is due to embeddings of the [REG] token and a masked output patch token.
        loc_scale
            a list of tuples with the mean and standard deviation of each time series.
        """
        if context_length is None:
            context_length = self.model_context_length

        if context_length > self.model_context_length:
            warnings.warn(
                f"The specified context_length {context_length} is greater than the model's default context length {self.model_context_length}. "
                f"Resetting context_length to {self.model_context_length}."
            )
            context_length = self.model_context_length

        test_dataset = LiteSpecFormerDataset.convert_inputs(
            inputs=inputs,
            context_length=context_length,
            prediction_length=0,
            batch_size=batch_size,
            output_patch_size=self.model_output_patch_size,
            mode=DatasetMode.TEST,
        )
        test_loader = DataLoader(
            test_dataset,
            batch_size=None,
            num_workers=0,
            pin_memory=self.model.device.type == "cuda",
            shuffle=False,
            drop_last=False,
        )

        all_embeds: list[torch.Tensor] = []
        all_loc_scales: list[tuple[torch.Tensor, torch.Tensor]] = []

        for batch in test_loader:
            assert batch["future_target"] is None
            batch_context = batch["context"]
            batch_target_idx_ranges = batch["target_idx_ranges"]

            encoder_outputs, (locs, scales), *_ = self.model.encode(
                context=batch_context.to(device=self.model.device, dtype=torch.float32)
            )
            batch_embeds = [
                encoder_outputs[0][start:end].cpu()
                for (start, end) in batch_target_idx_ranges
            ]
            batch_loc_scales = list(
                zip(
                    [locs[start:end].cpu() for (start, end) in batch_target_idx_ranges],
                    [
                        scales[start:end].cpu()
                        for (start, end) in batch_target_idx_ranges
                    ],
                )
            )
            all_embeds.extend(batch_embeds)
            all_loc_scales.extend(batch_loc_scales)

        return all_embeds, all_loc_scales

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        """
        Load the model, either from a local path, S3 prefix or from the HuggingFace Hub.
        Supports the same arguments as ``AutoConfig`` and ``AutoModel`` from ``transformers``.
        """

        # Check if the model is on S3 and cache it locally first
        # NOTE: Only base models (not LoRA adapters) are supported via S3
        if str(pretrained_model_name_or_path).startswith("s3://"):
            return BaseForecastPipeline.from_pretrained(
                pretrained_model_name_or_path, *args, **kwargs
            )

        # Check if the hub model_id or local path is a LoRA adapter
        if find_adapter_config_file(pretrained_model_name_or_path) is not None:
            if not is_peft_available():
                raise ImportError(
                    f"The model at {pretrained_model_name_or_path} is a `peft` adaptor, but `peft` is not available. "
                    f"Please install `peft` with `pip install peft` to use this model. "
                )
            from peft import AutoPeftModel

            model = AutoPeftModel.from_pretrained(
                pretrained_model_name_or_path, *args, **kwargs
            )
            model = model.merge_and_unload()
            return cls(model=model)

        # Handle the case for the base model
        config = AutoConfig.from_pretrained(
            pretrained_model_name_or_path, *args, **kwargs
        )

        assert hasattr(
            config, "litespecformer_pipeline_class"
        ), "Not a LiteSpecFormer config file"
        pipeline_class_name = getattr(
            config, "litespecformer_pipeline_class", "LiteSpecFormerPipeline"
        )

        model = LiteSpecFormerModel.from_pretrained(
            pretrained_model_name_or_path, *args, **kwargs
        )
        return cls(model=model)

    def save_pretrained(self, save_directory: str | Path, *args, **kwargs):
        """
        Save the underlying model to a local directory or to HuggingFace Hub.
        """
        self.model.save_pretrained(save_directory, *args, **kwargs)

    def calculate_metrics(
        self,
        contexts: Union[np.ndarray, torch.FloatTensor],
        targets: Union[np.ndarray, torch.FloatTensor],
        predictions: Union[np.ndarray, torch.FloatTensor],
    ) -> Dict[str, float]:
        """
        Compute evaluation metrics given the contexts, targets and predictions. This is a placeholder method and should be implemented based on the specific metrics you want to compute.

        Arguments
        ----------
        contexts: Union[np.ndarray, torch.FloatTensor]
            The historical context time series data used for making predictions, with shape (batch_size, num_variables, context_length).
        targets: Union[np.ndarray, torch.FloatTensor]
            The true values for the test set, with shape (batch_size, num_variables, prediction_length).
        predictions: Union[np.ndarray, torch.FloatTensor]
            The predicted values of LiteSpecFormer, expected to have shape (batch_size, num_variables, num_quantiles, prediction_length).

        Returns
        -------
        Dict[str, float]
            A dictionary containing the computed evaluation metrics, including:
            - "MSE": Mean Squared Error between the median predictions and the targets.
            - "MAE": Mean Absolute Error between the median predictions and the targets.
            - "RMSE": Root Mean Squared Error between the median predictions and the targets.
            - "MAPE": Mean Absolute Percentage Error between the median predictions and the targets.
            - "MASE": Mean Absolute Scaled Error, which is the mean absolute error of the predictions scaled by the mean absolute error of a naive forecast (e.g., using the last observed value as the forecast).
            - "MSPE": Mean Scaled Pinball Error, which is a common metric for evaluating quantile forecasts, calculated as the pinball loss of the predicted quantiles normalized by the pinball loss of a naive forecast (e.g., using the last observed value as the forecast).
            - "RSE": Relative Squared Error, calculated as the sum of squared errors of the predictions relative to the sum of squared errors of a naive forecast.
        """
        # Convert inputs to torch tensors if they are numpy arrays
        if isinstance(targets, np.ndarray):
            targets = torch.from_numpy(targets).float().to(self.model.device)
        if isinstance(predictions, np.ndarray):
            predictions = torch.from_numpy(predictions).float().to(self.model.device)

        # Check the shape of inputs data
        assert (
            targets.ndim == 3
        ), f"Expected targets to have 3 dimensions (batch_size, n_variates, prediction_length), but got {targets.shape}"
        assert (
            predictions.ndim == 4
        ), f"Expected predictions to have 4 dimensions (batch_size, n_variates, num_quantiles, prediction_length), but got {predictions.shape}"

        # Compute metrics
        return calculate_metrics(
            contexts=contexts, predictions=predictions, targets=targets
        )

    def evaluate_spectrum_library(
        self,
        dataset_name_or_path: str,
        seq_length: int,
        prediction_length: int,
        batch_size: int = 256,
        split_ratio: float = 0.6,
        accelerator: Accelerator | None = None,
        output_dir: str | None = None,
    ) -> Dict[str, float]:
        """
        该方法用于测试当前pipeline中频谱预测模型的零样本置信预测能力。
        其主要的测试方法和流程主要来自于: https://github.com/wwhenxuan/Spectrum-Prediction-Library

        Arguments
        ----------
        dataset_name_or_path: str
            Path to the spectrum library dataset.
        seq_length: int
            Length of the input sequence.
        prediction_length: int
            Length of the prediction sequence.
        batch_size: int, optional
            Batch size for evaluation (default is 256).
        split_ratio: float, optional
            Ratio for splitting the dataset (default is 0.6).
        accelerator: Accelerator | None, optional
            Accelerator for distributed evaluation (default is None).
        output_dir: str | None, optional
            Directory to save evaluation results (default is None).

        Returns
        -------
        Dict[str, float]
            A dictionary containing the computed evaluation metrics.
        """
        time_now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        setting = f"{time_now}_seq{seq_length}_pred{prediction_length}_batch{batch_size}_split{split_ratio}"

        if accelerator is None:
            # Create an Accelerator for distributed evaluation if one is not provided.
            # This allows the evaluation to be performed on multiple GPUs if available, which can speed up the process significantly for large datasets.
            accelerator = Accelerator()
        accelerator.print(
            f"[{time_now}] Starting evaluation on spectrum library with:",
            dataset_name_or_path,
        )
        accelerator.print(
            f"\tseq_length={seq_length}, prediction_length={prediction_length}, batch_size={batch_size}, split_ratio={split_ratio}"
        )

        accelerator.print("Loading model onto device:", accelerator.device)
        accelerator.print(
            Fore.RED + "Loading the Spectrum Library dataset" + Style.RESET_ALL,
            end=" -> ",
        )
        # load the dataset for evaluation
        dataset = SpectrumLibraryDataset(
            seq_length=seq_length,
            prediction_length=prediction_length,
            dataset_name_or_path=dataset_name_or_path,
            split_ratio=split_ratio,
            flag="test",
        )
        # create the dataloader for evaluation
        data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
        accelerator.print(Fore.GREEN + "Done" + Style.RESET_ALL)

        # Prepare the model and dataloader for distributed evaluation with accelerator
        accelerator.print(
            "Preparing the model and dataloader for distributed evaluation..."
        )
        self.model, data_loader = accelerator.prepare(self.model, data_loader)

        # Create lists to store contexts, targets and predictions for metric computation after the evaluation loop
        context_list, target_list, prediction_list = [], [], []

        accelerator.print(Fore.BLUE + "Starting evaluation loop" + Style.RESET_ALL)
        # Create a tqdm progress bar for the dataloader
        data_loader = tqdm(data_loader, desc="Evaluating")
        with torch.no_grad():
            # iterate through the dataloader and evaluate the model
            for batch_x, batch_y in data_loader:
                # [batch_size, num_variates, seq_length]
                context, target = batch_x.permute(0, 2, 1), batch_y.permute(0, 2, 1)
                # [batch_size * num_variates, seq_length] -> channel independent processing
                # context = context.reshape(-1, context.shape[-1])
                # target = target.reshape(-1, target.shape[-1])

                # Predict the future values using the model's predict method
                prediciton = self.predict(
                    inputs=context,
                    prediction_length=prediction_length,
                )
                # [batch_size, num_variates, num_quantiles, prediction_length]
                prediction = torch.concatenate(
                    [pred.unsqueeze(0) for pred in prediciton], dim=0
                )

                context_list.append(context.cpu())
                target_list.append(target.cpu())
                prediction_list.append(prediction.cpu())
        accelerator.print(Fore.BLUE + "Evaluation loop completed" + Style.RESET_ALL)

        # Compute evaluation metrics using the collected contexts, targets and predictions
        accelerator.print("Computing evaluation metrics...")
        contexts, targets, predictions = (
            torch.concatenate(context_list, dim=0).to(self.model.device),
            torch.concatenate(target_list, dim=0).to(self.model.device),
            torch.concatenate(prediction_list, dim=0).to(self.model.device),
        )
        computed_metrics = self.calculate_metrics(
            contexts=contexts, targets=targets, predictions=predictions
        )
        for key in computed_metrics.keys():
            computed_metrics[key] = float(computed_metrics[key])
            accelerator.print(f"\t{key}: {computed_metrics[key]}")

        mse, mae, rmse, mape, mase, mspe, rse = (
            computed_metrics.get("mse"),
            computed_metrics.get("mae"),
            computed_metrics.get("rmse"),
            computed_metrics.get("mape"),
            computed_metrics.get("mase"),
            computed_metrics.get("mspe"),
            computed_metrics.get("rse"),
        )

        # Create output directory if it doesn't exist and save the computed metrics
        if output_dir is not None:
            os.makedirs(output_dir, exist_ok=True)
        else:
            output_dir = "./results/zero_shot_evaluation"
            os.makedirs(output_dir, exist_ok=True)

        save_path = path.join(output_dir, setting + ".pth")
        accelerator.print(f"Saving outputs results to {save_path}", end=" -> ")
        if accelerator.is_main_process:
            torch.save(
                obj={
                    "context": contexts,
                    "target": targets,
                    "prediction": predictions,
                    "mse": mse,
                    "mae": mae,
                    "rmse": rmse,
                    "mape": mape,
                    "mase": mase,
                    "mspe": mspe,
                    "rse": rse,
                },
                f=save_path,
            )
        accelerator.print(Fore.GREEN + "Done" + Style.RESET_ALL)

        # Logging the results to csv file locally
        accelerator.print("Logging results to csv file", end=" -> ")
        csv_save_path = path.join(output_dir, "results.csv")
        logging_messages = {
            "time": time_now,
            "model": "LiteSpecFormer",
            "dataset": dataset_name_or_path,
            "seq_length": seq_length,
            "prediction_length": prediction_length,
            "mse": mse,
            "mae": mae,
            "rmse": rmse,
            "mape": mape,
            "mase": mase,
            "mspe": mspe,
            "rse": rse,
        }
        logging_results(
            accelerator=accelerator,
            logging_path=csv_save_path,
            headers=list(logging_messages.keys()),
            messages=logging_messages,
        )
        accelerator.print(Fore.GREEN + "Done" + Style.RESET_ALL)

        # Clear the model and dataloader from accelerator to free up memory before computing metrics
        accelerator.print(
            "Clearing model and dataloader from accelerator to free up memory",
            end=" -> ",
        )
        accelerator.clear(self.model, data_loader)
        accelerator.print(Fore.GREEN + "Done" + Style.RESET_ALL)

        return computed_metrics
