"""
Tests that the specified index column (a.k.a "index_col")
is properly handled or inferred during parsing for all of
the parsers defined in parsers.py
"""

from io import StringIO

import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm

xfail_pyarrow = pytest.mark.usefixtures("pyarrow_xfail")
skip_pyarrow = pytest.mark.usefixtures("pyarrow_skip")


@pytest.mark.parametrize("with_header", [True, False])
def test_index_col_named(all_parsers, with_header):
    parser = all_parsers
    no_header = """\
KORD1,19990127, 19:00:00, 18:56:00, 0.8100, 2.8100, 7.2000, 0.0000, 280.0000
KORD2,19990127, 20:00:00, 19:56:00, 0.0100, 2.2100, 7.2000, 0.0000, 260.0000
KORD3,19990127, 21:00:00, 20:56:00, -0.5900, 2.2100, 5.7000, 0.0000, 280.0000
KORD4,19990127, 21:00:00, 21:18:00, -0.9900, 2.0100, 3.6000, 0.0000, 270.0000
KORD5,19990127, 22:00:00, 21:56:00, -0.5900, 1.7100, 5.1000, 0.0000, 290.0000
KORD6,19990127, 23:00:00, 22:56:00, -0.5900, 1.7100, 4.6000, 0.0000, 280.0000"""
    header = "ID,date,NominalTime,ActualTime,TDew,TAir,Windspeed,Precip,WindDir\n"

    if with_header:
        data = header + no_header

        result = parser.read_csv(StringIO(data), index_col="ID")
        expected = parser.read_csv(StringIO(data), header=0).set_index("ID")
        tm.assert_frame_equal(result, expected)
    else:
        data = no_header
        msg = "Index ID invalid"

        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), index_col="ID")


def test_index_col_named2(all_parsers):
    parser = all_parsers
    data = """\
1,2,3,4,hello
5,6,7,8,world
9,10,11,12,foo
"""

    expected = pd.DataFrame(
        {"a": [1, 5, 9], "b": [2, 6, 10], "c": [3, 7, 11], "d": [4, 8, 12]},
        index=pd.Index(["hello", "world", "foo"], name="message"),
    )
    names = ["a", "b", "c", "d", "message"]

    result = parser.read_csv(StringIO(data), names=names, index_col=["message"])
    tm.assert_frame_equal(result, expected)


def test_index_col_is_true(all_parsers):
    # see gh-9798
    data = "a,b\n1,2"
    parser = all_parsers

    msg = "The value of index_col couldn't be 'True'"
    with pytest.raises(ValueError, match=msg):
        parser.read_csv(StringIO(data), index_col=True)


@skip_pyarrow  # CSV parse error: Expected 3 columns, got 4
def test_infer_index_col(all_parsers):
    data = """A,B,C
foo,1,2,3
bar,4,5,6
baz,7,8,9
"""
    parser = all_parsers
    result = parser.read_csv(StringIO(data))

    expected = pd.DataFrame(
        [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
        index=["foo", "bar", "baz"],
        columns=["A", "B", "C"],
    )
    tm.assert_frame_equal(result, expected)


@skip_pyarrow  # CSV parse error: Empty CSV file or block
@pytest.mark.parametrize(
    "index_col,kwargs",
    [
        (None, {"columns": ["x", "y", "z"]}),
        (False, {"columns": ["x", "y", "z"]}),
        (0, {"columns": ["y", "z"], "index": pd.Index([], name="x")}),
        (1, {"columns": ["x", "z"], "index": pd.Index([], name="y")}),
        ("x", {"columns": ["y", "z"], "index": pd.Index([], name="x")}),
        ("y", {"columns": ["x", "z"], "index": pd.Index([], name="y")}),
        (
            [0, 1],
            {
                "columns": ["z"],
                "index": pd.MultiIndex.from_arrays([[]] * 2, names=["x", "y"]),
            },
        ),
        (
            ["x", "y"],
            {
                "columns": ["z"],
                "index": pd.MultiIndex.from_arrays([[]] * 2, names=["x", "y"]),
            },
        ),
        (
            [1, 0],
            {
                "columns": ["z"],
                "index": pd.MultiIndex.from_arrays([[]] * 2, names=["y", "x"]),
            },
        ),
        (
            ["y", "x"],
            {
                "columns": ["z"],
                "index": pd.MultiIndex.from_arrays([[]] * 2, names=["y", "x"]),
            },
        ),
    ],
)
def test_index_col_empty_data(all_parsers, index_col, kwargs):
    data = "x,y,z"
    parser = all_parsers
    result = parser.read_csv(StringIO(data), index_col=index_col)

    expected = pd.DataFrame(**kwargs)
    tm.assert_frame_equal(result, expected)


@skip_pyarrow  # CSV parse error: Empty CSV file or block
def test_empty_with_index_col_false(all_parsers):
    # see gh-10413
    data = "x,y"
    parser = all_parsers
    result = parser.read_csv(StringIO(data), index_col=False)

    expected = pd.DataFrame(columns=["x", "y"])
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "index_names",
    [
        ["", ""],
        ["foo", ""],
        ["", "bar"],
        ["foo", "bar"],
        ["NotReallyUnnamed", "Unnamed: 0"],
    ],
)
def test_multi_index_naming(all_parsers, index_names):
    parser = all_parsers

    # We don't want empty index names being replaced with "Unnamed: 0"
    data = ",".join([*index_names, "col\na,c,1\na,d,2\nb,c,3\nb,d,4"])
    result = parser.read_csv(StringIO(data), index_col=[0, 1])

    expected = pd.DataFrame(
        {"col": [1, 2, 3, 4]},
        index=pd.MultiIndex.from_product([["a", "b"], ["c", "d"]]),
    )
    expected.index.names = [name if name else None for name in index_names]
    tm.assert_frame_equal(result, expected)


