"""
Test the Pandas in-memory adapter.
"""
import pandas as pd
import pytest

from shillelagh.adapters.memory.pandas import find_dataframe
from shillelagh.adapters.memory.pandas import PandasMemory
from shillelagh.backends.apsw.db import connect
from shillelagh.exceptions import ProgrammingError
from shillelagh.fields import Order
from shillelagh.filters import Equal
from shillelagh.filters import IsNotNull
from shillelagh.filters import IsNull
from shillelagh.filters import NotEqual
from shillelagh.filters import Operator


def test_pandas():
    """
    Test basic operations with a dataframe.
    """
    mydf = pd.DataFrame(  # noqa: F841  pylint: disable=unused-variable
        [
            {"index": 10, "temperature": 15.2, "site": "Diamond_St"},
            {"index": 11, "temperature": 13.1, "site": "Blacktail_Loop"},
            {"index": 12, "temperature": 13.3, "site": "Platinum_St"},
            {"index": 13, "temperature": 12.1, "site": "Kodiak_Trail"},
        ],
    )

    connection = connect(":memory:")
    cursor = connection.cursor()

    sql = "SELECT * FROM mydf"
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (10, 15.2, "Diamond_St"),
        (11, 13.1, "Blacktail_Loop"),
        (12, 13.3, "Platinum_St"),
        (13, 12.1, "Kodiak_Trail"),
    ]

    sql = "SELECT * FROM mydf WHERE temperature > 10 AND temperature < 10"
    cursor.execute(sql)
    assert cursor.fetchall() == []

    sql = "SELECT * FROM mydf WHERE temperature > 13 ORDER BY site DESC"
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (12, 13.3, "Platinum_St"),
        (10, 15.2, "Diamond_St"),
        (11, 13.1, "Blacktail_Loop"),
    ]

    sql = "SELECT * FROM mydf WHERE temperature < 13"
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (13, 12.1, "Kodiak_Trail"),
    ]

    sql = "SELECT * FROM mydf WHERE site = 'Platinum_St'"
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (12, 13.3, "Platinum_St"),
    ]

    sql = """INSERT INTO mydf ("index", temperature, site) VALUES (14, 10.1, 'New_Site')"""
    cursor.execute(sql)
    sql = 'SELECT * FROM mydf WHERE "index" > 11'
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (12, 13.3, "Platinum_St"),
        (13, 12.1, "Kodiak_Trail"),
        (14, 10.1, "New_Site"),
    ]

    sql = "DELETE FROM mydf WHERE site = 'New_Site'"
    cursor.execute(sql)
    sql = 'SELECT * FROM mydf WHERE "index" > 11'
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (12, 13.3, "Platinum_St"),
        (13, 12.1, "Kodiak_Trail"),
    ]

    sql = "UPDATE mydf SET temperature = temperature * 2"
    cursor.execute(sql)
    sql = "SELECT * FROM mydf"
    cursor.execute(sql)
    assert cursor.fetchall() == [
        (10, 30.4, "Diamond_St"),
        (11, 26.2, "Blacktail_Loop"),
        (12, 26.6, "Platinum_St"),
        (13, 24.2, "Kodiak_Trail"),
    ]

    sql = "UPDATE mydf SET temperature = temperature / 2"
    cursor.execute(sql)


