import os
import os.path as osp
from datetime import datetime

import numpy as np
import torch
from torch import nn, optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from utils.network_utils import load_checkpoint, save_checkpoint


class BaseTrainer:
    def __init__(self, config):
        self.config = config
        self.exp_name = self.config.get("exp_name", None)
        if self.exp_name is None:
            self.exp_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

        self.log_dir = osp.join(self.config["exp_dir"], self.exp_name, "logs")
        self.pth_dir = osp.join(self.config["exp_dir"], self.exp_name, "checkpoints")
        os.makedirs(self.log_dir, exist_ok=True)
        os.makedirs(self.pth_dir, exist_ok=True)

        self.writer = SummaryWriter(log_dir=self.log_dir)

        self.model = self._init_net()
        self.optimizer = self._init_optimizer()
        self.criterion = nn.CrossEntropyLoss().to(self.config["device"])

        self.train_loader, self.val_loader = self._init_dataloaders()

        pretrained_path = self.config.get("model_path", False)
        if pretrained_path:
            self.training_epoch, self.total_iter = load_checkpoint(
                pretrained_path, self.model, optimizer=self.optimizer,
            )

        else:
            self.training_epoch = 0
            self.total_iter = 0

        self.epochs = self.config.get("epochs", int(1e5))

    def _init_net(self):
        raise NotImplemented

    def _init_dataloaders(self):
        raise NotImplemented

    def _init_optimizer(self):
        optimizer = getattr(optim, self.config["optim"])(
            self.model.parameters(), **self.config["optim_params"]
        )
        return optimizer

    def train_epoch(self):
        self.model.train()
        total_loss = 0

        gt_class = np.empty(0)
        pd_class = np.empty(0)

        for i, batch in enumerate(self.train_loader):
            inputs = batch["image"].to(self.config["device"])
            targets = batch["class"].to(self.config["device"])

            predictions = self.model(inputs)
            loss = self.criterion(predictions, targets)

            classes = predictions.topk(k=1)[1].view(-1).cpu().numpy()

            gt_class = np.concatenate((gt_class, batch["class"].numpy()))
            pd_class = np.concatenate((pd_class, classes))

            total_loss += loss.item()

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            if (i + 1) % 10 == 0:
                print(
                    "\tIter [%d/%d] Loss: %.4f"
                    % (i + 1, len(self.train_loader), loss.item()),
                )

            self.writer.add_scalar(
                "Train loss (iterations)", loss.item(), self.total_iter,
            )
            self.total_iter += 1

        total_loss /= len(self.train_loader)
        class_accuracy = sum(pd_class == gt_class) / pd_class.shape[0]

        print("Train loss - {:4f}".format(total_loss))
        print("Train CLASS accuracy - {:4f}".format(class_accuracy))

        self.writer.add_scalar("Train loss (epochs)", total_loss, self.training_epoch)
        self.writer.add_scalar(
            "Train CLASS accuracy", class_accuracy, self.training_epoch,
        )

    def val(self):
        self.model.eval()
        total_loss = 0

        gt_class = np.empty(0)
        pd_class = np.empty(0)

        with torch.no_grad():
            for i, batch in tqdm(enumerate(self.val_loader)):
                inputs = batch["image"].to(self.config["device"])
                targets = batch["class"].to(self.config["device"])

                predictions = self.model(inputs)
                loss = self.criterion(predictions, targets)

                classes = predictions.topk(k=1)[1].view(-1).cpu().numpy()

                gt_class = np.concatenate((gt_class, batch["class"].numpy()))
                pd_class = np.concatenate((pd_class, classes))

                total_loss += loss.item()

        total_loss /= len(self.val_loader)
        class_accuracy = sum(pd_class == gt_class) / pd_class.shape[0]

        print("Validation loss - {:4f}".format(total_loss))
        print("Validation CLASS accuracy - {:4f}".format(class_accuracy))

        self.writer.add_scalar("Validation loss", total_loss, self.training_epoch)
        self.writer.add_scalar(
            "Validation CLASS accuracy", class_accuracy, self.training_epoch,
        )

    def loop(self):
        for epoch in range(self.training_epoch, self.epochs):
            print("Epoch - {}".format(self.training_epoch + 1))
            self.train_epoch()
            save_checkpoint(
                {
                    "state_dict": self.model.state_dict(),
                    "optimizer": self.optimizer.state_dict(),
                    "epoch": epoch,
                    "total_iter": self.total_iter,
                },
                osp.join(self.pth_dir, "{:0>8}.pth".format(epoch)),
            )
            self.val()

            self.training_epoch += 1