#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "Christian Heider Nielsen"
__all__ = []
__doc__ = r""" description """

import pickle
import shutil
import time
from pathlib import Path

import torch
from draugr import AverageMeter
from draugr.torch_utilities import (
    TensorBoardPytorchWriter,
)
from draugr.visualisation import progress_bar
from draugr.writers import MockWriter, Writer

# from tensorboard_logger import configure, log_value
from torch.nn import functional as F
from torch.optim.lr_scheduler import ReduceLROnPlateau
from warg import NOD, ensure_existence

from neodroidvision.classification.mechanims.attention.foveal.architecture.ram import (
    RecurrentAttention,
)
from neodroidvision.data.classification import MNISTDataset
from ram_params import get_ram_config

model_file_ending = ".model.tar"

__author__ = "Christian Heider Nielsen"
__doc__ = r""" description """


class Trainer:
    """A Recurrent Attention Model trainer.

    All hyperparameters are provided by the user in the
    config file."""

    def __init__(self, data_loader, **config):
        """
        Construct a new Trainer instance.

        Args:
            config: object containing command line arguments.
            data_loader: A data iterator."""
        config = NOD(**config)
        self.config = config

        if config.use_gpu and torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")

        # glimpse network params
        self.patch_size = config.patch_size
        self.glimpse_scale = config.glimpse_scale
        self.num_patches = config.num_patches
        self.loc_hidden = config.loc_hidden
        self.glimpse_hidden = config.glimpse_hidden

        # core network params
        self.num_glimpses = config.num_glimpses
        self.hidden_size = config.hidden_size

        # reinforce params
        self.std = config.std
        self.M = config.M

        # data params
        if config.is_train:
            self.train_loader = data_loader[0]
            self.valid_loader = data_loader[1]
            self.num_train = len(self.train_loader.sampler.indices)
            self.num_valid = len(self.valid_loader.sampler.indices)
        else:
            self.test_loader = data_loader
            self.num_test = len(self.test_loader.dataset)
        self.num_classes = 10
        self.num_channels = 1

        # training params
        self.epochs = config.epochs
        self.start_epoch = 0
        self.momentum = config.momentum
        self.lr = config.init_lr

        # misc params
        self.model_name = (
            f"ram_{config.num_glimpses}_{config.patch_size}x{config.patch_size}_"
            f"{config.glimpse_scale}"
        )
        self.best = config.best
        self.ckpt_dir = Path(config.ckpt_dir)
        self.logs_dir = Path(config.logs_dir)
        self.plot_dir = Path(config.plot_dir) / self.model_name
        self.best_valid_acc = 0.0
        self.counter = 0
        self.lr_patience = config.lr_patience
        self.train_patience = config.train_patience
        self.use_tensorboard = config.use_tensorboard
        self.resume = config.resume
        self.print_freq = config.print_freq
        self.plot_freq = config.plot_freq

        ensure_existence(self.ckpt_dir)
        ensure_existence(self.logs_dir)
        ensure_existence(self.plot_dir)

        # configure tensorboard logging
        """
if self.use_tensorboard:
tensorboard_dir = self.logs_dir / self.model_name
print(f"[*] Saving tensorboard logs to {tensorboard_dir}")
if not os.path.exists(tensorboard_dir):
os.makedirs(tensorboard_dir)
configure(tensorboard_dir)
"""

        self.model = RecurrentAttention(
            self.patch_size,
            self.num_patches,
            self.glimpse_scale,
            self.num_channels,
            self.loc_hidden,
            self.glimpse_hidden,
            self.std,
            self.hidden_size,
            self.num_classes,
        ).to(self.device)

        self.optimiser = torch.optim.Adam(
            self.model.parameters(), lr=self.config.init_lr
        )
        self.scheduler = ReduceLROnPlateau(
            self.optimiser, "min", patience=self.lr_patience
        )

    def reset(self):
        """

        Args:
          self:

        Returns:

        """
        h_t = torch.zeros(
            self.batch_size,
            self.hidden_size,
            dtype=torch.float,
            device=self.device,
            requires_grad=True,
        )
        l_t = torch.FloatTensor(self.batch_size, 2).uniform_(-1, 1).to(self.device)
        l_t.requires_grad = True

        return h_t, l_t

    def train(self, *, writer):
        """Train the model on the training set.

        A checkpoint of the model is saved after each epoch
        and if the validation accuracy is improved upon,
        a separate ckpt is created for use on the test set."""
        # load the most recent checkpoint
        if self.resume:
            self.load_checkpoint(best=False)

        print(
            f"\n[*] Train on {self.num_train} samples, validate on {self.num_valid} samples"
        )

        for epoch in range(self.start_epoch, self.epochs):

            print(
                f"\nEpoch: {epoch + 1}/{self.epochs} - LR: {self.optimiser.param_groups[0]['lr']:.6f}"
            )

            # train for 1 epoch
            train_loss, train_acc = self.train_one_epoch(epoch, writer=writer)

            # evaluate on validation set
            valid_loss, valid_acc = self.validate(epoch, writer=writer)

            # # reduce lr if validation loss plateaus
            self.scheduler.step(-valid_acc)

            is_best = valid_acc > self.best_valid_acc
            msg1 = "train loss: {:.3f} - train acc: {:.3f} "
            msg2 = "- val loss: {:.3f} - val acc: {:.3f} - val err: {:.3f}"
            if is_best:
                self.counter = 0
                msg2 += " [*]"
            print(
                (msg1 + msg2).format(
                    train_loss, train_acc, valid_loss, valid_acc, 100 - valid_acc
                )
            )

            # check for improvement
            if not is_best:
                self.counter += 1
            if self.counter > self.train_patience:
                print("[!] No improvement in a while, stopping training.")
                return
            self.best_valid_acc = max(valid_acc, self.best_valid_acc)
            self.save_checkpoint(
                {
                    "epoch": epoch + 1,
                    "model_state": self.model.state_dict(),
                    "optim_state": self.optimiser.state_dict(),
                    "best_valid_acc": self.best_valid_acc,
                },
                is_best,
            )

    def train_one_epoch(self, epoch, *, writer: Writer = MockWriter()):
        """
        Train the model for 1 epoch of the training set.

        An epoch corresponds to one full pass through the entire
        training set in successive mini-batches.

        This is used by train() and should not be called manually."""
        self.model.train()
        batch_time = AverageMeter()
        losses = AverageMeter()
        accs = AverageMeter()

        tic = time.time()
        with progress_bar(total=self.num_train) as pbar:
            for i, (x, y) in enumerate(self.train_loader):
                self.optimiser.zero_grad()

                x, y = x.to(self.device), y.to(self.device)

                plot = False
                if (epoch % self.plot_freq == 0) and (i == 0):
                    plot = True

                # initialize location vector and hidden state
                self.batch_size = x.shape[0]
                h_t, l_t = self.reset()

                # save images
                imgs = [x[0:9]]

                # extract the glimpses
                locs = []
                log_pi = []
                baselines = []
                for t in range(self.num_glimpses - 1):
                    # forward pass through model
                    h_t, l_t, b_t, p = self.model(x, l_t, h_t)

                    # store
                    locs.append(l_t[0:9])
                    baselines.append(b_t)
                    log_pi.append(p)

                # last iteration
                h_t, l_t, b_t, log_probas, p = self.model(x, l_t, h_t, last=True)
                log_pi.append(p)
                baselines.append(b_t)
                locs.append(l_t[0:9])

                # convert list to tensors and reshape
                baselines = torch.stack(baselines).transpose(1, 0)
                log_pi = torch.stack(log_pi).transpose(1, 0)

                # calculate reward
                predicted = torch.max(log_probas, 1)[1]
                R = (predicted.detach() == y).float()
                R = R.unsqueeze(1).repeat(1, self.num_glimpses)

                # compute losses for differentiable modules
                loss_action = F.nll_loss(log_probas, y)
                loss_baseline = F.mse_loss(baselines, R)

                # compute reinforce loss
                # summed over timesteps and averaged across batch
                adjusted_reward = R - baselines.detach()
                loss_reinforce = torch.sum(-log_pi * adjusted_reward, dim=1)
                loss_reinforce = torch.mean(loss_reinforce, dim=0)

                # sum up into a hybrid loss
                loss = loss_action + loss_baseline + loss_reinforce * 0.01

                # compute accuracy
                correct = (predicted == y).float()
                acc = 100 * (correct.sum() / len(y))

                # store
                losses.update(loss.item(), x.size()[0])
                accs.update(acc.item(), x.size()[0])

                # compute gradients and update SGD
                loss.backward()
                self.optimiser.step()

                # measure elapsed time
                toc = time.time()
                batch_time.update(toc - tic)

                pbar.set_description(
                    (
                        f"{(toc - tic):.1f}s - loss: {loss.item():.3f} - acc: {acc.item():.3f}"
                    )
                )
                pbar.update(self.batch_size)

                # dump the glimpses and locs
                if plot:
                    imgs = [g.cpu().data.numpy().squeeze() for g in imgs]
                    locs = [l.cpu().data.numpy() for l in locs]
                    pickle.dump(
                        imgs, open(str(self.plot_dir / f"g_{epoch + 1}.p"), "wb")
                    )
                    pickle.dump(
                        locs, open(str(self.plot_dir / f"l_{epoch + 1}.p"), "wb")
                    )

                # log to tensorboard
                if self.use_tensorboard:
                    iteration = epoch * len(self.train_loader) + i
                    writer.scalar("train_loss", losses.avg, iteration)
                    writer.scalar("train_acc", accs.avg, iteration)

            return losses.avg, accs.avg

    @torch.no_grad()
    def validate(self, epoch, *, writer: Writer = MockWriter()):
        """Evaluate the RAM model on the validation set."""
        losses = AverageMeter()
        accs = AverageMeter()

        for i, (x, y) in enumerate(self.valid_loader):
            x, y = x.to(self.device), y.to(self.device)

            # duplicate M times
            x = x.repeat(self.M, 1, 1, 1)

            # initialize location vector and hidden state
            self.batch_size = x.shape[0]
            h_t, l_t = self.reset()

            # extract the glimpses
            log_pi = []
            baselines = []
            for t in range(self.num_glimpses - 1):
                # forward pass through model
                h_t, l_t, b_t, p = self.model(x, l_t, h_t)

                # store
                baselines.append(b_t)
                log_pi.append(p)

            # last iteration
            h_t, l_t, b_t, log_probas, p = self.model(x, l_t, h_t, last=True)
            log_pi.append(p)
            baselines.append(b_t)

            # convert list to tensors and reshape
            baselines = torch.stack(baselines).transpose(1, 0)
            log_pi = torch.stack(log_pi).transpose(1, 0)

            # average
            log_probas = log_probas.view(self.M, -1, log_probas.shape[-1])
            log_probas = torch.mean(log_probas, dim=0)

            baselines = baselines.contiguous().view(self.M, -1, baselines.shape[-1])
            baselines = torch.mean(baselines, dim=0)

            log_pi = log_pi.contiguous().view(self.M, -1, log_pi.shape[-1])
            log_pi = torch.mean(log_pi, dim=0)

            # calculate reward
            predicted = torch.max(log_probas, 1)[1]
            R = (predicted.detach() == y).float()
            R = R.unsqueeze(1).repeat(1, self.num_glimpses)

            # compute losses for differentiable modules
            loss_action = F.nll_loss(log_probas, y)
            loss_baseline = F.mse_loss(baselines, R)

            # compute reinforce loss
            adjusted_reward = R - baselines.detach()
            loss_reinforce = torch.sum(-log_pi * adjusted_reward, dim=1)
            loss_reinforce = torch.mean(loss_reinforce, dim=0)

            # sum up into a hybrid loss
            loss = loss_action + loss_baseline + loss_reinforce * 0.01

            # compute accuracy
            correct = (predicted == y).float()
            acc = 100 * (correct.sum() / len(y))

            # store
            losses.update(loss.item(), x.size()[0])
            accs.update(acc.item(), x.size()[0])

            # log to tensorboard
            if self.use_tensorboard:
                iteration = epoch * len(self.valid_loader) + i
                writer.scalar("valid_loss", losses.avg, iteration)
                writer.scalar("valid_acc", accs.avg, iteration)

        return losses.avg, accs.avg

    @torch.no_grad()
    def test(self, *, writer):
        """Test the RAM model.

        This function should only be called at the very
        end once the model has finished training."""
        correct = 0

        # load the best checkpoint
        self.load_checkpoint(best=self.best)

        for i, (x, y) in enumerate(self.test_loader):
            x, y = x.to(self.device), y.to(self.device)

            # duplicate M times
            x = x.repeat(self.M, 1, 1, 1)

            # initialize location vector and hidden state
            self.batch_size = x.shape[0]
            h_t, l_t = self.reset()

            # extract the glimpses
            for t in range(self.num_glimpses - 1):
                # forward pass through model
                h_t, l_t, b_t, p = self.model(x, l_t, h_t)

            # last iteration
            h_t, l_t, b_t, log_probas, p = self.model(x, l_t, h_t, last=True)

            log_probas = log_probas.view(self.M, -1, log_probas.shape[-1])
            log_probas = torch.mean(log_probas, dim=0)

            pred = log_probas.data.max(1, keepdim=True)[1]
            correct += pred.eq(y.data.view_as(pred)).cpu().sum()

        perc = (100.0 * correct) / self.num_test
        error = 100 - perc
        print(f"[*] Test Acc: {correct}/{self.num_test} ({perc:.2f}% - {error:.2f}%)")

    def save_checkpoint(self, state, is_best):
        """Saves a checkpoint of the model.

        If this model has reached the best validation accuracy thus
        far, a seperate file with the suffix `best` is created."""
        ckpt_path = str(
            (self.ckpt_dir / f"{self.model_name}_checkpoint").with_suffix(
                model_file_ending
            )
        )
        torch.save(state, ckpt_path)
        if is_best:
            shutil.copyfile(
                ckpt_path,
                str(
                    (self.ckpt_dir / f"{self.model_name}_model_best").with_suffix(
                        model_file_ending
                    )
                ),
            )

    def load_checkpoint(self, best=False):
        """Load the best copy of a model.

        This is useful for 2 cases:
        - Resuming training with the most recent model checkpoint.
        - Loading the best validation model to evaluate on the test data.

        Args:
            best: if set to True, loads the best model.
                Use this if you want to evaluate your model
                on the test data. Else, set to False in which
                case the most recent version of the checkpoint
                is used."""
        print(f"[*] Loading model from {self.ckpt_dir}")

        filename = f"{self.model_name}_checkpoint"
        if best:
            filename = f"{self.model_name}_model_best"
        saved = (self.ckpt_dir / filename).with_suffix(model_file_ending)
        if saved.exists():
            ckpt = torch.load(str(saved))

            # load variables from checkpoint
            self.start_epoch = ckpt["epoch"]
            self.best_valid_acc = ckpt["best_valid_acc"]
            self.model.load_state_dict(ckpt["model_state"])
            self.optimiser.load_state_dict(ckpt["optim_state"])

            if best:
                print(
                    f"[*] Loaded {filename} checkpoint @ epoch {ckpt['epoch']} with "
                    f"best valid acc of {ckpt['best_valid_acc']:.3f}"
                )
            else:
                print(f"[*] Loaded {filename} checkpoint @ epoch {ckpt['epoch']}")


def main(**config):
    """

    Args:
      config:
    """

    config = NOD(**config)

    torch.manual_seed(config.random_seed)  # ensure reproducibility
    kwargs = {}
    if config.use_gpu:
        torch.cuda.manual_seed(config.random_seed)
        kwargs["num_workers"] = 1
        kwargs["pin_memory"] = True

    if config.is_train:  # instantiate data loaders
        data_loader = MNISTDataset.get_train_valid_loader(
            config.data_dir,
            batch_size=config.batch_size,
            random_seed=config.random_seed,
            valid_size=config.valid_size,
            shuffle=config.shuffle,
            **kwargs,
        )
    else:
        data_loader = MNISTDataset.get_test_loader(
            config.data_dir, config.batch_size, **kwargs
        )

    trainer = Trainer(data_loader, **config)

    with TensorBoardPytorchWriter() as writer:
        if config.is_train:
            trainer.train(writer=writer)
        else:  # or load a pretrained model and test
            trainer.test(writer=writer)


if __name__ == "__main__":
    config = get_ram_config()
    config.is_train = False
    # config.is_train = True

    main(**config)
