"""
Tests that skipped rows are properly handled during
parsing for all of the parsers defined in parsers.py
"""

from datetime import datetime
from io import StringIO

import numpy as np
import pytest

from pandas.errors import (
    EmptyDataError,
    ParserError,
)

import pandas as pd
import pandas._testing as tm


@pytest.mark.parametrize("skiprows", [list(range(6)), 6])
def test_skip_rows_bug(all_parsers, skiprows):
    # see gh-505
    parser = all_parsers
    text = """#foo,a,b,c
#foo,a,b,c
#foo,a,b,c
#foo,a,b,c
#foo,a,b,c
#foo,a,b,c
1/1/2000,1.,2.,3.
1/2/2000,4,5,6
1/3/2000,7,8,9
"""
    if parser.engine == "pyarrow" and not isinstance(skiprows, int):
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(
                StringIO(text),
                skiprows=skiprows,
                header=None,
                index_col=0,
                parse_dates=True,
            )
        return

    result = parser.read_csv(
        StringIO(text), skiprows=skiprows, header=None, index_col=0, parse_dates=True
    )
    index = pd.Index(
        [datetime(2000, 1, 1), datetime(2000, 1, 2), datetime(2000, 1, 3)],
        dtype="M8[us]",
        name=0,
    )

    expected = pd.DataFrame(
        np.arange(1.0, 10.0).reshape((3, 3)), columns=[1, 2, 3], index=index
    )
    tm.assert_frame_equal(result, expected)


def test_deep_skip_rows(all_parsers):
    # see gh-4382
    parser = all_parsers
    data = "a,b,c\n" + "\n".join(
        [",".join([str(i), str(i + 1), str(i + 2)]) for i in range(10)]
    )
    condensed_data = "a,b,c\n" + "\n".join(
        [",".join([str(i), str(i + 1), str(i + 2)]) for i in [0, 1, 2, 3, 4, 6, 8, 9]]
    )

    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), skiprows=[6, 8])
        return

    result = parser.read_csv(StringIO(data), skiprows=[6, 8])
    condensed_result = parser.read_csv(StringIO(condensed_data))
    tm.assert_frame_equal(result, condensed_result)


def test_skip_rows_blank(all_parsers):
    # see gh-9832
    parser = all_parsers
    text = """#foo,a,b,c
#foo,a,b,c

#foo,a,b,c
#foo,a,b,c

1/1/2000,1.,2.,3.
1/2/2000,4,5,6
1/3/2000,7,8,9
"""
    data = parser.read_csv(
        StringIO(text), skiprows=6, header=None, index_col=0, parse_dates=True
    )
    index = pd.Index(
        [datetime(2000, 1, 1), datetime(2000, 1, 2), datetime(2000, 1, 3)],
        dtype="M8[us]",
        name=0,
    )

    expected = pd.DataFrame(
        np.arange(1.0, 10.0).reshape((3, 3)), columns=[1, 2, 3], index=index
    )
    tm.assert_frame_equal(data, expected)


@pytest.mark.parametrize(
    "data,kwargs,expected",
    [
        (
            """id,text,num_lines
1,"line 11
line 12",2
2,"line 21
line 22",2
3,"line 31",1""",
            {"skiprows": [1]},
            pd.DataFrame(
                [[2, "line 21\nline 22", 2], [3, "line 31", 1]],
                columns=["id", "text", "num_lines"],
            ),
        ),
        (
            "a,b,c\n~a\n b~,~e\n d~,~f\n f~\n1,2,~12\n 13\n 14~",
            {"quotechar": "~", "skiprows": [2]},
            pd.DataFrame([["a\n b", "e\n d", "f\n f"]], columns=["a", "b", "c"]),
        ),
        (
            (
                "Text,url\n~example\n "
                "sentence\n one~,url1\n~"
                "example\n sentence\n two~,url2\n~"
                "example\n sentence\n three~,url3"
            ),
            {"quotechar": "~", "skiprows": [1, 3]},
            pd.DataFrame(
                [["example\n sentence\n two", "url2"]], columns=["Text", "url"]
            ),
        ),
    ],
)
def test_skip_row_with_newline(all_parsers, data, kwargs, expected):
    # see gh-12775 and gh-10911
    parser = all_parsers
    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), **kwargs)
        return

    result = parser.read_csv(StringIO(data), **kwargs)
    tm.assert_frame_equal(result, expected)


def test_skip_row_with_quote(all_parsers):
    # see gh-12775 and gh-10911
    parser = all_parsers
    data = """id,text,num_lines
1,"line '11' line 12",2
2,"line '21' line 22",2
3,"line '31' line 32",1"""

    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), skiprows=[1])
        return

    exp_data = [[2, "line '21' line 22", 2], [3, "line '31' line 32", 1]]
    expected = pd.DataFrame(exp_data, columns=["id", "text", "num_lines"])

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


