import pytest
from test_tc.algorithms.models import *  

# --------- ProphetModel Tests --------- #
## test ProphetModel __init__ ##
# The __init__ test also the method make_grid_params 
@pytest.fixture
def sample_params():
    return {'weekly_seasonality' : ['True'],
            'yearly_seasonality' : ['True']
}

def test_prophet_model_init(sample_params):
    config_path = "./config/config_check"
    max_na_ratio = 0.3
    model = ProphetModel(sample_params, config_path, max_na_ratio)

    assert model.config_path == config_path
    assert model.max_na_ratio == max_na_ratio
    assert model.grid_params is not None
    assert model.grid_params == {'weekly_seasonality' : [True], 
                                 'yearly_seasonality' : [True],
                                 'interval_width' : [0.95]}

##  test ProphetModel fit method ##
def create_data():
    unique_ids_train = ['a'] * 20 + ['a/b'] * 20 + ['a/c'] * 20 + ['a/d'] * 20
    unique_ids_val = ['a'] * 10 + ['a/b'] * 10 + ['a/c'] * 10 + ['a/d'] * 10

    # Creazione delle date con ripetizioni ogni 20 giorni per train_data e ogni 10 giorni per val_data
    ds_train = pd.date_range(start='2023-01-01', periods=20, freq='20D').repeat(4)
    ds_val = pd.date_range(start='2023-01-06', periods=10, freq='10D').repeat(4)

    y_train = list(range(1, 61))
    y_train += [0]*20
    y_val = list(range(61, 91))
    y_val += [0]*10

    # Creazione dei DataFrame train_data e val_data
    train_data = pd.DataFrame({'unique_id': unique_ids_train, 'ds': ds_train, 'y': y_train})
    val_data = pd.DataFrame({'unique_id': unique_ids_val, 'ds': ds_val, 'y': y_val})

    # Reorder columns
    train_data = train_data[['unique_id', 'ds', 'y']]
    val_data = val_data[['unique_id', 'ds', 'y']]
    return train_data, val_data

def test_prophet_model_fit():
    config_path = "./config/config_check"
    max_na_ratio = 0.3
    params = {'weekly_seasonality' : ['True'],
            'yearly_seasonality' : ['True']}
    reference_params = {'weekly_seasonality' : True, 
                                 'yearly_seasonality' : True,
                                 'interval_width' : 0.95}
    model = ProphetModel(params, config_path, max_na_ratio)

    train_data, val_data = create_data()

    # Specificare il numero di ripetizioni desiderate
    num_repetitions = 20

    # Creazione dei DataFrame con più righe
    train_data = pd.concat([train_data] * num_repetitions)
    val_data = pd.concat([val_data] * num_repetitions)

    train_data = train_data.set_index('unique_id')
    val_data = val_data.set_index('unique_id')
    train_data = train_data.sort_index()
    val_data = val_data.sort_index()

    # Call the fit method
    model.fit(train_data, val_data)

    # Check if models_dict and models_best_params are populated
    assert model.models_dict != {}
    assert model.models_best_params != {}
    for k, best_param in model.models_best_params.items():
        if k != 'a/d':
            assert best_param == reference_params
        else:
            assert best_param == {}

## test ProphetModel refit method ##
def test_prophet_model_refit():
    # first fit
    config_path = "./config/config_check"
    max_na_ratio = 0.3
    params = {'weekly_seasonality' : ['True'],
            'yearly_seasonality' : ['True']}
    reference_params = {'weekly_seasonality' : True, 
                                 'yearly_seasonality' : True,
                                 'interval_width' : 0.95}
    model = ProphetModel(params, config_path, max_na_ratio)

    train_data, val_data = create_data()

    # Specificare il numero di ripetizioni desiderate
    num_repetitions = 20

    # Creazione dei DataFrame con più righe
    train_data = pd.concat([train_data] * num_repetitions)
    val_data = pd.concat([val_data] * num_repetitions)

    train_data = train_data.set_index('unique_id')
    val_data = val_data.set_index('unique_id')
    train_data = train_data.sort_index()
    val_data = val_data.sort_index()

    # Call the fit method
    model.fit(train_data, val_data)

    # Then refit
    model.refit(train_data)

    # Check if models_dict and models_best_params are populated
    assert model.models_dict != {}
    assert model.models_best_params != {}
    for k, best_param in model.models_best_params.items():
        if k != 'a/d':
            assert best_param == reference_params
        else:
            assert best_param == {}

