#|default_exp basics

#|export
from fastcore.imports import *
import builtins,types
import pprint
try: from types import UnionType
except ImportError: UnionType = None

#|hide
from __future__ import annotations
from fastcore.test import *
from nbdev.showdoc import *
from fastcore.nb_imports import *

#|export
defaults = SimpleNamespace()

#|export
def ifnone(a, b):
    "`b` if `a` is None else `a`"
    return b if a is None else a

test_eq(ifnone(None,1), 1)
test_eq(ifnone(2   ,1), 2)

#|export
def maybe_attr(o, attr):
    "`getattr(o,attr,o)`"
    return getattr(o,attr,o)

class myobj: myattr='foo'

test_eq(maybe_attr(myobj, 'myattr'), 'foo')
test_eq(maybe_attr(myobj, 'another_attr'), myobj)

#|export
def basic_repr(flds=None):
    "Minimal `__repr__`"
    if isinstance(flds, str): flds = re.split(', *', flds)
    flds = list(flds or [])
    def _f(self):
        res = f'{type(self).__module__}.{type(self).__name__}'
        if not flds: return f'<{res}>'
        sig = ', '.join(f'{o}={getattr(self,o)!r}' for o in flds)
        return f'{res}({sig})'
    return _f

#|export
def is_array(x):
    "`True` if `x` supports `__array__` or `iloc`"
    return hasattr(x,'__array__') or hasattr(x,'iloc')

#|export
def listify(o=None, *rest, use_list=False, match=None):
    "Convert `o` to a `list`"
    if rest: o = (o,)+rest
    if use_list: res = list(o)
    elif o is None: res = []
    elif isinstance(o, list): res = o
    elif isinstance(o, str) or is_array(o): res = [o]
    elif is_iter(o): res = list(o)
    else: res = [o]
    if match is not None:
        if is_coll(match): match = len(match)
        if len(res)==1: res = res*match
        else: assert len(res)==match, 'Match length mismatch'
    return res

test_eq(listify('hi'), ['hi'])
test_eq(listify(array(1)), [array(1)])
test_eq(listify(1), [1])
test_eq(listify([1,2]), [1,2])
test_eq(listify(range(3)), [0,1,2])
test_eq(listify(None), [])
test_eq(listify(1,2), [1,2])

gen = (o for o in range(3))
test_eq(listify(gen), [0,1,2])

test_eq(listify(1,match=3), [1,1,1])

test_eq(listify(1,match=range(3)), [1,1,1])

test_eq(listify([1,1,1],match=3), [1,1,1])
test_fail(lambda: listify([1,1],match=3))

#|export
def tuplify(o, use_list=False, match=None):
    "Make `o` a tuple"
    return tuple(listify(o, use_list=use_list, match=match))

test_eq(tuplify(None),())
test_eq(tuplify([1,2,3]),(1,2,3))
test_eq(tuplify(1,match=[1,2,3]),(1,1,1))

#|export
def true(x):
    "Test whether `x` is truthy; collections with >0 elements are considered `True`"
    try: return bool(len(x))
    except: return bool(x)

#|export
class NullType:
    "An object that is `False` and can be called, chained, and indexed"
    def __getattr__(self,*args):return null
    def __call__(self,*args, **kwargs):return null
    def __getitem__(self, *args):return null
    def __bool__(self): return False

null = NullType()

#|export
def tonull(x):
    "Convert `None` to `null`"
    return null if x is None else x

#|export
def get_class(nm, *fld_names, sup=None, doc=None, funcs=None, **flds):
    "Dynamically create a class, optionally inheriting from `sup`, containing `fld_names`"
    attrs = {}
    for f in fld_names: attrs[f] = None
    for f in listify(funcs): attrs[f.__name__] = f
    for k,v in flds.items(): attrs[k] = v
    sup = ifnone(sup, ())
    if not isinstance(sup, tuple): sup=(sup,)

    def _init(self, *args, **kwargs):
        for i,v in enumerate(args): setattr(self, list(attrs.keys())[i], v)
        for k,v in kwargs.items(): setattr(self,k,v)

    all_flds = [*fld_names,*flds.keys()]
    def _eq(self,b):
        return all([getattr(self,k)==getattr(b,k) for k in all_flds])

    if not sup: attrs['__repr__'] = basic_repr(all_flds)
    attrs['__init__'] = _init
    attrs['__eq__'] = _eq
    res = type(nm, sup, attrs)
    if doc is not None: res.__doc__ = doc
    return res