@pytest.mark.parametrize(
    "data,exp_data",
    [
        (
            """id,text,num_lines
1,"line \n'11' line 12",2
2,"line \n'21' line 22",2
3,"line \n'31' line 32",1""",
            [[2, "line \n'21' line 22", 2], [3, "line \n'31' line 32", 1]],
        ),
        (
            """id,text,num_lines
1,"line '11\n' line 12",2
2,"line '21\n' line 22",2
3,"line '31\n' line 32",1""",
            [[2, "line '21\n' line 22", 2], [3, "line '31\n' line 32", 1]],
        ),
        (
            """id,text,num_lines
1,"line '11\n' \r\tline 12",2
2,"line '21\n' \r\tline 22",2
3,"line '31\n' \r\tline 32",1""",
            [[2, "line '21\n' \r\tline 22", 2], [3, "line '31\n' \r\tline 32", 1]],
        ),
    ],
)
def test_skip_row_with_newline_and_quote(all_parsers, data, exp_data):
    # see gh-12775 and gh-10911
    parser = all_parsers
    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), skiprows=[1])
        return

    result = parser.read_csv(StringIO(data), skiprows=[1])

    expected = pd.DataFrame(exp_data, columns=["id", "text", "num_lines"])
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "lineterminator",
    ["\n", "\r\n", "\r"],  # "LF"  # "CRLF"  # "CR"
)
def test_skiprows_lineterminator(all_parsers, lineterminator, request):
    # see gh-9079
    parser = all_parsers
    data = "\n".join(
        [
            "SMOSMANIA ThetaProbe-ML2X ",
            "2007/01/01 01:00   0.2140 U M ",
            "2007/01/01 02:00   0.2141 M O ",
            "2007/01/01 04:00   0.2142 D M ",
        ]
    )
    expected = pd.DataFrame(
        [
            ["2007/01/01", "01:00", 0.2140, "U", "M"],
            ["2007/01/01", "02:00", 0.2141, "M", "O"],
            ["2007/01/01", "04:00", 0.2142, "D", "M"],
        ],
        columns=["date", "time", "var", "flag", "oflag"],
    )

    if parser.engine == "python" and lineterminator == "\r":
        mark = pytest.mark.xfail(reason="'CR' not respect with the Python parser yet")
        request.applymarker(mark)

    data = data.replace("\n", lineterminator)

    if parser.engine == "pyarrow":
        msg = "the 'pyarrow' engine does not support separators > 1 char"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(
                StringIO(data),
                skiprows=1,
                sep=r"\s+",
                names=["date", "time", "var", "flag", "oflag"],
            )
        return

    result = parser.read_csv(
        StringIO(data),
        skiprows=1,
        sep=r"\s+",
        names=["date", "time", "var", "flag", "oflag"],
    )
    tm.assert_frame_equal(result, expected)


def test_skiprows_infield_quote(all_parsers):
    # see gh-14459
    parser = all_parsers
    data = 'a"\nb"\na\n1'
    expected = pd.DataFrame({"a": [1]})

    result = parser.read_csv(StringIO(data), skiprows=2)
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "kwargs,expected",
    [
        ({}, "1"),
        ({"header": 0, "names": ["foo"]}, "foo"),
    ],
)
def test_skip_rows_callable(all_parsers, kwargs, expected):
    parser = all_parsers
    data = "a\n1\n2\n3\n4\n5"

    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), skiprows=lambda x: x % 2 == 0, **kwargs)
        return

    result = parser.read_csv(StringIO(data), skiprows=lambda x: x % 2 == 0, **kwargs)
    expected = pd.DataFrame({expected: [3, 5]})
    tm.assert_frame_equal(result, expected)


def test_skip_rows_callable_not_in(all_parsers):
    parser = all_parsers
    data = "0,a\n1,b\n2,c\n3,d\n4,e"
    expected = pd.DataFrame([[1, "b"], [3, "d"]])

    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(
                StringIO(data), header=None, skiprows=lambda x: x not in [1, 3]
            )
        return

    result = parser.read_csv(
        StringIO(data), header=None, skiprows=lambda x: x not in [1, 3]
    )
    tm.assert_frame_equal(result, expected)


def test_skip_rows_skip_all(all_parsers):
    parser = all_parsers
    data = "a\n1\n2\n3\n4\n5"
    msg = "No columns to parse from file"

    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), skiprows=lambda x: True)
        return

    with pytest.raises(EmptyDataError, match=msg):
        parser.read_csv(StringIO(data), skiprows=lambda x: True)


def test_skip_rows_bad_callable(all_parsers):
    msg = "by zero"
    parser = all_parsers
    data = "a\n1\n2\n3\n4\n5"

    if parser.engine == "pyarrow":
        msg = "skiprows argument must be an integer when using engine='pyarrow'"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), skiprows=lambda x: 1 / 0)
        return

    with pytest.raises(ZeroDivisionError, match=msg):
        parser.read_csv(StringIO(data), skiprows=lambda x: 1 / 0)