def test_adapter(mocker):
    """
    Test additional operations on the adapter.
    """
    mydf = pd.DataFrame(
        [
            {"index": 10, "temperature": 15.2, "site": "Diamond_St"},
            {"index": 11, "temperature": 13.1, "site": "Blacktail_Loop"},
            {"index": 12, "temperature": 13.3, "site": "Platinum_St"},
            {"index": 13, "temperature": 12.1, "site": "Kodiak_Trail"},
        ],
    )

    mock_find_dataframe = mocker.patch(
        "shillelagh.adapters.memory.pandas.find_dataframe",
    )
    mock_find_dataframe.return_value = mydf

    adapter = PandasMemory("mydf")

    assert list(adapter.get_data({"index": Equal(10)}, [])) == [
        {"rowid": 0, "index": 10, "site": "Diamond_St", "temperature": 15.2},
    ]

    assert list(adapter.get_data({"index": NotEqual(10)}, [])) == [
        {"rowid": 1, "index": 11, "temperature": 13.1, "site": "Blacktail_Loop"},
        {"rowid": 2, "index": 12, "temperature": 13.3, "site": "Platinum_St"},
        {"rowid": 3, "index": 13, "temperature": 12.1, "site": "Kodiak_Trail"},
    ]

    adapter.update_data(
        3,
        {"rowid": 5, "index": 13, "temperature": 12.1, "site": "Kodiak_Trail"},
    )
    assert list(adapter.get_data({}, [])) == [
        {"index": 10, "rowid": 0, "site": "Diamond_St", "temperature": 15.2},
        {"index": 11, "rowid": 1, "site": "Blacktail_Loop", "temperature": 13.1},
        {"index": 12, "rowid": 2, "site": "Platinum_St", "temperature": 13.3},
        {"index": 13, "rowid": 5, "site": "Kodiak_Trail", "temperature": 12.1},
    ]

    adapter.insert_data(
        {"rowid": 7, "index": 14, "temperature": 10.1, "site": "New_Site"},
    )
    assert list(adapter.get_data({}, [])) == [
        {"index": 10, "rowid": 0, "site": "Diamond_St", "temperature": 15.2},
        {"index": 11, "rowid": 1, "site": "Blacktail_Loop", "temperature": 13.1},
        {"index": 12, "rowid": 2, "site": "Platinum_St", "temperature": 13.3},
        {"index": 13, "rowid": 5, "site": "Kodiak_Trail", "temperature": 12.1},
        {"index": 14, "rowid": 7, "site": "New_Site", "temperature": 10.1},
    ]

    adapter.update_data(
        5,
        {"rowid": 3, "index": 13, "temperature": 12.1, "site": "Kodiak_Trail"},
    )
    adapter.delete_data(7)
    assert list(adapter.get_data({}, [])) == [
        {"index": 10, "rowid": 0, "site": "Diamond_St", "temperature": 15.2},
        {"index": 11, "rowid": 1, "site": "Blacktail_Loop", "temperature": 13.1},
        {"index": 12, "rowid": 2, "site": "Platinum_St", "temperature": 13.3},
        {"index": 13, "rowid": 3, "site": "Kodiak_Trail", "temperature": 12.1},
    ]

    with pytest.raises(ProgrammingError) as excinfo:
        list(adapter.get_data({"a": [1, 2, 3]}, []))
    assert str(excinfo.value) == "Invalid filter: [1, 2, 3]"


def test_adapter_nulls(mocker):
    """
    Test operations with nulls on the adapter.
    """
    mydf = pd.DataFrame([{"a": None, "b": 10}, {"a": 20, "b": None}])

    mock_find_dataframe = mocker.patch(
        "shillelagh.adapters.memory.pandas.find_dataframe",
    )
    mock_find_dataframe.return_value = mydf

    adapter = PandasMemory("mydf")

    rows = list(adapter.get_data({"a": IsNull()}, []))
    assert len(rows) == 1
    row = rows[0]
    assert row["a"] != row["a"]  # NaN
    assert row["b"] == 10.0
    assert row["rowid"] == 0

    rows = list(adapter.get_data({"a": IsNotNull()}, []))
    assert len(rows) == 1
    row = rows[0]
    assert row["a"] == 20.0
    assert row["b"] != row["b"]  # NaN
    assert row["rowid"] == 1


outer_df = pd.DataFrame()


def test_find_dataframe():
    """
    Test that we can find the dataframe.
    """
    mydf = pd.DataFrame()

    def inner_scope() -> pd.DataFrame:
        return find_dataframe("mydf")

    assert inner_scope() is mydf

    assert find_dataframe("outer_df") is outer_df

    assert find_dataframe("foobar") is None
    with pytest.raises(ProgrammingError) as excinfo:
        PandasMemory("foobar")
    assert str(excinfo.value) == "Could not find dataframe"


def test_get_cost(mocker):
    """
    Test cost estimation.
    """
    mydf = pd.DataFrame(
        [
            {"index": 10, "temperature": 15.2, "site": "Diamond_St"},
            {"index": 11, "temperature": 13.1, "site": "Blacktail_Loop"},
            {"index": 12, "temperature": 13.3, "site": "Platinum_St"},
            {"index": 13, "temperature": 12.1, "site": "Kodiak_Trail"},
        ],
    )

    mock_find_dataframe = mocker.patch(
        "shillelagh.adapters.memory.pandas.find_dataframe",
    )
    mock_find_dataframe.return_value = mydf

    adapter = PandasMemory(mydf)
    assert adapter.get_cost([], []) == 0

    # linear filtering cost
    assert adapter.get_cost([("one", Operator.EQ)], []) == 1000
    assert adapter.get_cost([("one", Operator.EQ), ("two", Operator.EQ)], []) == 2000

    # linear sorting cost
    assert (
        adapter.get_cost(
            [("one", Operator.EQ), ("two", Operator.EQ)],
            [("one", Order.ASCENDING)],
        )
        == 11965
    )
    assert (
        adapter.get_cost(
            [("one", Operator.EQ), ("two", Operator.EQ)],
            [("one", Order.ASCENDING), ("two", Order.DESCENDING)],
        )
        == 21931
    )