#|export
def mk_class(nm, *fld_names, sup=None, doc=None, funcs=None, mod=None, **flds):
    "Create a class using `get_class` and add to the caller's module"
    if mod is None: mod = sys._getframe(1).f_locals
    res = get_class(nm, *fld_names, sup=sup, doc=doc, funcs=funcs, **flds)
    mod[nm] = res

mk_class('_t', a=1, sup=dict)
t = _t()
test_eq(t.a, 1)
assert(isinstance(t,dict))

#|export
def wrap_class(nm, *fld_names, sup=None, doc=None, funcs=None, **flds):
    "Decorator: makes function a method of a new class `nm` passing parameters to `mk_class`"
    def _inner(f):
        mk_class(nm, *fld_names, sup=sup, doc=doc, funcs=listify(funcs)+[f], mod=f.__globals__, **flds)
        return f
    return _inner

@wrap_class('_t', a=2)
def bar(self,x): return x+1

t = _t()
test_eq(t.a, 2)
test_eq(t.bar(3), 4)

#|export
class ignore_exceptions:
    "Context manager to ignore exceptions"
    def __enter__(self): pass
    def __exit__(self, *args): return True

with ignore_exceptions(): 
    # Exception will be ignored
    raise Exception

#|export
def exec_local(code, var_name):
    "Call `exec` on `code` and return the var `var_name"
    loc = {}
    exec(code, globals(), loc)
    return loc[var_name]

test_eq(exec_local("a=1", "a"), 1)

#|export
def risinstance(types, obj=None):
    "Curried `isinstance` but with args reversed"
    types = tuplify(types)
    if obj is None: return partial(risinstance,types)
    if any(isinstance(t,str) for t in types):
        return any(t.__name__ in types for t in type(obj).__mro__)
    return isinstance(obj, types)

assert risinstance(int, 1)
assert not risinstance(str, 0)
assert risinstance(int)(1)

assert risinstance(('str','int'), 'a')
assert risinstance('str', 'a')
assert not risinstance('int', 'a')

noop()
test_eq(noop(1),1)

class _t: foo=noops
test_eq(_t().foo(1),1)

#|export
#|hide
class _InfMeta(type):
    @property
    def count(self): return itertools.count()
    @property
    def zeros(self): return itertools.cycle([0])
    @property
    def ones(self):  return itertools.cycle([1])
    @property
    def nones(self): return itertools.cycle([None])

#|export
class Inf(metaclass=_InfMeta):
    "Infinite lists"
    pass

show_doc(Inf);

test_eq([o for i,o in zip(range(5), Inf.count)],
        [0, 1, 2, 3, 4])

test_eq([o for i,o in zip(range(5), Inf.zeros)],
        [0]*5)

test_eq([o for i,o in zip(range(5), Inf.ones)],
        [1]*5)

test_eq([o for i,o in zip(range(5), Inf.nones)],
        [None]*5)

#|export
_dumobj = object()
def _oper(op,a,b=_dumobj): return (lambda o:op(o,a)) if b is _dumobj else op(a,b)

def _mk_op(nm, mod):
    "Create an operator using `oper` and add to the caller's module"
    op = getattr(operator,nm)
    def _inner(a, b=_dumobj): return _oper(op, a,b)
    _inner.__name__ = _inner.__qualname__ = nm
    _inner.__doc__ = f'Same as `operator.{nm}`, or returns partial if 1 arg'
    mod[nm] = _inner

#|export
def in_(x, a):
    "`True` if `x in a`"
    return x in a

operator.in_ = in_

#|export
_all_ = ['lt','gt','le','ge','eq','ne','add','sub','mul','truediv','is_','is_not','in_']

#|export
for op in ['lt','gt','le','ge','eq','ne','add','sub','mul','truediv','is_','is_not','in_']: _mk_op(op, globals())

# test if element is in another
assert in_('c', ('b', 'c', 'a'))
assert in_(4, [2,3,4,5])
assert in_('t', 'fastai')
test_fail(in_('h', 'fastai'))

