# -*- coding: utf-8 -*-
import logging
import collections
import ast
import _ast
import string
import numpy as np
import math
import sys
import six

logger = logging.getLogger("expr")
logger.setLevel(logging.ERROR)


class Base(object):
    def __neg__(self):
        return Neg(self)

    def __add__(self, rhs):
        return Add(self, rhs)

    def __radd__(self, lhs):
        return Add(lhs, self)

    def __sub__(self, rhs):
        return Sub(self, rhs)

    def __rsub__(self, lhs):
        return Sub(lhs, self)

    def __mul__(self, rhs):
        return Mul(self, rhs)

    def __rmul__(self, lhs):
        return Mul(lhs, self)

    def __div__(self, rhs):
        return Div(self, rhs)

    def __pow__(self, rhs):
        return Pow(self, rhs)

    def walk(self, f):
        return f(self)


class Sqrt(Base):
    def __init__(self, value):
        self.value = check_const(value)

    def __repr__(self):
        return "sqrt(" + repr(self.value) + ")"


class Const(Base):
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return repr(self.value)

    def c_code(self):
        return repr(self.value)


def check_const(value):
    if isinstance(value, (int, float)):
        return Const(value)
    else:
        return value


class Neg(Base):
    def __init__(self, obj):
        self.obj = check_const(obj)

    def __repr__(self):
        return "(-" + repr(self.obj) + ")"

    def walk(self, f):
        # logger.debug("walk neg")
        return f(Neg(self.obj.walk(f)))


class Add(Base):
    def __init__(self, lhs, rhs):
        self.lhs = check_const(lhs)
        self.rhs = check_const(rhs)

    def __repr__(self):
        return "(" + repr(self.lhs) + " + " + repr(self.rhs) + ")"

    def walk(self, f):
        return f(Add(self.lhs.walk(f), self.rhs.walk(f)))

    def c_code(self):
        return "(" + self.lhs.c_code() + " + " + self.rhs.c_code() + ")"


class Sub(Base):
    def __init__(self, lhs, rhs):
        self.lhs = check_const(lhs)
        self.rhs = check_const(rhs)

    def __repr__(self):
        return "(" + repr(self.lhs) + " - " + repr(self.rhs) + ")"

    def walk(self, f):
        return f(Sub(self.lhs.walk(f), self.rhs.walk(f)))


class Div(Base):
    def __init__(self, lhs, rhs):
        self.lhs = check_const(lhs)
        self.rhs = check_const(rhs)

    def __repr__(self):
        return "(" + repr(self.lhs) + " / " + repr(self.rhs) + ")"

    def walk(self, f):
        return f(Div(self.lhs.walk(f), self.rhs.walk(f)))


class Mul(Base):
    def __init__(self, lhs, rhs):
        self.lhs = check_const(lhs)
        self.rhs = check_const(rhs)

    def __repr__(self):
        return "(" + repr(self.lhs) + " * " + repr(self.rhs) + ")"

    def c_code(self):
        return "(" + self.lhs.c_code() + " * " + self.rhs.c_code() + ")"

    def walk(self, f):
        return f(Mul(self.lhs.walk(f), self.rhs.walk(f)))


class Pow(Base):
    def __init__(self, lhs, rhs):
        self.lhs = check_const(lhs)
        self.rhs = check_const(rhs)

    def __repr__(self):
        return "(" + repr(self.lhs) + " ** " + repr(self.rhs) + ")"

    def walk(self, f):
        return f(Pow(self.lhs.walk(f), self.rhs.walk(f)))


class Function(Base):
    def __init__(self, var, args):
        self.var = var
        self.args = args
        # print args

    def __repr__(self):
        return self.var.name + "(" + ", ".join([repr(k) for k in self.args]) + ")"

    def walk(self, f):
        args = [arg.walk(f) for arg in self.args]
        return f(Function(self.var.walk(f), args))


