import ast
from ast import (
    Assign,
    arg,
    FunctionDef,
    Module,
    Store,
    Subscript,
    Name,
    Load,
    Index,
    Num,
)
from ast import arguments as ast_arguments

from dolang.grammar import tree_to_ast
from dolang.factory import FlatFunctionFactory
from dolang.symbolic import parse_string
from dolang.codegen import to_source

import sys


def compile_factory(fff: FlatFunctionFactory):
    arguments = [*fff.arguments.keys()]
    funname = fff.funname

    unpacking = []
    for i, (arg_group_name, arg_group) in enumerate(fff.arguments.items()):
        for pos, sym in enumerate(arg_group):
            rhs = Subscript(
                value=Name(id=arg_group_name, ctx=Load()),
                slice=Index(Num(pos)),
                ctx=Load(),
            )
            val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs)
            unpacking.append(val)

    body = []

    for k, neq in fff.preamble.items():
        tree = parse_string(neq)
        val = tree_to_ast(tree).value
        line = Assign(targets=[Name(id=k, ctx=Store())], value=val)
        body.append(line)

    for n, (k, neq) in enumerate(fff.content.items()):
        tree = parse_string(neq)
        val = tree_to_ast(tree).value
        line = Assign(targets=[Name(id=k, ctx=Store())], value=val)
        body.append(line)
    #
    for n, (lhs, neq) in enumerate(fff.content.items()):
        line = Assign(
            targets=[
                Subscript(
                    value=Name(id="out", ctx=Load()), slice=Index(Num(n)), ctx=Store()
                )
            ],
            value=Name(id=lhs, ctx=Load()),
        )
        body.append(line)

    if sys.version_info >= (3, 8, 0):
        f = FunctionDef(
            name=funname,
            args=ast_arguments(
                posonlyargs=[],
                args=[arg(arg=a) for a in arguments] + [arg(arg="out")],
                vararg=None,
                kwarg=None,
                kwonlyargs=[],
                kw_defaults=[],
                defaults=[],
            ),
            body=unpacking + body,
            decorator_list=[],
        )
        mod = Module(body=[f], type_ignores=[])
    else:
        f = FunctionDef(
            name=funname,
            args=ast_arguments(
                args=[arg(arg=a) for a in arguments] + [arg(arg="out")],
                vararg=None,
                kwarg=None,
                kwonlyargs=[],
                kw_defaults=[],
                defaults=[],
            ),
            body=unpacking + body,
            decorator_list=[],
        )
        mod = Module(body=[f])

    mmod = ast.fix_missing_locations(mod)
    return mmod


def make_method_from_factory(
    fff: FlatFunctionFactory, vectorize=True, use_file=False, debug=False, compile=True
):
    mod = compile_factory(fff)

    if debug:
        from .codegen import to_source

        print(to_source(mod))

    if vectorize:
        coredims = [len(v) for k, v in fff.arguments.items()]
        signature = str.join(",", ["(n_{})".format(d) for d in coredims])
        n_out = len(fff.content)
        if n_out in coredims:
            signature += "->(n_{})".format(n_out)
            # ftylist = float64[:](*([float64[:]] * len(coredims)))
            fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
        else:
            signature += ",(n_{})".format(n_out)
            fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
    else:
        signature = None

    fun = eval_ast(mod)

    if compile:
        from numba import jit, guvectorize

        jfun = jit(fun, nopython=True)
        if vectorize:
            gufun = guvectorize([fty], signature, target="parallel", nopython=True)(fun)
            return jfun, gufun
    else:
        jfun = fun

    return jfun


def eval_ast(mod):
    context = {}

    import numpy

    context["inf"] = numpy.inf
    context["exp"] = numpy.exp
    context["log"] = numpy.log
    context["sin"] = numpy.sin
    context["cos"] = numpy.cos
    context["tan"] = numpy.tan
    context["tanh"] = numpy.tanh
    context["atan"] = numpy.arctan
    context["atanh"] = numpy.arctanh

    context["maximum"] = numpy.maximum
    context["minimum"] = numpy.minimum
    context["max"] = numpy.maximum
    context["min"] = numpy.minimum
    context["abs"] = numpy.abs

    name = mod.body[0].name
    mod = ast.fix_missing_locations(mod)
    code = compile(mod, "<string>", "exec")
    exec(code, context, context)
    fun = context[name]

    return fun
