from pywr.dataframe_tools import align_and_resample_dataframe, ResamplingError
import numpy as np
import pandas as pd
import pytest


def make_df(freq, start="2015-01-01", end="2015-12-31"):
    # Daily time-step
    index = pd.period_range(start, end, freq=freq)
    series = pd.DataFrame(np.arange(len(index), dtype=np.float64), index=index)
    return series


def make_model_index(freq, start="2015-01-01", end="2015-12-31"):
    return pd.period_range(start, end, freq=freq)


class TestDownSampling:
    """Test for down-sampling a dataframe to lower frequency model time-step."""

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_daily_to_monthly(self, resample_func):
        """Test daily data to monthly model time-step."""
        input_df = make_df("D")
        model_index = make_model_index("M")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(
            input_resampled, input_resampled.resample("M").agg(resample_func)
        )

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_daily_to_weekly(self, resample_func):
        """Test daily to weekly model time-step."""
        input_df = make_df("D")
        model_index = make_model_index("W")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(
            input_resampled, input_resampled.resample("W").agg(resample_func)
        )

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_daily_to_7daily(self, resample_func):
        """Test daily to 7-day model time-step."""
        input_df = make_df("D")
        model_index = make_model_index("7D")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(
            input_resampled, input_resampled.resample("7D").agg(resample_func)
        )

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_misaligned_daily_to_7daily(self, resample_func):
        """Test daily to 7-day model time-step."""
        input_df = make_df("D", start="2014-12-20")
        model_index = make_model_index("7D")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(
            input_resampled, input_resampled.resample("7D").agg(resample_func)
        )

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_weekly_to_monthly(self, resample_func):
        """Test weekly to monthly model time-step."""
        input_df = make_df("W")
        model_index = make_model_index("M")

        with pytest.raises(ResamplingError):
            input_resampled = align_and_resample_dataframe(
                input_df, model_index, resample_func=resample_func
            )

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_weekly_to_two_weekly(self, resample_func):
        """Test weekly to two-weekly model time-step."""
        input_df = make_df("W")
        model_index = make_model_index("2W")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(
            input_resampled, input_resampled.resample("2W").agg(resample_func)
        )


class TestUpSampling:
    """Test for up-sampling a dataframe to higher frequency model time-step."""

    def test_monthly_to_daily(self):
        """Test monthly data to daily model time-step."""
        input_df = make_df("M")
        model_index = make_model_index("D")

        input_resampled = align_and_resample_dataframe(input_df, model_index)

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    def test_weekly_to_daily(self):
        """Test weekly aligned to daily model time-step."""
        # Note that we have to choose a week starting on Wednesday because 2015-01-01 is a Wednesday.
        input_df = make_df("W-WED")
        model_index = make_model_index("D")

        input_resampled = align_and_resample_dataframe(input_df, model_index)

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    def test_7daily_to_daily(self):
        """Test weekly aligned to daily model time-step."""
        # Note that we have to choose a week starting on Wednesday because 2015-01-01 is a Wednesday.
        input_df = make_df("7D")
        model_index = make_model_index("D")

        input_resampled = align_and_resample_dataframe(input_df, model_index)

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    def test_nonaligned_weekly_to_daily(self):
        """Test weekly non-aligned data to daily model time-step."""
        input_df = make_df("W")
        model_index = make_model_index("D")

        input_resampled = align_and_resample_dataframe(input_df, model_index)

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_weekly_to_7daily(self, resample_func):
        """Test weekly to 7-day model time-step."""
        # Note that we have to choose a week starting on Wednesday because 2015-01-01 is a Wednesday.
        input_df = make_df("W-WED")
        model_index = make_model_index("7D")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]
        expected_df = expected_df.resample("7D").agg(resample_func)

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    @pytest.mark.parametrize("resample_func", ["mean", "max"])
    def test_nonaligned_weekly_to_7daily(self, resample_func):
        """Test invalid weekly data to 7-day model time-step."""
        # Because the weeks do not align with the 7-day time-step it is not possible to align this data.
        input_df = make_df("W")
        model_index = make_model_index("7D")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func=resample_func
        )

        pd.testing.assert_index_equal(input_resampled.index, model_index)

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]
        expected_df = expected_df.resample("7D").agg(resample_func)

        pd.testing.assert_frame_equal(input_resampled, expected_df)

    def test_annual_to_daily(self):
        input_df = make_df("A", "2010-01-01", "2020-01-01")
        model_index = make_model_index("D", "2010-01-01", "2020-01-01")

        input_resampled = align_and_resample_dataframe(input_df, model_index)

        expected_df = input_df.resample("D").ffill()
        expected_df = expected_df["2010-01-01":"2020-01-01"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    def test_annual_to_monthly(self):
        input_df = make_df("A", "2010-01-01", "2020-01-01")
        model_index = make_model_index("M", "2010-01-01", "2020-01-01")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func="ffill"
        )

        expected_df = input_df.resample("M").ffill()
        expected_df = expected_df["2010-01-01":"2020-01-01"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)

    def test_two_weekly_to_weekly(self):
        """Test weekly to two-weekly model time-step."""
        input_df = make_df("2W")
        model_index = make_model_index("W")

        input_resampled = align_and_resample_dataframe(
            input_df, model_index, resample_func="ffill"
        )

        expected_df = input_df.resample("W").ffill()
        expected_df = expected_df["2015-01-01":"2015-12-31"]

        pd.testing.assert_index_equal(input_resampled.index, model_index)
        pd.testing.assert_frame_equal(input_resampled, expected_df)