def test_skip_rows_and_n_rows(all_parsers):
    # GH#44021
    data = """a,b
1,a
2,b
3,c
4,d
5,e
6,f
7,g
8,h
"""
    parser = all_parsers
    if parser.engine == "pyarrow":
        msg = "The 'nrows' option is not supported with the 'pyarrow' engine"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(StringIO(data), nrows=5, skiprows=[2, 4, 6])
        return

    result = parser.read_csv(StringIO(data), nrows=5, skiprows=[2, 4, 6])
    expected = pd.DataFrame({"a": [1, 3, 5, 7, 8], "b": ["a", "c", "e", "g", "h"]})
    tm.assert_frame_equal(result, expected)


def test_skip_rows_with_chunks(all_parsers):
    # GH 55677
    data = """col_a
10
20
30
40
50
60
70
80
90
100
"""
    parser = all_parsers
    if parser.engine == "pyarrow":
        msg = "The 'chunksize' option is not supported with the 'pyarrow' engine"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(
                StringIO(data),
                engine=parser,
                skiprows=lambda x: x in [1, 4, 5],
                chunksize=4,
            )
        return

    reader = parser.read_csv(
        StringIO(data), engine=parser, skiprows=lambda x: x in [1, 4, 5], chunksize=4
    )
    df1 = next(reader)
    df2 = next(reader)

    tm.assert_frame_equal(df1, pd.DataFrame({"col_a": [20, 30, 60, 70]}))
    tm.assert_frame_equal(df2, pd.DataFrame({"col_a": [80, 90, 100]}, index=[4, 5, 6]))


DATA_WITH_PREAMBLE = "j0,j0b\nj1,j1b\nj2,j2b\na,b\n1,2\n3,4\n"


@pytest.mark.parametrize(
    "skiprows,header", [(3, "infer"), (3, 0), (2, 1), (1, 2), (3, [0])]
)
def test_skiprows_with_header(all_parsers, skiprows, header):
    # GH#48507 the pyarrow engine dropped skiprows whenever header was not None
    parser = all_parsers

    result = parser.read_csv(
        StringIO(DATA_WITH_PREAMBLE), skiprows=skiprows, header=header
    )
    expected = pd.DataFrame({"a": [1, 3], "b": [2, 4]})
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("skiprows,header", [(3, 0), (2, 1), (1, 2)])
def test_skiprows_with_header_and_names(all_parsers, skiprows, header):
    # GH#48507
    parser = all_parsers

    result = parser.read_csv(
        StringIO(DATA_WITH_PREAMBLE), skiprows=skiprows, header=header, names=["x", "y"]
    )
    expected = pd.DataFrame({"x": [1, 3], "y": [2, 4]})
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("kwargs", [{"usecols": ["b"]}, {"index_col": "a"}])
def test_skiprows_with_header_column_selection(all_parsers, kwargs):
    # GH#48507 usecols/index_col were resolved against the wrong header row
    parser = all_parsers

    result = parser.read_csv(StringIO(DATA_WITH_PREAMBLE), skiprows=3, **kwargs)
    expected = pd.DataFrame({"b": [2, 4]})
    if "index_col" in kwargs:
        expected.index = pd.Index([1, 3], name="a")
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("header", ["infer", None])
def test_skiprows_negative(all_parsers, header):
    # GH#48507 a negative skiprows skips nothing, as range(-1) is empty for the
    # c and python engines; pyarrow must not be handed it
    parser = all_parsers
    data = "a,b\n1,2\n"

    result = parser.read_csv(StringIO(data), skiprows=-1, header=header)
    expected = parser.read_csv(StringIO(data), header=header)
    tm.assert_frame_equal(result, expected)


def test_skiprows_past_end_of_file(all_parsers):
    # GH#48507 skiprows reaches pyarrow's skip_rows on the default header path,
    # where skipping past the end of the file raises instead of EmptyDataError
    parser = all_parsers
    data = "a,b\n1,2\n"

    if parser.engine == "pyarrow":
        with pytest.raises(ParserError, match="Could not skip initial 5 rows"):
            parser.read_csv(StringIO(data), skiprows=5)
        return

    with pytest.raises(EmptyDataError, match="No columns to parse from file"):
        parser.read_csv(StringIO(data), skiprows=5)


def test_skiprows_to_final_line_without_newline(all_parsers):
    # GH#48507 pyarrow's skip_rows cannot skip past a final line lacking a
    # trailing newline; skip_rows_after_names covers only the header row
    parser = all_parsers
    data = "j0\na,b"

    if parser.engine == "pyarrow":
        with pytest.raises(EmptyDataError, match="No columns to parse from file"):
            parser.read_csv(StringIO(data), skiprows=1)
        return

    result = parser.read_csv(StringIO(data), skiprows=1)
    expected = pd.DataFrame(columns=["a", "b"], dtype=object)
    tm.assert_frame_equal(result, expected)
