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

__all__ = []

# Cell
from ..soft_dependencies import SoftDependencies
if not SoftDependencies.check()['light']:
    raise ImportError("The fastai module is not installed.")

# Cell
from fastai.learner import _ConstantFunc
from fastai.vision.all import *
from fastai.tabular.all import *

# Cell
@typedispatch
def _fully_decode(dl:TfmdDL, inps, outs):
    "Attempt to fully decode the `inp"
    inps = dl.decode(inps)
    dec = []
    for d in inps:
        for e in d:
            e = retain_type(e, typ=type(d))
            dec.append(e)
    outs.insert(len(outs), dec)
    return outs

@typedispatch
def _fully_decode(dl:TabDataLoader, inps, outs):
    "Attempt to fully decode the `inp"
    for i in range(dl.n_inp):
        inps[i] = torch.cat(inps[i], dim=0)
    dec = dl.decode(inps)
    outs.insert(len(outs), dec)
    return outs

def _decode_loss(vocab, dec_out, outs):
    i2c, dec2c = {}, []
    try:
        for i, vocab in enumerate(list(vocab)):
                i2c[i] = vocab
        for i, item in enumerate(dec_out):
            dec2c.append(i2c[int(item)])
        dec_out = dec2c
        outs.insert(0, dec_out)
    except:
        outs.insert(0, dec_out)
    return outs

# Cell
#@delegates(GatherPredsCallback.__init__)
#@patch
#def get_preds(self:Learner, ds_idx=1, dl=None, with_input=False, with_decoded=False, with_loss=False, raw=False, act=None,
#                inner=False, reorder=True, cbs=None, **kwargs):
#    if dl is None: dl = self.dls[ds_idx].new(shuffled=False, drop_last=False)
#    if reorder and hasattr(dl, 'get_idxs'):
#       idxs = dl.get_idxs()
#        dl = dl.new(get_idxs = _ConstantFunc(idxs))
#    cb = GatherPredsCallback(with_input=with_input, with_loss=with_loss, **kwargs)
#    ctx_mgrs = self.validation_context(cbs=L(cbs)+[cb], inner=inner)
#    if with_loss: ctx_mgrs.append(self.loss_not_reduced())
#    with ContextManagers(ctx_mgrs):
#        self._do_epoch_validate(dl=dl)
#        if act is None: act = getattr(self.loss_func, 'activation', noop)
#        res = cb.all_tensors()
#        pred_i = 1 if with_input else 0
#        if res[pred_i] is not None:
#            res[pred_i] = act(res[pred_i]) if not raw else res[pred_i]
#            if with_decoded:
#                res.insert(pred_i+2, getattr(self.loss_func, 'decodes', noop)(res[pred_i]))
#        if reorder and hasattr(dl, 'get_idxs'): res = nested_reorder(res, tensor(idxs).argsort())
#        if with_decoded:
#            if hasattr(self.dls, 'categorize'): res[pred_i+1] = apply(getattr(self.dls, 'categorize', noop).decode, [*res[pred_i+2]])
#            if hasattr(self.dls, 'multi_categorize'): res[pred_i+1] = apply(getattr(self.dls, 'multi_categorize', noop).decode, [*res[pred_i+2]])
#        return tuple(res)
#    self._end_cleanup()

# Cell
@patch
def predict(x:Learner, item, with_input=False, rm_type_tfms=None):
        dl = x.dls.test_dl([item], rm_type_tfms=rm_type_tfms, num_workers=0)
        res = x.get_preds(dl=dl, with_input=with_input, with_decoded=True)
        return res

# Cell
@patch
def predict(x:TabularLearner, row, with_input=False, rm_type_tfms=None):
        dl = x.dls.test_dl(row.to_frame().T)
        dl.dataset.conts = dl.dataset.conts.astype(np.float32)
        res = x.get_preds(dl=dl, with_decoded=True, with_input=with_input)
        return res
