# AUTOGENERATED! DO NOT EDIT! File to edit: 01_error.fastai.ipynb (unless otherwise specified).

__all__ = ['loss_func_error', 'callback_error', 'catch_pred_errors', 'catch_loss_errors']

# Cell
from .torch import layer_error, device_error

from fastai.data.all import *
from fastai.optimizer import *
from fastai.learner import *
from fastai.callback.training import ShortEpochCallback
from fastai.torch_core import default_device


from fastcore.basics import patch
from fastcore.meta import delegates

# Cell
def loss_func_error(e:Exception, learn) -> Exception:
    """
    Error that should be run when there is an issue when working with the loss function

    Raises with a message stating the shapes of the inputs and targs, and the error
    """
    err = f'There was an issue with calculating the loss with `{getattr(learn.loss_func, "__name__", learn.loss_func)}`'
    err += f'\n\nPrediction shape(s): {[p.shape for p in listify(learn.pred)]}'
    err += f'\nLabel Shape(s): {[y.shape for y in learn.yb]}'
    err += f'\nError: {e.args[0]}'
    e.args = [err]
    raise e

# Cell
def callback_error(e:Exception, cb:str, event_name:str) -> Exception:
    """
    Raises an error from when a Callback event failed, showing what event, the name of the Callback and the trace
    """
    e.args = [f"Exception raised in the {cb} Callback during {event_name}:\n\n{e.args[0]}"]
    raise e

# Cell
def catch_pred_errors(e:Exception, model) -> Exception:
    "Catches any errors relating to prediction that are either related to the device or model layers. Else raise `e`"
    if "Input type" in e.args[0]: device_error(e, 'Input', 'Model weights')
    elif "Expected" in e.args[0]: layer_error(e, model)
    else: raise e # anything else

# Cell
def catch_loss_errors(e:Exception, learn):
    "Catches any errors that occur with the loss function and its calculation"
    if "Input type" in e.args[0]: device_error(e, 'Model prediction', 'Truths')
    else: loss_func_error(e, learn)

# Cell
@patch
def sanity_check(self:Learner, show_table=False):
    "Performs a short epoch and uses all the callbacks in `self.cbs` on the CPU to ensure nothing is broken"
    device = getattr(self.dls, 'device', default_device())
    if hasattr(self.dls, 'device'):
        self.dls.device = 'cpu'
    else:
        # Using raw torch
        self.model.to('cpu')
    self.save('tmp')
    cbs = [ShortEpochCallback(short_valid=False)]
    if show_table:
        with self.no_bar(), self.no_logging():
            self.fit(1, cbs=cbs)
    else:
        self.fit(1, cbs=cbs)
    if hasattr(self.dls, 'device'):
        self.dls.device = device
    else:
        self.model.to(device)
    self.load('tmp')

# Cell
@patch
@delegates(Learner.sanity_check)
def __init__(self:Learner, dls, model, loss_func=None, opt_func=Adam, lr=defaults.lr, splitter=trainable_params, cbs=None,
                 metrics=None, path=None, model_dir='models', wd=None, wd_bn_bias=False, train_bn=True,
                 moms=(0.95,0.85,0.95), sanity_check=False, **kwargs):
    "Group together a `model`, some `dls` and a `loss_func` to handle training, potentially run a sanity check"
    path = Path(path) if path is not None else getattr(dls, 'path', Path('.'))
    if loss_func is None:
        loss_func = getattr(dls.train_ds, 'loss_func', None)
        assert loss_func is not None, "Could not infer loss function from the data, please pass a loss function."
    self.dls,self.model = dls,model
    store_attr(but='dls,model,cbs')
    self.training,self.create_mbar,self.logger,self.opt,self.cbs = False,True,print,None,L()
    self.add_cbs(L(defaults.callbacks)+L(cbs))
    self("after_create")
    if sanity_check: self.sanity_check(**kwargs)

# Cell
@patch
def _do_one_batch(self:Learner):
    try:
        self.pred = self.model(*self.xb)
    except RuntimeError as e:
        catch_pred_errors(e, self.model)
    self('after_pred')
    if len(self.yb):
        try:
            self.loss_grad = self.loss_func(self.pred, *self.yb)
        except Exception as e:
            catch_loss_errors(e, self)
        self.loss = self.loss_grad.clone()
    self('after_loss')
    if not self.training or not len(self.yb): return
    self('before_backward')
    self.loss_grad.backward()
    self._with_events(self.opt.step, 'step', CancelStepException)
    self.opt.zero_grad()

# Cell
@patch
def _call_one(self:Learner, event_name):
    if not hasattr(event, event_name): raise Exception(f'missing {event_name}')
    for cb in self.cbs.sorted('order'):
        try:
            cb(event_name)
        except Exception as e:
            callback_error(e, cb.__repr__(), event_name)