import numpy as np
import pandas as pd
import inspect
from .parameters import default_data
from numba import jit, vectorize, guvectorize
from functools import wraps
from six import StringIO
import ast
import toolz


class GetReturnNode(ast.NodeVisitor):
    """
    A Visitor to get the return tuple names from a calc-style function
    """
    def visit_Return(self, node):
        if isinstance(node.value, ast.Tuple):
            return [e.id for e in node.value.elts]
        else:
            return [node.value.id]


def dataframe_guvectorize(dtype_args, dtype_sig):
    """
    Extracts numpy arrays from caller arguments and passes them
    to guvectorized numba functions
    """
    def make_wrapper(func):
        vecd_f = guvectorize(dtype_args, dtype_sig)(func)

        @wraps(func)
        def wrapper(*args, **kwargs):
            # np_arrays = [getattr(args[0], i).values for i in theargs]
            arrays = [arg.values for arg in args]
            ans = vecd_f(*arrays)
            return ans
        return wrapper
    return make_wrapper


def dataframe_vectorize(dtype_args):
    """
    Extracts numpy arrays from caller arguments and passes them
    to vectorized numba functions
    """
    def make_wrapper(func):
        vecd_f = vectorize(dtype_args)(func)

        @wraps(func)
        def wrapper(*args, **kwargs):
            arrays = [arg.values for arg in args]
            ans = vecd_f(*arrays)
            return ans
        return wrapper
    return make_wrapper


def dataframe_wrap_guvectorize(dtype_args, dtype_sig):
    """
    Extracts particular numpy arrays from caller argments and passes
    them to guvectorize. Goes one step further than dataframe_guvectorize
    by looking for the column names in the dataframe and just extracting those
    """
    def make_wrapper(func):
        theargs = inspect.getargspec(func).args
        vecd_f = guvectorize(dtype_args, dtype_sig)(func)

        def wrapper(*args, **kwargs):
            np_arrays = [getattr(args[0], i).values for i in theargs]
            ans = vecd_f(*np_arrays)
            return ans
        return wrapper
    return make_wrapper


def create_apply_function_string(sigout, sigin, parameters):
    """
    Create a string for a function of the form::

        def ap_fuc(x_0, x_1, x_2, ...):
            for i in range(len(x_0)):
                x_0[i], ... = jitted_f(x_j[i], ...)
            return x_0[i], ...

    where the specific args to jitted_f and the number of
    values to return is determined by sigout and sigin

    Parameters
    ----------
    sigout: iterable of the out arguments

    sigin: iterable of the in arguments

    parameters: iterable of which of the args (from in_args) are parameter
                variables (as opposed to column records). This influences
                how we construct the '_apply' function

    Returns
    -------
    a String representing the function
    """

    s = StringIO()
    total_len = len(sigout) + len(sigin)
    out_args = ["x_" + str(i) for i in range(0, len(sigout))]
    in_args = ["x_" + str(i) for i in range(len(sigout), total_len)]

    s.write("def ap_func({0}):\n".format(",".join(out_args + in_args)))
    s.write("  for i in range(len(x_0)):\n")

    out_index = [x + "[i]" for x in out_args]
    in_index = []
    for arg, _var in zip(in_args, sigin):
        in_index.append(arg + "[i]" if _var not in parameters else arg)
    s.write("    " + ",".join(out_index) + " = ")
    s.write("jitted_f(" + ",".join(in_index) + ")\n")
    s.write("  return " + ",".join(out_args) + "\n")

    return s.getvalue()


