#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import absolute_import, unicode_literals

import unittest
from datetime import date, datetime
from decimal import Decimal

import numpy as np

from pyathenajdbc.util import as_pandas
from tests import WithConnect
from tests.util import with_cursor


class TestUtil(unittest.TestCase, WithConnect):
    @with_cursor
    def test_as_pandas(self, cursor):
        cursor.execute(
            """
        SELECT
          col_boolean
          ,col_tinyint
          ,col_smallint
          ,col_int
          ,col_bigint
          ,col_float
          ,col_double
          ,col_string
          ,col_timestamp
          ,CAST(col_timestamp AS time) AS col_time
          ,col_date
          ,col_binary
          ,col_array
          ,CAST(col_array AS json) AS col_array_json
          ,col_map
          ,CAST(col_map AS json) AS col_map_json
          ,col_struct
          ,col_decimal
        FROM one_row_complex
        """
        )
        df = as_pandas(cursor)
        rows = [
            tuple(
                [
                    row["col_boolean"],
                    row["col_tinyint"],
                    row["col_smallint"],
                    row["col_int"],
                    row["col_bigint"],
                    row["col_float"],
                    row["col_double"],
                    row["col_string"],
                    row["col_timestamp"],
                    row["col_time"],
                    row["col_date"],
                    row["col_binary"],
                    row["col_array"],
                    row["col_array_json"],
                    row["col_map"],
                    row["col_map_json"],
                    row["col_struct"],
                    row["col_decimal"],
                ]
            )
            for _, row in df.iterrows()
        ]
        expected = [
            (
                True,
                127,
                32767,
                2147483647,
                9223372036854775807,
                0.5,
                0.25,
                "a string",
                datetime(2017, 1, 1, 0, 0, 0),
                datetime(2017, 1, 1, 0, 0, 0).time(),
                date(2017, 1, 2),
                b"123",
                "[1, 2]",
                [1, 2],
                "{1=2, 3=4}",
                {"1": 2, "3": 4},
                "{a=1, b=2}",
                Decimal("0.1"),
            )
        ]
        self.assertEqual(rows, expected)

    @with_cursor
    def test_as_pandas_integer_na_values(self, cursor):
        cursor.execute(
            """
            SELECT * FROM integer_na_values
            """
        )
        df = as_pandas(cursor, coerce_float=True)
        rows = [tuple([row["a"], row["b"]]) for _, row in df.iterrows()]
        # TODO AssertionError: Lists differ:
        #  [(1.0, 2.0), (1.0, nan), (nan, nan)] != [(1.0, 2.0), (1.0, nan), (nan, nan)]
        # self.assertEqual(rows, [
        #     (1.0, 2.0),
        #     (1.0, np.nan),
        #     (np.nan, np.nan),
        # ])
        np.testing.assert_array_equal(rows, [(1, 2), (1, np.nan), (np.nan, np.nan)])

    @with_cursor
    def test_as_pandas_boolean_na_values(self, cursor):
        cursor.execute(
            """
            SELECT * FROM boolean_na_values
            """
        )
        df = as_pandas(cursor)
        rows = [tuple([row["a"], row["b"]]) for _, row in df.iterrows()]
        self.assertEqual(rows, [(True, False), (False, None), (None, None)])

@@ -196,16 +196,36 @@
 py as np
+%0Aimport pandas as pd
 %0A%0Afrom p
@@ -687,59 +687,8 @@
 amp%0A
-          ,CAST(col_timestamp AS time) AS col_time%0A
     
@@ -761,117 +761,15 @@
    ,
-CAST(col_array AS json) AS col_array_json%0A          ,col_map%0A          ,CAST(col_map AS json) AS col_map_json
+col_map
 %0A   
@@ -1304,45 +1304,8 @@
 %22%5D,%0A
-                    row%5B%22col_time%22%5D,%0A
     
@@ -1447,95 +1447,11 @@
 col_
-array_json%22%5D,%0A                    row%5B%22col_map%22%5D,%0A                    row%5B%22col_map_json
+map
 %22%5D,%0A
@@ -1868,63 +1868,20 @@
     
-datetime(2017, 1, 1, 0, 0, 0),%0A                datetime
+pd.Timestamp
 (201
@@ -1901,15 +1901,8 @@
 , 0)
-.time()
 ,%0A  
@@ -1974,22 +1974,20 @@
        %22
-%5B
 1, 2
-%5D
 %22,%0A     
@@ -2001,78 +2001,20 @@
     
-%5B1, 2%5D,%0A                %22%7B1=2, 3=4%7D%22,%0A                %7B%221%22: 2, %223%22: 4%7D
+%22%7B1=2, 3=4%7D%22
 ,%0A  

Remove columns of type that JDBC does not support