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

# %% auto 0
__all__ = ['run_cbs', 'CancelFitException', 'CancelBatchException', 'CancelEpochException', 'Learner']

# %% ../nbs/05_learner.ipynb 3
import torch
from contextlib import contextmanager
from operator import attrgetter
from torch import optim
import torch.nn.functional as F
from functools import partial

from .callbacks import Callback
from .model import get_seg_model, get_model_name

# %% ../nbs/05_learner.ipynb 4
def run_cbs(cbs, method_nm, learn=None):
    for cb in sorted(cbs, key=attrgetter('order')):
        method = getattr(cb, method_nm, None)
        
        if method is not None: 
            method(learn)

# %% ../nbs/05_learner.ipynb 7
class CancelFitException(Exception): pass
class CancelBatchException(Exception): pass
class CancelEpochException(Exception): pass

# %% ../nbs/05_learner.ipynb 8
class Learner():
    def __init__(self, model, dls, cbs, loss_func=F.cross_entropy, lr=0.1, opt_func=optim.SGD):
        self.model = model
        self.dls = dls
        self.loss_func = loss_func
        self.opt_func = opt_func
        self.cbs = cbs
        self.n_epochs = 1
        self.epochs = None
        self.opt = None
        self.dl = None
        self.lr = lr

    def one_batch(self):
        self.callback('before_preds')
        self.preds = self.model(self.batch[0])
        self.callback('after_preds')
        self.loss = self.loss_func(self.preds, self.batch[1])
        self.callback('after_loss')
        if self.model.training:
            self.loss.backward()
            self.opt.step()
            self.opt.zero_grad()
        self.callback('after_backprop')

    def one_epoch(self, train):
        self.model.train(train)
        self.dl = self.dls.train if train else self.dls.valid
        try:
            self.callback('before_epoch')
            for self.iter,self.batch in enumerate(self.dl):
                try:
                    self.callback('before_batch')
                    self.one_batch()
                    self.callback('after_batch')
                except CancelBatchException: pass
            self.callback('after_epoch')
        except CancelEpochException: pass
    
    def fit(self, n_epochs):
        self.n_epochs = n_epochs
        self.epochs = range(n_epochs)
        self.opt = self.opt_func(self.model.parameters(), self.lr)
        try:
            self.callback('before_fit')
            for self.epoch in self.epochs:
                self.callback('before_full_epoch')
                self.one_epoch(True)
                self.one_epoch(False)
                self.callback('after_full_epoch')
            self.callback('after_fit')
        except CancelFitException: pass

    def callback(self, method_nm): run_cbs(self.cbs, method_nm, self)