def create_toplevel_function_string(args_out, args_in, pm_or_pf,
                                    kwargs_for_func={}):
    """
    Create a string for a function of the form::

        def hl_func(x_0, x_1, x_2, ...):
            outputs = (...) = calc_func(...)
            header = [...]
            return DataFrame(data, columns=header)

    where the specific args to jitted_f and the number of
    values to return is destermined by sigout and sigin

    Parameters
    ----------
    args_out: iterable of the out arguments

    args_in: iterable of the in arguments

    pm_or_pf: iterable of strings for object that holds each arg

    kwargs_for_func: dictionary of keyword args for the function

    Returns
    -------
    a String representing the function
    """

    s = StringIO()

    s.write("def hl_func(pm, pf")

    if kwargs_for_func:
        kwargs = ",".join(str(k) + "=" + str(v) for k, v in
                          kwargs_for_func.items())
        s.write(", " + kwargs + " ")

    s.write("):\n")
    s.write("    from pandas import DataFrame\n")
    s.write("    import numpy as np\n")
    s.write("    outputs = \\\n")
    outs = []
    for arg in kwargs_for_func:
        args_in.remove(arg)

    for p, attr in zip(pm_or_pf, args_out + args_in):
        outs.append(p + "." + attr + ", ")
    outs = [m_or_f + "." + arg for m_or_f, arg in zip(pm_or_pf, args_out)]
    s.write("        (" + ", ".join(outs) + ") = \\\n")
    s.write("        " + "applied_f(")
    for p, attr in zip(pm_or_pf, args_out + args_in):
        s.write(p + "." + attr + ", ")
    for arg in kwargs_for_func:
        s.write(arg + ", ")
    s.write(")\n")

    s.write("    header = [")
    col_headers = ["'" + out + "'" for out in args_out]
    s.write(", ".join(col_headers))
    s.write("]\n")
    if len(args_out) == 1:
        s.write("    return DataFrame(data=outputs,"
                "columns=header)")
    else:
        s.write("    return DataFrame(data=np.column_stack("
                "outputs),columns=header)")

    return s.getvalue()


def make_apply_function(func, out_args, in_args, parameters, do_jit=True,
                        **kwargs):
    """
    Takes a '_calc' function and creates the necessary Python code for an
    _apply style function. Will also jit the function if desired

    Parameters
    ----------
    func: the 'calc' style function

    out_args: list of out arguments for the apply function

    in_args: list of in arguments for the apply function

    parameters: iterable of which of the args (from in_args) are parameter
                variables (as opposed to column records). This influences
                how we construct the '_apply' function

    do_jit: Bool, if True, jit the resulting apply function

    Returns
    -------
    '_apply' style function

    """
    jitted_f = jit(**kwargs)(func)
    apfunc = create_apply_function_string(out_args, in_args, parameters)

    func_code = compile(apfunc, "<string>", "exec")
    fakeglobals = {}
    eval(func_code, {"jitted_f": jitted_f}, fakeglobals)
    if do_jit:
        return jit(**kwargs)(fakeglobals['ap_func'])
    else:
        return fakeglobals['ap_func']


def apply_jit(dtype_sig_out, dtype_sig_in, parameters=None, **kwargs):
    """
    make a decorator that takes in a _calc-style function, handle
    the apply step
    """
    if not parameters:
        parameters = []

    def make_wrapper(func):
        theargs = inspect.getargspec(func).args
        jitted_f = jit(**kwargs)(func)

        jitted_apply = make_apply_function(func, dtype_sig_out,
                                           dtype_sig_in, parameters,
                                           **kwargs)

        def wrapper(*args, **kwargs):
            in_arrays = []
            out_arrays = []
            for farg in theargs:
                if hasattr(args[0], farg):
                    in_arrays.append(getattr(args[0], farg))
                else:
                    in_arrays.append(getattr(args[1], farg))

            for farg in dtype_sig_out:
                if hasattr(args[0], farg):
                    out_arrays.append(getattr(args[0], farg))
                else:
                    out_arrays.append(getattr(args[1], farg))

            final_array = out_arrays + in_arrays
            ans = jitted_apply(*final_array)
            return ans

        return wrapper
    return make_wrapper