class Slice(Base):
    def __init__(self, var, args):
        self.var = var
        self.args = args
        try:
            len(self.args)
        except:
            self.args = (self.args,)

    def tovar(self):
        parts = [self.var.name]
        for arg in self.args:
            if isinstance(arg, (int, float)):
                parts.append(repr(arg).replace("-", "_min_"))
            elif isinstance(arg, Var):
                parts.append(arg.name)
            elif arg is None:
                parts.append("None")
            else:
                s = ""
                s += repr(arg.start)
                s += "_"
                s += repr(arg.stop)
                s += "_"
                s += repr(arg.step)
                parts.append(s)
        return "__".join(parts)

    def __repr__(self):
        parts = []
        # logger.debug("args: %r" % self.args)
        for arg in self.args:
            if isinstance(arg, (int, float)):
                parts.append(repr(arg))
            elif isinstance(arg, Var):
                parts.append(arg.name)
            elif arg is None:
                parts.append("")
            else:
                s = ""
                if arg.start is not None:
                    s += repr(arg.start)
                s += ":"
                if arg.stop is not None:
                    s += repr(arg.stop)
                s += ":"
                if arg.step is not None:
                    s += repr(arg.step)
                if s == "::":
                    s = ":"
                parts.append(s)
        return self.var.name + "[" + ",".join(parts) + "]"


class Var(Base):
    def __init__(self, name):
        self.name = name

    def __getitem__(self, args):
        logger.debug("[%s] Var.__getitem__(%r)" % (self.name, args,))
        return Slice(self, args)

    def __call__(self, *args):
        logger.debug("[%s] Var.__call__(%r)" % (self.name, args,))
        return Function(self, args)

    def __str__(self):
        return "Var(%r)" % self.name

    def __repr__(self):
        return self.name

    def c_code(self):
        return self.name

    def walk(self, f):
        return f(self)


class Scope(object):
    def __init__(self, name):
        self.name = name
        self.resolved = collections.OrderedDict()

    def write(self, f):
        for block_index, (key, value) in enumerate(self.resolved.items()):
            # this assumes the var is part of a 'block'
            f.write("double %s = blocks[%i][index];\n" % (key, block_index))

    def __getitem__(self, name):
        logger.debug("[%s] Scope.__getitem__(%r)" % (self.name, name,))
        if name not in self.resolved:
            self.resolved[name] = Var(name)
        else:
            logger.debug("[%s] Scope.__getitem__(%r) (resolved again)" % (self.name, name,))
        return self.resolved[name]

        # return Scope(name)


def distance(args):
    if len(args) == 1:
        result = args[0]
    else:
        result = Add(args[0]**2, args[1]**2)
        for arg in args[2:]:
            result = Add(result, arg**2)
        result = Sqrt(result)
    return result


macros = {"l2": distance}


def translate(expression, replacements={}):
    scope = Scope(None)
    newvars = {}
    used_vars = set()
    result = eval(expression, {}, scope)
    # for key, value in replacements:
    translated_replacements = {}
    for i, (key, expression) in enumerate(replacements.items()):
        current_replacements = collections.OrderedDict(list(replacements.items())[:i])
        # print "*" * 70
        translated_replacements[key] = translate(expression, current_replacements)[0]
        # print "TRANSLATE:", expression, translated_replacements[key], current_replacements

    def slice_to_var(slice):
        newname = slice.tovar()
        logger.debug("replace %s by %s" % (slice.var.name, newname))
        newvars[newname] = slice
        return Var(newname)

    def walker(expr):
        # print ">>>>", expr, type(expr)
        if isinstance(expr, Slice):
            return slice_to_var(expr)
        elif isinstance(expr, Function):
            # logger.debug("expr: " + expr.var.name)
            if expr.var.name in macros:
                return macros[expr.var.name](expr.args)
            else:
                return expr
        elif isinstance(expr, Var):
            # logger.debug("var: " + expr.name)
            used_vars.add(expr.name)
            if expr.name in translated_replacements:
                return translated_replacements[expr.name]
            else:
                return expr
        else:
            return expr
    # print "walking.."
    newresult = result.walk(walker)
    return newresult, newvars  # , used_vars


class MathExpression(object):
    def __init__(self, expression_string, macros):
        self.expression_string = expression_string
        self.scope = Scope()
        self.macros = macros
        for macro in self.macros:
            scope.add_macro(macro)
        self.expression = eval(self.expression_string)

    def evaluate(self, values):
        return eval(self.expression, math.__dict__, values)

    def evaluate_numpy(self, values):
        return eval(self.expression, np.__dict__, values)

    def variables(self):
        eval(self.expression, np.__dict__, scope)


