from querybuilder.fields import (
    RankField, RowNumberField, DenseRankField, PercentRankField, CumeDistField, NTileField, LagField,
    LeadField, FirstValueField, LastValueField, NthValueField, NumStdDevField
)
from querybuilder.query import QueryWindow, Query
from querybuilder.tests.models import Order
from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str


class QueryWindowTest(QueryTestCase):
    def test_query_window(self):
        query_window = QueryWindow()
        query_str = query_window.get_sql()
        expected_query = 'OVER ()'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_query_window_partition(self):
        query_window = QueryWindow().partition_by('field_one')
        query_str = query_window.get_sql()
        expected_query = 'OVER (PARTITION BY field_one)'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_query_window_order(self):
        query_window = QueryWindow().order_by('field_one')
        query_str = query_window.get_sql()
        expected_query = 'OVER (ORDER BY field_one ASC)'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_query_window_partition_order(self):
        query_window = QueryWindow().partition_by(
            'field_one'
        ).order_by(
            'field_one'
        )
        query_str = query_window.get_sql()
        expected_query = 'OVER (PARTITION BY field_one ORDER BY field_one ASC)'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_query_window_partition_order_many(self):
        query_window = QueryWindow().partition_by(
            'field_one'
        ).partition_by(
            'field_two'
        ).order_by(
            'field_one'
        ).order_by(
            '-field_two'
        )
        query_str = query_window.get_sql()
        expected_query = 'OVER (PARTITION BY field_one, field_two ORDER BY field_one ASC, field_two DESC)'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))


