# Copyright 2020 The OATomobile Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Trains the deep imitative model on expert demostrations."""

import os
from typing import Mapping

import torch
import torch.distributions as D
import torch.optim as optim
import tqdm
from absl import app
from absl import flags
from absl import logging
import numpy as np

logging.set_verbosity(logging.DEBUG)
FLAGS = flags.FLAGS
flags.DEFINE_string(
    name="dataset_dir",
    default=None,
    help="The full path to the processed dataset.",
)
flags.DEFINE_string(
    name="output_dir",
    default=None,
    help="The full path to the output directory (for logs, ckpts).",
)
flags.DEFINE_integer(
    name="batch_size",
    default=512,
    help="The batch size used for training the neural network.",
)
flags.DEFINE_integer(
    name="num_epochs",
    default=None,
    help="The number of training epochs for the neural network.",
)
flags.DEFINE_integer(
    name="save_model_frequency",
    default=4,
    help="The number epochs between saves of the model.",
)
flags.DEFINE_float(
    name="learning_rate",
    default=1e-4,
    help="The ADAM learning rate.",
)
flags.DEFINE_integer(
    name="num_timesteps_to_keep",
    default=4,
    help="The numbers of time-steps to keep from the target, with downsampling.",
)
flags.DEFINE_float(
    name="weight_decay",
    default=0.0,
    help="The L2 penalty (regularization) coefficient.",
)
flags.DEFINE_bool(
    name="clip_gradients",
    default=False,
    help="If True it clips the gradients norm to 1.0.",
)
flags.DEFINE_integer(
    name="bootstrap_id",
    default=-1,
    help="Bootstrap id",
)
flags.DEFINE_integer(
    name="pretrain_epoch",
    default=-1,
    help="Pretrain Epoch",
)
flags.DEFINE_integer(
    name="gpu",
    default=0,
    help="GPU id",
)
flags.DEFINE_integer(
    name="num_workers",
    default=0,
    help="num_workers",
)
flags.DEFINE_integer(
    name="val_num_workers",
    default=0,
    help="val_num_workers",
)
flags.DEFINE_integer(
    name="validation_interval",
    default=1,
    help="validation_interval",
)
flags.DEFINE_integer(
    name="latent_dim",
    default=128,
    help="Dimension of latent space.",
)
flags.DEFINE_integer(
    name="in_channels",
    default=3,
    help="Channel Dimensions.",
)
flags.DEFINE_integer(
    name="beta",
    default=4,
    help="Beta Value.",
)