def test_multi_index_naming_not_all_at_beginning(all_parsers):
    parser = all_parsers
    data = ",Unnamed: 2,\na,c,1\na,d,2\nb,c,3\nb,d,4"
    result = parser.read_csv(StringIO(data), index_col=[0, 2])

    expected = pd.DataFrame(
        {"Unnamed: 2": ["c", "d", "c", "d"]},
        index=pd.MultiIndex(
            levels=[["a", "b"], [1, 2, 3, 4]], codes=[[0, 0, 1, 1], [0, 1, 2, 3]]
        ),
    )
    tm.assert_frame_equal(result, expected)


def test_no_multi_index_level_names_empty(temp_file, all_parsers):
    # GH 10984
    parser = all_parsers
    midx = pd.MultiIndex.from_tuples([("A", 1, 2), ("A", 1, 2), ("B", 1, 2)])
    expected = pd.DataFrame(
        np.random.default_rng(2).standard_normal((3, 3)),
        index=midx,
        columns=["x", "y", "z"],
    )
    expected.to_csv(temp_file)
    result = parser.read_csv(temp_file, index_col=[0, 1, 2])
    tm.assert_frame_equal(result, expected)


@xfail_pyarrow  # TypeError: an integer is required
def test_header_with_index_col(all_parsers):
    # GH 33476
    parser = all_parsers
    data = """
I11,A,A
I12,B,B
I2,1,3
"""
    midx = pd.MultiIndex.from_tuples([("A", "B"), ("A", "B.1")], names=["I11", "I12"])
    idx = pd.Index(["I2"])
    expected = pd.DataFrame([[1, 3]], index=idx, columns=midx)

    result = parser.read_csv(StringIO(data), index_col=0, header=[0, 1])
    tm.assert_frame_equal(result, expected)

    col_idx = pd.Index(["A", "A.1"])
    idx = pd.Index(["I12", "I2"], name="I11")
    expected = pd.DataFrame([["B", "B"], ["1", "3"]], index=idx, columns=col_idx)

    result = parser.read_csv(StringIO(data), index_col="I11", header=0)
    tm.assert_frame_equal(result, expected)


@pytest.mark.slow
def test_index_col_large_csv(temp_file, all_parsers, monkeypatch):
    # https://github.com/pandas-dev/pandas/issues/37094
    parser = all_parsers

    ARR_LEN = 100
    df = pd.DataFrame(
        {
            "a": range(ARR_LEN + 1),
            "b": np.random.default_rng(2).standard_normal(ARR_LEN + 1),
        }
    )

    df.to_csv(temp_file, index=False)
    with monkeypatch.context() as m:
        m.setattr("pandas.core.algorithms._MINIMUM_COMP_ARR_LEN", ARR_LEN)
        result = parser.read_csv(temp_file, index_col=[0])

    tm.assert_frame_equal(result, df.set_index("a"))


def test_index_col_multiindex_columns_no_data(all_parsers):
    # GH#38292
    parser = all_parsers
    if parser.engine == "pyarrow":
        with pytest.raises(ValueError, match="does not support a list of integers"):
            parser.read_csv(
                StringIO("a0,a1,a2\nb0,b1,b2\n"), header=[0, 1], index_col=0
            )
        return

    result = parser.read_csv(
        StringIO("a0,a1,a2\nb0,b1,b2\n"), header=[0, 1], index_col=0
    )
    expected = pd.DataFrame(
        [],
        index=pd.Index([]),
        columns=pd.MultiIndex.from_arrays(
            [["a1", "a2"], ["b1", "b2"]], names=["a0", "b0"]
        ),
    )
    tm.assert_frame_equal(result, expected)


# header=[0] is a singleton, so it behaves like header=0; the empty-frame
# index dtype still differs for the pyarrow engine (float64 vs object)
@xfail_pyarrow  # AssertionError: DataFrame.index are different
def test_index_col_header_no_data(all_parsers):
    # GH#38292
    parser = all_parsers
    result = parser.read_csv(StringIO("a0,a1,a2\n"), header=[0], index_col=0)
    expected = pd.DataFrame(
        [],
        columns=["a1", "a2"],
        index=pd.Index([], name="a0"),
    )
    tm.assert_frame_equal(result, expected)


