from __future__ import print_function, division, absolute_import, unicode_literals

import io
import pytest

from dask.dataframe.io.sql import read_sql_table
from dask.utils import tmpfile
from dask.dataframe.utils import assert_eq

pd = pytest.importorskip("pandas")
dd = pytest.importorskip("dask.dataframe")
pytest.importorskip("sqlalchemy")
pytest.importorskip("sqlite3")
np = pytest.importorskip("numpy")


data = """
name,number,age,negish
Alice,0,33,-5
Bob,1,40,-3
Chris,2,22,3
Dora,3,16,5
Edith,4,53,0
Francis,5,30,0
Garreth,6,20,0
"""

df = pd.read_csv(io.StringIO(data), index_col="number")


@pytest.yield_fixture
def db():
    with tmpfile() as f:
        uri = "sqlite:///%s" % f
        df.to_sql("test", uri, index=True, if_exists="replace")
        yield uri


def test_empty(db):
    from sqlalchemy import create_engine, MetaData, Table, Column, Integer

    with tmpfile() as f:
        uri = "sqlite:///%s" % f
        metadata = MetaData()
        engine = create_engine(uri)
        table = Table(
            "empty_table",
            metadata,
            Column("id", Integer, primary_key=True),
            Column("col2", Integer),
        )
        metadata.create_all(engine)

        dask_df = read_sql_table(table.name, uri, index_col="id", npartitions=1)
        assert dask_df.index.name == "id"
        assert dask_df.col2.dtype == np.dtype("int64")
        pd_dataframe = dask_df.compute()
        assert pd_dataframe.empty is True


@pytest.mark.skip(
    reason="Requires a postgres server. Sqlite does not support multiple schemas."
)
def test_empty_other_schema():
    from sqlalchemy import create_engine, MetaData, Table, Column, Integer, event, DDL

    # Database configurations.
    pg_host = "localhost"
    pg_port = "5432"
    pg_user = "user"
    pg_pass = "pass"
    pg_db = "db"
    db_url = "postgresql://%s:%s@%s:%s/%s" % (pg_user, pg_pass, pg_host, pg_port, pg_db)

    # Create an empty table in a different schema.
    table_name = "empty_table"
    schema_name = "other_schema"
    engine = create_engine(db_url)
    metadata = MetaData()
    table = Table(
        table_name,
        metadata,
        Column("id", Integer, primary_key=True),
        Column("col2", Integer),
        schema=schema_name,
    )
    # Create the schema and the table.
    event.listen(
        metadata, "before_create", DDL("CREATE SCHEMA IF NOT EXISTS %s" % schema_name)
    )
    metadata.create_all(engine)

    # Read the empty table from the other schema.
    dask_df = read_sql_table(
        table.name, db_url, index_col="id", schema=table.schema, npartitions=1
    )

    # Validate that the retrieved table is empty.
    assert dask_df.index.name == "id"
    assert dask_df.col2.dtype == np.dtype("int64")
    pd_dataframe = dask_df.compute()
    assert pd_dataframe.empty is True

    # Drop the schema and the table.
    engine.execute("DROP SCHEMA IF EXISTS %s CASCADE" % schema_name)


def test_needs_rational(db):
    import datetime

    now = datetime.datetime.now()
    d = datetime.timedelta(seconds=1)
    df = pd.DataFrame(
        {
            "a": list("ghjkl"),
            "b": [now + i * d for i in range(5)],
            "c": [True, True, False, True, True],
        }
    )
    df = df.append(
        [
            {"a": "x", "b": now + d * 1000, "c": None},
            {"a": None, "b": now + d * 1001, "c": None},
        ]
    )
    with tmpfile() as f:
        uri = "sqlite:///%s" % f
        df.to_sql("test", uri, index=False, if_exists="replace")

        # one partition contains NULL
        data = read_sql_table("test", uri, npartitions=2, index_col="b")
        df2 = df.set_index("b")
        assert_eq(data, df2.astype({"c": bool}))  # bools are coerced

        # one partition contains NULL, but big enough head
        data = read_sql_table("test", uri, npartitions=2, index_col="b", head_rows=12)
        df2 = df.set_index("b")
        assert_eq(data, df2)

        # empty partitions
        data = read_sql_table("test", uri, npartitions=20, index_col="b")
        part = data.get_partition(12).compute()
        assert part.dtypes.tolist() == ["O", bool]
        assert part.empty
        df2 = df.set_index("b")
        assert_eq(data, df2.astype({"c": bool}))

        # explicit meta
        data = read_sql_table("test", uri, npartitions=2, index_col="b", meta=df2[:0])
        part = data.get_partition(1).compute()
        assert part.dtypes.tolist() == ["O", "O"]
        df2 = df.set_index("b")
        assert_eq(data, df2)


def test_simple(db):
    # single chunk
    data = read_sql_table("test", db, npartitions=2, index_col="number").compute()
    assert (data.name == df.name).all()
    assert data.index.name == "number"
    assert_eq(data, df)


