from __future__ import absolute_import, division, print_function

import math
import itertools
import operator
import pytest
from datetime import datetime, date
import datashape
from collections import Iterator, Iterable

import blaze
from blaze.compute.python import (nunique, mean, rrowfunc, rowfunc,
                                  reduce_by_funcs, optimize)
from blaze import dshape
from blaze.compute.core import compute, compute_up, pre_compute
from blaze.expr import (symbol, by, merge, join, count, distinct,
                        Apply, sum, min, max, any, summary,
                        count, std, head, transform)
import numpy as np

from blaze import cos, sin
from blaze.compatibility import builtins
from blaze.utils import raises


t = symbol('t', 'var * {name: string, amount: int, id: int}')


data = [['Alice', 100, 1],
        ['Bob', 200, 2],
        ['Alice', 50, 3]]


tbig = symbol('tbig', 'var * {name: string, sex: string[1], amount: int, id: int}')


databig = [['Alice', 'F', 100, 1],
           ['Alice', 'F', 100, 3],
           ['Drew', 'F', 100, 4],
           ['Drew', 'M', 100, 5],
           ['Drew', 'M', 200, 5]]


def test_dispatched_rowfunc():
    cw = optimize(t['amount'] + 100, [])
    assert rowfunc(t)(t) == t
    assert rowfunc(cw)(('Alice', 100, 1)) == 200


def test_reduce_by_funcs():
    e = summary(number=t.id.max(), sum=t.amount.sum())
    b = by(t, e)
    assert reduce_by_funcs(b)[2]([1,2,3], [4,5,6]) == (1, 7)


def test_symbol():
    assert compute(t, data) == data


def test_projection():
    assert list(compute(t['name'], data)) == [x[0] for x in data]


def test_eq():
    assert list(compute(t['amount'] == 100, data)) == [x[1] == 100 for x in data]


def test_selection():
    assert list(compute(t[t['amount'] == 0], data)) == \
                [x for x in data if x[1] == 0]
    assert list(compute(t[t['amount'] > 150], data)) == \
                [x for x in data if x[1] > 150]


def test_arithmetic():
    assert list(compute(t['amount'] + t['id'], data)) == \
                [b + c for a, b, c, in data]
    assert list(compute(t['amount'] * t['id'], data)) == \
                [b * c for a, b, c, in data]
    assert list(compute(t['amount'] % t['id'], data)) == \
                [b % c for a, b, c, in data]


def test_unary_ops():
    for op in ('cos', 'sin', 'exp', 'ceil', 'floor', 'trunc', 'isnan'):
        f = getattr(blaze, op)
        pyf = getattr(math, op)
        result = list(compute(f(t['amount']), data))
        assert result == [pyf(x[1]) for x in data]


def test_neg():
    expr = optimize(-t.amount, [])
    assert list(compute(expr, data)) == [-x[1] for x in data]


def test_reductions():
    assert compute(sum(t['amount']), data) == 100 + 200 + 50
    assert compute(min(t['amount']), data) == 50
    assert compute(max(t['amount']), data) == 200
    assert compute(nunique(t['amount']), data) == 3
    assert compute(nunique(t['name']), data) == 2
    assert compute(count(t['amount']), data) == 3
    assert compute(any(t['amount'] > 150), data) is True
    assert compute(any(t['amount'] > 250), data) is False
    assert compute(t.amount[0], data) == 100