# use in_ as a partial
assert in_('fastai')('t')
assert in_([2,3,4,5])(4)
test_fail(in_('fastai')('h'))

#|export
def ret_true(*args, **kwargs):
    "Predicate: always `True`"
    return True

assert ret_true(1,2,3)
assert ret_true(False)

#|export
def ret_false(*args, **kwargs):
    "Predicate: always `False`"
    return False

#|export
def stop(e=StopIteration):
    "Raises exception `e` (by default `StopException`)"
    raise e

#|export
def gen(func, seq, cond=ret_true):
    "Like `(func(o) for o in seq if cond(func(o)))` but handles `StopIteration`"
    return itertools.takewhile(cond, map(func,seq))

test_eq(gen(noop, Inf.count, lt(5)),
        range(5))
test_eq(gen(operator.neg, Inf.count, gt(-5)),
        [0,-1,-2,-3,-4])
test_eq(gen(lambda o:o if o<5 else stop(), Inf.count),
        range(5))

#|export
def chunked(it, chunk_sz=None, drop_last=False, n_chunks=None):
    "Return batches from iterator `it` of size `chunk_sz` (or return `n_chunks` total)"
    assert bool(chunk_sz) ^ bool(n_chunks)
    if n_chunks: chunk_sz = max(math.ceil(len(it)/n_chunks), 1)
    if not isinstance(it, Iterator): it = iter(it)
    while True:
        res = list(itertools.islice(it, chunk_sz))
        if res and (len(res)==chunk_sz or not drop_last): yield res
        if len(res)<chunk_sz: return

t = list(range(10))
test_eq(chunked(t,3),      [[0,1,2], [3,4,5], [6,7,8], [9]])
test_eq(chunked(t,3,True), [[0,1,2], [3,4,5], [6,7,8],    ])

t = map(lambda o:stop() if o==6 else o, Inf.count)
test_eq(chunked(t,3), [[0, 1, 2], [3, 4, 5]])
t = map(lambda o:stop() if o==7 else o, Inf.count)
test_eq(chunked(t,3), [[0, 1, 2], [3, 4, 5], [6]])

t = np.arange(10)
test_eq(chunked(t,3),      [[0,1,2], [3,4,5], [6,7,8], [9]])
test_eq(chunked(t,3,True), [[0,1,2], [3,4,5], [6,7,8],    ])

test_eq(chunked([], 3),          [])
test_eq(chunked([], n_chunks=3), [])

#|export
def otherwise(x, tst, y):
    "`y if tst(x) else x`"
    return y if tst(x) else x

test_eq(otherwise(2+1, gt(3), 4), 3)
test_eq(otherwise(2+1, gt(2), 4), 4)

#|export
def custom_dir(c, add):
    "Implement custom `__dir__`, adding `add` to `cls`"
    return object.__dir__(c) + listify(add)

class _T: 
    def f(): pass

s = custom_dir(_T(), add=['foo', 'bar'])
assert {'foo', 'bar', 'f'}.issubset(s)

#|export
class AttrDict(dict):
    "`dict` subclass that also provides access to keys as attrs"
    def __getattr__(self,k): return self[k] if k in self else stop(AttributeError(k))
    def __setattr__(self, k, v): (self.__setitem__,super().__setattr__)[k[0]=='_'](k,v)
    def __dir__(self): return super().__dir__() + list(self.keys())
    def _repr_markdown_(self): return f'```json\n{pprint.pformat(self, indent=2)}\n```'
    def copy(self): return AttrDict(**self)

d = AttrDict(a=1,b="two")
test_eq(d.a, 1)
test_eq(d['b'], 'two')
test_eq(d.get('c','nope'), 'nope')
d.b = 2
test_eq(d.b, 2)
test_eq(d['b'], 2)
d['b'] = 3
test_eq(d['b'], 3)
test_eq(d.b, 3)
assert 'a' in dir(d)