class WindowFunctionTest(QueryTestCase):
    def test_rank_no_over(self):
        query = Query().from_table(
            table=Order,
            fields=[
                RankField()
            ]
        )
        query_str = query.get_sql()
        expected_query = 'SELECT RANK() AS "rank" FROM querybuilder_tests_order'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_rank_over(self):
        query = Query().from_table(
            table=Order,
            fields=[
                RankField(
                    over=QueryWindow()
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = 'SELECT RANK() OVER () AS "rank" FROM querybuilder_tests_order'
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_rank_over_order(self):
        query = Query().from_table(
            table=Order,
            fields=[
                'id',
                RankField(
                    over=QueryWindow().order_by(
                        'id'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.id, RANK() OVER (ORDER BY id ASC) AS "rank" FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_rank_over_partition(self):
        query = Query().from_table(
            table=Order,
            fields=[
                'id',
                RankField(
                    over=QueryWindow().partition_by(
                        'account_id'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.id, RANK() OVER (PARTITION BY account_id) AS "rank" FROM '
            'querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_row_number(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                RowNumberField(
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        ).order_by(
            'row_number'
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'ROW_NUMBER() OVER (ORDER BY margin DESC) AS "row_number" '
            'FROM querybuilder_tests_order '
            'ORDER BY row_number '
            'ASC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_rank(self):
        query = Query().from_table(
            table=Order,
            fields=[
                'id',
                RankField(
                    over=QueryWindow().partition_by(
                        'account_id'
                    ).order_by(
                        'id'
                    )
                )
            ]
        ).order_by(
            '-rank'
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.id, '
            'RANK() OVER (PARTITION BY account_id ORDER BY id ASC) AS "rank" '
            'FROM querybuilder_tests_order '
            'ORDER BY rank '
            'DESC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_dense_rank(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                DenseRankField(
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        ).order_by(
            'dense_rank'
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'DENSE_RANK() OVER (ORDER BY margin DESC) AS "dense_rank" '
            'FROM querybuilder_tests_order '
            'ORDER BY dense_rank '
            'ASC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_rank_percent(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                PercentRankField(
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        ).order_by(
            'percent_rank'
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'PERCENT_RANK() OVER (ORDER BY margin DESC) AS "percent_rank" '
            'FROM querybuilder_tests_order '
            'ORDER BY percent_rank '
            'ASC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_cume_dist(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                CumeDistField(
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        ).order_by(
            'cume_dist'
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'CUME_DIST() OVER (ORDER BY margin DESC) AS "cume_dist" '
            'FROM querybuilder_tests_order '
            'ORDER BY cume_dist '
            'ASC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_ntile(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                NTileField(
                    num_buckets=2,
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        ).order_by(
            'ntile'
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'NTILE(2) OVER (ORDER BY margin DESC) AS "ntile" '
            'FROM querybuilder_tests_order '
            'ORDER BY ntile '
            'ASC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_lag(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                LagField(
                    'margin',
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'LAG(querybuilder_tests_order.margin, 1) OVER (ORDER BY margin DESC) AS "margin_lag" '
            'FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_lag_default(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                LagField(
                    'margin',
                    default=0,
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'LAG(querybuilder_tests_order.margin, 1, \'0\') OVER (ORDER BY margin DESC) AS "margin_lag" '
            'FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_lead(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                LeadField(
                    'margin',
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'LEAD(querybuilder_tests_order.margin, 1) OVER (ORDER BY margin DESC) AS "margin_lead" '
            'FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_first_value(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                FirstValueField(
                    'margin',
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'FIRST_VALUE(querybuilder_tests_order.margin) OVER (ORDER BY margin DESC) AS "margin_first_value" '
            'FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_last_value(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                LastValueField(
                    'margin',
                    over=QueryWindow().order_by(
                        'margin'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'LAST_VALUE(querybuilder_tests_order.margin) OVER (ORDER BY margin ASC) AS "margin_last_value" '
            'FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_nth_value(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                NthValueField(
                    'margin',
                    n=2,
                    over=QueryWindow().order_by(
                        '-margin'
                    )
                )
            ]
        )
        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            'NTH_VALUE(querybuilder_tests_order.margin, 2) OVER (ORDER BY margin DESC) AS "margin_nth_value" '
            'FROM querybuilder_tests_order'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

    def test_num_stddev(self):
        query = Query().from_table(
            table=Order,
            fields=[
                '*',
                NumStdDevField(
                    'margin',
                    over=QueryWindow()
                )
            ]
        ).order_by(
            '-margin_num_stddev'
        )

        query_str = query.get_sql()
        expected_query = (
            'SELECT querybuilder_tests_order.*, '
            '(CASE WHEN (STDDEV(querybuilder_tests_order.margin) OVER ()) <> 0 '
            'THEN ((querybuilder_tests_order.margin - ('
            'AVG(querybuilder_tests_order.margin) OVER ())) / (STDDEV(querybuilder_tests_order.margin) OVER ())) '
            'ELSE 0 '
            'END) '
            'AS "margin_num_stddev" '
            'FROM querybuilder_tests_order '
            'ORDER BY margin_num_stddev '
            'DESC'
        )
        self.assertEqual(query_str, expected_query, get_comparison_str(query_str, expected_query))

[(12, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (15, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (21, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (27, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (37, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (51, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (64, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (77, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (95, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (114, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (138, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (164, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (188, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (212, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (236, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (261, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (282, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (304, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (325, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (346, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (367, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (389, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (418, 'querybuilder.tests.query_tests.get_comparison_str', 'get_comparison_str', (['query_str', 'expected_query'], {}), False, 'from querybuilder.tests.query_tests import QueryTestCase, get_comparison_str\n'), (18, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (24, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (56, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (67, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (80, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (98, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (264, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (285, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (307, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (328, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (349, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (370, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (59, 'querybuilder.fields.RankField', 'RankField', ([], {}), False, 'from querybuilder.fields import RankField, RowNumberField, DenseRankField, PercentRankField, CumeDistField, NTileField, LagField, LeadField, FirstValueField, LastValueField, NthValueField, NumStdDevField\n'), (30, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (117, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (141, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (167, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (191, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (215, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (239, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (392, 'querybuilder.query.Query', 'Query', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (71, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (85, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (103, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (270, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (292, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (313, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (334, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (355, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (377, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (398, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (40, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (122, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (172, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (196, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (220, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (245, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n'), (146, 'querybuilder.query.QueryWindow', 'QueryWindow', ([], {}), False, 'from querybuilder.query import QueryWindow, Query\n')]