from typing import Union, Tuple, List, Dict, Any

import argparse
import os
from os import path
import time
import csv
import sys

import numpy as np
import pandas as pd

from accelerate import Accelerator
from colorama import Fore, Style

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler

from tqdm import tqdm

from litespecformer import (
    LiteSpecFormerPipeline,
    LiteSpecFormerModel,
    LiteSpecFormerConfig,
)
from litespecformer.metrics import (
    mean_squared_error,
    mean_absolute_error,
    root_mean_squared_error,
    mean_absolute_percentage_error,
    mean_squared_percentage_error,
    root_squared_error,
    mean_absolute_scaled_error,
)


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

    def __init__(self, config, flag: str = "train") -> None:
        super().__init__()
        # The configuration parameters for the dataset
        self.config = config

        # The length of the input sequence and the prediction sequence
        self.seq_len = config.seq_len
        self.pred_len = config.pred_len

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

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

        # The path to load the csv data
        self.data_path = config.data_path
        print(f"Loading data from {self.data_path}...")

        # Load the data from the csv file
        self.data, self.data_x, self.data_y = self.load_data(self.data_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._n_vars = self.data_x.shape[1]

    def load_data(self, data_path: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Load the dataset from the given path.
        """
        if not path.exists(data_path):
            raise FileNotFoundError(f"Dataset not found at {data_path}")

        # Get the data from the csv file
        data = pd.read_csv(data_path).iloc[:, 1:].values

        # 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.config.split_ratio)

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

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

        return data, data_x, data_y

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

    def __len__(self) -> int:
        """Get the number of samples in the dataset."""
        return len(self.data_x) - self.seq_len - self.pred_len + 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_len

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

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


def calculate_metrics(
    contexts: torch.FloatTensor,
    predictions: torch.FloatTensor,
    targets: torch.FloatTensor,
    medium_index: int = 5,
) -> Dict[str, torch.FloatTensor]:
    """
    Calculates a set of common regression metrics between predictions and targets.
    Includes Mean Squared Error (MSE), Mean Absolute Error (MAE), Root Mean Squared Error (RMSE),
    Mean Absolute Percentage Error (MAPE), Mean Squared Percentage Error (MSPE), and Root Squared Error (RSE).

    Arguments
    ----------
    contexts
        The historical time series data used for training, with shape [batch_size, num_variables, context_length].
    predictions
        The predicted values of LiteSpecFormer, expected to have shape [batch_size, num_variables, num_quantiles, prediction_length].
    targets
        The true values for the test set, expected to have shape [batch_size, num_variables, prediction_length].
    medium_index
        The index of the median quantile (0.5) in the predictions, used for loss calculation and evaluation (default is 5 for 11 quantiles).

    Returns
    -------
    A dictionary containing the calculated metrics:
    - "mse": Mean Squared Error
    - "mae": Mean Absolute Error
    - "rmse": Root Mean Squared Error
    - "mape": Mean Absolute Percentage Error
    - "mase": Mean Absolute Scaled Error
    - "mspe": Mean Squared Percentage Error
    - "rse": Root Squared Error
    """
    print("contexts shape:", contexts.shape)
    print("predictions shape:", predictions.shape)
    print("targets shape:", targets.shape)

    # Input validation
    batch_size, num_variables, context_length = contexts.shape
    prediction_length = predictions.shape[1]

    return {
        "mse": mean_squared_error(predictions, targets),
        "mae": mean_absolute_error(predictions, targets),
        "rmse": root_mean_squared_error(predictions, targets),
        "mape": mean_absolute_percentage_error(predictions, targets),
        "mase": mean_absolute_scaled_error(
            contexts=torch.reshape(
                contexts,
                shape=(batch_size * num_variables, context_length),
            ),
            targets=torch.reshape(
                targets.permute(0, 2, 1),
                shape=(batch_size * num_variables, prediction_length),
            ),
            predictions=torch.reshape(
                predictions.permute(0, 2, 1),
                shape=(batch_size * num_variables, prediction_length),
            ),
        ),
        "mspe": mean_squared_percentage_error(predictions, targets),
        "rse": root_squared_error(predictions, targets),
    }


def logging_results(
    accelerator: Accelerator,
    logging_path: str,
    headers: List[str],
    messages: Dict[str, Union[str, float]],
) -> None:
    """
    将训练结果记录到CSV文件中，支持多卡训练场景

    Args:
        accelerator: Accelerate库的Accelerator对象，用于多卡同步和主进程判断
        logging_path: CSV文件的保存路径（包含文件名）
        headers: CSV文件的表头列表
        messages: 要记录的结果字典，键需与headers对应
    """
    # 只在主进程执行文件写入操作，避免多进程冲突
    if accelerator.is_main_process:
        # 检查message的键是否与headers匹配
        message_keys = set(messages.keys())
        header_set = set(headers)
        if not message_keys.issubset(header_set):
            missing_keys = message_keys - header_set
            raise ValueError(f"Message中包含表头不存在的键: {missing_keys}")

        # 确保目录存在
        os.makedirs(path.dirname(logging_path), exist_ok=True)

        # 判断文件是否存在
        file_exists = path.exists(logging_path)

        # 写入CSV文件
        with open(logging_path, "a", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=headers)

            # 如果文件不存在，先写入表头
            if not file_exists:
                writer.writeheader()

            # 按表头顺序写入数据（确保顺序一致）
            # 过滤并按headers顺序整理数据
            row_data = {header: messages.get(header, "") for header in headers}
            writer.writerow(row_data)

        # 同步所有进程，确保文件写入完成后其他进程再继续
        accelerator.wait_for_everyone()


def print_configs(
    accelerator: "Accelerator",  # 加引号避免未导入时的类型提示报错，实际使用时可移除引号
    time_now: str,
    config: Dict[str, Any],
    experiment_name: str = "Spectrum Prediction",
    model_name: str = "SpectrumTime",
    dataset: str = "spec4",
    mode: str = "supervised",
    print_separator: bool = True,
) -> None:
    """
    Prints key configuration parameters for deep learning experiments, optimized for distributed training scenarios.
    This print configs function from wwhenxuan for the `Spectrum Prediction Library`.

    Args:
    -----
        accelerator: Accelerator instance from Hugging Face Accelerate library, used for unified printing in distributed training.
        time_now: Formatted string of the experiment start time (e.g., "2026-01-13 14:25:30").
        config: Dictionary containing all critical experiment parameters (e.g., lr, batch_size, epochs, etc.).
        experiment_name: Name of the experiment for title display (default: "Spectrum Prediction").
        model_name: Name of the model being used (default: "SpectrumTime").
        dataset: Name of the dataset for the experiment (default: "spec4").
        mode: Training mode (e.g., "supervised", "unsupervised", default: "supervised").
        print_separator: Whether to print separator lines for better readability (default: True).
    """
    # Build the content to be printed
    output_lines = []

    # Title and time information
    if print_separator:
        separator = "=" * 60
        output_lines.append(separator)
    output_lines.append(f"📋 {experiment_name} - Experiment Parameters")
    output_lines.append(f"🤖 {model_name} model ( {mode} mode ) on {dataset} dataset.")
    output_lines.append(f"⏰ Experiment Start Time: {time_now}")
    if print_separator:
        output_lines.append("-" * 60)

    # Calculate the maximum length of parameter names for aligned display
    max_key_len = max(len(str(key)) for key in config.keys()) if config else 0
    # Print all parameters with aligned formatting
    for key, value in config.items():
        # Format: parameter name left-justified, value left-aligned for readability
        line = f"🔧 {key.ljust(max_key_len)} : {value}"
        output_lines.append(line)

    if print_separator:
        output_lines.append(separator)

    # Concatenate all lines into a single text block
    output_text = "\n".join(output_lines)

    # Print to console via accelerator (supports distributed training)
    accelerator.print(output_text)


class BaseExperiment(object):
    """
    The base class for spectrum prediction experiment.
    This module will be inherited later for
    - supervised training,
    - zero-shot generalization, and
    - fine-tuning of pre-trained models, respectively.`
    """

    def __init__(
        self, configs, model_config, accelerator: Accelerator, setting: str
    ) -> None:
        # The configs and the accelerator for model training
        self.configs = configs
        self.model_config = model_config
        self.accelerator = accelerator
        self.setting = setting

        # Build the deep learning model
        self.model_name = configs.model
        self.model = self.build_model(name=self.model_name)

        # Whether to train or test the model with zero-shot forecasting
        self.mode = configs.mode

        # The hyper-parameters for data loading
        self.batch_size = configs.batch_size
        self.shuffle = configs.shuffle

        # The checkpoint directory
        self.checkpoint_dir = configs.checkpoint
        self.checkpoint_path = path.join(self.checkpoint_dir, setting)
        os.makedirs(self.checkpoint_path, exist_ok=True)

        # The file path of the training dataset for supervised learning
        # Or the testing dataset for unsupervised learning (zero-shot learning)
        self.data_path = configs.data_path

        # The root path of the data for model fine-tuning or large scale pre-training
        self.root_path = configs.root_path

    @property
    def model_dict(self) -> dict:
        """The dictionary of available models."""
        # 这里返回的模型需要进行修改
        return {"LiteSpecFormer": LiteSpecFormerModel}

    def build_model(self, name: str = "SpectrumTime") -> nn.Module:
        """Build the model for training."""
        # Check if the model name is valid
        assert name in self.model_dict, f"Model {name} is not supported."

        self.accelerator.print(f"Building the model: {name}", end=" -> ")

        # Create the model for experiment
        model = self.model_dict[name](self.model_config)

        self.accelerator.print(Fore.GREEN + "Done!" + Style.RESET_ALL)

        return model

    def _load_data(self, flag: str = "train") -> DataLoader:
        """
        Load the data for training or testing.
        The `flag` parameter controls whether a training or testing set is created.

        TODO: This interface module may be adjusted in the future to adapt to different spectrum datasets.
        These datasets may have different formats.
        """
        assert flag in ["train", "test"], "flag must be either 'train' or 'test'!"

        # Create the dataset for Spectrum data
        dataset = SpectrumDataset(self.configs, flag=flag)

        # Print dataset information
        self.accelerator.print(
            f"{flag.capitalize()} dataset loaded with {len(dataset)} samples.",
            end=" -> ",
        )

        # Create the data loader
        data_loader = DataLoader(
            dataset,
            batch_size=self.configs.batch_size,
            shuffle=(flag == "train"),
            drop_last=False,  # Do not drop the last batch
        )
        self.accelerator.print(Fore.GREEN + "Done!" + Style.RESET_ALL)

        return data_loader

    def get_trainable_params(self) -> List[nn.Parameter]:
        """Obtain trainable parameters of the model in experiment."""
        assert hasattr(
            self, "model"
        ), "The model has not been built yet. Please call the build_model method first."
        return [p for p in self.model.parameters() if p.requires_grad]

    def get_num_trainable_params(self) -> int:
        """Obtain the number of trainable parameters"""
        return sum(p.numel() for p in self.get_trainable_params())

    def get_learning_rate(self, optimizer: torch.optim.Optimizer) -> float:
        """Get the current learning rate of the optimizer"""
        return optimizer.param_groups[0]["lr"]

    def print_start_message(self, time_now: str) -> None:
        """Prints a message indicating the start of the experiment."""
        print_configs(
            accelerator=self.accelerator,
            time_now=time_now,
            config={  # The main configs parameters to be printed
                "seq_len": self.configs.seq_len,
                "pred_len": self.configs.pred_len,
                "batch_size": self.batch_size,
                "learning_rate": self.configs.learning_rate,
            },
            experiment_name="Spectrum Prediction",
            model_name=self.model_name,
            dataset=self.configs.dataset_name,
            mode=self.mode,
            print_separator=True,
        )

    def test(self, data_loader: DataLoader) -> Tuple[float, float, float]:
        """
        A unified interface for spectrum prediction testing.

        To ensure the reliability of the experiment,
        all base classes used for verifying results uniformly call this method to obtain the results.

        The created `data_loader` object in PyTorch needs to be passed in.
        """

        # Only load the checkpoint for supervised and fine-tuning
        if self.mode in ["supervised", "fine-tuning"]:
            # loading the model parameters from checkpoint
            self.accelerator.print(
                "🚀 "
                + Fore.BLUE
                + "Loading the best model for testing..."
                + Style.RESET_ALL
            )
            self.accelerator.load_state(self.checkpoint_path)

        # Create the lists to store predictions and targets
        inputs = []
        predictions = []
        targets = []
        time_list = []

        # Set the model to evaluation mode
        self.model.model.eval()

        with torch.no_grad():
            for batch_x, batch_y in tqdm(data_loader):
                # 切换输入数据的形式将其转换为通道独立的样式
                batch_x = batch_x.permute(0, 2, 1)  # (batch_size, n_vars, seq_len)
                # batch_y = batch_y.permute(0, 2, 1)  # (batch_size, n_vars, pred_len)

                batch_size, n_vars, seq_len = batch_x.shape

                # batch_x = torch.reshape(batch_x, (batch_size * n_vars, seq_len))

                time_start = time.time()
                outputs = self.model.predict(
                    batch_x, prediction_length=self.configs.pred_len
                )

                outputs = torch.concat(
                    [output[:, 5, :].unsqueeze(dim=0) for output in outputs], dim=0
                )
                outputs = outputs.permute(0, 2, 1)  # (batch_size, pred_len, n_vars)

                time_end = time.time()

                time_list.append(time_end - time_start)

                inputs.append(batch_x)
                predictions.append(outputs)
                targets.append(batch_y)

        inputs = torch.concatenate(inputs, axis=0).to(self.accelerator.device)
        predictions = torch.concatenate(predictions, axis=0).to(self.accelerator.device)
        targets = torch.concatenate(targets, axis=0).to(self.accelerator.device)

        # Calculate the evaluation metrics
        metrics = self.metrics(inputs, predictions, targets)

        # Save the test results to the folder
        self.save_results(
            inputs=inputs, predictions=predictions, targets=targets, metrics=metrics
        )

        return metrics, np.mean(time_list)

    def metrics(
        self,
        inputs: torch.FloatTensor,
        outputs: torch.FloatTensor,
        labels: torch.FloatTensor,
    ) -> Dict[str, float]:
        """
        Calculate evaluation indicators.

        The model's performance metrics are calculated based on its output and the true labels.

        For the prediction task, we used two metrics:
        mean squared error (MSE) and mean absolute error (MAE).
        """
        self.accelerator.print(
            "🧪 " + Fore.RED + "Calculating test metrics..." + Style.RESET_ALL
        )

        # Calculate the metrics
        metrics = calculate_metrics(
            contexts=inputs, predictions=outputs, targets=labels
        )

        for metric_name in metrics.keys():
            metrics[metric_name] = float(
                metrics[metric_name]
            )  # Convert to Python float

            self.accelerator.print(
                f"🎖️ {metric_name}: {Fore.GREEN}{metrics[metric_name]}{Style.RESET_ALL}"
            )

        return metrics

    def save_results(
        self,
        inputs: torch.FloatTensor,
        predictions: torch.FloatTensor,
        targets: torch.FloatTensor,
        metrics: Dict[str, float],
    ) -> None:
        """Save experimental results to a specified path."""
        results_path = self.checkpoint_path + "/results.pth"

        # Wait for all processes to finish before saving
        self.accelerator.wait_for_everyone()

        # Only main process saves the results
        if self.accelerator.is_main_process:
            # Save the model inputs, predictions, targets, MSE, and MAE
            torch.save(
                {
                    "inputs": inputs,
                    "predictions": predictions,
                    "targets": targets,
                    "mse": metrics["mse"],
                    "mae": metrics["mae"],
                    "rmse": metrics["rmse"],
                    "mape": metrics["mape"],
                    "mase": metrics["mase"],
                    "mspe": metrics["mspe"],
                    "rse": metrics["rse"],
                },
                results_path,
            )
        self.accelerator.print("Test results saved to " + results_path)

    def logging(
        self,
        time_now: str,
        metrics: Dict[str, float],
        time_mean: float,
    ) -> None:
        """Record the experimental results to a CSV file."""

        # Log the experiment results
        # Wait for all processes to finish before saving
        self.accelerator.wait_for_everyone()

        # Only save main process saves the model
        if self.accelerator.is_main_process:
            messages = {
                "Timestamp": time_now,
                "Dataset": self.configs.dataset_name,
                "Model": self.model_name,
                "MSE": metrics["mse"],
                "MAE": metrics["mae"],
                "RMSE": metrics["rmse"],
                "MAPE": metrics["mape"],
                "MASE": metrics["mase"],
                "MSPE": metrics["mspe"],
                "RSE": metrics["rse"],
                "Time": time_mean,
                "split_ratio": self.configs.split_ratio,
                "seq_len": self.configs.seq_len,
                "pred_len": self.configs.pred_len,
                "setting": self.setting,
            }

            # logging the results to csv file
            logging_results(
                accelerator=self.accelerator,
                logging_path="./results.csv",
                headers=self.logging_headers,
                messages=messages,
            )

            # Print the state
            self.accelerator.print(
                Fore.GREEN + "The results logging is done." + Fore.RESET
            )

    @property
    def logging_headers(self) -> List[str]:
        """The Log header for the results of the csv."""
        return [
            "Timestamp",
            "Dataset",
            "Model",
            "Mode",  # The mode for the model
            "MSE",
            "MAE",
            "RMSE",
            "MAPE",
            "MASE",
            "MSPE",
            "RSE",
            "Time",
            "split_ratio",
            "seq_len",
            "pred_len",
            "Loss_function",
            "Optimizer",
            "Learning_rate",
            "Batch_size",
            "setting",
            "seed",
            "path",  # The path to the saved checkpoint and results
        ]

    def _check_loss(self, loss: torch.FloatTensor) -> None:
        """Check the training and validation losses to avoid gradient explosion."""
        if not torch.isfinite(loss):
            self.accelerator.print(
                Fore.RED
                + "WARNING: non-finite loss, ending training!"
                + Style.RESET_ALL
            )
            sys.exit(1)


class ZeroShotExperiment(BaseExperiment):
    """The epxeriment interface for the zero-shot validation of the time series foundation models."""

    def __init__(
        self,
        configs,
        model_config: LiteSpecFormerConfig,
        accelerator: Accelerator,
        setting: str,
    ) -> None:
        super().__init__(
            configs, model_config=model_config, accelerator=accelerator, setting=setting
        )

        # The only thing can do is to load the testing data loader
        self.test_loader = self._load_data(flag="test")

    def run(self, time_now: str) -> None:
        """Run the zero-shot validation experiment."""

        self.print_start_message(time_now=time_now)

        # Prepare the model and data loader for evaluation
        self.model, test_loader = self.accelerator.prepare(self.model, self.test_loader)

        # Create the pipeline for the model
        self.model = LiteSpecFormerPipeline(model=self.model)

        self.accelerator.load_state(self.configs.model_params)

        # Run the zero-shot validation
        metrics, time_mean = self.test(data_loader=test_loader)

        # Logging the results
        self.logging(time_now=time_now, metrics=metrics, time_mean=time_mean)


parser = argparse.ArgumentParser(description="The configuration for LiteSpecFormer")

parser.add_argument(
    "--model", type=str, default="LiteSpecFormer", help="The model name."
)
parser.add_argument(
    "--mode",
    type=str,
    default="zero-shot",
    choices=["supervised", "zero-shot", "fine-tuning"],
    help="Mode of operation: train, zero-shot prediction, or fine-tuning.",
)
parser.add_argument(
    "--dataset_name",
    type=str,
    default="hello",
    choices=["spec4", "spec25", "spec151"],
    help="The name of dataset to be used.",
)
parser.add_argument(
    "--data_path",
    type=str,
    default="./dataset/hello.csv",
    help="The path of the training and testing dataset for supervised learning.",
)
parser.add_argument(
    "--root_path",
    type=str,
    default="./dataset/",
    help="The root path of spectrum datasets mainly used to load all the data for pre-training.",
)
parser.add_argument(
    "--checkpoint",
    type=str,
    default="./results",
    help="The directory to save checkpoints.",
)
parser.add_argument("--model_params", type=str, default="./checkpoints", help="")
parser.add_argument(
    "--seq_len",
    type=int,
    default=96,
    choices=[96, 192, 336, 512],
    help="Input sequence length or the look-back windows.",
)
parser.add_argument(
    "--pred_len",
    type=int,
    default=96,
    choices=[12, 16, 24, 48, 96, 192, 336, 720],
    help="The length of model to be predicted.",
)
parser.add_argument(
    "--split_ratio",
    type=float,
    default=0.6,
    help="The ratio to split the trainning and testing dataset.",
)
# training hyper-parameters
parser.add_argument(
    "--batch_size", type=int, default=32, help="The batch size of training."
)
parser.add_argument(
    "--shuffle",
    type=bool,
    default=True,
    help="Whether to shuffle the training dataset.",
)
parser.add_argument(
    "--learning_rate", type=float, default=0.001, help="The learning rate of optimizer."
)
# Whether to scaling the time series data
parser.add_argument(
    "--scale", type=bool, default=True, help="Whether to standard the training data."
)

args = parser.parse_args()

if __name__ == "__main__":
    time_now = time.strftime(f"%Y-%m-%d_%H:%M:%S", time.localtime())

    experiment = ZeroShotExperiment(
        configs=args,
        model_config=LiteSpecFormerConfig(),
        accelerator=Accelerator(),
        setting=f"{time_now}_{args.dataset_name}_{args.seq_len}_{args.pred_len}",
    )

    experiment.run(time_now=time_now)
