#AUTOGENERATED! DO NOT EDIT! File to edit: dev/02_data_pipeline.ipynb (unless otherwise specified).

__all__ = ['get_func', 'show_title', 'Func', 'Sig', 'SelfFunc', 'Self', 'positional_annotations', 'Transform',
           'compose_tfms', 'TupleOfTypes', 'can_show', 'Pipeline', 'TfmdList', 'Tuplify']

from ..imports import *
from ..test import *
from ..core import *
from ..notebook.showdoc import show_doc

def get_func(t, name, *args, **kwargs):
    "Get the `t.name` (potentially partial-ized with `args` and `kwargs`) or `noop` if not defined"
    f = getattr(t, name, noop)
    return f if not (args or kwargs) else partial(f, *args, **kwargs)

def show_title(o, ax=None, ctx=None):
    "Set title of `ax` to `o`, or print `o` if `ax` is `None`"
    ax = ifnone(ax,ctx)
    if ax is None: print(o)
    else: ax.set_title(o)

class Func():
    "Basic wrapper around a `name` with `args` and `kwargs` to call on a given type"
    def __init__(self, name, *args, **kwargs): self.name,self.args,self.kwargs = name,args,kwargs
    def __repr__(self): return f'sig: {self.name}({self.args}, {self.kwargs})'
    def _get(self, t): return get_func(t, self.name, *self.args, **self.kwargs)
    def __call__(self,t): return L(t).mapped(self._get) if is_listy(t) else self._get(t)

class _Sig():
    def __getattr__(self,k):
        def _inner(*args, **kwargs): return Func(k, *args, **kwargs)
        return _inner

Sig = _Sig()

class SelfFunc():
    "Search for `name` attribute and call it with `args` and `kwargs` on any object it's passed."
    def __init__(self, nm, *args, **kwargs): self.nm,self.args,self.kwargs = nm,args,kwargs
    def __repr__(self): return f'self: {self.nm}({self.args}, {self.kwargs})'
    def __call__(self, o):
        if not is_listy(o): return getattr(o,self.nm)(*self.args, **self.kwargs)
        else: return [getattr(o_,self.nm)(*self.args, **self.kwargs) for o_ in o]

class _SelfFunc():
    def __getattr__(self,k):
        def _inner(*args, **kwargs): return SelfFunc(k, *args, **kwargs)
        return _inner

Self = _SelfFunc()

def positional_annotations(f):
    "Get list of annotated types for all positional params, or None if no annotation"
    sig = inspect.signature(f)
    return [p.annotation if p.annotation != inspect._empty else None
            for p in sig.parameters.values() if p.default == inspect._empty and p.kind != inspect._VAR_KEYWORD]

from multimethod import multimeta,DispatchError

class Transform(metaclass=multimeta):
    "A function that `encodes` if `filt` matches, and optionally `decodes`"
    order,add_before_setup,filt,t = 0,False,None,None
    def __init__(self,encodes=None,decodes=None):
        self.encodes = getattr(self, 'encodes', noop) if encodes is None else encodes
        self.decodes = getattr(self, 'decodes', noop) if decodes is None else decodes

    def _apply(self, fs, x, filt):
        if self.filt is not None and self.filt!=filt: return x
        if self.t:
            gs = self._get_func(fs, self.t)
            if is_listy(self.t) and len(positional_annotations(gs)) != len(self.t):
                gs = [self._get_func(fs,t_) for t_ in self.t]
                if len(gs) == 1: gs = gs[0]
        else: gs=fs
        if is_listy(gs): return tuple(f(x_) for f,x_ in zip(gs,x))
        return gs(*L(x))

    def _get_func(self,f,t):
        idx = (object,) + tuple(t) if is_listy(t) else (object,t)
        try: f = f.__func__[idx]
        except DispatchError: return noop
        return partial(f,self)

    def accept_types(self, t): self.t = t
        # We can't create encodes/decodes here since patching might change things later
        # So we call _get_func in _apply instead

    def __call__(self, x, filt=None): return self._apply(self.encodes, x, filt)
    def decode  (self, x, filt=None): return self._apply(self.decodes, x, filt)
    def __getitem__(self, x): return self(x) # So it can be used as a `Dataset`

add_docs(Transform,
         __call__="Dispatch and apply the proper encodes to `x` if `filt` matches",
         decode="Dispatch and apply the proper decodes to `x` if `filt` matches",
         accept_types="Indicate the type of input received by the transform is `t`")

def compose_tfms(x, tfms, func_nm='__call__', reverse=False, **kwargs):
    "Apply all `func_nm` attribute of `tfms` on `x`, maybe in `reverse` order"
    if reverse: tfms = reversed(tfms)
    for tfm in tfms: x = getattr(tfm,func_nm,noop)(x, **kwargs)
    return x

def _get_ret(func):
    "Get the return annotation of `func`"
    ann = getattr(func,'__annotations__', None)
    if not ann: return None
    typ = ann.get('return')
    return list(typ.__args__) if getattr(typ, '_name', '')=='Tuple' else typ

class TupleOfTypes():
    "Class that handles the show method of a tuple of types"
    def __init__(self, ts): self.ts = ts
    def show(self, o, ctx=None, **kwargs):
        for p,t in zip(o,self.ts): ctx = t.show(p, ctx=ctx, **kwargs)
        return ctx

def can_show(t):
    "Check if `t` can show itself and wrap tuples in `TupleOfTypes`"
    if is_listy(t):
        res = [can_show(t_) for t_ in t]
        return TupleOfTypes(res) if len([o for o in res if o is None]) == 0 else None
    return t if hasattr(t, 'show') else None