def iterate_jit(parameters=None, **kwargs):
    """
    make a decorator that takes in a _calc-style function, create a
    function that handles the "high-level" function and the "_apply"
    style function

    Note: perhaps a better "bigger picture" description of what this does?
    """
    if not parameters:
        parameters = []

    def make_wrapper(func):

        # Step 1. Wrap this function in apply_jit
        # from apply_jit

        # Get the input arguments from the function
        in_args = inspect.getargspec(func).args
        try:
            jit_args = inspect.getargspec(jit).args + ['nopython']
        except TypeError:
            #print ("This should only be seen in RTD, if not install numba!")
            return func

        kwargs_for_func = toolz.keyfilter(in_args.__contains__, kwargs)
        kwargs_for_jit = toolz.keyfilter(jit_args.__contains__, kwargs)

        # Any name that is a taxcalc parameter (or the special case 'puf'
        # Boolean flag is given special treatment. Identify those names here
        allowed_parameters = list(default_data(metadata=True).keys())
        allowed_parameters += list(arg[1:] for arg in  default_data(metadata=True).keys())
        allowed_parameters.append("puf")
        additional_parameters = [arg for arg in in_args if arg in allowed_parameters]
        additional_parameters += parameters
        # Remote duplicates
        all_parameters = list(set(additional_parameters))

        src = inspect.getsourcelines(func)[0]

        # Discover the return arguments by walking
        # the AST of the function
        all_returned_vals = []
        gnr = GetReturnNode()
        all_out_args = None
        for node in ast.walk(ast.parse(''.join(src))):
            all_out_args = gnr.visit(node)
            if all_out_args:
                break

        if not all_out_args:
            raise ValueError("Can't find return statement in function!")

        # Now create the apply jitted function
        applied_jitted_f = make_apply_function(func,
                                               list(reversed(all_out_args)),
                                               in_args,
                                               parameters=all_parameters,
                                               do_jit=True,
                                               **kwargs_for_jit)

        def wrapper(*args, **kwargs):
            in_arrays = []
            out_arrays = []
            pm_or_pf = []
            for farg in all_out_args + in_args:
                if hasattr(args[0], farg):
                    in_arrays.append(getattr(args[0], farg))
                    pm_or_pf.append("pm")
                elif hasattr(args[1], farg):
                    in_arrays.append(getattr(args[1], farg))
                    pm_or_pf.append("pf")
                elif farg not in kwargs_for_func:
                    raise ValueError("Unknown arg: " + farg)

            # Create the high level function
            high_level_func = create_toplevel_function_string(all_out_args,
                                                              list(in_args),
                                                              pm_or_pf,
                                                              kwargs_for_func)
            func_code = compile(high_level_func, "<string>", "exec")
            fakeglobals = {}
            eval(func_code, {"applied_f": applied_jitted_f}, fakeglobals)
            high_level_fn = fakeglobals['hl_func']

            ans = high_level_fn(*args, **kwargs)
            return ans

        return wrapper

    return make_wrapper

import numpy as np
import pandas as pd
import inspect
from .parameters import default_data
from numba import jit, vectorize, guvectorize
from functools import wraps
from six import StringIO
import ast
import toolz


class GetReturnNode(ast.NodeVisitor):
    """
    A Visitor to get the return tuple names from a calc-style function
    """
    def visit_Return(self, node):
        if isinstance(node.value, ast.Tuple):
            return [e.id for e in node.value.elts]
        else:
            return [node.value.id]


def dataframe_guvectorize(dtype_args, dtype_sig):
    """
    Extracts numpy arrays from caller arguments and passes them
    to guvectorized numba functions
    """
    def make_wrapper(func):
        vecd_f = guvectorize(dtype_args, dtype_sig)(func)

        @wraps(func)
        def wrapper(*args, **kwargs):
            # np_arrays = [getattr(args[0], i).values for i in theargs]
            arrays = [arg.values for arg in args]
            ans = vecd_f(*arrays)
            return ans
        return wrapper
    return make_wrapper


def dataframe_vectorize(dtype_args):
    """
    Extracts numpy arrays from caller arguments and passes them
    to vectorized numba functions
    """
    def make_wrapper(func):
        vecd_f = vectorize(dtype_args)(func)

        @wraps(func)
        def wrapper(*args, **kwargs):
            arrays = [arg.values for arg in args]
            ans = vecd_f(*arrays)
            return ans
        return wrapper
    return make_wrapper


