import re
import inspect


def formula_str(fit_func, params, num_format_or_precision=2):
    """
    Returns a string showing what a fit functions' formula is, with params injected.
    The normal use is when you have a fit function, say:
        def func(x, a, b):
            return a + b * x)
    which you then fit to data
        params, pcov = curve_fit(func, xdata, ydata, p0=(1.0, 1.0))
    to get some numerical params, say:
        params == [1.07815647e-06,  1.28497311e+00]

    Then if you call
        formula_str(func, params)
    you'll get
        1.07e-06 + 1.28e+00 * x

    You can control the appearance of the numerical params in the formula using the num_format_or_precision argument.

    Note that to work, the formula of the fit function has to fit completely in the return line of the function.
    """
    if not isinstance(num_format_or_precision, str):
        num_format_or_precision = "{:." + str(num_format_or_precision) + "e}"

    # get the param names from the code of the function
    param_names = re.compile('def [^(]+\(([^)]+)\)').findall(inspect.getsource(fit_func))[0].split(', ')[1:]

    # get the formula string from the code of the function
    formula_str = re.compile('return ([^\n]*)').findall(inspect.getsource(fit_func))[0]

    # replace the param values in the formula string to get the result
    rep = dict((re.escape(k), num_format_or_precision.format(v)) for k, v in zip(param_names, params))
    pattern = re.compile("|".join(list(rep.keys())))
    formula_str_with_nums = pattern.sub(lambda m: rep[re.escape(m.group(0))], formula_str)

    return formula_str_with_nums