## test ProphetModel predict method ##
# The function checks also the method prepare_output
def test_prophet_model_predict():
    # first fit
    config_path = "./config/config_check"
    max_na_ratio = 0.3
    params = {'weekly_seasonality' : ['True'],
            'yearly_seasonality' : ['True']}
   
    model = ProphetModel(params, config_path, max_na_ratio)

    train_data, val_data = create_data()

    # Specificare il numero di ripetizioni desiderate
    num_repetitions = 20

    # Creazione dei DataFrame con più righe
    train_data = pd.concat([train_data] * num_repetitions)
    val_data = pd.concat([val_data] * num_repetitions)

    train_data = train_data.set_index('unique_id')
    val_data = val_data.set_index('unique_id')
    train_data = train_data.sort_index()
    val_data = val_data.sort_index()

    # Call the fit method
    model.fit(train_data, val_data)

    # Then predict
    output = model.predict(val_data)

    # Assert output columns
    assert np.all(output.columns == ['Timestamp', 'Id_Pred' , 'Pred_Mean'  ,'Sigma',  'Pi_Lower_95' , 'Pi_Upper_95'])
    
    # Assert Id_Pred consistency
    assert np.all(output['Id_Pred'].unique() == ['a', 'a/b', 'a/c', 'a/d'])

    # Assert predictions for 'a/d'
    for col in ['Pred_Mean', 'Sigma', 'Pi_Lower_95', 'Pi_Upper_95']:
        assert  np.all(output[output['Id_Pred'] == 'a/d'][col] == 0)

## test ProphetModel score method ## 
def test_prophet_model_score():
    # Obtain predictions
    config_path = "./config/config_check"
    max_na_ratio = 0.3
    params = {'weekly_seasonality' : ['True'],
            'yearly_seasonality' : ['True']}
   
    model = ProphetModel(params, config_path, max_na_ratio)

    train_data, val_data = create_data()

    # Specificare il numero di ripetizioni desiderate
    num_repetitions = 20

    # Creazione dei DataFrame con più righe
    train_data = pd.concat([train_data] * num_repetitions)
    val_data = pd.concat([val_data] * num_repetitions)

    train_data = train_data.set_index('unique_id')
    val_data = val_data.set_index('unique_id')
    train_data = train_data.sort_index()
    val_data = val_data.sort_index()

    # Call the fit method
    model.fit(train_data, val_data)

    # Then predict
    df_pred = model.predict(val_data)

    # Then score
    score = model.score(val_data, df_pred)
    
    # Assert score columns
    assert np.all(score.columns == ['Id_Pred', 'MAE', 'MAPE', 'RMSE', 'MSE', 'R2', 'Percentage Coverage'])


# --------- HierarchicalReconcileMethod Tests --------- #
## test HierarchicalReconcileMethod __init__ ##
def sample_data():
    model_name = "test_model"
    S = pd.DataFrame({
        "column1": [1, 2, 3, 4, None],
        "column2": [10, 20, None, 40, 50]
    })
    tags = {"tag1": ["column1", "column2"]}
    reconciler = None
    max_na_ratio = 0.5

    return model_name, S, tags, reconciler, max_na_ratio

def test_init_with_none_reconciler():
    model_name, S, tags, reconciler, max_na_ratio = sample_data()
    model = HierarchicalReconcileModel(model_name, S, tags, reconciler, max_na_ratio)

    assert model.model_name == model_name
    assert model.S.equals(S)
    assert model.tags == tags
    assert model.max_na_ratio == max_na_ratio

def test_init_with_custom_reconciler():
    model_name, S, tags, reconciler, max_na_ratio = sample_data()
    reconciler = BottomUp()
    model = HierarchicalReconcileModel(model_name, S, tags, reconciler, max_na_ratio)

    assert model.model_name == model_name
    assert model.S.equals(S)
    assert model.tags == tags
    assert len(model.reconciler) == 1
    assert model.max_na_ratio == max_na_ratio


## test HierarchicalReconcileMethod prepare_dataframe_to_hf ##
def prepare_data():
    X_pred = pd.DataFrame({
        'Timestamp': ['2023-07-20', '2023-07-21'],
        'Pred_Mean': [10, 15],
        'Id_Pred': [1, 2]
    })
    X_true = pd.DataFrame({
        'Timestamp': ['2023-07-20', '2023-07-21'],
        'y': [8, 12]
    })

    return X_pred, X_true