def dataframe_wrap_guvectorize(dtype_args, dtype_sig):
    """
    Extracts particular numpy arrays from caller argments and passes
    them to guvectorize. Goes one step further than dataframe_guvectorize
    by looking for the column names in the dataframe and just extracting those
    """
    def make_wrapper(func):
        theargs = inspect.getargspec(func).args
        vecd_f = guvectorize(dtype_args, dtype_sig)(func)

        def wrapper(*args, **kwargs):
            np_arrays = [getattr(args[0], i).values for i in theargs]
            ans = vecd_f(*np_arrays)
            return ans
        return wrapper
    return make_wrapper


def create_apply_function_string(sigout, sigin, parameters):
    """
    Create a string for a function of the form::

        def ap_fuc(x_0, x_1, x_2, ...):
            for i in range(len(x_0)):
                x_0[i], ... = jitted_f(x_j[i], ...)
            return x_0[i], ...

    where the specific args to jitted_f and the number of
    values to return is determined by sigout and sigin

    Parameters
    ----------
    sigout: iterable of the out arguments

    sigin: iterable of the in arguments

    parameters: iterable of which of the args (from in_args) are parameter
                variables (as opposed to column records). This influences
                how we construct the '_apply' function

    Returns
    -------
    a String representing the function
    """

    s = StringIO()
    total_len = len(sigout) + len(sigin)
    out_args = ["x_" + str(i) for i in range(0, len(sigout))]
    in_args = ["x_" + str(i) for i in range(len(sigout), total_len)]

    s.write("def ap_func({0}):\n".format(",".join(out_args + in_args)))
    s.write("  for i in range(len(x_0)):\n")

    out_index = [x + "[i]" for x in out_args]
    in_index = []
    for arg, _var in zip(in_args, sigin):
        in_index.append(arg + "[i]" if _var not in parameters else arg)
    s.write("    " + ",".join(out_index) + " = ")
    s.write("jitted_f(" + ",".join(in_index) + ")\n")
    s.write("  return " + ",".join(out_args) + "\n")

    return s.getvalue()


def create_toplevel_function_string(args_out, args_in, pm_or_pf,
                                    kwargs_for_func={}):
    """
    Create a string for a function of the form::

        def hl_func(x_0, x_1, x_2, ...):
            outputs = (...) = calc_func(...)
            header = [...]
            return DataFrame(data, columns=header)

    where the specific args to jitted_f and the number of
    values to return is destermined by sigout and sigin

    Parameters
    ----------
    args_out: iterable of the out arguments

    args_in: iterable of the in arguments

    pm_or_pf: iterable of strings for object that holds each arg

    kwargs_for_func: dictionary of keyword args for the function

    Returns
    -------
    a String representing the function
    """

    s = StringIO()

    s.write("def hl_func(pm, pf")

    if kwargs_for_func:
        kwargs = ",".join(str(k) + "=" + str(v) for k, v in
                          kwargs_for_func.items())
        s.write(", " + kwargs + " ")

    s.write("):\n")
    s.write("    from pandas import DataFrame\n")
    s.write("    import numpy as np\n")
    s.write("    outputs = \\\n")
    outs = []
    for arg in kwargs_for_func:
        args_in.remove(arg)

    for p, attr in zip(pm_or_pf, args_out + args_in):
        outs.append(p + "." + attr + ", ")
    outs = [m_or_f + "." + arg for m_or_f, arg in zip(pm_or_pf, args_out)]
    s.write("        (" + ", ".join(outs) + ") = \\\n")
    s.write("        " + "applied_f(")
    for p, attr in zip(pm_or_pf, args_out + args_in):
        s.write(p + "." + attr + ", ")
    for arg in kwargs_for_func:
        s.write(arg + ", ")
    s.write(")\n")

    s.write("    header = [")
    col_headers = ["'" + out + "'" for out in args_out]
    s.write(", ".join(col_headers))
    s.write("]\n")
    if len(args_out) == 1:
        s.write("    return DataFrame(data=outputs,"
                "columns=header)")
    else:
        s.write("    return DataFrame(data=np.column_stack("
                "outputs),columns=header)")

    return s.getvalue()


