# Unit test exog_long_to_dict
# ==============================================================================
import pandas as pd
import numpy as np
import pytest
from ..preprocessing import exog_long_to_dict

# Fixtures
n_exog_A = 10
n_exog_B = 5
n_exog_C = 15
index_exog_A = pd.date_range("2020-01-01", periods=n_exog_A, freq="D")
index_exog_B = pd.date_range("2020-01-01", periods=n_exog_B, freq="D")
index_exog_C = pd.date_range("2020-01-01", periods=n_exog_C, freq="D")
exog_A = pd.DataFrame(
    {
        "series_id": "A",
        "datetime": index_exog_A,
        "exog_1": np.arange(n_exog_A),
    }
)
exog_B = pd.DataFrame(
    {
        "series_id": "B",
        "datetime": index_exog_B,
        "exog_1": np.arange(n_exog_B),
        "exog_2": "b",
    }
)
exog_C = pd.DataFrame(
    {
        "series_id": "C",
        "datetime": index_exog_C,
        "exog_1": np.arange(n_exog_C),
        "exog_3": 1,
    }
)
exog_long = pd.concat([exog_A, exog_B, exog_C], axis=0)


def test_check_output_series_long_to_dict_dropna_False():
    """
    Check output of exog_long_to_dict with dropna=False.
    """
    exog_columns = pd.DataFrame(
        columns=["exog_1", "exog_2", "exog_3"],
    )
    exog_columns = pd.DataFrame(
        columns=["exog_1", "exog_2", "exog_3"],
    ).astype({"exog_1": int, "exog_2": str, "exog_3": float})
    expected = {
        "A": pd.concat(
            [
                exog_A.set_index("datetime").asfreq("D").drop(columns="series_id"),
                exog_columns,
            ],
            axis=0,
        )[["exog_1", "exog_2", "exog_3"]],
        "B": pd.concat(
            [
                exog_B.set_index("datetime").asfreq("D").drop(columns="series_id"),
                exog_columns,
            ],
            axis=0,
        )[["exog_1", "exog_2", "exog_3"]],
        "C": pd.concat(
            [
                exog_C.set_index("datetime").asfreq("D").drop(columns="series_id"),
                exog_columns,
            ],
            axis=0,
        )[["exog_1", "exog_2", "exog_3"]],
    }
    results = exog_long_to_dict(
        data=exog_long,
        series_id="series_id",
        index="datetime",
        freq="D",
        dropna=False,
    )

    for k in expected.keys():
        pd.testing.assert_frame_equal(results[k], expected[k])


def test_check_output_series_long_to_dict_dropna_True():
    """
    Check output of series_long_to_dict with dropna=True.
    """

    expected = {
        "A": exog_A.set_index("datetime").asfreq("D").drop(columns="series_id"),
        "B": exog_B.set_index("datetime").asfreq("D").drop(columns="series_id"),
        "C": exog_C.set_index("datetime").asfreq("D").drop(columns="series_id"),
    }

    for k in expected.keys():
        expected[k].index.name = None

    results = exog_long_to_dict(
        data=exog_long,
        series_id="series_id",
        index="datetime",
        freq="D",
        dropna=True,
    )

    for k in expected.keys():
        pd.testing.assert_frame_equal(results[k], expected[k], check_dtype=False)


def test_TypeError_when_data_is_not_dataframe():
    """
    Raise TypeError if data is not a pandas DataFrame.
    """
    err_msg = "`data` must be a pandas DataFrame."
    with pytest.raises(TypeError, match=err_msg):
        exog_long_to_dict(
            data="not_a_dataframe",
            series_id="series_id",
            index="datetime",
            freq="D",
        )


def test_ValueError_when_series_id_not_in_data():
    """
    Raise ValueError if series_id is not in data.
    """
    series_id = "series_id_not_in_data"
    err_msg = f"Column '{series_id}' not found in `data`."
    with pytest.raises(ValueError, match=err_msg):
        exog_long_to_dict(
            data=exog_long,
            series_id=series_id,
            index="datetime",
            freq="D",
        )


def test_ValueError_when_index_not_in_data():
    """
    Raise ValueError if index is not in data.
    """
    index = "series_id_not_in_data"
    err_msg = f"Column '{index}' not found in `data`."
    with pytest.raises(ValueError, match=err_msg):
        exog_long_to_dict(
            data=exog_long,
            series_id="series_id",
            index=index,
            freq="D",
        )