def test_prepare_dataframe_to_hf_with_X_true():
    X_pred, X_true = prepare_data()
    model_name, S, tags, reconciler, max_na_ratio = sample_data()
    model = HierarchicalReconcileModel(model_name, S, tags, reconciler, max_na_ratio)

    result = model.prepare_dataframe_to_hf(X_pred, model_name, X_true)

    expected_columns = ['ds', 'Test_Model', 'y']
    assert result.columns.tolist() == expected_columns
    assert result.index.tolist() == [1, 2]
    assert result['Test_Model'].tolist() == [10, 15]
    assert result['y'].tolist() == [8, 12]

def test_prepare_dataframe_to_hf_without_X_true():
    X_pred, _ = prepare_data()
    model_name, S, tags, reconciler, max_na_ratio = sample_data()
    model = HierarchicalReconcileModel(model_name, S, tags, reconciler, max_na_ratio)

    result = model.prepare_dataframe_to_hf(X_pred, model_name)

    expected_columns = ['ds', 'Test_Model']
    assert result.columns.tolist() == expected_columns
    assert result.index.tolist() == [1, 2]
    assert result['Test_Model'].tolist() == [10, 15]


## test HierarchicalReconcileMethod fit ##
def test_HierarchicalReconcileMethod_fit():
    X_pred, X_true = prepare_data()
    model_name, S, tags, reconciler, max_na_ratio = sample_data()
    model = HierarchicalReconcileModel(model_name, S, tags, reconciler, max_na_ratio)

    updated_self = model.fit(X_true, X_pred, X_pred)

    assert updated_self.X.columns.tolist() == [ 'ds', 'Test_Model', 'y']  
    assert updated_self.X_hat.columns.tolist() == [ 'ds', 'Test_Model']  


## test HierarchicalReconcileMethod predict ##
# The test checks also the prepare_output method
from hierarchicalforecast.utils import aggregate
import warnings

def prepare_data_reconcile():
    df = pd.DataFrame(columns = ['level0',	'level1',	'ds',	'y'])
    df['ds'] = ['2023-07-20', '2023-07-21', '2023-07-22']*2
    df['y'] = 1
    df['level0'] = 'a'
    df['level1'] =  ['b', 'b', 'b', 'c', 'c', 'c']
    spec = [['level0'], ['level0', 'level1'],]

    Y_df, S_df, tags = aggregate(df=df, spec=spec)
    Y_df = Y_df.reset_index()
    return Y_df, S_df, tags 

def prepare_input_reconcile():
    X_pred = pd.DataFrame({
        'Timestamp': ['2023-07-20', '2023-07-21', '2023-07-22', '2023-07-23',
                      '2023-07-24', '2023-07-25', '2023-07-26', '2023-07-27']*3,
        'Pred_Mean': [10] * 8 + [4] * 8 + [2] * 8,
        'Id_Pred': ['a'] * 8 + ['a/b'] * 8 + ['a/c'] * 8,
        'Pi_Upper_95': [11] * 8 + [5] * 8 + [3] * 8,
        'Pi_Lower_95': [9] * 8 + [3] * 8 + [2] * 8

    })

    X_true = pd.DataFrame({
        'ds': ['2023-07-20', '2023-07-21', '2023-07-22', '2023-07-23',
                      '2023-07-24', '2023-07-25', '2023-07-26', '2023-07-27']*3,
        'y': [2] * 8 + [1] * 8 + [1] * 8,
        'unique_id':  ['a'] * 8 + ['a/b'] * 8 + ['a/c'] * 8
    }).set_index('unique_id')

    X_future_pred = pd.DataFrame({
        'Timestamp': ['2023-07-28', '2023-07-29', '2023-07-30']*3,
        'Pred_Mean': [10, 10, 10, 4, 4, 2, 5, 5, 2],
        'Id_Pred': ['a', 'a', 'a', 'a/b', 'a/b', 'a/b', 'a/c', 'a/c', 'a/c']
    })

    return X_true, X_pred, X_future_pred

