#!/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)

[((29, 27, 29, 35), 'agate.data_types.Number', 'Number', ({}, {}), '()', False, 'from agate.data_types import Number, Text\n'), ((30, 25, 30, 31), 'agate.data_types.Text', 'Text', ({}, {}), '()', False, 'from agate.data_types import Number, Text\n'), ((36, 16, 36, 70), 'agate.Table', 'Table', ({(36, 22, 36, 31): 'self.rows', (36, 33, 36, 50): 'self.column_names', (36, 52, 36, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((53, 16, 53, 70), 'agate.Table', 'Table', ({(53, 22, 53, 31): 'self.rows', (53, 33, 53, 50): 'self.column_names', (53, 52, 53, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((68, 16, 68, 70), 'agate.Table', 'Table', ({(68, 22, 68, 31): 'self.rows', (68, 33, 68, 50): 'self.column_names', (68, 52, 68, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((83, 16, 83, 70), 'agate.Table', 'Table', ({(83, 22, 83, 31): 'self.rows', (83, 33, 83, 50): 'self.column_names', (83, 52, 83, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((89, 16, 89, 70), 'agate.Table', 'Table', ({(89, 22, 89, 31): 'self.rows', (89, 33, 89, 50): 'self.column_names', (89, 52, 89, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((102, 16, 102, 70), 'agate.Table', 'Table', ({(102, 22, 102, 31): 'self.rows', (102, 33, 102, 50): 'self.column_names', (102, 52, 102, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((118, 16, 118, 70), 'agate.Table', 'Table', ({(118, 22, 118, 31): 'self.rows', (118, 33, 118, 50): 'self.column_names', (118, 52, 118, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((134, 16, 134, 70), 'agate.Table', 'Table', ({(134, 22, 134, 31): 'self.rows', (134, 33, 134, 50): 'self.column_names', (134, 52, 134, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((158, 16, 158, 70), 'agate.Table', 'Table', ({(158, 22, 158, 31): 'self.rows', (158, 33, 158, 50): 'self.column_names', (158, 52, 158, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((175, 16, 175, 70), 'agate.Table', 'Table', ({(175, 22, 175, 31): 'self.rows', (175, 33, 175, 50): 'self.column_names', (175, 52, 175, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((191, 16, 191, 70), 'agate.Table', 'Table', ({(191, 22, 191, 31): 'self.rows', (191, 33, 191, 50): 'self.column_names', (191, 52, 191, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((207, 16, 207, 70), 'agate.Table', 'Table', ({(207, 22, 207, 31): 'self.rows', (207, 33, 207, 50): 'self.column_names', (207, 52, 207, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((223, 16, 223, 70), 'agate.Table', 'Table', ({(223, 22, 223, 31): 'self.rows', (223, 33, 223, 50): 'self.column_names', (223, 52, 223, 69): 'self.column_types'}, {}), '(self.rows, self.column_names, self.column_types)', False, 'from agate import Table\n'), ((120, 52, 120, 62), 'agate.aggregations.Sum', 'Sum', ({(120, 56, 120, 61): '"""age"""'}, {}), "('age')", False, 'from agate.aggregations import Sum\n'), ((193, 56, 193, 72), 'agate.computations.Percent', 'Percent', ({(193, 64, 193, 71): '"""Count"""'}, {}), "('Count')", False, 'from agate.computations import Percent\n'), ((198, 21, 198, 32), 'decimal.Decimal', 'Decimal', ({(198, 29, 198, 31): '(50)'}, {}), '(50)', False, 'from decimal import Decimal\n'), ((199, 23, 199, 34), 'decimal.Decimal', 'Decimal', ({(199, 31, 199, 33): '(50)'}, {}), '(50)', False, 'from decimal import Decimal\n'), ((209, 65, 209, 81), 'agate.computations.Percent', 'Percent', ({(209, 73, 209, 80): '"""Count"""'}, {}), "('Count')", False, 'from agate.computations import Percent\n'), ((214, 21, 214, 32), 'decimal.Decimal', 'Decimal', ({(214, 29, 214, 31): '(50)'}, {}), '(50)', False, 'from decimal import Decimal\n'), ((225, 65, 225, 90), 'agate.computations.Percent', 'Percent', (), '', False, 'from agate.computations import Percent\n'), ((215, 49, 215, 61), 'decimal.Decimal', 'Decimal', ({(215, 57, 215, 60): '(100)'}, {}), '(100)', False, 'from decimal import Decimal\n'), ((215, 89, 215, 101), 'decimal.Decimal', 'Decimal', ({(215, 97, 215, 100): '(100)'}, {}), '(100)', False, 'from decimal import Decimal\n'), ((230, 47, 230, 59), 'decimal.Decimal', 'Decimal', ({(230, 55, 230, 58): '(100)'}, {}), '(100)', False, 'from decimal import Decimal\n'), ((231, 49, 231, 61), 'decimal.Decimal', 'Decimal', ({(231, 57, 231, 60): '(100)'}, {}), '(100)', False, 'from decimal import Decimal\n'), ((231, 89, 231, 101), 'decimal.Decimal', 'Decimal', ({(231, 97, 231, 100): '(100)'}, {}), '(100)', False, 'from decimal import Decimal\n'), ((215, 23, 215, 33), 'decimal.Decimal', 'Decimal', ({(215, 31, 215, 32): '(1)'}, {}), '(1)', False, 'from decimal import Decimal\n'), ((215, 36, 215, 46), 'decimal.Decimal', 'Decimal', ({(215, 44, 215, 45): '(6)'}, {}), '(6)', False, 'from decimal import Decimal\n'), ((215, 63, 215, 73), 'decimal.Decimal', 'Decimal', ({(215, 71, 215, 72): '(1)'}, {}), '(1)', False, 'from decimal import Decimal\n'), ((215, 76, 215, 86), 'decimal.Decimal', 'Decimal', ({(215, 84, 215, 85): '(3)'}, {}), '(3)', False, 'from decimal import Decimal\n'), ((230, 21, 230, 31), 'decimal.Decimal', 'Decimal', ({(230, 29, 230, 30): '(3)'}, {}), '(3)', False, 'from decimal import Decimal\n'), ((230, 34, 230, 44), 'decimal.Decimal', 'Decimal', ({(230, 42, 230, 43): '(8)'}, {}), '(8)', False, 'from decimal import Decimal\n'), ((231, 23, 231, 33), 'decimal.Decimal', 'Decimal', ({(231, 31, 231, 32): '(1)'}, {}), '(1)', False, 'from decimal import Decimal\n'), ((231, 36, 231, 46), 'decimal.Decimal', 'Decimal', ({(231, 44, 231, 45): '(8)'}, {}), '(8)', False, 'from decimal import Decimal\n'), ((231, 63, 231, 73), 'decimal.Decimal', 'Decimal', ({(231, 71, 231, 72): '(2)'}, {}), '(2)', False, 'from decimal import Decimal\n'), ((231, 76, 231, 86), 'decimal.Decimal', 'Decimal', ({(231, 84, 231, 85): '(8)'}, {}), '(8)', False, 'from decimal import Decimal\n')]