def make_apply_function(func, out_args, in_args, parameters, do_jit=True,
                        **kwargs):
    """
    Takes a '_calc' function and creates the necessary Python code for an
    _apply style function. Will also jit the function if desired

    Parameters
    ----------
    func: the 'calc' style function

    out_args: list of out arguments for the apply function

    in_args: list of in arguments for the apply function

    parameters: iterable of which of the args (from in_args) are parameter
                variables (as opposed to column records). This influences
                how we construct the '_apply' function

    do_jit: Bool, if True, jit the resulting apply function

    Returns
    -------
    '_apply' style function

    """
    jitted_f = jit(**kwargs)(func)
    apfunc = create_apply_function_string(out_args, in_args, parameters)

    func_code = compile(apfunc, "<string>", "exec")
    fakeglobals = {}
    eval(func_code, {"jitted_f": jitted_f}, fakeglobals)
    if do_jit:
        return jit(**kwargs)(fakeglobals['ap_func'])
    else:
        return fakeglobals['ap_func']


def apply_jit(dtype_sig_out, dtype_sig_in, parameters=None, **kwargs):
    """
    make a decorator that takes in a _calc-style function, handle
    the apply step
    """
    if not parameters:
        parameters = []

    def make_wrapper(func):
        theargs = inspect.getargspec(func).args
        jitted_f = jit(**kwargs)(func)

        jitted_apply = make_apply_function(func, dtype_sig_out,
                                           dtype_sig_in, parameters,
                                           **kwargs)

        def wrapper(*args, **kwargs):
            in_arrays = []
            out_arrays = []
            for farg in theargs:
                if hasattr(args[0], farg):
                    in_arrays.append(getattr(args[0], farg))
                else:
                    in_arrays.append(getattr(args[1], farg))

            for farg in dtype_sig_out:
                if hasattr(args[0], farg):
                    out_arrays.append(getattr(args[0], farg))
                else:
                    out_arrays.append(getattr(args[1], farg))

            final_array = out_arrays + in_arrays
            ans = jitted_apply(*final_array)
            return ans

        return wrapper
    return make_wrapper


def iterate_jit(parameters=None, **kwargs):
    """
    make a decorator that takes in a _calc-style function, create a
    function that handles the "high-level" function and the "_apply"
    style function

    Note: perhaps a better "bigger picture" description of what this does?
    """
    if not parameters:
        parameters = []

    def make_wrapper(func):

        # Step 1. Wrap this function in apply_jit
        # from apply_jit

        # Get the input arguments from the function
        in_args = inspect.getargspec(func).args
        try:
            jit_args = inspect.getargspec(jit).args + ['nopython']
        except TypeError:
            #print ("This should only be seen in RTD, if not install numba!")
            return func

        kwargs_for_func = toolz.keyfilter(in_args.__contains__, kwargs)
        kwargs_for_jit = toolz.keyfilter(jit_args.__contains__, kwargs)

        # Any name that is a taxcalc parameter (or the special case 'puf'
        # Boolean flag is given special treatment. Identify those names here
        allowed_parameters = list(default_data(metadata=True).keys())
        allowed_parameters += list(arg[1:] for arg in  default_data(metadata=True).keys())
        allowed_parameters.append("puf")
        additional_parameters = [arg for arg in in_args if arg in allowed_parameters]
        additional_parameters += parameters
        # Remote duplicates
        all_parameters = list(set(additional_parameters))

        src = inspect.getsourcelines(func)[0]

        # Discover the return arguments by walking
        # the AST of the function
        all_returned_vals = []
        gnr = GetReturnNode()
        all_out_args = None
        for node in ast.walk(ast.parse(''.join(src))):
            all_out_args = gnr.visit(node)
            if all_out_args:
                break

        if not all_out_args:
            raise ValueError("Can't find return statement in function!")

        # Now create the apply jitted function
        applied_jitted_f = make_apply_function(func,
                                               list(reversed(all_out_args)),
                                               in_args,
                                               parameters=all_parameters,
                                               do_jit=True,
                                               **kwargs_for_jit)

        def wrapper(*args, **kwargs):
            in_arrays = []
            out_arrays = []
            pm_or_pf = []
            for farg in all_out_args + in_args:
                if hasattr(args[0], farg):
                    in_arrays.append(getattr(args[0], farg))
                    pm_or_pf.append("pm")
                elif hasattr(args[1], farg):
                    in_arrays.append(getattr(args[1], farg))
                    pm_or_pf.append("pf")
                elif farg not in kwargs_for_func:
                    raise ValueError("Unknown arg: " + farg)

            # Create the high level function
            high_level_func = create_toplevel_function_string(all_out_args,
                                                              list(in_args),
                                                              pm_or_pf,
                                                              kwargs_for_func)
            func_code = compile(high_level_func, "<string>", "exec")
            fakeglobals = {}
            eval(func_code, {"applied_f": applied_jitted_f}, fakeglobals)
            high_level_fn = fakeglobals['hl_func']

            ans = high_level_fn(*args, **kwargs)
            return ans

        return wrapper

    return make_wrapper