def test_HierarchicalReconcileMethod_predict():
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        _, S_df, tags = prepare_data_reconcile()
        X_true, X_pred, X_future_pred = prepare_input_reconcile()
        h_model = HierarchicalReconcileModel(model_name = 'test_model', S = S_df, tags = tags,
                                            reconciler = [TopDown(method = 'average_proportions')])
        h_model.fit(X_true, X_pred, X_future_pred)

        Y_rec = h_model.predict(as_output = True)

    # Assert if Y_rec is consistent with as_output = True
    assert Y_rec.columns.tolist() == ['Timestamp', 'Id_Pred', 'Pred_Mean', 'Sigma', 'Pi_Lower_95','Pi_Upper_95']
    # Assert if the reconciliation is logically correct
    for ds in Y_rec['Timestamp'].unique():
        check_df = Y_rec[Y_rec['Timestamp'] == ds].reset_index()
        assert check_df['Pred_Mean'][0] == check_df['Pred_Mean'][1:].sum()


## test HierarchicalReconcileMethod score ##
def test_HierarchicalReconcileMethod_score():
    import math
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        _, S_df, tags = prepare_data_reconcile()
        X_true, X_pred, X_future_pred = prepare_input_reconcile()
        h_model = HierarchicalReconcileModel(model_name = 'test_model', S = S_df, tags = tags,
                                            reconciler = [TopDown(method = 'average_proportions')])
        h_model.fit(X_true, X_pred, X_future_pred)

        # Check score NOT equal values
        # just check the columns
        Y_rec = h_model.predict(as_output = True)
        X_future_true = X_future_pred
        X_future_true.rename(columns = {'Id_Pred': 'unique_id'}, inplace= True)
        X_future_true = X_future_true.set_index('unique_id')
        X_future_true.columns = ['ds', 'y'] 
        _, scores = h_model.score(X_true = X_future_true.reset_index(), X_pred = Y_rec, check_model_pred = Y_rec, hierarchical_output = False)

        assert scores.columns.tolist() == ['Id_Pred', 'MAE', 'MAPE', 'RMSE', 'MSE', 'R2', 'Percentage Coverage']

        # Check score equal values
        # check the expected results
        X_future_true['y'] = list(Y_rec['Pred_Mean'])
        scores_aggr, _ = h_model.score(X_true = X_future_true.reset_index(), X_pred = Y_rec, check_model_pred = Y_rec, hierarchical_output = False)
        expected_scores = [0.0, 0.0, 0.0, 0.0, np.nan, 0.0]
        for i, (x, y) in enumerate(zip(scores_aggr.values, expected_scores)):
            if math.isnan(x):
                assert math.isnan(y)
            else:
                assert x == y

## test HierarchicalReconcileMethod best_reconciler ##
def test_HierarchicalReconcileMethod_best_reconciler():
    with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            _, S_df, tags = prepare_data_reconcile()
            X_true, X_pred, X_future_pred = prepare_input_reconcile()
            h_model = HierarchicalReconcileModel(model_name = 'test_model', S = S_df, tags = tags,
                                                reconciler = [TopDown(method = 'average_proportions')])
            h_model.fit(X_true, X_pred, X_future_pred)

            X_future_true = X_future_pred.copy()
            X_future_true.rename(columns = {'Id_Pred': 'unique_id'}, inplace= True)
            X_future_true = X_future_true.set_index('unique_id')
            X_future_true.columns = ['ds', 'y']

            best_reconciler = h_model.best_reconciler(X_true=X_true, X_pred=X_pred, X_future_true = X_future_true, X_future_pred=X_future_pred)
            reconciler_string = str(best_reconciler).split(' ')[0]
            top_down_string = reconciler_string.split('.')[2]
            assert top_down_string == 'TopDown'


def test_boosting_reconciler_init():
    model = BoostingReconcileModel()
    assert isinstance(model.cat_features, list)
    assert isinstance(model.xgb_params, dict)
    assert isinstance(model.preprocess_params, dict)
    
    xgb = {"learning_rate": 0.2, "n_iterations": 4}
    preproc = {"father_target": True, "father_preds": False}
    cat_features = ["dayofweek", "month"]
    model = BoostingReconcileModel(
        params_xgb=xgb,
        params_preprocess=preproc,
        cat_features=cat_features
        )
    assert model.cat_features == cat_features
    assert model.xgb_params == xgb
    assert model.preprocess_params == preproc