def test_npartitions(db):
    data = read_sql_table(
        "test", db, columns=list(df.columns), npartitions=2, index_col="number"
    )
    assert len(data.divisions) == 3
    assert (data.name.compute() == df.name).all()
    data = read_sql_table(
        "test", db, columns=["name"], npartitions=6, index_col="number"
    )
    assert_eq(data, df[["name"]])
    data = read_sql_table(
        "test",
        db,
        columns=list(df.columns),
        bytes_per_chunk=2 ** 30,
        index_col="number",
    )
    assert data.npartitions == 1
    assert (data.name.compute() == df.name).all()

    data_1 = read_sql_table(
        "test",
        db,
        columns=list(df.columns),
        bytes_per_chunk=2 ** 30,
        index_col="number",
        head_rows=1,
    )
    assert data_1.npartitions == 1
    assert (data_1.name.compute() == df.name).all()

    data = read_sql_table(
        "test",
        db,
        columns=list(df.columns),
        bytes_per_chunk=250,
        index_col="number",
        head_rows=1,
    )
    assert data.npartitions == 2


def test_divisions(db):
    data = read_sql_table(
        "test", db, columns=["name"], divisions=[0, 2, 4], index_col="number"
    )
    assert data.divisions == (0, 2, 4)
    assert data.index.max().compute() == 4
    assert_eq(data, df[["name"]][df.index <= 4])


def test_division_or_partition(db):
    with pytest.raises(TypeError):
        read_sql_table(
            "test",
            db,
            columns=["name"],
            index_col="number",
            divisions=[0, 2, 4],
            npartitions=3,
        )

    out = read_sql_table("test", db, index_col="number", bytes_per_chunk=100)
    m = out.map_partitions(
        lambda d: d.memory_usage(deep=True, index=True).sum()
    ).compute()
    assert (50 < m).all() and (m < 200).all()
    assert_eq(out, df)


def test_range(db):
    data = read_sql_table("test", db, npartitions=2, index_col="number", limits=[1, 4])
    assert data.index.min().compute() == 1
    assert data.index.max().compute() == 4


def test_datetimes():
    import datetime

    now = datetime.datetime.now()
    d = datetime.timedelta(seconds=1)
    df = pd.DataFrame(
        {"a": list("ghjkl"), "b": [now + i * d for i in range(2, -3, -1)]}
    )
    with tmpfile() as f:
        uri = "sqlite:///%s" % f
        df.to_sql("test", uri, index=False, if_exists="replace")
        data = read_sql_table("test", uri, npartitions=2, index_col="b")
        assert data.index.dtype.kind == "M"
        assert data.divisions[0] == df.b.min()
        df2 = df.set_index("b")
        assert_eq(data.map_partitions(lambda x: x.sort_index()), df2.sort_index())


def test_with_func(db):
    from sqlalchemy import sql

    index = sql.func.abs(sql.column("negish")).label("abs")

    # function for the index, get all columns
    data = read_sql_table("test", db, npartitions=2, index_col=index)
    assert data.divisions[0] == 0
    part = data.get_partition(0).compute()
    assert (part.index == 0).all()

    # now an arith op for one column too; it's name will be 'age'
    data = read_sql_table(
        "test", db, npartitions=2, index_col=index, columns=[index, -sql.column("age")]
    )
    assert (data.age.compute() < 0).all()

    # a column that would have no name, give it a label
    index = (-sql.column("negish")).label("index")
    data = read_sql_table(
        "test", db, npartitions=2, index_col=index, columns=["negish", "age"]
    )
    d = data.compute()
    assert (-d.index == d["negish"]).all()


def test_no_nameless_index(db):
    from sqlalchemy import sql

    index = -sql.column("negish")
    with pytest.raises(ValueError):
        read_sql_table(
            "test", db, npartitions=2, index_col=index, columns=["negish", "age", index]
        )

    index = sql.func.abs(sql.column("negish"))

    # function for the index, get all columns
    with pytest.raises(ValueError):
        read_sql_table("test", db, npartitions=2, index_col=index)


def test_select_from_select(db):
    from sqlalchemy import sql

    s1 = sql.select([sql.column("number"), sql.column("name")]).select_from(
        sql.table("test")
    )
    out = read_sql_table(s1, db, npartitions=2, index_col="number")
    assert_eq(out, df[["name"]])


def test_extra_connection_engine_keywords(capsys, db):
    data = read_sql_table(
        "test", db, npartitions=2, index_col="number", engine_kwargs={"echo": False}
    ).compute()
    # no captured message from the stdout with the echo=False parameter (this is the default)
    out, err = capsys.readouterr()
    assert "SELECT" not in out
    assert_eq(data, df)
    # with the echo=True sqlalchemy parameter, you should get all SQL queries in the stdout
    data = read_sql_table(
        "test", db, npartitions=2, index_col="number", engine_kwargs={"echo": True}
    ).compute()
    out, err = capsys.readouterr()
    assert "SELECT" in out
    assert_eq(data, df)


def test_no_character_index_without_divisions(db):

    # attempt to read the sql table with a character index and no divisions
    with pytest.raises(TypeError):
        read_sql_table("test", db, npartitions=2, index_col="name", divisions=None)
