# -*- coding: utf-8 -*-
from __future__ import print_function, division, absolute_import
import ast
from functools import reduce

import numba
from numba import *
from numba import error
from numba import visitors, nodes
from numba import function_util
from numba.symtab import Variable
from numba.typesystem import is_obj
from numba import pyconsts

logger = logging.getLogger(__name__)

opmap = {
    ast.Eq        : pyconsts.Py_EQ,
    ast.NotEq     : pyconsts.Py_NE,
    ast.Lt        : pyconsts.Py_LT,
    ast.LtE       : pyconsts.Py_LE,
    ast.Gt        : pyconsts.Py_GT,
    ast.GtE       : pyconsts.Py_GE,
}

def build_boolop(right, left):
    node = ast.BoolOp(ast.And(), [left, right])
    return nodes.typednode(node, bool_)

def extract(complex_node):
    complex_node = nodes.CloneableNode(complex_node)

    real = nodes.ComplexAttributeNode(complex_node, 'real')
    imag = nodes.ComplexAttributeNode(complex_node.clone, 'imag')

    return real, imag

def compare(lhs, rhs):
    result = ast.Compare(lhs, [ast.Eq()], [rhs])
    return nodes.typednode(result, bool_)

class SpecializeComparisons(visitors.NumbaTransformer):
    """
    Rewrite cascaded ast.Compare nodes to a sequence of boolean operations
    ANDed together:

        a < b < c

    becomes

        a < b and b < c
    """

    def single_compare(self, node):
        rhs = node.comparators[0]

        if is_obj(node.left.type):
            node = self.single_compare_objects(node)

        elif node.left.type.is_pointer and rhs.type.is_pointer:
            # Coerce pointers to integer values before comparing
            node.left = nodes.CoercionNode(node.left, Py_uintptr_t)
            node.comparators = [nodes.CoercionNode(rhs, Py_uintptr_t)]

        elif node.left.type.is_complex and rhs.type.is_complex:
            real1, imag1 = extract(node.left)
            real2, imag2 = extract(rhs)
            lhs = compare(real1, real2)
            rhs = compare(imag1, imag2)
            result = ast.BoolOp(ast.And(), [lhs, rhs])
            node = nodes.typednode(result, bool_)

        return node

    def single_compare_objects(self, node):
        op = type(node.ops[0])
        if op not in opmap:
            raise error.NumbaError(
                    node, "%s comparisons not yet implemented" % (op,))

        # Build arguments for PyObject_RichCompareBool
        operator = nodes.const(opmap[op], int_)
        args = [node.left, node.comparators[0], operator]

        # Call PyObject_RichCompareBool
        compare = function_util.external_call(self.context,
                                              self.llvm_module,
                                              'PyObject_RichCompare',
                                              args=args)

        # Coerce int result to bool
        return nodes.CoercionNode(compare, node.type)

    def visit_Compare(self, node):
        "Reduce cascaded comparisons into single comparisons"

        # Process children
        self.generic_visit(node)

        compare_nodes = []
        comparators = [nodes.CloneableNode(c) for c in node.comparators]

        if len(node.comparators) > 1:
            if node.type.is_array:
                raise error.NumbaError(
                        node, "Cannot determine truth value of boolean array "
                              "(use any or all)")

        # Build comparison nodes
        left = node.left
        for op, right in zip(node.ops, comparators):
            node = ast.Compare(left=left, ops=[op], comparators=[right])

            # Set result type of comparison:
            #     bool array of array comparison
            #     bool otherwise

            if left.type.is_array or right.type.is_array:
                # array < x -> Array(bool_, array.ndim)
                result_type = self.env.crnt.typesystem.promote(
                    left.type, right.type)
            else:
                result_type = bool_

            nodes.typednode(node, result_type)

            # Handle comparisons specially based on their types
            node = self.single_compare(node)
            compare_nodes.append(node)

            left = right.clone

        # AND the comparisons together
        node = reduce(build_boolop, reversed(compare_nodes))

        return node

@@ -974,24 +974,28 @@
 compare(lhs,
+ op,
  rhs):%0A    r
@@ -1024,16 +1024,10 @@
 s, %5B
-ast.Eq()
+op
 %5D, %5B
@@ -1879,16 +1879,84 @@
 ct(rhs)%0A
+            op = type(node.ops%5B0%5D)%0A            if op == ast.Eq:%0A    
         
@@ -1983,16 +1983,26 @@
 al1,
+ ast.Eq(),
  real2)%0A
     
@@ -1997,16 +1997,20 @@
  real2)%0A
+    
         
@@ -2037,16 +2037,26 @@
 ag1,
+ ast.Eq(),
  imag2)%0A
     
@@ -2051,16 +2051,20 @@
  imag2)%0A
+    
         
@@ -2106,24 +2106,398 @@
 %5Blhs, rhs%5D)%0A
+            elif op == ast.NotEq:%0A                lhs = compare(real1, ast.NotEq(), real2)%0A                rhs = compare(imag1, ast.NotEq(), imag2)%0A                result = ast.BoolOp(ast.Or(), %5Blhs, rhs%5D)%0A            else:%0A                raise NotImplementedError(%22ordered comparisons are not %22%0A                                          %22implemented for complex numbers%22)%0A
             

implement != for complex numbers (fix issue 256)