"""
Exploring NePS Compatibility with PyTorch Lightning
=======================================================

1. Introduction:
----------------
Welcome to this tutorial on utilizing NePS-generated files and directories
in conjunction with PyTorch Lightning.

2. Setup:
---------
Ensure you have the necessary dependencies installed. You can install the 'NePS'
package by executing the following command:

```bash
pip install neural-pipeline-search
```

Additionally, note that 'NePS' does not include 'torchvision' as a dependency.
You can install it with this command:

```bash
pip install torchvision==0.14
```

Make sure to download the torchvision version that is compatible with your
pytorch version. More info on this link:

https://pypi.org/project/torchvision/

Additionally, you will need to install the PyTorch Lightning package. This can
be achieved with the following command:

```bash
pip install lightning
```

These dependencies ensure you have everything you need for this tutorial.
"""
import argparse
import glob
import logging
import random
import time
from pathlib import Path
from typing import Tuple

import lightning as L
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import TensorBoardLogger
from torch.utils.data import SubsetRandomSampler
from torch.utils.data.dataloader import DataLoader
from torchmetrics import Accuracy
from torchvision.datasets import MNIST
from torchvision.transforms import transforms

import neps
from neps.utils.common import get_initial_directory, load_lightning_checkpoint

#############################################################
# Definig the seeds for reproducibility


def set_seed(seed=123):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)


#############################################################
# Define the lightning model


