from __future__ import annotations

import logging
from pathlib import Path
from timeit import default_timer

import pandas as pd
import torch

from mlpforecast.forecaster.utils import format_target, get_latest_checkpoint
from mlpforecast.metrics.deterministic import evaluate_point_forecast

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("Forecaster")
import lightning as pl
from lightning.pytorch import loggers
from lightning.pytorch.callbacks import (
    EarlyStopping,
    LearningRateMonitor,
    ModelCheckpoint,
    RichProgressBar,
    TQDMProgressBar,
)
from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBarTheme
from optuna_integration.pytorch_lightning import PyTorchLightningPruningCallback

from mlpforecast.data.loader import TimeseriesDataModule


class PytorchForecast:
    """
    PytorchForecast class for setting up and managing the training process of a PyTorch model using PyTorch Lightning.

    Attributes:
        exp_name (str): The name of the experiment.
        file_name (str): The name of the file to save the logs and model checkpoints.
        seed (int): The seed for random number generation to ensure reproducibility.
        root_dir (str): The root directory for saving logs and checkpoints.
        trial (optuna.trial.Trial): Optuna trial for hyperparameter optimization.
        metric (str): The metric to monitor for early stopping and model checkpointing.
        max_epochs (int): The maximum number of training epochs.
        wandb (bool): Flag to use Wandb logger instead of TensorBoard logger.
        rich_progress_bar (bool): Flag to use rich progress bar for training visualization.
    """

    def __init__(
        self,
        exp_name="Tanesco",
        file_name=None,
        seed=42,
        root_dir="../",
        trial=None,
        metric="val_mae",
        max_epochs=10,
        wandb=False,
        model_type="MLPF",
        rich_progress_bar=False,
        gradient_clip_val=10,
    ):
        """
        Initializes the PytorchForecast class with the given parameters.

        Args:
            exp_name (str): The name of the experiment. Defaults to "Tanesco".
            file_name (str): The name of the file to save the logs and model checkpoints. Defaults to None.
            seed (int): The seed for random number generation. Defaults to 42.
            root_dir (str): The root directory for saving logs and checkpoints. Defaults to "../".
            trial (optuna.trial.Trial): Optuna trial for hyperparameter optimization. Defaults to None.
            metric (str): The metric to monitor for early stopping and model checkpointing. Defaults to "val_mae".
            max_epochs (int): The maximum number of training epochs. Defaults to 10.
            wandb (bool): Flag to use Wandb logger instead of TensorBoard logger. Defaults to False.
            rich_progress_bar (bool): Flag to use rich progress bar for training visualization. Defaults to True.
        """
        super().__init__()

        self.exp_name = exp_name
        self.file_name = file_name
        self.seed = seed
        self.root_dir = root_dir
        self.trial = trial
        self.metric = metric
        self.max_epochs = max_epochs
        self.wandb = wandb
        self.model_type = model_type
        self.rich_progress_bar = rich_progress_bar
        self.model_type = "default_model"  # Update with actual model type
        self.model = None
        self.train_df=None
        self.datamodule = None
        self.gradient_clip_val = gradient_clip_val
        self._create_folder()


    def _create_folder(self):
        """
        Creates directories for storing the results, logs, and figures generated by the model.
        """
        self.results_path = Path(
            f"{self.root_dir}/results/{self.exp_name}/{self.model_type}/"
        )

        self.logs = Path(f"{self.root_dir}/logs/{self.exp_name}/{self.model_type}/")
        self.figures = Path(
            f"{self.root_dir}/figures/{self.exp_name}/{self.model_type}/"
        )
        self.figures.mkdir(parents=True, exist_ok=True)
        self.logs.mkdir(parents=True, exist_ok=True)
        self.results_path.mkdir(parents=True, exist_ok=True)
        if self.file_name is not None:
            self.checkpoints = Path(
                f"{self.root_dir}/checkpoints/{self.exp_name}/{self.model_type}/{self.file_name}"
            )
        else:
            self.checkpoints = Path(
                f"{self.root_dir}/checkpoints/{self.exp_name}/{self.model_type}"
            )
        self.checkpoints.mkdir(parents=True, exist_ok=True)


    def _set_up_trainer(self):
        """
        Set up the PyTorch Lightning trainer with the necessary configurations and callbacks.

        This method initializes the logger, sets up early stopping, model checkpointing,
        learning rate monitoring, and progress bar settings. It creates a trainer instance
        that will be used for model training.

        Raises:
            ValueError: If the logger type is unsupported.
        """

        # Initialize the callback list
        callback = []

        # Set random seed for reproducibility
        pl.seed_everything(self.seed, workers=True)

        # Initialize the logger
        if not self.wandb:
            self.logger = loggers.TensorBoardLogger(
                save_dir=self.logs,
                version=(self.file_name if self.file_name is not None else 0),
            )
        else:
            self.logger = loggers.WandbLogger(
                save_dir=self.logs,
                name=self.file_name,
                project=self.exp_name,
                log_model="all",
            )

        # Set up early stopping callback
        if self.trial is not None:
            early_stopping = PyTorchLightningPruningCallback(
                self.trial, monitor=self.metric
            )
            callback.append(early_stopping)
        else:
            early_stopping = EarlyStopping(
                monitor=self.metric,
                min_delta=0.0,
                patience=int(self.max_epochs * 0.25),
                verbose=False,
                mode="min",
                check_on_train_epoch_end=True,
            )
            callback.append(early_stopping)

        # Create checkpoints directory if it doesn't exist
        self.checkpoints.mkdir(parents=True, exist_ok=True)

        # Set up model checkpointing
        checkpoint_callback = ModelCheckpoint(
            dirpath=self.checkpoints,
            monitor=self.metric,
            mode="min",
            save_top_k=1,
            filename="{epoch:02d}",
        )
        callback.append(checkpoint_callback)

        lr_logger = LearningRateMonitor()
        callback.append(lr_logger)

        if self.rich_progress_bar:
            progress_bar = RichProgressBar(
                theme=RichProgressBarTheme(
                    description="green_yellow",
                    progress_bar="green1",
                    progress_bar_finished="green1",
                    progress_bar_pulse="#6206E0",
                    batch_progress="green_yellow",
                    time="grey82",
                    processing_speed="grey82",
                    metrics="grey82",
                )
            )
        else:
            progress_bar = TQDMProgressBar()

        callback.append(progress_bar)

        self.trainer = pl.Trainer(
            logger=self.logger,
            gradient_clip_val=self.gradient_clip_val,
            max_epochs=self.max_epochs,
            callbacks=callback,
            accelerator="auto",
            devices=1,
        )


    def fit(
        self,
        train_df,
        val_df=None,
        train_ratio=0.80,
        drop_last=True,
        num_worker=1,
        batch_size=64,
        pin_memory=True,
    ):
        """
        Fit the model using the provided training DataFrame.

        Args:
            train_df (pd.DataFrame): The training data.
            val_df (pd.DataFrame, optional): The validation data. If None, will split train_df based on train_ratio.
            train_ratio (float, optional): Proportion of data to use for training. Default is 0.80.
            drop_last (bool, optional): Whether to drop the last incomplete batch. Default is True.
            num_worker (int, optional): Number of workers for data loading. Default is 1.
            batch_size (int, optional): Size of each batch for training. Default is 64.
            pin_memory (bool, optional): \
                If True, the data loader will copy Tensors into CUDA pinned memory. Default is True.

        Returns:
            (float): The training wall time or cost metric based on the training configuration.

        Raises:
            ValueError: If the model instance is not initialized.
        """

        # Check if the model instance is initialized
        if self.model is None:
            raise ValueError("Model instance is empty")

        # Set the checkpoint path for the model
        self.model.checkpoint_path = self.checkpoints

        # Fit the data pipeline using the training DataFrame
        self.model.data_pipeline.fit(train_df.copy())

        # If no validation DataFrame is provided, split the training DataFrame

        if val_df is None and train_ratio > 0.0:
            self.train_df = train_df.copy()
            train_df, val_df = (
                train_df.iloc[: int(train_ratio * len(train_df))],
                train_df.iloc[int(train_ratio * len(train_df)) :],
            )
            self.metric = f"val_{self.model.hparams['metric']}"
            val_feature, val_target = self.model.data_pipeline.transform(val_df)

        else:
            val_feature, val_target = None, None
            self.metric = f"train_{self.model.hparams['metric']}"
            self.train_df = pd.concat([train_df.copy(), val_df.copy()], axis=0)
        # Transform the training data into features and targets
        train_feature, train_target = self.model.data_pipeline.transform(train_df)

        # Initialize the TimeseriesDataModule
        self.datamodule = TimeseriesDataModule(
            train_inputs=train_feature,
            train_targets=train_target,
            val_inputs=val_feature,
            val_targets=val_target,
            drop_last=drop_last,
            num_worker=num_worker,
            batch_size=batch_size,
            pin_memory=pin_memory,
        )

        # Sort the training DataFrame by the specified date column
        self.train_df = self.train_df.sort_values(
            by=self.model.data_pipeline.date_column
        )

        # Select the last 'max_data_drop' rows plus 'input_window_size' rows
        data_drop = self.model.data_pipeline.max_data_drop
        input_window = self.model.data_pipeline.input_window_size

        self.train_df = self.train_df.iloc[-(data_drop + input_window) :]

        # Set up the trainer
        self._set_up_trainer()
        start_time = default_timer()
        logger.info("""---------------Training started ---------------------------""")

        # Train the model and log training duration

        self.trainer.fit(
            self.model,
            self.datamodule.train_dataloader(),
            self.datamodule.val_dataloader(),
            ckpt_path=get_latest_checkpoint(self.checkpoints),
        )

        self.train_walltime = default_timer() - start_time
        logging.info(f"Training complete after {self.train_walltime / 60:.2f} minutes")

        if self.trial is not None:
            # Make predictions and compute the cost metric for hyper-param optimisation
            self.predict(val_df)
            cost = self.metrics.groupby("target")["MAE"].mean().iloc[-1].round(2)

            return cost  # Return the computed cost metric
        else:
            return self.train_walltime


    def load_and_prepare_data(self, test_df: pd.DataFrame, daily_feature: str):
        """
        Loads the checkpoint and prepares the ground truth data.
        
        Args:
            test_df (pd.DataFrame): The test DataFrame containing the input features for prediction.
            daily_feature (str): The daily feature to use in the model.

        Returns:
            ground_truth (pd.DataFrame): A DataFrame containing the ground truth data.
        """
        self.load_checkpoint()
        self.model.data_pipeline.daily_features = daily_feature

        # Prepare ground truth data
        ground_truth = test_df.iloc[self.model.data_pipeline.max_data_drop :].copy()
        ground_truth[self.model.data_pipeline.date_column] = pd.to_numeric(
            ground_truth[self.model.data_pipeline.date_column]
        )
        return ground_truth


    def perform_prediction(self, test_df: pd.DataFrame):
        """
        Performs the model prediction.
        
        Args:
            test_df (pd.DataFrame): The test DataFrame containing the input features for prediction.

        Returns:
            (dict): A dictionary containing the forecasted values.
        
        """
        features, _ = self.model.data_pipeline.transform(test_df.copy())
        features = torch.FloatTensor(features.copy())
        self.model.to(features.device)
        self.model.eval()
        return self.model.forecast(features)


    def create_results_df(
        self, time_stamp, ground_truth, predictions, target_series, date_column
    ):
        """
        Creates a DataFrame with the timestamp index and populates it with ground truth and forecasted values.
        
        Args:
            time_stamp (np.array): The timestamp index.
            ground_truth (np.array): The ground truth values.
            predictions (np.array): The forecasted values.
            target_series (list): A list of target series names.
            date_column (str): The name of the date column.

        Returns:
            results_df (pd.DataFrame): A DataFrame containing the ground truth and forecasted values, indexed by timestamp.
        
        """
        results_df = pd.DataFrame(index=pd.to_datetime(time_stamp.flatten(), unit="ns"))
        results_df.index.name = date_column

        for target_idx, target_name in enumerate(target_series):
            results_df[target_name] = ground_truth[:, :, target_idx].flatten()
            results_df[f"{target_name}_forecast"] = predictions[
                :, :, target_idx
            ].flatten()

        return results_df


    def evaluate_point_forecast(self, ground_truth, pred, time_stamp):
        """
        Evaluates the point forecast.

        Args:
            ground_truth (np.array): The ground truth values.
            pred (np.array): The forecasted values.
            time_stamp (np.array): The timestamp index.

        Returns:
            (dict): A dictionary containing the evaluation metrics.
        
        """
        return evaluate_point_forecast({
            "true": ground_truth,
            "loc": pred,
            "index": time_stamp,
            "targets": self.model.data_pipeline.target_series,
        })


    def predict(self, test_df=None, covariate_df=None, daily_feature=True):
        """
        Perform prediction on the test DataFrame and return a DataFrame with ground truth and forecasted values.

        Args:
            test_df (pd.DataFrame): The test DataFrame containing the input features for prediction.
            daily_feature (bool): Flag indicating whether daily features are used in the model. Default is True.

        Returns:
            results_df (pd.DataFrame): A DataFrame containing the ground truth and forecasted values, indexed by timestamp.
        """
        if (test_df is not None) and (self.train_df is not None):
            test_df = pd.concat([self.train_df, test_df], axis=0)
        else:
            assert ValueError("test_df can not be None")
        test_df = test_df.sort_values(by=self.model.data_pipeline.date_column)

        ground_truth = self.load_and_prepare_data(test_df, daily_feature)
        time_stamp = ground_truth[[self.model.data_pipeline.date_column]].values
        ground_truth = ground_truth[self.model.data_pipeline.target_series].values

        time_stamp = format_target(
            time_stamp,
            self.model.hparams["input_window_size"],
            self.model.hparams["forecast_horizon"],
            daily_feature=self.model.data_pipeline.daily_features,
        )
        ground_truth = format_target(
            ground_truth,
            self.model.hparams["input_window_size"],
            self.model.hparams["forecast_horizon"],
            daily_feature=self.model.data_pipeline.daily_features,
        )

        start_time = default_timer()
        pred = self.perform_prediction(test_df)
        self.test_walltime = default_timer() - start_time

        # Inverse transform predictions
        scaler = self.model.data_pipeline.data_pipeline.named_steps["scaling"]
        target_scaler = scaler.named_transformers_["target_scaler"]

        N, T, C = pred["pred"].size()
        pred["pred"] = target_scaler.inverse_transform(
            pred["pred"].numpy().reshape(N * T, C)
        )
        pred["pred"] = pred["pred"].reshape(N, T, C)

        # Evaluate point forecast
        self.metrics = self.evaluate_point_forecast(
            ground_truth, pred["pred"], time_stamp
        )
        self.metrics["test-time"] = self.test_walltime
        self.metrics["Model"] = self.model_type.upper()

        # Assert that the prediction and ground truth shapes are the same
        assert (
            pred["pred"].shape == ground_truth.shape
        ), "Shape mismatch: pred['pred'] and ground_truth must have the same shape."

        # Create results DataFrame
        results_df = self.create_results_df(
            time_stamp,
            ground_truth,
            pred["pred"],
            self.model.data_pipeline.target_series,
            self.model.data_pipeline.date_column,
        )
        results_df["Model"] = self.model_type.upper()

        return results_df
