"""
Pipeline Tests

Description: Pipeline tests ensures that the functionalities of the feature processing pipelines
work as expected to avoid feeding models incorrect data.


"""

import pytest
import numpy as np
import pandas as pd
from absl import logging
from forecasting_toolkit.feature_preprocessing.timeseries.pipelines import (
    feature_preprocessing_pipeline_factory_v1
)

@pytest.fixture()
def date_range():
    return ['1/1/2018','1/08/2018']

def test_onehot_encoding_of_categorical_columns(date_range):
    """
    Tests the pipeline that encodes categorical features
    """
    categories = ['audio', 'video']

    date_range = pd.date_range(start=date_range[0], end=date_range[1])
    N = len(date_range)
    test_df = pd.DataFrame({
        'date': date_range,
        'streams': [np.random.randint(0, 10e3) for _ in range(N)],
        'stream_type': [np.random.choice(categories) for _ in range(N)]
    })
    
    # feature preprocessing pipelines
    pipeline_fn = feature_preprocessing_pipeline_factory_v1(
                    categorical_cols=["stream_type"], 
                    numerical_float_cols=[],
                    bool_cols=[],
                    time_cols=[],
                    text_cols=[]
    )
    transformed_test_arr = pipeline_fn.fit_transform(test_df)
    
    # transformed test
    logging.debug(transformed_test_arr)
    assert transformed_test_arr.shape == (N, 2)


def test_float_standardization(date_range):
    """
    Tests the pipeline for float standardization
    """
    categories = ['audio', 'video']

    date_range = pd.date_range(start=date_range[0], end=date_range[1])
    N = len(date_range)
    test_df = pd.DataFrame({
        'date': date_range,
        'streams': [np.random.randint(0, 10e3) for _ in range(N)],
        'stream_type': [np.random.choice(categories) for _ in range(N)]
    })
    mean_streams = test_df["streams"].mean()
    std_streams = test_df["streams"].std()
    test_df["standardized_streams"] = (test_df["streams"] - mean_streams)/ std_streams
    
    # feature preprocessing pipelines
    pipeline_fn = feature_preprocessing_pipeline_factory_v1(
                    categorical_cols=[], 
                    numerical_float_cols=["streams"],
                    bool_cols=[],
                    time_cols=[],
                    text_cols=[]
    )
    transformed_test_arr = pipeline_fn.fit_transform(test_df)
    
    # transformed test
    assert transformed_test_arr.shape == (N, 1)
    assert np.allclose(a=transformed_test_arr.reshape(-1), b=test_df["standardized_streams"].values, atol=0.5)

def test_float_standardization(date_range):
    """
    Tests the pipeline that encodes categorical features
    """
    categories = ['audio', 'video']

    date_range = pd.date_range(start=date_range[0], end=date_range[1])
    N = len(date_range)
    test_df = pd.DataFrame({
        'date': date_range,
        'streams': [np.random.randint(0, 10e3) for _ in range(N)],
        'stream_type': [np.random.choice(categories) for _ in range(N)]
    })
    mean_streams = test_df["streams"].mean()
    std_streams = test_df["streams"].std()
    test_df["standardized_streams"] = (test_df["streams"] - mean_streams)/ std_streams
    
    # feature preprocessing pipelines
    pipeline_fn = feature_preprocessing_pipeline_factory_v1(
                    categorical_cols=[], 
                    numerical_float_cols=["streams"],
                    bool_cols=[],
                    time_cols=[],
                    text_cols=[]
    )
    transformed_test_arr = pipeline_fn.fit_transform(test_df)
    
    # transformed test
    assert transformed_test_arr.shape == (N, 1)
    assert np.allclose(a=transformed_test_arr.reshape(-1), b=test_df["standardized_streams"].values, atol=0.5)