class LitMNIST(L.LightningModule):
    def __init__(
        self,
        configuration: dict,
        n_train: int = 8192,
        n_valid: int = 1024,
    ):
        super().__init__()

        # Initialize the model's hyperparameters with the configuration
        self.save_hyperparameters(configuration)

        self.n_train = n_train
        self.n_valid = n_valid

        # Define data transformation and loss function
        self.transform = transforms.ToTensor()
        self.criterion = nn.NLLLoss()

        # Define the model's architecture
        self.linear1 = nn.Linear(in_features=784, out_features=392)
        self.linear2 = nn.Linear(in_features=392, out_features=196)
        self.linear3 = nn.Linear(in_features=196, out_features=10)

        # Define PyTorch Lightning metrics for training, validation, and testing
        metric = Accuracy(task="multiclass", num_classes=10)
        self.train_accuracy = metric.clone()
        self.val_accuracy = metric.clone()
        self.test_accuracy = metric.clone()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Forward pass function
        x = x.view(x.size(0), -1)
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = self.linear3(x)

        return F.log_softmax(x, dim=1)

    def common_step(
        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Perform a forward pass and compute loss, predictions, and get the ground
        truth labels for a batch of data.
        """
        x, y = batch
        logits = self.forward(x)
        loss = self.criterion(logits, y)
        preds = torch.argmax(logits, dim=1)

        return loss, preds, y

    def training_step(
        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
    ) -> float:
        loss, preds, y = self.common_step(batch, batch_idx)
        self.train_accuracy.update(preds, y)

        self.log_dict(
            {"train_loss": loss, "train_acc": self.val_accuracy.compute()},
            on_epoch=True,
            on_step=False,
            prog_bar=True,
        )

        return loss

    def validation_step(
        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
    ) -> None:
        loss, preds, y = self.common_step(batch, batch_idx)
        self.val_accuracy.update(preds, y)

        self.log_dict(
            {"val_loss": loss, "val_acc": self.val_accuracy.compute()},
            on_epoch=True,
            on_step=False,
            prog_bar=True,
        )

    def test_step(
        self, batch: Tuple[torch.Tensor, torch.Tensor], batch_idx: int
    ) -> None:
        _, preds, y = self.common_step(batch, batch_idx)
        self.test_accuracy.update(preds, y)

        self.log(name="test_acc", value=self.test_accuracy.compute())

    def configure_optimizers(self) -> torch.optim.Optimizer:
        # Configure and return the optimizer based on the configuration
        if self.hparams.optimizer == "Adam":
            optimizer = torch.optim.Adam(
                self.parameters(),
                lr=self.hparams.lr,
                weight_decay=self.hparams.weight_decay,
            )
        elif self.hparams.optimizer == "SGD":
            optimizer = torch.optim.SGD(
                self.parameters(),
                lr=self.hparams.lr,
                weight_decay=self.hparams.weight_decay,
            )
        else:
            raise ValueError(
                "The optimizer choices is not one of the available optimizers"
            )
        return optimizer

    def on_train_end(self):
        # Get the metric at the end of the training and log it with respect to
        # it's hyperparameters
        val_acc_metric = {
            "val_accuracy": self.val_accuracy.compute(),
        }

        # Log hyperparameters
        self.logger.log_hyperparams(self.hparams, metrics=val_acc_metric)

    def prepare_data(self) -> None:
        # Downloading the dataste if not already downloaded
        MNIST(self.hparams.data_dir, train=True, download=True)
        MNIST(self.hparams.data_dir, train=False, download=True)

    def setup(self, stage: str = None) -> None:
        # Assign train/val datasets for use in dataloaders
        if stage == "fit" or stage is None:
            self.mnist_full = MNIST(
                self.hparams.data_dir, train=True, transform=self.transform
            )

            # Create random subsets of the training dataset for validation.
            self.train_sampler = SubsetRandomSampler(range(self.n_train))
            self.val_sampler = SubsetRandomSampler(
                range(self.n_train, self.n_train + self.n_valid)
            )

        # Assign test dataset for use in dataloader
        if stage == "test" or stage is None:
            self.mnist_test = MNIST(
                self.hparams.data_dir, train=False, transform=self.transform
            )

    def train_dataloader(self) -> DataLoader:
        return DataLoader(
            self.mnist_full,
            batch_size=self.hparams.batch_size,
            sampler=self.train_sampler,
            num_workers=16,
        )

    def val_dataloader(self) -> DataLoader:
        return DataLoader(
            self.mnist_full,
            batch_size=self.hparams.batch_size,
            sampler=self.val_sampler,
            num_workers=16,
        )

    def test_dataloader(self) -> DataLoader:
        return DataLoader(
            self.mnist_test,
            batch_size=self.hparams.batch_size,
            num_workers=16,
        )


#############################################################
# Define search space


def search_space() -> dict:
    # Define a dictionary to represent the hyperparameter search space
    space = dict(
        data_dir=neps.ConstantParameter("./data"),
        batch_size=neps.ConstantParameter(64),
        lr=neps.FloatParameter(lower=1e-5, upper=1e-2, log=True, default=1e-3),
        weight_decay=neps.FloatParameter(
            lower=1e-5, upper=1e-3, log=True, default=5e-4
        ),
        optimizer=neps.CategoricalParameter(choices=["Adam", "SGD"], default="Adam"),
        epochs=neps.IntegerParameter(lower=1, upper=9, log=False, is_fidelity=True),
    )
    return space


#############################################################
# Define the run pipeline function


def run_pipeline(pipeline_directory, previous_pipeline_directory, **config) -> dict:
    # Initialize the first directory to store the event and checkpoints files
    init_dir = get_initial_directory(pipeline_directory)
    checkpoint_dir = init_dir / "checkpoints"

    # Initialize the model and checkpoint dir
    model = LitMNIST(config)

    # Create the TensorBoard logger for logging
    logger = TensorBoardLogger(
        save_dir=init_dir, name="data", version="logs", default_hp_metric=False
    )

    # Add checkpoints at the end of training
    checkpoint_callback = ModelCheckpoint(
        dirpath=checkpoint_dir,
        filename="{epoch}-{val_loss:.2f}",
    )

    # Use this function to load the previous checkpoint if it exists
    checkpoint_path, checkpoint = load_lightning_checkpoint(
        previous_pipeline_directory=previous_pipeline_directory,
        checkpoint_dir=checkpoint_dir,
    )

    if checkpoint is None:
        previously_spent_epochs = 0
    else:
        previously_spent_epochs = checkpoint["epoch"]

    # Create a PyTorch Lightning Trainer
    epochs = config["epochs"]

    trainer = L.Trainer(
        logger=logger,
        max_epochs=epochs,
        callbacks=[checkpoint_callback],
    )

    # Train the model and retrieve training/validation metrics
    if checkpoint_path:
        trainer.fit(model, ckpt_path=checkpoint_path)
    else:
        trainer.fit(model)

    train_accuracy = trainer.logged_metrics.get("train_acc", None)
    val_loss = trainer.logged_metrics.get("val_loss", None)
    val_accuracy = trainer.logged_metrics.get("val_acc", None)

    # Test the model and retrieve test metrics
    trainer.test(model)

    test_accuracy = trainer.logged_metrics.get("test_acc", None)

    return {
        "loss": val_loss,
        "cost": epochs - previously_spent_epochs,
        "info_dict": {
            "train_accuracy": train_accuracy,
            "val_accuracy": val_accuracy,
            "test_accuracy": test_accuracy,
        },
    }


if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--max_evaluations_total",
        type=int,
        default=15,
        help="Number of different configurations to train",
    )
    args = parser.parse_args()

    # Initialize the logger and record start time
    start_time = time.time()
    set_seed(112)
    logging.basicConfig(level=logging.INFO)

    # Run NePS with specified parameters
    neps.run(
        run_pipeline=run_pipeline,
        pipeline_space=search_space(),
        root_directory="results/hyperband",
        max_evaluations_total=args.max_evaluations_total,
        searcher="hyperband",
    )

    # Record the end time and calculate execution time
    end_time = time.time()
    execution_time = end_time - start_time

    # Log the execution time
    logging.info(f"Execution time: {execution_time} seconds")