def test_multiindex_columns_no_data(all_parsers):
    # GH#38292
    parser = all_parsers
    if parser.engine == "pyarrow":
        with pytest.raises(ValueError, match="does not support a list of integers"):
            parser.read_csv(StringIO("a0,a1,a2\nb0,b1,b2\n"), header=[0, 1])
        return

    result = parser.read_csv(StringIO("a0,a1,a2\nb0,b1,b2\n"), header=[0, 1])
    expected = pd.DataFrame(
        [], columns=pd.MultiIndex.from_arrays([["a0", "a1", "a2"], ["b0", "b1", "b2"]])
    )
    tm.assert_frame_equal(result, expected)


def test_multiindex_columns_index_col_with_data(all_parsers):
    # GH#38292
    parser = all_parsers
    if parser.engine == "pyarrow":
        with pytest.raises(ValueError, match="does not support a list of integers"):
            parser.read_csv(
                StringIO("a0,a1,a2\nb0,b1,b2\ndata,data,data"),
                header=[0, 1],
                index_col=0,
            )
        return

    result = parser.read_csv(
        StringIO("a0,a1,a2\nb0,b1,b2\ndata,data,data"), header=[0, 1], index_col=0
    )
    expected = pd.DataFrame(
        [["data", "data"]],
        columns=pd.MultiIndex.from_arrays(
            [["a1", "a2"], ["b1", "b2"]], names=["a0", "b0"]
        ),
        index=pd.Index(["data"]),
    )
    tm.assert_frame_equal(result, expected)


@skip_pyarrow  # CSV parse error: Empty CSV file or block
def test_infer_types_boolean_sum(all_parsers):
    # GH#44079
    parser = all_parsers
    result = parser.read_csv(
        StringIO("0,1"),
        names=["a", "b"],
        index_col=["a"],
        dtype={"a": "UInt8"},
    )
    expected = pd.DataFrame(
        data={
            "a": [
                0,
            ],
            "b": [1],
        }
    ).set_index("a")
    # Not checking index type now, because the C parser will return a
    # index column of dtype 'object', and the Python parser will return a
    # index column of dtype 'int64'.
    tm.assert_frame_equal(result, expected, check_index_type=False)


@pytest.mark.parametrize("dtype, val", [(object, "01"), ("int64", 1)])
def test_specify_dtype_for_index_col(all_parsers, dtype, val):
    # GH#9435, GH#57666
    data = "a,b\n01,2"
    parser = all_parsers
    result = parser.read_csv(StringIO(data), index_col="a", dtype={"a": dtype})
    expected = pd.DataFrame({"b": [2]}, index=pd.Index([val], name="a", dtype=dtype))
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("dtype", [str, object])
def test_scalar_dtype_for_index_col(all_parsers, dtype, using_infer_string, request):
    # GH#45801 a scalar dtype applies to the index column too, like the c
    #  engine, except when it resolves to object (dtype=object always;
    #  dtype=str with infer_string disabled), in which case index columns
    #  keep their inferred values
    parser = all_parsers
    data = "a,b\n1,x\n2,y"
    apply_to_index = dtype is str and using_infer_string
    if parser.engine == "python" and apply_to_index:
        request.applymarker(
            pytest.mark.xfail(
                reason="python engine does not apply scalar dtype to index_col"
            )
        )
    result = parser.read_csv(StringIO(data), index_col="a", dtype=dtype)
    if apply_to_index:
        expected_index = pd.Index(["1", "2"], name="a", dtype=dtype)
    else:
        expected_index = pd.Index([1, 2], name="a")
    expected = pd.DataFrame({"b": ["x", "y"]}, index=expected_index, dtype=dtype)
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("dtype", ["float64", "Int64"])
def test_scalar_numeric_dtype_for_index_col(all_parsers, dtype, request):
    # GH#45801 a scalar dtype applies to the index column too, like the c
    #  engine
    parser = all_parsers
    data = "a,b\n1,2\n3,4"
    if parser.engine == "python":
        request.applymarker(
            pytest.mark.xfail(
                reason="python engine does not apply scalar dtype to index_col"
            )
        )
    result = parser.read_csv(StringIO(data), index_col="a", dtype=dtype)
    expected = pd.DataFrame(
        {"b": [2, 4]}, index=pd.Index([1, 3], name="a"), dtype=dtype
    )
    expected.index = expected.index.astype(dtype)
    tm.assert_frame_equal(result, expected)


def test_multiindex_columns_not_leading_index_col(all_parsers):
    # GH#38549
    parser = all_parsers
    data = """a,b,c,d
e,f,g,h
x,y,1,2
"""
    if parser.engine == "pyarrow":
        with pytest.raises(ValueError, match="does not support a list of integers"):
            parser.read_csv(StringIO(data), header=[0, 1], index_col=1)
        return

    result = parser.read_csv(
        StringIO(data),
        header=[0, 1],
        index_col=1,
    )
    cols = pd.MultiIndex.from_tuples(
        [("a", "e"), ("c", "g"), ("d", "h")], names=["b", "f"]
    )
    expected = pd.DataFrame([["x", 1, 2]], columns=cols, index=["y"])
    tm.assert_frame_equal(result, expected)