#!/usr/bin/env python
# -*- coding: utf8 -*-

import sys

try:
    from cdecimal import Decimal
except ImportError:  # pragma: no cover
    from decimal import Decimal

from agate import Table
from agate.aggregations import Sum
from agate.computations import Percent
from agate.data_types import Number, Text
from agate.testcase import AgateTestCase


class TestPivot(AgateTestCase):
    def setUp(self):
        self.rows = (
            ('joe', 'white', 'male', 20, 'blue'),
            ('jane', 'white', 'female', 20, 'blue'),
            ('josh', 'black', 'male', 20, 'blue'),
            ('jim', 'latino', 'male', 25, 'blue'),
            ('julia', 'white', 'female', 25, 'green'),
            ('joan', 'asian', 'female', 25, 'green')
        )

        self.number_type = Number()
        self.text_type = Text()

        self.column_names = ['name', 'race', 'gender', 'age', 'color']
        self.column_types = [self.text_type, self.text_type, self.text_type, self.number_type, self.text_type]

    def test_pivot(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('race', 'gender')

        pivot_rows = (
            ('white', 1, 2),
            ('black', 1, 0),
            ('latino', 1, 0),
            ('asian', 0, 1)
        )

        self.assertColumnNames(pivot_table, ['race', 'male', 'female'])
        self.assertRowNames(pivot_table, ['white', 'black', 'latino', 'asian'])
        self.assertColumnTypes(pivot_table, [Text, Number, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_by_lambda(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot(lambda r: r['gender'])

        pivot_rows = (
            ('male', 3),
            ('female', 3)
        )

        self.assertColumnNames(pivot_table, ['group', 'Count'])
        self.assertRowNames(pivot_table, ['male', 'female'])
        self.assertColumnTypes(pivot_table, [Text, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_by_lambda_group_name(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot(lambda r: r['gender'], key_name='gender')

        pivot_rows = (
            ('male', 3),
            ('female', 3)
        )

        self.assertColumnNames(pivot_table, ['gender', 'Count'])
        self.assertRowNames(pivot_table, ['male', 'female'])
        self.assertColumnTypes(pivot_table, [Text, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_by_lambda_group_name_sequence_invalid(self):
        table = Table(self.rows, self.column_names, self.column_types)

        with self.assertRaises(ValueError):
            table.pivot(['race', 'gender'], key_name='foo')

    def test_pivot_no_key(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot(pivot='gender')

        pivot_rows = (
            (3, 3),
        )

        self.assertColumnNames(pivot_table, ['male', 'female'])
        self.assertColumnTypes(pivot_table, [Number, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_no_pivot(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('race')

        pivot_rows = (
            ('white', 3),
            ('black', 1),
            ('latino', 1),
            ('asian', 1)
        )

        self.assertColumnNames(pivot_table, ['race', 'Count'])
        self.assertColumnTypes(pivot_table, [Text, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_sum(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('race', 'gender', Sum('age'))

        pivot_rows = (
            ('white', 20, 45),
            ('black', 20, 0),
            ('latino', 25, 0),
            ('asian', 0, 25)
        )

        self.assertColumnNames(pivot_table, ['race', 'male', 'female'])
        self.assertColumnTypes(pivot_table, [Text, Number, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_multiple_keys(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot(['race', 'gender'], 'age')

        pivot_rows = (
            ('white', 'male', 1, 0),
            ('white', 'female', 1, 1),
            ('black', 'male', 1, 0),
            ('latino', 'male', 0, 1),
            ('asian', 'female', 0, 1),
        )

        self.assertRows(pivot_table, pivot_rows)
        self.assertColumnNames(pivot_table, ['race', 'gender', '20', '25'])
        self.assertRowNames(pivot_table, [
            ('white', 'male'),
            ('white', 'female'),
            ('black', 'male'),
            ('latino', 'male'),
            ('asian', 'female'),
        ])
        self.assertColumnTypes(pivot_table, [Text, Text, Number, Number])

    def test_pivot_multiple_keys_no_pivot(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot(['race', 'gender'])

        pivot_rows = (
            ('white', 'male', 1),
            ('white', 'female', 2),
            ('black', 'male', 1),
            ('latino', 'male', 1),
            ('asian', 'female', 1),
        )

        self.assertRows(pivot_table, pivot_rows)
        self.assertColumnNames(pivot_table, ['race', 'gender', 'Count'])
        self.assertColumnTypes(pivot_table, [Text, Text, Number])

    def test_pivot_default_value(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('race', 'gender', default_value=None)

        pivot_rows = (
            ('white', 1, 2),
            ('black', 1, None),
            ('latino', 1, None),
            ('asian', None, 1)
        )

        self.assertColumnNames(pivot_table, ['race', 'male', 'female'])
        self.assertColumnTypes(pivot_table, [Text, Number, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_compute(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('gender', computation=Percent('Count'))

        pivot_table.print_table(output=sys.stdout)

        pivot_rows = (
            ('male', Decimal(50)),
            ('female', Decimal(50)),
        )

        self.assertColumnNames(pivot_table, ['gender', 'Percent'])
        self.assertColumnTypes(pivot_table, [Text, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_compute_pivots(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('gender', 'color', computation=Percent('Count'))

        pivot_table.print_table(output=sys.stdout)

        pivot_rows = (
            ('male', Decimal(50), 0),
            ('female', Decimal(1) / Decimal(6) * Decimal(100), Decimal(1) / Decimal(3) * Decimal(100)),
        )

        self.assertColumnNames(pivot_table, ['gender', 'blue', 'green'])
        self.assertColumnTypes(pivot_table, [Text, Number, Number])
        self.assertRows(pivot_table, pivot_rows)

    def test_pivot_compute_kwargs(self):
        table = Table(self.rows, self.column_names, self.column_types)

        pivot_table = table.pivot('gender', 'color', computation=Percent('Count', total=8))

        pivot_table.print_table(output=sys.stdout)

        pivot_rows = (
            ('male', Decimal(3) / Decimal(8) * Decimal(100), 0),
            ('female', Decimal(1) / Decimal(8) * Decimal(100), Decimal(2) / Decimal(8) * Decimal(100)),
        )

        self.assertColumnNames(pivot_table, ['gender', 'blue', 'green'])
        self.assertColumnTypes(pivot_table, [Text, Number, Number])
        self.assertRows(pivot_table, pivot_rows)

[((778, 786), 'agate.data_types.Number', 'Number', ([], {}), '()\n', (784, 786), False, 'from agate.data_types import Number, Text\n'), ((812, 818), 'agate.data_types.Text', 'Text', ([], {}), '()\n', (816, 818), False, 'from agate.data_types import Number, Text\n'), ((1045, 1099), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (1050, 1099), False, 'from agate import Table\n'), ((1626, 1680), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (1631, 1680), False, 'from agate import Table\n'), ((2123, 2177), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (2128, 2177), False, 'from agate import Table\n'), ((2657, 2711), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (2662, 2711), False, 'from agate import Table\n'), ((2867, 2921), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (2872, 2921), False, 'from agate import Table\n'), ((3255, 3309), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (3260, 3309), False, 'from agate import Table\n'), ((3711, 3765), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (3716, 3765), False, 'from agate import Table\n'), ((4233, 4287), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (4238, 4287), False, 'from agate import Table\n'), ((5054, 5108), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (5059, 5108), False, 'from agate import Table\n'), ((5619, 5673), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (5624, 5673), False, 'from agate import Table\n'), ((6147, 6201), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (6152, 6201), False, 'from agate import Table\n'), ((6670, 6724), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (6675, 6724), False, 'from agate import Table\n'), ((7286, 7340), 'agate.Table', 'Table', (['self.rows', 'self.column_names', 'self.column_types'], {}), '(self.rows, self.column_names, self.column_types)\n', (7291, 7340), False, 'from agate import Table\n'), ((3819, 3829), 'agate.aggregations.Sum', 'Sum', (['"""age"""'], {}), "('age')\n", (3822, 3829), False, 'from agate.aggregations import Sum\n'), ((6259, 6275), 'agate.computations.Percent', 'Percent', (['"""Count"""'], {}), "('Count')\n", (6266, 6275), False, 'from agate.computations import Percent\n'), ((6374, 6385), 'decimal.Decimal', 'Decimal', (['(50)'], {}), '(50)\n', (6381, 6385), False, 'from decimal import Decimal\n'), ((6411, 6422), 'decimal.Decimal', 'Decimal', (['(50)'], {}), '(50)\n', (6418, 6422), False, 'from decimal import Decimal\n'), ((6791, 6807), 'agate.computations.Percent', 'Percent', (['"""Count"""'], {}), "('Count')\n", (6798, 6807), False, 'from agate.computations import Percent\n'), ((6906, 6917), 'decimal.Decimal', 'Decimal', (['(50)'], {}), '(50)\n', (6913, 6917), False, 'from decimal import Decimal\n'), ((7407, 7432), 'agate.computations.Percent', 'Percent', (['"""Count"""'], {'total': '(8)'}), "('Count', total=8)\n", (7414, 7432), False, 'from agate.computations import Percent\n'), ((6972, 6984), 'decimal.Decimal', 'Decimal', (['(100)'], {}), '(100)\n', (6979, 6984), False, 'from decimal import Decimal\n'), ((7012, 7024), 'decimal.Decimal', 'Decimal', (['(100)'], {}), '(100)\n', (7019, 7024), False, 'from decimal import Decimal\n'), ((7557, 7569), 'decimal.Decimal', 'Decimal', (['(100)'], {}), '(100)\n', (7564, 7569), False, 'from decimal import Decimal\n'), ((7624, 7636), 'decimal.Decimal', 'Decimal', (['(100)'], {}), '(100)\n', (7631, 7636), False, 'from decimal import Decimal\n'), ((7664, 7676), 'decimal.Decimal', 'Decimal', (['(100)'], {}), '(100)\n', (7671, 7676), False, 'from decimal import Decimal\n'), ((6946, 6956), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (6953, 6956), False, 'from decimal import Decimal\n'), ((6959, 6969), 'decimal.Decimal', 'Decimal', (['(6)'], {}), '(6)\n', (6966, 6969), False, 'from decimal import Decimal\n'), ((6986, 6996), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (6993, 6996), False, 'from decimal import Decimal\n'), ((6999, 7009), 'decimal.Decimal', 'Decimal', (['(3)'], {}), '(3)\n', (7006, 7009), False, 'from decimal import Decimal\n'), ((7531, 7541), 'decimal.Decimal', 'Decimal', (['(3)'], {}), '(3)\n', (7538, 7541), False, 'from decimal import Decimal\n'), ((7544, 7554), 'decimal.Decimal', 'Decimal', (['(8)'], {}), '(8)\n', (7551, 7554), False, 'from decimal import Decimal\n'), ((7598, 7608), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (7605, 7608), False, 'from decimal import Decimal\n'), ((7611, 7621), 'decimal.Decimal', 'Decimal', (['(8)'], {}), '(8)\n', (7618, 7621), False, 'from decimal import Decimal\n'), ((7638, 7648), 'decimal.Decimal', 'Decimal', (['(2)'], {}), '(2)\n', (7645, 7648), False, 'from decimal import Decimal\n'), ((7651, 7661), 'decimal.Decimal', 'Decimal', (['(8)'], {}), '(8)\n', (7658, 7661), False, 'from decimal import Decimal\n')]