def main(argv):
    # Debugging purposes.
    logging.debug(argv)
    logging.debug(FLAGS)

    # Parses command line arguments.
    bootstrap_id = FLAGS.bootstrap_id
    dataset_dir = FLAGS.dataset_dir
    output_dir = FLAGS.output_dir + '_bootstrap{}'.format(bootstrap_id) if bootstrap_id != -1 else FLAGS.output_dir
    batch_size = FLAGS.batch_size
    num_epochs = FLAGS.num_epochs
    learning_rate = FLAGS.learning_rate
    save_model_frequency = FLAGS.save_model_frequency
    num_timesteps_to_keep = FLAGS.num_timesteps_to_keep
    weight_decay = FLAGS.weight_decay
    clip_gradients = FLAGS.clip_gradients
    pretrain_epoch = FLAGS.pretrain_epoch
    num_workers = FLAGS.num_workers
    val_num_workers = FLAGS.val_num_workers
    validation_interval = FLAGS.validation_interval
    latent_dim = FLAGS.latent_dim
    in_channels = FLAGS.in_channels
    beta = FLAGS.beta

    print('GPU: ' + str(FLAGS.gpu))
    os.environ["CUDA_VISIBLE_DEVICES"] = str(FLAGS.gpu)
    from oatomobile.baselines.torch.cim.perception.model import BetaVAE
    from oatomobile.baselines.torch.logging import Checkpointer
    from oatomobile.baselines.torch.logging import TensorBoardWriter
    from oatomobile.baselines.torch.typing import ArrayLike
    from oatomobile.datasets.carla import CARLADataset

    noise_level = 1e-2

    # Determines device, accelerator.
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # pylint: disable=no-member

    # Creates the necessary output directory.
    os.makedirs(output_dir, exist_ok=True)
    log_dir = os.path.join(output_dir, "logs")
    os.makedirs(log_dir, exist_ok=True)
    ckpt_dir = os.path.join(output_dir, "ckpts")
    os.makedirs(ckpt_dir, exist_ok=True)

    # Initializes the model and its optimizer.
    output_shape = [num_timesteps_to_keep, 2]
    model = BetaVAE(latent_dim, in_channels, beta).to(device)
    optimizer = optim.Adam(
        model.parameters(),
        lr=learning_rate,
        weight_decay=weight_decay,
    )
    writer = TensorBoardWriter(log_dir=log_dir)
    checkpointer = Checkpointer(model=model, ckpt_dir=ckpt_dir)

    if pretrain_epoch != -1:
        print('Loading the parameters of epoch {}.'.format(pretrain_epoch))
        checkpointer.load(pretrain_epoch)

    def transform(batch: Mapping[str, ArrayLike]) -> Mapping[str, torch.Tensor]:
        """Preprocesses a batch for the model.

        Args:
          batch: (keyword arguments) The raw batch variables.

        Returns:
          The processed batch.
        """
        # Sends tensors to `device`.
        batch = {key: tensor.to(device) for (key, tensor) in batch.items()}
        # Preprocesses batch for the model.
        batch = model.transform(batch)
        return batch

    # Setups the dataset and the dataloader.
    modalities = (
        "bird_view_camera_cityscapes",
    )

    bootstrap_ids = None
    if bootstrap_id != -1:
        bootstrap_ids = np.load(os.path.join(dataset_dir, 'bootstrap{}.npy'.format(bootstrap_id)), allow_pickle=True)

    dataset_train = CARLADataset.as_torch(
        dataset_dir=os.path.join(dataset_dir, "train"),
        modalities=modalities,
        bootstrap_ids=bootstrap_ids
    )
    dataloader_train = torch.utils.data.DataLoader(
        dataset_train,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
    )

    if validation_interval > 0:
        dataset_val = CARLADataset.as_torch(
            dataset_dir=os.path.join(dataset_dir, "val"),
            modalities=modalities,
        )
        dataloader_val = torch.utils.data.DataLoader(
            dataset_val,
            batch_size=batch_size * 5,
            shuffle=True,
            num_workers=val_num_workers,
        )

    # Theoretical limit of NLL.
    nll_limit = -torch.sum(  # pylint: disable=no-member
        D.MultivariateNormal(
            loc=torch.zeros(output_shape[-2] * output_shape[-1]),  # pylint: disable=no-member
            scale_tril=torch.eye(output_shape[-2] * output_shape[-1]) *  # pylint: disable=no-member
                       noise_level,  # pylint: disable=no-member
        ).log_prob(torch.zeros(output_shape[-2] * output_shape[-1])))  # pylint: disable=no-member

    def train_step(
            model: BetaVAE,
            optimizer: optim.Optimizer,
            batch: Mapping[str, torch.Tensor],
            clip: bool = False,
    ) -> torch.Tensor:
        """Performs a single gradient-descent optimisation step."""
        # Resets optimizer's gradients.
        optimizer.zero_grad()
        recons, input, mu, log_var = model(batch['visual_features'])
        losses = model.loss(recons, input, mu, log_var)
        loss = losses['loss']
        loss.backward()

        if clip:
            torch.nn.utils.clip_grad_norm(model.parameters(), 1.0)

        # Performs a gradient descent step.
        optimizer.step()

        return losses

    def train_epoch(
            model: BetaVAE,
            optimizer: optim.Optimizer,
            dataloader: torch.utils.data.DataLoader,
    ) -> torch.Tensor:
        """Performs an epoch of gradient descent optimization on `dataloader`."""
        model.train()
        losses = {'loss': 0.0, 'Reconstruction_Loss': 0.0, 'KLD': 0.0}
        with tqdm.tqdm(dataloader) as pbar:
            for batch in pbar:
                # Prepares the batch.
                batch = transform(batch)
                # Performs a gradient-descent step.
                step_losses = train_step(model, optimizer, batch, clip=clip_gradients)
                losses = {key: losses.get(key, 0) + step_losses.get(key, 0) for key in set(losses) | set(step_losses)}

        losses.update((key, val / len(dataloader)) for key, val in losses.items())
        return losses

    def evaluate_step(
            model: BetaVAE,
            batch: Mapping[str, torch.Tensor],
    ) -> torch.Tensor:
        """Evaluates `model` on a `batch`."""
        # Forward pass from the model.
        recons, input, mu, log_var = model(batch['visual_features'])
        losses = model.loss(recons, input, mu, log_var)
        loss = losses['loss']

        return loss

    def evaluate_epoch(
            model: BetaVAE,
            dataloader: torch.utils.data.DataLoader,
    ) -> torch.Tensor:
        """Performs an evaluation of the `model` on the `dataloader."""
        model.eval()
        loss = 0.0
        with tqdm.tqdm(dataloader) as pbar:
            for batch in pbar:
                # Prepares the batch.
                batch = transform(batch)
                # Accumulates loss in dataset.
                with torch.no_grad():
                    loss += evaluate_step(model, batch)
        return loss / len(dataloader)

    def write(
            writer: TensorBoardWriter,
            split: str,
            loss: torch.Tensor,
            epoch: int,
    ) -> None:
        """Visualises model performance on `TensorBoard`."""
        writer.log(
            split=split,
            loss=loss.detach().cpu().numpy().item(),
            global_step=epoch,
            overhead_features=None,
            predictions=None,
            ground_truth=None
        )

    with tqdm.tqdm(range(pretrain_epoch + 1, num_epochs)) as pbar_epoch:
        for epoch in pbar_epoch:
            # Trains model on whole training dataset, and writes on `TensorBoard`.
            losses_train = train_epoch(model, optimizer, dataloader_train)
            write(writer, "train", losses_train['loss'], epoch)

            # Evaluates model on whole validation dataset, and writes on `TensorBoard`.
            loss_val = None
            if validation_interval > 0 and epoch % validation_interval == 0:
                loss_val = evaluate_epoch(model, dataloader_val)
                write(writer, "val", loss_val, epoch)

            # Checkpoints model weights.
            if epoch % save_model_frequency == 0:
                checkpointer.save(epoch)

            # Updates progress bar description.
            pbar_epoch.set_description(
                "Epoch {} -> TL: {:.2f} | Rec: {:.2f} | KL: {:.2f} | VL: {:.2f}".format(
                    epoch,
                    losses_train['loss'].detach().cpu().numpy().item(),
                    losses_train['Reconstruction_Loss'].detach().cpu().numpy().item(),
                    losses_train['KLD'].detach().cpu().numpy().item(),
                    loss_val.detach().cpu().numpy().item() if loss_val is not None else 0,
                ))


if __name__ == "__main__":
    flags.mark_flag_as_required("dataset_dir")
    flags.mark_flag_as_required("output_dir")
    flags.mark_flag_as_required("num_epochs")
    app.run(main)