#|export
def get_annotations_ex(obj, *, globals=None, locals=None):
    "Backport of py3.10 `get_annotations` that returns globals/locals"
    if isinstance(obj, type):
        obj_dict = getattr(obj, '__dict__', None)
        if obj_dict and hasattr(obj_dict, 'get'):
            ann = obj_dict.get('__annotations__', None)
            if isinstance(ann, types.GetSetDescriptorType): ann = None
        else: ann = None

        obj_globals = None
        module_name = getattr(obj, '__module__', None)
        if module_name:
            module = sys.modules.get(module_name, None)
            if module: obj_globals = getattr(module, '__dict__', None)
        obj_locals = dict(vars(obj))
        unwrap = obj
    elif isinstance(obj, types.ModuleType):
        ann = getattr(obj, '__annotations__', None)
        obj_globals = getattr(obj, '__dict__')
        obj_locals,unwrap = None,None
    elif callable(obj):
        ann = getattr(obj, '__annotations__', None)
        obj_globals = getattr(obj, '__globals__', None)
        obj_locals,unwrap = None,obj
    else: raise TypeError(f"{obj!r} is not a module, class, or callable.")

    if ann is None: ann = {}
    if not isinstance(ann, dict): raise ValueError(f"{obj!r}.__annotations__ is neither a dict nor None")
    if not ann: ann = {}

    if unwrap is not None:
        while True:
            if hasattr(unwrap, '__wrapped__'):
                unwrap = unwrap.__wrapped__
                continue
            if isinstance(unwrap, functools.partial):
                unwrap = unwrap.func
                continue
            break
        if hasattr(unwrap, "__globals__"): obj_globals = unwrap.__globals__

    if globals is None: globals = obj_globals
    if locals is None: locals = obj_locals

    return dict(ann), globals, locals

#|export
def eval_type(t, glb, loc):
    "`eval` a type or collection of types, if needed, for annotations in py3.10+"
    if isinstance(t,str):
        if '|' in t: return Union[eval_type(tuple(t.split('|')), glb, loc)]
        return eval(t, glb, loc)
    if isinstance(t,(tuple,list)): return type(t)([eval_type(c, glb, loc) for c in t])
    return t

#|export
def _eval_type(t, glb, loc):
    res = eval_type(t, glb, loc)
    return NoneType if res is None else res

def type_hints(f):
    "Like `typing.get_type_hints` but returns `{}` if not allowed type"
    if not isinstance(f, typing._allowed_types): return {}
    ann,glb,loc = get_annotations_ex(f)
    return {k:_eval_type(v,glb,loc) for k,v in ann.items()}

def f(a:int)->bool: ... # a function with type hints (allowed)
exp = {'a':int,'return':bool}
test_eq(type_hints(f), typing.get_type_hints(f))
test_eq(type_hints(f), exp)

class _T:
    def __init__(self, a:int=0)->bool: ...
assert not type_hints(_T)

#|export
def annotations(o):
    "Annotations for `o`, or `type(o)`"
    res = {}
    if not o: return res
    res = type_hints(o)
    if not res: res = type_hints(getattr(o,'__init__',None))
    if not res: res = type_hints(type(o))
    return res

for o in _T,_T(),_T.__init__,f: test_eq(annotations(o), exp)
assert not annotations(int)
assert not annotations(print)

#|export
def anno_ret(func):
    "Get the return annotation of `func`"
    return annotations(func).get('return', None) if func else None

def f(x) -> float: return x
test_eq(anno_ret(f), float)

def f(x) -> typing.Tuple[float,float]: return x
test_eq(anno_ret(f), typing.Tuple[float,float])

def f(x) -> None: return x

test_eq(anno_ret(f), NoneType)
assert anno_ret(f) is not None # returns NoneType instead of None

def f(x): return x

test_eq(anno_ret(f), None)
test_eq(anno_ret(None), None) # instead of passing in a func, pass in None

#|export
def _ispy3_10(): return sys.version_info.major >=3 and sys.version_info.minor >=10

def signature_ex(obj, eval_str:bool=False):
    "Backport of `inspect.signature(..., eval_str=True` to <py310"
    from inspect import Signature, Parameter, signature

    def _eval_param(ann, k, v):
        if k not in ann: return v
        return Parameter(v.name, v.kind, annotation=ann[k], default=v.default)

    if not eval_str: return signature(obj)
    if _ispy3_10(): return signature(obj, eval_str=eval_str)
    sig = signature(obj)
    if sig is None: return None
    ann = type_hints(obj)
    params = [_eval_param(ann,k,v) for k,v in sig.parameters.items()]
    return Signature(params, return_annotation=sig.return_annotation)