valid_binary_operators = [_ast.Add, _ast.Sub, _ast.Mult, _ast.Pow, _ast.Div, _ast.FloorDiv, _ast.BitAnd, _ast.BitOr, _ast.BitXor, _ast.Mod]
valid_compare_operators = [_ast.Lt, _ast.LtE, _ast.Gt, _ast.GtE, _ast.Eq, _ast.NotEq]
valid_unary_operators = [_ast.USub, _ast.UAdd]
valid_id_characters = string.ascii_letters + string.digits + "_"
valid_functions = "sin cos".split()


def math_parse(expression, macros=[]):
    # TODO: validate macros?
    node = ast.parse(expression)
    if len(node.body) != 1:
        raise ValueError("expected one expression, got %r" % len(node.body))
    expr = node.body[0]
    if not isinstance(expr, _ast.Expr):
        raise ValueError("expected an expression got a %r" % type(node.body))

    validate_expression(expr.value)
    return MathExpression(expression, macros)


last_func = None


def validate_expression(expr, variable_set, function_set=[], names=None):
    global last_func
    names = names if names is not None else []
    if isinstance(expr, six.string_types):
        node = ast.parse(expr)
        if len(node.body) != 1:
            raise ValueError("expected one expression, got %r" % len(node.body))
        first_expr = node.body[0]
        if not isinstance(first_expr, _ast.Expr):
            raise ValueError("expected an expression got a %r" % type(node.body))
        validate_expression(first_expr.value, variable_set, function_set, names)
    elif isinstance(expr, _ast.BinOp):
        if expr.op.__class__ in valid_binary_operators:
            validate_expression(expr.right, variable_set, function_set, names)
            validate_expression(expr.left, variable_set, function_set, names)
        else:
            raise ValueError("Binary operator not allowed: %r" % expr.op)
    elif isinstance(expr, _ast.UnaryOp):
        if expr.op.__class__ in valid_unary_operators:
            validate_expression(expr.operand, variable_set, function_set, names)
        else:
            raise ValueError("Unary operator not allowed: %r" % expr.op)
    elif isinstance(expr, _ast.Name):
        validate_id(expr.id)
        if expr.id not in variable_set:
            raise NameError("variable %r is not defined (available are: %s)" % (expr.id, ", ".join(list(variable_set))))
        names.append(expr.id)
    elif isinstance(expr, _ast.Num):
        pass  # numbers are fine
    elif isinstance(expr, _ast.Call):
        validate_func(expr.func, function_set)
        last_func = expr
        for arg in expr.args:
            validate_expression(arg, variable_set, function_set, names)
        for arg in expr.keywords:
            validate_expression(arg, variable_set, function_set, names)
    elif isinstance(expr, _ast.Compare):
        validate_expression(expr.left, variable_set, function_set, names)
        for op in expr.ops:
            if op.__class__ not in valid_compare_operators:
                raise ValueError("Compare operator not allowed: %r" % op)
        for comparator in expr.comparators:
            validate_expression(comparator, variable_set, function_set, names)
    elif isinstance(expr, _ast.keyword):
        validate_expression(expr.value, variable_set, function_set, names)
    elif isinstance(expr, _ast.Subscript):
        validate_expression(expr.value, variable_set, function_set, names)
        if isinstance(expr.slice.value, _ast.Num):
            pass  # numbers are fine
        else:
            raise ValueError("Only subscript/slices with numbers allowed, not: %r" % expr.slice.value)
    else:
        last_func = expr
        raise ValueError("Unknown expression type: %r" % type(expr))


def validate_func(name, function_set):
    if name.id not in function_set:
        raise NameError("function %r is not defined" % name.id)
    # if name.id not in valid_functions:
    # raise ValueError, "invalid or unknown function %r" % name.id


def validate_id(id):
    for char in id:
        if char not in valid_id_characters:
            raise ValueError("invalid character %r in id %r" % (char, id))


expressions = [expr for expr in """
a * b
a + b
sin(a)
a**b - sin(a**2, b--a)
""".split("\n") if expr]

