import base64
from argparse import ArgumentParser

import mlflow.pytorch
import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.metrics.functional import accuracy
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms

from config import CONDA_ENV


class MNISTDataModule(pl.LightningDataModule):
    def __init__(self, **kwargs):
        """
        Initialization of inherited lightning data module
        """
        super(MNISTDataModule, self).__init__()
        self.df_train = None
        self.df_val = None
        self.df_test = None
        self.train_data_loader = None
        self.val_data_loader = None
        self.test_data_loader = None
        self.args = kwargs

        # transforms for images
        self.transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        )

    def prepare_data(self):
        """
        Implementation of abstract class
        """
        pass

    def setup(self, stage=None):
        """
        Downloads the data, parse it and split the data into train, test, validation data
        :param stage: Stage - training or testing
        """

        self.df_train = datasets.MNIST(
            "dataset", download=True, train=True, transform=self.transform
        )
        self.df_train, self.df_val = random_split(self.df_train, [55000, 5000])
        self.df_test = datasets.MNIST(
            "dataset", download=True, train=False, transform=self.transform
        )

    def create_data_loader(self, df):
        """
        Generic data loader function
        :param df: Input tensor
        :return: Returns the constructed dataloader
        """
        return DataLoader(
            df, batch_size=self.args["batch_size"], num_workers=self.args["num_workers"],
        )

    def train_dataloader(self):
        """
        :return: output - Train data loader for the given input
        """
        return self.create_data_loader(self.df_train)

    def val_dataloader(self):
        """
        :return: output - Validation data loader for the given input
        """
        return self.create_data_loader(self.df_val)

    def test_dataloader(self):
        """
        :return: output - Test data loader for the given input
        """
        return self.create_data_loader(self.df_test)


class LightningMNISTClassifier(pl.LightningModule):
    def __init__(self, **kwargs):
        """
        Initializes the network
        """
        super(LightningMNISTClassifier, self).__init__()

        # mnist images are (1, 28, 28) (channels, width, height)
        self.optimizer = None
        self.scheduler = None
        self.layer_1 = torch.nn.Linear(28 * 28, 128)
        self.layer_2 = torch.nn.Linear(128, 256)
        self.layer_3 = torch.nn.Linear(256, 10)
        self.args = kwargs

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument(
            "--batch_size",
            type=int,
            default=64,
            metavar="N",
            help="input batch size for training (default: 64)",
        )
        parser.add_argument(
            "--num_workers",
            type=int,
            default=3,
            metavar="N",
            help="number of workers (default: 3)",
        )
        parser.add_argument(
            "--lr", type=float, default=0.001, metavar="LR", help="learning rate (default: 0.001)",
        )
        return parser

    def forward(self, x):
        """
        :param x: Input data
        :return: output - mnist digit label for the input image
        """
        batch_size = x.size()[0]

        # (b, 1, 28, 28) -> (b, 1*28*28)
        x = x.view(batch_size, -1)

        # layer 1 (b, 1*28*28) -> (b, 128)
        x = self.layer_1(x)
        x = torch.relu(x)

        # layer 2 (b, 128) -> (b, 256)
        x = self.layer_2(x)
        x = torch.relu(x)

        # layer 3 (b, 256) -> (b, 10)
        x = self.layer_3(x)

        # probability distribution over labels
        x = torch.log_softmax(x, dim=1)

        return x

    def cross_entropy_loss(self, logits, labels):
        """
        Initializes the loss function
        :return: output - Initialized cross entropy loss function
        """
        return F.nll_loss(logits, labels)

    def training_step(self, train_batch, batch_idx):
        """
        Training the data as batches and returns training loss on each batch
        :param train_batch: Batch data
        :param batch_idx: Batch indices
        :return: output - Training loss
        """
        x, y = train_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        return {"loss": loss}

    def validation_step(self, val_batch, batch_idx):
        """
        Performs validation of data in batches
        :param val_batch: Batch data
        :param batch_idx: Batch indices
        :return: output - valid step loss
        """
        x, y = val_batch
        logits = self.forward(x)
        loss = self.cross_entropy_loss(logits, y)
        return {"val_loss": loss}

    def validation_epoch_end(self, outputs):
        """
        Computes average validation accuracy
        :param outputs: outputs after every epoch end
        :return: output - average valid loss
        """
        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        self.log("val_loss", avg_loss, sync_dist=True)

    def test_step(self, test_batch, batch_idx):
        """
        Performs test and computes the accuracy of the model
        :param test_batch: Batch data
        :param batch_idx: Batch indices
        :return: output - Testing accuracy
        """
        x, y = test_batch
        output = self.forward(x)
        _, y_hat = torch.max(output, dim=1)
        test_acc = accuracy(y_hat.cpu(), y.cpu())
        return {"test_acc": test_acc}

    def test_epoch_end(self, outputs):
        """
        Computes average test accuracy score
        :param outputs: outputs after every epoch end
        :return: output - average test loss
        """
        avg_test_acc = torch.stack([x["test_acc"] for x in outputs]).mean()
        self.log("avg_test_acc", avg_test_acc)

    def prepare_data(self):
        """
        Prepares the data for training and prediction
        """
        return {}

    def configure_optimizers(self):
        """
        Initializes the optimizer and learning rate scheduler
        :return: output - Initialized optimizer and scheduler
        """
        self.optimizer = torch.optim.Adam(self.parameters(), lr=self.args["lr"])
        self.scheduler = {
            "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer, mode="min", factor=0.2, patience=2, min_lr=1e-6, verbose=True,
            ),
            "monitor": "val_loss",
        }
        return [self.optimizer], [self.scheduler]