#|export
def union2tuple(t):
    if (getattr(t, '__origin__', None) is Union
        or (UnionType and isinstance(t, UnionType))): return t.__args__
    return t

test_eq(union2tuple(Union[int,str]), (int,str))
test_eq(union2tuple(int), int)
test_eq(union2tuple(Tuple[int,str]), Tuple[int,str])
test_eq(union2tuple((int,str)), (int,str))
if UnionType: test_eq(union2tuple(int|str), (int,str))

#|export
def argnames(f, frame=False):
    "Names of arguments to function or frame `f`"
    code = getattr(f, 'f_code' if frame else '__code__')
    return code.co_varnames[:code.co_argcount+code.co_kwonlyargcount]

test_eq(argnames(f), ['x'])

#|export
def with_cast(f):
    "Decorator which uses any parameter annotations as preprocessing functions"
    anno, out_anno, params = annotations(f), anno_ret(f), argnames(f)
    c_out = ifnone(out_anno, noop)
    defaults = dict(zip(reversed(params), reversed(f.__defaults__ or {})))
    @functools.wraps(f)
    def _inner(*args, **kwargs):
        args = list(args)
        for i,v in enumerate(params):
            if v in anno:
                c = anno[v]
                if v in kwargs: kwargs[v] = c(kwargs[v])
                elif i<len(args): args[i] = c(args[i])
                elif v in defaults: kwargs[v] = c(defaults[v])
        return c_out(f(*args, **kwargs))
    return _inner

@with_cast
def _f(a, b:Path, c:str='', d=0): return (a,b,c,d)

test_eq(_f(1, '.', 3), (1,Path('.'),'3',0))
test_eq(_f(1, '.'), (1,Path('.'),'',0))

@with_cast
def _g(a:int=0)->str: return a

test_eq(_g(4.0), '4')
test_eq(_g(4.4), '4')
test_eq(_g(2), '2')

#|export
def _store_attr(self, anno, **attrs):
    stored = getattr(self, '__stored_args__', None)
    for n,v in attrs.items():
        if n in anno: v = anno[n](v)
        setattr(self, n, v)
        if stored is not None: stored[n] = v

#|export
def store_attr(names=None, self=None, but='', cast=False, store_args=None, **attrs):
    "Store params named in comma-separated `names` from calling context into attrs in `self`"
    fr = sys._getframe(1)
    args = argnames(fr, True)
    if self: args = ('self', *args)
    else: self = fr.f_locals[args[0]]
    if store_args is None: store_args = not hasattr(self,'__slots__')
    if store_args and not hasattr(self, '__stored_args__'): self.__stored_args__ = {}
    anno = annotations(self) if cast else {}
    if names and isinstance(names,str): names = re.split(', *', names)
    ns = names if names is not None else getattr(self, '__slots__', args[1:])
    added = {n:fr.f_locals[n] for n in ns}
    attrs = {**attrs, **added}
    if isinstance(but,str): but = re.split(', *', but)
    attrs = {k:v for k,v in attrs.items() if k not in but}
    return _store_attr(self, anno, **attrs)

class T:
    def __init__(self, a,b,c): self.a,self.b,self.c = a,b,c

class T:
    def __init__(self, a,b,c): store_attr('a,b,c', self)

t = T(1,c=2,b=3)
assert t.a==1 and t.b==3 and t.c==2

test_eq(t.__stored_args__, {'a':1, 'b':3, 'c':2})

class T:
    def __init__(self, a,b,c:str): store_attr('a,b,c')

t = T(1,c=2,b=3)
assert t.a==1 and t.b==3 and t.c==2

#|hide
class _T:
    def __init__(self, a,b):
        c = 2
        store_attr('a,b,c')

t = _T(1,b=3)
assert t.a==1 and t.b==3 and t.c==2

class T:
    def __init__(self, a:listify, b, c:str): store_attr('a,b,c', cast=True)

t = T(1,c=2,b=3)
assert t.a==[1] and t.b==3 and t.c=='2'

class T2(T):
    def __init__(self, d, **kwargs):
        super().__init__(**kwargs)
        store_attr('d')

t = T2(d=1,a=2,b=3,c=4)
assert t.a==2 and t.b==3 and t.c==4 and t.d==1

class T:
    def __init__(self, a,b,c): store_attr()

