import duckdb
import pandas as pd
import numpy
import pytest


def check_category_equal(category):
    df_in = pd.DataFrame(
        {
            'x': pd.Categorical(category, ordered=True),
        }
    )
    df_out = duckdb.query_df(df_in, "data", "SELECT * FROM data").df()
    assert df_in.equals(df_out)


def check_result_list(category, res):
    for i in range(len(category)):
        assert category[i][0] == res[i]


def check_create_table(category):
    conn = duckdb.connect()

    conn.execute("PRAGMA enable_verification")
    df_in = pd.DataFrame({'x': pd.Categorical(category, ordered=True), 'y': pd.Categorical(category, ordered=True)})

    df_out = duckdb.query_df(df_in, "data", "SELECT * FROM data").df()
    assert df_in.equals(df_out)

    conn.execute("CREATE TABLE t1 AS SELECT * FROM df_in")
    conn.execute("CREATE TABLE t2 AS SELECT * FROM df_in")

    # Check fetchall
    res = conn.execute("SELECT t1.x FROM t1").fetchall()
    check_result_list(res, category)

    # Do a insert to trigger string -> cat
    conn.execute("INSERT INTO t1 VALUES ('2','2')")

    res = conn.execute("SELECT x FROM t1 where x = '1'").fetchall()
    assert res == [('1',)]

    res = conn.execute("SELECT t1.x FROM t1 inner join t2 on (t1.x = t2.x)").fetchall()
    assert res == conn.execute("SELECT x FROM t1").fetchall()

    res = conn.execute("SELECT t1.x FROM t1 inner join t2 on (t1.x = t2.y)").fetchall()
    assert res == conn.execute("SELECT x FROM t1").fetchall()

    assert res == conn.execute("SELECT x FROM t1").fetchall()
    # Triggering the cast with ENUM as a src
    conn.execute("ALTER TABLE t1 ALTER x SET DATA TYPE VARCHAR")
    # We should be able to drop the table without any dependencies
    conn.execute("DROP TABLE t1")


class TestCategory(object):
    def test_category_simple(self, duckdb_cursor):
        df_in = pd.DataFrame({'float': [1.0, 2.0, 1.0], 'int': pd.Series([1, 2, 1], dtype="category")})

        df_out = duckdb.query_df(df_in, "data", "SELECT * FROM data").df()
        print(duckdb.query_df(df_in, "data", "SELECT * FROM data").fetchall())
        print(df_out['int'])
        assert numpy.all(df_out['float'] == numpy.array([1.0, 2.0, 1.0]))
        assert numpy.all(df_out['int'] == numpy.array([1, 2, 1]))

    def test_category_nulls(self, duckdb_cursor):
        df_in = pd.DataFrame({'int': pd.Series([1, 2, None], dtype="category")})
        df_out = duckdb.query_df(df_in, "data", "SELECT * FROM data").df()
        print(duckdb.query_df(df_in, "data", "SELECT * FROM data").fetchall())
        assert df_out['int'][0] == 1
        assert df_out['int'][1] == 2
        assert numpy.isnan(df_out['int'][2])

    def test_category_string(self, duckdb_cursor):
        check_category_equal(['foo', 'bla', 'zoo', 'foo', 'foo', 'bla'])

    def test_category_string_null(self, duckdb_cursor):
        check_category_equal(['foo', 'bla', None, 'zoo', 'foo', 'foo', None, 'bla'])

    def test_category_string_null_bug_4747(self, duckdb_cursor):
        check_category_equal([str(i) for i in range(160)] + [None])

    def test_categorical_fetchall(self, duckdb_cursor):
        df_in = pd.DataFrame(
            {
                'x': pd.Categorical(['foo', 'bla', None, 'zoo', 'foo', 'foo', None, 'bla'], ordered=True),
            }
        )
        assert duckdb.query_df(df_in, "data", "SELECT * FROM data").fetchall() == [
            ('foo',),
            ('bla',),
            (None,),
            ('zoo',),
            ('foo',),
            ('foo',),
            (None,),
            ('bla',),
        ]

    def test_category_string_uint8(self, duckdb_cursor):
        category = []
        for i in range(10):
            category.append(str(i))
        check_create_table(category)

    def test_category_fetch_df_chunk(self, duckdb_cursor):
        con = duckdb.connect()
        categories = ['foo', 'bla', None, 'zoo', 'foo', 'foo', None, 'bla']
        result = categories * 256
        categories = result * 2
        df_result = pd.DataFrame(
            {
                'x': pd.Categorical(result, ordered=True),
            }
        )
        df_in = pd.DataFrame(
            {
                'x': pd.Categorical(categories, ordered=True),
            }
        )
        con.register("data", df_in)

        query = con.execute("SELECT * FROM data")
        cur_chunk = query.fetch_df_chunk()
        assert cur_chunk.equals(df_result)

        cur_chunk = query.fetch_df_chunk()
        assert cur_chunk.equals(df_result)

        cur_chunk = query.fetch_df_chunk()
        assert cur_chunk.empty

    def test_category_mix(self, duckdb_cursor):
        df_in = pd.DataFrame(
            {
                'float': [1.0, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 0.0],
                'x': pd.Categorical(['foo', 'bla', None, 'zoo', 'foo', 'foo', None, 'bla'], ordered=True),
            }
        )

        df_out = duckdb.query_df(df_in, "data", "SELECT * FROM data").df()
        assert df_out.equals(df_in)
