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

__all__ = ['device_error', 'hook_fn', 'PreHook', 'ForwardHooks', 'hook_outputs', 'layer_error']

# Cell
import torch
import re
from fastai.callback.hook import Hook
from fastai.torch_core import to_detach

from fastcore.basics import store_attr

# Cell
def device_error(e:Exception, a:str, b:str) -> Exception:
    """
    Verbose error for if `a` and `b` are on different devices
    Should be used when checking if a model is on the same device, or two tensors
    """
    inp, weight, _ = e.args[0].replace('( ', '').split(')')
    inp = inp.replace('Input type', f'{a} has type: \t\t')
    weight = weight.replace(' and weight type', f'{b} have type: \t')
    err = f'Mismatch between weight types\n\n{inp})\n{weight})\n\nBoth should be the same.'
    e.args = [err]
    raise e

# Cell
def hook_fn(m, i):
    "Simple hook fn to return the layer"
    return m

# Cell
class PreHook(Hook):
    "Creates and registers a hook on `m` with `hook_func` as a forward pre_hook"
    def __init__(self, m, hook_func, is_forward=True, detach=True, cpu=False, gather=False):
        store_attr('hook_func,detach,cpu,gather')
        f = m.register_forward_pre_hook if is_forward else m.register_backward_pre_hook
        self.hook = f(self.hook_fn)
        self.stored,self.removed = None, False

    def hook_fn(self, module, inp):
        "Applies `hook_fn` to `module` and `inp`"
        if self.detach:
            inp = to_detach(inp, cpu=self.cpu, gather=self.gather)
        self.stored = self.hook_func(module, inp)

# Cell
class ForwardHooks():
    "Create several forward-hooks on the modules in `ms` with `hook_func`"
    def __init__(self, ms, hook_func, is_forward=True, detach=True, cpu=False):
        self.hooks = [PreHook(m, hook_func, is_forward, detach, cpu) for m in ms]

# Cell
def hook_outputs(modules, detach=True, cpu=False, grad=False):
    "Return `Hooks` that store activations of all `modules` in `self.stored`"
    return ForwardHooks(modules, hook_fn, detach=detach, cpu=cpu, is_forward=not grad)

# Cell
def layer_error(e:Exception, model, inp) -> Exception:
    """
    Verbose error for when there is a size mismatch between some input and the model.
    `model` should be any torch model
    `inp` is the input that went to the model
    """
    args = e.args[0].replace("Expected", "Model expected")
    hooks = hook_outputs(model)
    try:
        _ = model(inp)
    except:
        pass
    finally:
        layer = [hook.stored for hook in hooks.hooks if hook.stored is not None][-1]
        [h.remove() for h in hooks.hooks]
        e.args = [f'Size mismatch between input tensors and what the model expects\n{"-"*76}\nLayer: {layer}\nError: > {args}']
        raise e