A Visitor to get the return tuple names from a calc-style function Extracts numpy arrays from caller arguments and passes them     to guvectorized numba functions # np_arrays = [getattr(args[0], i).values for i in theargs] Extracts numpy arrays from caller arguments and passes them     to vectorized numba functions Extracts particular numpy arrays from caller argments and passes     them to guvectorize. Goes one step further than dataframe_guvectorize     by looking for the column names in the dataframe and just extracting those Create a string for a function of the form::          def ap_fuc(x_0, x_1, x_2, ...):             for i in range(len(x_0)):                 x_0[i], ... = jitted_f(x_j[i], ...)             return x_0[i], ...      where the specific args to jitted_f and the number of     values to return is determined by sigout and sigin      Parameters     ----------     sigout: iterable of the out arguments      sigin: iterable of the in arguments      parameters: iterable of which of the args (from in_args) are parameter                 variables (as opposed to column records). This influences                 how we construct the '_apply' function      Returns     -------     a String representing the function Create a string for a function of the form::          def hl_func(x_0, x_1, x_2, ...):             outputs = (...) = calc_func(...)             header = [...]             return DataFrame(data, columns=header)      where the specific args to jitted_f and the number of     values to return is destermined by sigout and sigin      Parameters     ----------     args_out: iterable of the out arguments      args_in: iterable of the in arguments      pm_or_pf: iterable of strings for object that holds each arg      kwargs_for_func: dictionary of keyword args for the function      Returns     -------     a String representing the function Takes a '_calc' function and creates the necessary Python code for an     _apply style function. Will also jit the function if desired      Parameters     ----------     func: the 'calc' style function      out_args: list of out arguments for the apply function      in_args: list of in arguments for the apply function      parameters: iterable of which of the args (from in_args) are parameter                 variables (as opposed to column records). This influences                 how we construct the '_apply' function      do_jit: Bool, if True, jit the resulting apply function      Returns     -------     '_apply' style function make a decorator that takes in a _calc-style function, handle     the apply step make a decorator that takes in a _calc-style function, create a     function that handles the "high-level" function and the "_apply"     style function      Note: perhaps a better "bigger picture" description of what this does? # Step 1. Wrap this function in apply_jit # from apply_jit # Get the input arguments from the function #print ("This should only be seen in RTD, if not install numba!") # Any name that is a taxcalc parameter (or the special case 'puf' # Boolean flag is given special treatment. Identify those names here # Remote duplicates # Discover the return arguments by walking # the AST of the function # Now create the apply jitted function # Create the high level function