# Unit test predict_dist ForecasterRecursiveMultiSeries
# ==============================================================================
import re
import pytest
import pandas as pd
from ....recursive import ForecasterRecursiveMultiSeries
from scipy.stats import norm
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import OneHotEncoder
from sklearn.linear_model import LinearRegression

# Fixtures
from .fixtures_forecaster_recursive_multiseries import (
    series_dict_range,
    exog_wide_range,
    exog_pred_wide_range
)

transformer_exog = ColumnTransformer(
                       [('scale', StandardScaler(), ['exog_1']),
                        ('onehot', OneHotEncoder(), ['exog_2'])],
                       remainder = 'passthrough',
                       verbose_feature_names_out = False
                   )


def test_predict_dist_TypeError_when_distribution_object_is_not_valid():
    """
    Test TypeError is raise in predict_dist when `distribution` is not a valid
    probability distribution object from scipy.stats.
    """
    forecaster = ForecasterRecursiveMultiSeries(
                     regressor          = LinearRegression(),
                     lags               = 3,
                     transformer_series = StandardScaler(),
                     transformer_exog   = transformer_exog,
                 )
    forecaster.fit(
        series=series_dict_range, exog=exog_wide_range, store_in_sample_residuals=True
    )
    
    class CustomObject:  # pragma: no cover
        pass
    
    err_msg = re.escape(
        "`distribution` must be a valid probability distribution object "
        "from scipy.stats, with methods `_pdf` and `fit`."
    )
    with pytest.raises(TypeError, match = err_msg):
        forecaster.predict_dist(
            steps                   = 2,
            exog                    = exog_pred_wide_range,
            distribution            = CustomObject(),
            n_boot                  = 4,
            use_in_sample_residuals = True,
            use_binned_residuals    = False
        )


@pytest.mark.parametrize("level", 
                         ['l1', ['l1']], 
                         ids=lambda lvl: f'level: {lvl}')
def test_predict_dist_output_when_forecaster_is_LinearRegression_steps_is_2_in_sample_residuals_True_exog_and_transformer(level):
    """
    Test output of predict_dist when regressor is LinearRegression,
    2 steps are predicted, using in-sample residuals, exog is included and both
    inputs are transformed. Single level.
    """
    forecaster = ForecasterRecursiveMultiSeries(
                     regressor          = LinearRegression(),
                     lags               = 3,
                     transformer_series = StandardScaler(),
                     transformer_exog   = transformer_exog,
                 )

    forecaster.fit(
        series=series_dict_range, exog=exog_wide_range, store_in_sample_residuals=True
    )
    results = forecaster.predict_dist(
                  steps                   = 2,
                  distribution            = norm,
                  levels                  = level,
                  exog                    = exog_pred_wide_range,
                  n_boot                  = 4,
                  use_in_sample_residuals = True,
                  use_binned_residuals    = False,
                  suppress_warnings       = True
              )

    expected = pd.DataFrame(
        {
            "level": ["l1", "l1"],
            "loc": [0.3071804611998517, 0.33695529446570166],
            "scale": [0.1435578185939024, 0.21900963160500286],
        },
        index=pd.RangeIndex(start=50, stop=52),
    )

    pd.testing.assert_frame_equal(results, expected)


@pytest.mark.parametrize("levels", 
                         [['l1', 'l2'], None], 
                         ids=lambda lvl: f'levels: {lvl}')
def test_predict_dist_output_when_forecaster_is_LinearRegression_steps_is_2_in_sample_residuals_False_exog_and_transformer(levels):
    """
    Test output of predict_dist when regressor is LinearRegression,
    2 steps are predicted, using out-sample residuals, exog is included and both
    inputs are transformed. Multiple levels.
    """
    forecaster = ForecasterRecursiveMultiSeries(
                     regressor          = LinearRegression(),
                     lags               = 3,
                     transformer_series = StandardScaler(),
                     transformer_exog   = transformer_exog,
                 )

    forecaster.fit(
        series=series_dict_range, exog=exog_wide_range, store_in_sample_residuals=True
    )
    forecaster.out_sample_residuals_ = forecaster.in_sample_residuals_
    results = forecaster.predict_dist(
                  steps                   = 2,
                  distribution            = norm,
                  levels                  = levels,
                  exog                    = exog_pred_wide_range,
                  n_boot                  = 4,
                  use_in_sample_residuals = False,
                  use_binned_residuals    = False
              )

    expected = pd.DataFrame(
        {
            "level": ["l1", "l2", "l1", "l2"],
            "loc": [
                0.3071804611998517,
                0.5812127014174517,
                0.33695529446570166,
                0.1296887412374588,
            ],
            "scale": [
                0.1435578185939024,
                0.317378881479421,
                0.21900963160500286,
                0.06418038020863531,
            ],
        },
        index=pd.Index([50, 50, 51, 51]),
    )

    pd.testing.assert_frame_equal(results, expected)
