# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['History', 'Module']

# Cell
#export
import torch
import torch.nn as nn
from tqdm.auto import tqdm

# Internal Cell
#exporti
def get_pbar_description_from_batch_metrics(batch_metrics, prefix=""):
    description = ""
    for name, value in batch_metrics.items():
        description += f'{prefix}{name}: {value:.2f} '
    return description[:-1]

# Cell
class History(dict):
    """
    This object inherits from base `dict` to provide a `History` object similar to Keras',
    allowing the automatic logging of the loss and the different metrics during training.
    It's automatically used in `.fit()` (as it is in Keras).
    """
    def log_dict(self, data, prefix=""):
        """
        Parameters
        ----------
        data: dict
        """
        for name, value in data.items():
            name = prefix+name
            if name in self.keys():
                self[name].append(value)
            else:
                self[name] = [value]

    def aggregate(self, agg_fn=lambda x: sum(x)/len(x)):
        """
        Aggregates the stored values using the designed aggregation function.

        Parameters
        ----------
        agg_fn: function
            Mean by default.

        Returns
        -------
        agg_data: dict
            Aggregated data.
        """
        return {name:agg_fn(values) for name, values in self.items()}

# Cell
class Module(nn.Module):
    """
    Modification of PyTorch base `nn.Module` to provide a basic
    predefined training loop with logging and a Keras-like interface
    to be able to customize the training.
    This Module implements as well as a `compile` method and an `evaluate` one.
    All is done to obtain a behaviour as similar to Keras as possible.
    """
    def __init__(self, **kwargs):
        super(Module, self).__init__(**kwargs)

    def train_step(self, batch):
        inputs, labels = batch
        self.optimizer.zero_grad()
        outputs = self(inputs)
        loss = self.loss_fn(outputs, labels)
        loss.backward()
        self.optimizer.step()

        ## Obtain metrics if needed
        if self.metrics is not None:
            metrics = self.metrics(outputs, labels)
            metrics = {name:value.item() for name, value in metrics.items()}
            metrics['Loss'] = loss.item()
        else:
            metrics = {'Loss':loss.item()}
        return metrics

    def validation_step(self, batch):
        inputs, labels = batch
        outputs = self(inputs)
        loss = self.loss_fn(outputs, labels)

        ## Obtain metrics if needed
        if self.metrics is not None:
            metrics = self.metrics(outputs, labels)
            metrics = {name:value.item() for name, value in metrics.items()}
            metrics['Loss'] = loss.item()
        else:
            metrics = {'Loss':loss.item()}
        return metrics

    def fit(self, trainloader, epochs, validationloader=None):
        history_epoch = History()
        for epoch in tqdm(range(epochs), desc='Epochs', position=0):
            self.train()
            pbar = tqdm(enumerate(trainloader), total=len(trainloader), position=1, leave=False)
            history_batch = History()
            for batch_idx, batch in pbar:
                batch_metrics = self.train_step(batch)
                history_batch.log_dict(batch_metrics)
                pbar.set_description(get_pbar_description_from_batch_metrics(batch_metrics))
            if validationloader is not None:
                self.eval()
                pbar = tqdm(enumerate(validationloader), total=len(validationloader), position=2, leave=False)
                for batch_idx, batch in pbar:
                    with torch.no_grad():
                        batch_metrics = self.validation_step(batch)
                    history_batch.log_dict(batch_metrics, prefix='Val_')
                    pbar.set_description(get_pbar_description_from_batch_metrics(batch_metrics, 'Val_'))
            self.metrics.reset()
            history_epoch.log_dict(history_batch.aggregate())
        return history_epoch

    def compile(self, loss=None, optimizer=None, metrics=None):
        """
        metrics: torchmetrics.MetricCollection
        """
        self.loss_fn = loss
        self.optimizer = optimizer
        self.metrics = metrics

    def evaluate(self, dataloader):
        self.eval()
        for i, batch in tqdm(enumerate(dataloader), total=len(dataloader)):
            if i == 0:
                results = self.validation_step(batch)
                results = {name:[value] for name, value in results.items()}
            else:
                result = self.validation_step(batch)
                for name, value in result.items():
                    results[name].append(value)
        results = {name:sum(values)/len(values) for name,values in results.items()}
        return results