if __name__ == "__main__":
    values = dict(a=1, b=2, pi=math.pi)
    for expression in expressions:
        ex = math_parse(expression)
        print((expression, "=", ex.evaluate(values)))
    values = dict(a=np.array([1, 2]), b=np.array([2, 3]), pi=math.pi)
    for expression in expressions:
        ex = math_parse(expression)
        print(expression, "=", ex.evaluate_numpy(values))
    import types
    import numba

    numba.jit
    types.FunctionType()
    sys.exit(0)
    import vaex as vx
    from setuptools import setup, Extension
    ds = vx.example()
    scope = Scope(ds)
    result = eval("x + y*3", {}, scope)
    import StringIO
    f = StringIO.StringIO()
    f.write("// this is generated code =====\n")
    scope.write(f)
    f.write("values[0] = %s;" % result.c_code())
    f.write("// end of generated code =====\n")
    print(f.getvalue())
    import os
    import platform
    from distutils.sysconfig import get_python_inc

    base_path = os.path.expanduser("~/.vaex/ext")
    # build_path = os.path.expanduser("~/.vaex/ext/build")
    # build_path = os.path.expanduser("~/.vaex/ext/temp")
    if not os.path.exists(base_path):
        os.makedirs(base_path)

    filename = os.path.join(base_path, "code.cpp")
    template = file(os.path.expanduser("src/vaex/vaexfast_template.cpp")).read()
    with open(filename, "w") as file:
        file.write(template.replace("{code_nd}", f.getvalue()))
    try:
        import numpy
        numdir = os.path.dirname(numpy.__file__)
    except:
        numdir = None

    if numdir is None:
        print("numpy not found, cannot install")

    include_dirs = []
    library_dirs = []
    libraries = []
    defines = []
    if "darwin" in platform.system().lower():
        extra_compile_args = ["-mfpmath=sse", "-O3", "-funroll-loops"]
    else:
        # extra_compile_args = ["-mfpmath=sse", "-msse4", "-Ofast", "-flto", "-march=native", "-funroll-loops"]
        extra_compile_args = ["-mfpmath=sse", "-msse4", "-Ofast", "-flto", "-funroll-loops"]
        # extra_compile_args = ["-mfpmath=sse", "-O3", "-funroll-loops"]
        extra_compile_args = ["-mfpmath=sse", "-msse4a", "-O3", "-funroll-loops"]
    extra_compile_args.extend(["-std=c++0x"])

    include_dirs.append(os.path.join(get_python_inc(plat_specific=1), "numpy"))
    if numdir is not None:
        include_dirs.append(os.path.join(numdir, "core", "include"))

    ex = Extension("vaexfast_user", [filename],
                   include_dirs=include_dirs,
                   library_dirs=library_dirs,
                   libraries=libraries,
                   define_macros=defines,
                   extra_compile_args=extra_compile_args
                   )
    from setuptools import setup, Extension
    setup(script_args=["build_ext", "--build-lib=" + base_path], ext_modules=[ex])

    if 0:
        logger.setLevel(logging.DEBUG)

        expr1 = "log(a[:,index]) + index"
        expr1 = "log(a[index]) + index"
        expr1 = "log(a[-1]) - index + 1 + 1**2 + index**index + d(a,b) + 1.2"
        expr1 = "-2.5 * log10(SPECTROFLUX[:,2])"
        scope = Scope(None)
        result = eval(expr1, {}, scope)
        print(("expression", expr1))
        print(("should equal", repr(result)))
        print("translating")
        print()
        newexpr, vars = translate(expr1)
        print((newexpr, vars))

        expr = "xc + yc"
        newexpr, vars = translate(expr, {"xc": "x+1", "yc": "sqrt(x**2+y**2)"})
        print((newexpr, vars))

        expr = "(xc + yc)/r"
        replacements = collections.OrderedDict()
        replacements["xc"] = "x+1"
        replacements["yc"] = "y-1"
        replacements["r"] = "sqrt(xc**2+yc**2)"
        print((list(replacements.items())))
        newexpr, vars = translate(expr, replacements)
        print((newexpr, vars))
