# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/Callbacks.ipynb.

# %% auto 0
__all__ = ['Callback', 'MetricsCB', 'DeviceCB', 'get_mem', 'MemoryCB', 'SaveModelCB']

# %% ../nbs/Callbacks.ipynb 2
import torch
from .utility import to_cpu, to_device
from copy import copy
from torcheval.metrics import Mean
from torch.utils.tensorboard import SummaryWriter

# %% ../nbs/Callbacks.ipynb 3
class Callback(): order = 0

# %% ../nbs/Callbacks.ipynb 4
class MetricsCB(Callback):
    def __init__(self, device=None, use_tensorboard=False, *ms, **metrics):
        for o in ms: metrics[type(o).__name__] = o
        self.metrics = metrics
        self.all_metrics = copy(metrics)
        self.all_metrics['loss'] = self.loss = Mean(device=device)
        self.use_tensorboard = use_tensorboard

    def _log(self, d): print(d)
    def before_fit(self, learn): 
        learn.metrics = self
        if self.use_tensorboard:
            self.writer = SummaryWriter()
            
    def before_epoch(self, learn): [o.reset() for o in self.all_metrics.values()]

    def after_epoch(self, learn):
        log = {k:f'{v.compute():.3f}' for k,v in self.all_metrics.items()}
        
        if learn.model.training:
            epoch_type = 'train'  
        else: 
            epoch_type = 'valid'
            learn.current_metric_values = log
            offset = 0.5
        log['epoch'] = learn.epoch
        log['train'] = epoch_type
        self._log(log)
        if self.use_tensorboard:
            for k, v in self.all_metrics.items():
                self.writer.add_scalars('data/'+ k, {epoch_type:v.compute()}, learn.epoch)
            
    def after_batch(self, learn):
        #x,y,*_ = to_cpu(learn.batch)
        x,y,*_ = learn.batch
        #for m in self.metrics.values(): m.update(to_cpu(learn.preds), y)
        #loss = to_cpu(learn.loss)
        for m in self.metrics.values(): m.update(learn.preds, y)
        
        loss = learn.loss
        self.loss.update(loss, weight=len(x))
        
    def after_fit(self, learn):
        if self.use_tensorboard:
            self.writer.flush()
            self.writer.close()

# %% ../nbs/Callbacks.ipynb 5
class DeviceCB(Callback):
    def __init__(self, device=None): 
        if not device:
            device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.device = device
    def before_fit(self, learn):
        if hasattr(learn.model, 'to'): 
            learn.model.to(self.device)
    def before_batch(self, learn): 
        learn.batch = to_device(learn.batch, device=self.device)    

# %% ../nbs/Callbacks.ipynb 6
def get_mem():
    return torch.cuda.memory_allocated()/1024/1024

class MemoryCB(Callback):
    def before_fit(self, learn):
        print(f'memory befor fit {get_mem()}')
    def before_batch(self, learn): 
        print(f'memory before batch {get_mem()}')
    def after_batch(self, learn): 
        print(f'memory after batch {get_mem()}')
    def before_epoch(self, learn): 
        print(f'memory before epoch {get_mem()}')
    def after_epoch(self, learn): 
        print(f'memory after epoch {get_mem()}')
    def after_loss(self, learn): 
        print(f'memory after loss {get_mem()}')
    def after_preds(self, learn): 
        print(f'memory after preds {get_mem()}')
    def before_preds(self, learn): 
        print(f'memory before preds{get_mem()}')
    def after_backprop(self, learn): 
        print(f'memory after backprop {get_mem()}')

# %% ../nbs/Callbacks.ipynb 7
from .model import save_model_checkpoint, create_checkpoint_data

class SaveModelCB(Callback):
    order = 9
    
    def __init__(self, metric, metric_limit, model_path, model_name, selected_classes, model_def):
        self.metric = metric
        self.metric_limit = int(1000 * metric_limit)
        self.best_metric_value = 0.0
        self.model_path = model_path
        self.model_name = model_name
        self.selected_classes = selected_classes
        self.model_def = model_def
        self.best_model_name = model_name + '_best_model.pth'
    
    def after_epoch(self, learn):
        if learn.model.training: return #  only save after validation epoch
        values = learn.current_metric_values
        if not values: return
        
        try:    
            current_metric_value = float(values[self.metric])
        except:
            print('metric as save criterion not in metrics')
            return 
        epoch = values['epoch']
        loss = values['loss']

        scaled_metric = int(1000 * current_metric_value)
        if self.best_metric_value < current_metric_value:
            self.best_metric_value = current_metric_value
            checkpoint_data = create_checkpoint_data(model_def = self.model_def,
                                                     select_classes=self.selected_classes,
                                                     epoch = epoch,
                                                     loss = loss,
                                                     metric = current_metric_value)
            
            save_model_checkpoint(learn.model, learn.opt, self.model_path/self.best_model_name, **checkpoint_data)
            print(f'Best model saved with {self.metric} = {current_metric_value:.3}')
            if scaled_metric > self.metric_limit:
                numbered_model_name = self.model_name + '_' + self.metric + f'{scaled_metric:03}' + '.pth'
                save_model_checkpoint(learn.model, learn.opt, self.model_path / numbered_model_name, **checkpoint_data)
                print(f'Numbered Model saved with {self.metric} = {current_metric_value:.3}')