class MNISTClassifierWrapper(mlflow.pyfunc.PythonModel):
    """
    This wrapper is needed for deploying model on mlflow server
    The server API can accept data only in the pandas Dataframe tabular format
    where each row represents a single data point and each column stands for a feature.
    Thus we need to encode our images so that they could be stored in a single Dataframe row.

    This object should only be used if you're up to deploy the trained model on the mlflow server.
    """

    def __init__(self, map_location=None, input_dim=(28, 28, 1)):
        super(MNISTClassifierWrapper, self).__init__()
        self.map_location = map_location
        self.input_dim = input_dim
        self.transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        )

    def load_context(self, context):
        self._detector = LightningMNISTClassifier.load_from_checkpoint(context.artifacts['model_state'],
                                                                       map_location=self.map_location)
        self._detector.eval()

    def predict(self, context, model_input):
        def decode_img(x):
            r = base64.decodebytes(bytes(x, encoding='utf-8'))
            q = np.frombuffer(r, dtype='uint8').reshape(self.input_dim)
            q = np.array(q)
            return q

        images = [self.transform(decode_img(x)).to(self.map_location) for x in model_input['image'].values]

        input = torch.cat(images, dim=0).float()
        with torch.no_grad():
            output = torch.argmax(self._detector(input), 1)
            return output.cpu().numpy().reshape((-1))

    @staticmethod
    def export_best_model(trainer, map_location='cpu'):
        mlflow.pyfunc.log_model('model',
                                python_model=MNISTClassifierWrapper(map_location),
                                conda_env=CONDA_ENV,
                                artifacts={'model_state': trainer.checkpoint_callback.best_model_path},
                                code_path=['config.py',
                                           'mnist_classifier.py'])


if __name__ == "__main__":
    parser = ArgumentParser(description="PyTorch autolog Mnist Example")

    parser = pl.Trainer.add_argparse_args(parent_parser=parser)
    parser = LightningMNISTClassifier.add_model_specific_args(parent_parser=parser)

    # Turn on autologging
    mlflow.pytorch.autolog(log_every_n_epoch=2)

    args = parser.parse_args()
    dict_args = vars(args)

    if "accelerator" in dict_args:
        if dict_args["accelerator"] == "None":
            dict_args["accelerator"] = None

    dm = MNISTDataModule(**dict_args)
    dm.prepare_data()
    dm.setup(stage="fit")

    # All manual logging and model saving should be done inside the current mlflow experiment context.
    # Otherwise the autologged metrics and checkpoints will get lost.
    with mlflow.start_run():
        model = LightningMNISTClassifier(**dict_args)
        trainer = pl.Trainer.from_argparse_args(args)
        trainer.fit(model, dm)
        trainer.test()

        # By default the mlflow autologger stores checkpoints in the native pytorch format.
        # We need to manually save the model inside our wrapper for deploying to mlflow server.
        MNISTClassifierWrapper.export_best_model(trainer, 'cpu')