t = T(1,c=2,b=3)
assert t.a==1 and t.b==3 and t.c==2

class T4(T):
    def __init__(self, d, **kwargs):
        super().__init__(**kwargs)
        store_attr()

t = T4(4, a=1,c=2,b=3)
assert t.a==1 and t.b==3 and t.c==2 and t.d==4

class T4:
    def __init__(self, *, a: int, b: float = 1):
        store_attr()
        
t = T4(a=3)
assert t.a==3 and t.b==1
t = T4(a=3, b=2)
assert t.a==3 and t.b==2

#|hide
# ensure that subclasses work with or without `store_attr`
class T4(T):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        store_attr()

t = T4(a=1,c=2,b=3)
assert t.a==1 and t.b==3 and t.c==2

class T4(T): pass

t = T4(a=1,c=2,b=3)
assert t.a==1 and t.b==3 and t.c==2

#|hide
#ensure that kwargs work with names==None
class T:
    def __init__(self, a,b,c,**kwargs): store_attr(**kwargs)

t = T(1,c=2,b=3,d=4,e=-1)
assert t.a==1 and t.b==3 and t.c==2 and t.d==4 and t.e==-1

#|hide
#ensure that kwargs work with names==''
class T:
    def __init__(self, a, **kwargs):
        self.a = a+1
        store_attr('', **kwargs)

t = T(a=1, d=4)
test_eq(t.a, 2)
test_eq(t.d, 4)

class T:
    def __init__(self, a,b,c): store_attr(but='a')

t = T(1,c=2,b=3)
assert t.b==3 and t.c==2
assert not hasattr(t,'a')

class T:
    def __init__(self): store_attr(a=1)

t = T()
assert t.a==1

def create_T(a, b):
    t = SimpleNamespace()
    store_attr(self=t)
    return t

t = create_T(a=1, b=2)
assert t.a==1 and t.b==2

#|export
def attrdict(o, *ks, default=None):
    "Dict from each `k` in `ks` to `getattr(o,k)`"
    return {k:getattr(o, k, default) for k in ks}

class T:
    def __init__(self, a,b,c): store_attr()

t = T(1,c=2,b=3)
test_eq(attrdict(t,'b','c'), {'b':3, 'c':2})

#|export
def properties(cls, *ps):
    "Change attrs in `cls` with names in `ps` to properties"
    for p in ps: setattr(cls,p,property(getattr(cls,p)))

class T:
    def a(self): return 1
    def b(self): return 2
properties(T,'a')

test_eq(T().a,1)
test_eq(T().b(),2)

#|export
_c2w_re = re.compile(r'((?<=[a-z])[A-Z]|(?<!\A)[A-Z](?=[a-z]))')
_camel_re1 = re.compile('(.)([A-Z][a-z]+)')
_camel_re2 = re.compile('([a-z0-9])([A-Z])')

#|export
def camel2words(s, space=' '):
    "Convert CamelCase to 'spaced words'"
    return re.sub(_c2w_re, rf'{space}\1', s)

test_eq(camel2words('ClassAreCamel'), 'Class Are Camel')

#|export
def camel2snake(name):
    "Convert CamelCase to snake_case"
    s1   = re.sub(_camel_re1, r'\1_\2', name)
    return re.sub(_camel_re2, r'\1_\2', s1).lower()

test_eq(camel2snake('ClassAreCamel'), 'class_are_camel')
test_eq(camel2snake('Already_Snake'), 'already__snake')

#|export
def snake2camel(s):
    "Convert snake_case to CamelCase"
    return ''.join(s.title().split('_'))

test_eq(snake2camel('a_b_cc'), 'ABCc')

#|export
def class2attr(self, cls_name):
    "Return the snake-cased name of the class; strip ending `cls_name` if it exists."
    return camel2snake(re.sub(rf'{cls_name}$', '', self.__class__.__name__) or cls_name.lower())

class Parent:
    @property
    def name(self): return class2attr(self, 'Parent')

class ChildOfParent(Parent): pass
class ParentChildOf(Parent): pass

p = Parent()
cp = ChildOfParent()
cp2 = ParentChildOf()

test_eq(p.name, 'parent')
test_eq(cp.name, 'child_of')
test_eq(cp2.name, 'parent_child_of')