import os

import os.path as osp

from datetime import datetime

import numpy as np

import torch

from tqdm import tqdm

from utils.network_utils import load_checkpoint

class BaseRunner:

    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.res_dir = osp.join(self.config["exp_dir"], self.exp_name, "results")

        os.makedirs(self.res_dir, exist_ok=True)

        self.model = self._init_net()

        self.inference_loader = self._init_dataloader()

        pretrained_path = self.config.get("model_path", False)

        if pretrained_path:

            load_checkpoint(pretrained_path, self.model)

        else:

            raise Exception(

                "model_path doesnt't exist in config. Please specify checkpoint path",

            )

    def _init_net(self):

        raise NotImplemented

    def _init_dataloader(self):

        raise NotImplemented

    def inference(self):

        self.model.eval()

        gt_class = np.empty(0)

        pd_class = np.empty(0)

        with torch.no_grad():

            for i, batch in tqdm(enumerate(self.inference_loader)):

                inputs = batch["image"].to(self.config["device"])

                predictions = self.model(inputs)

                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))

        class_accuracy = sum(pd_class == gt_class) / pd_class.shape[0]

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

        pd_class = pd_class.astype(int)

        np.savetxt(osp.join(self.res_dir, "predictions.txt"), pd_class)