import pandas

from hts import HTSRegressor
from hts.core.result import HTSResult


def test_instantiate_regressor():
    ht = HTSRegressor()
    assert ht.model == 'prophet'
    assert isinstance(ht.hts_result, HTSResult)
    assert isinstance(ht.hts_result.residuals, dict)


def test_fit_regressor_uv(load_df_and_hier_uv):
    hierarchical_sine_data, sine_hier = load_df_and_hier_uv
    ht = HTSRegressor(model='holt_winters', revision_method='OLS')
    ht.fit(df=hierarchical_sine_data, nodes=sine_hier)
    assert isinstance(ht.hts_result.errors, dict)
    assert isinstance(ht.hts_result.models, dict)
    assert 'total' in ht.hts_result.models


def test_predict_regressor_uv(load_df_and_hier_uv):
    hierarchical_sine_data, sine_hier = load_df_and_hier_uv
    hsd = hierarchical_sine_data.head(200)

    for model in ['holt_winters', 'auto_arima', 'sarimax']:
        ht = HTSRegressor(model=model, revision_method='OLS')

        ht.fit(df=hsd, nodes=sine_hier)
        preds = ht.predict(steps_ahead=10)

        assert isinstance(preds, pandas.DataFrame)

        # base + steps ahead
        assert len(preds) == len(hsd) + 10
        assert len(ht.hts_result.forecasts['total']) == len(hsd) + 10
        assert 'total' in ht.hts_result.errors
        assert 'total' in ht.hts_result.forecasts
        assert 'total' in ht.hts_result.residuals

        assert 'a' in ht.hts_result.errors
        assert 'a' in ht.hts_result.forecasts
        assert 'a' in ht.hts_result.residuals