class Pipeline():
    "A pipeline of composed (for encode/decode) transforms, setup with types"
    def __init__(self, funcs=None, t=None):
        if isinstance(funcs, Pipeline): funcs = funcs.raws
        self.raws,self.fs,self.t_show = L(funcs),[],None
        if len(self.raws) == 0: self.final_t = t
        else:
            for i,f in enumerate(self.raws.sorted(key='order')):
                if not isinstance(f,Transform): f = Transform(f)
                f.accept_types(t)
                self.fs.append(f)
                if self.t_show is None: self.t_idx,self.t_show = i,can_show(t)
                t = _get_ret(f.encodes) or t
            if self.t_show is None: self.t_idx,self.t_show = i+1,can_show(t)
            self.final_t = t

    def new(self, t=None): return Pipeline(self, t)
    def __repr__(self): return f"Pipeline over {self.fs}"

    def setup(self, items=None):
        tfms,self.fs = self.fs,[]
        for t in tfms:
            if t.add_before_setup:     self.fs.append(t)
            if hasattr(t, 'setup'):    t.setup(items)
            if not t.add_before_setup: self.fs.append(t)

    def __call__(self, o, **kwargs): return compose_tfms(o, self.fs)
    def decode  (self, i, **kwargs): return compose_tfms(i, self.fs, func_nm='decode', reverse=True)
    #def __getitem__(self, x): return self(x)
    #def decode_at(self, idx): return self.decode(self[idx])
    #def show_at(self, idx):   return self.show(self[idx])

    def show(self, o, ctx=None, **kwargs):
        if self.t_show is None: return self.decode(o)
        o = compose_tfms(o, self.fs[self.t_idx:], func_nm='decode', reverse=True)
        return self.t_show.show(o, ctx=ctx, **kwargs)

add_docs(Pipeline,
         __call__="Compose `__call__` of all `tfms` on `o`",
         decode="Compose `decode` of all `tfms` on `i`",
         new="Create a new `Pipeline`with the same `tfms` and a new initial `t`",
         show="Show item `o`",
         setup="Go through the transforms in order and call their potential setup on `items`")

@docs
class TfmdList(GetAttr):
    "A `Pipeline` of `tfms` applied to a collection of `items`"
    _xtra = 'decode __call__ show'.split()

    def __init__(self, items, tfms, do_setup=True):
        self.items = L(items)
        self.default = self.tfms = Pipeline(tfms)
        if do_setup: self.setup()

    def __getitem__(self, i):
        "Transformed item(s) at `i`"
        its = self.items[i]
        return its.mapped(self.tfms) if is_iter(i) else self.tfms(its)

    def decode_batch(self, b, **kwargs):
        "Decode `b`, a list of lists of pipeline outputs (i.e. output of a `DataLoader`)"
        transp = L(zip(*L(b)))
        return transp.mapped(self.decode, **kwargs).zipped()

    def setup(self): self.tfms.setup(self)
    def subset(self, idxs): return self.__class__(self.items[idxs], self.tfms, do_setup=False)
    def decode_at(self, idx, filt=None): return self.decode(self[idx], filt=filt)
    def show_at(self, idx): return self.show(self[idx])
    def __eq__(self, b): return all_equal(self, b)
    def __len__(self): return len(self.items)
    def __iter__(self): return (self[i] for i in range_of(self))
    def __repr__(self): return f"{self.__class__.__name__}: {self.items}\ntfms - {self.tfms}"

    _docs = dict(setup="Transform setup with self",
                 decode_at="Decoded item at `idx`",
                 decode_batch="Decode the items in `b",
                 show_at="Show item at `idx`",
                 subset="New `TfmdList` that only includes items at `idxs`")

def _fake_annote(t):
    def _inner(x): return x
    if is_listy(t):
        new_t = Tuple[int]
        new_t.__args__ = tuple(t)
        t = new_t
    setattr(_inner, '__annotations__', {'return': t})
    return _inner

class Tuplify(Transform):
    "Create tuple by applying each of of `tfms` to the passed `o`"
    add_before_setup=True
    def __init__(self, tfms=None):
        if tfms is None: tfms = [None]
        self.activ,self.tfms = None,L(tfms).mapped(Pipeline)
        self.encodes = _fake_annote([tfm.final_t for tfm in self.tfms])

    def encodes(self, o):
        return [t(o) for t in self.tfms] if self.activ is None else self.tfms[self.activ](o)

    def __call__(self, o, *args, **kwargs):
        "List of output of each of `tfms` on `o`"
        if self.activ is not None: return self.tfms[self.activ](o, *args, **kwargs)
        return [t(o, *args, **kwargs) for t in self.tfms]

    def decode(self, o, **kwargs): return [t.decode(p, **kwargs) for p,t in zip(o,self.tfms)]
    def __repr__(self): return f'Tuplify({self.tfms})'

    def setup(self, o=None):
        "Setup each of `tfms` independently"
        for i,tfm in enumerate(self.tfms):
            self.activ = i
            tfm.setup(o)
        self.activ=None

    @classmethod
    def piped(cls, tfms=None, tuple_tfms=None):
        "Return a `Pipeline` that combines `Tuplify` over `tfms`, optionally followed by any `tuple_tfms`"
        return Pipeline(cls(tfms) + L(tuple_tfms))

    xt,yt = add_props(lambda i,x:x.tfms[i])