# -*- coding: utf-8 -*-
"""
Some decorators useful for data analysis functions
"""
import numpy as np

def multi_input(f):
    """Allow a function to also be applied to each element in a dictionary"""
    def wrapper(data, *args, **kwargs):
        if type(data) is dict:
            output_data = {}
            for name in data:
                output_data[name] = f(data[name], *args, **kwargs)
            if all(x is None for x in output_data.values()):
                return 
            else:
                return output_data
        else:
            return f(data, *args, **kwargs)
    return wrapper

def av_output(f):
    """Allow running a function multiple times returning the average output"""
    def wrapper(average=1, *args, **kwargs):
        data_av = f(*args, **kwargs)
        try:
            if not isinstance(data_av, np.ndarray):
                data_av = np.array(data_av)
        except:
            pass
        for i in range(average - 1):
            delta = f(*args, **kwargs)
            try:
                if not isinstance(delta, np.ndarray):
                    delta = np.array(delta)
            except:
                pass
            data_av += delta
        if not isinstance(data_av, tuple):
            data_av /= average
        else:
            data_av = [d / average for d in data_av]
        return data_av
    return wrapper