# pylint: disable=missing-docstring
# pylint: disable=invalid-name
import pandas as pd
import numpy as np
from doors.features import (
    categorical_to_frequency,
    grouped_lagged_decay,
    days_to_first_event,
    grouped_days_since_result,
    ema,
    grouped_ema,
)


def test_categorical_to_frequency():
    df = pd.DataFrame({
        'cat': [1, 1, 2, 3, 4, 4, 4],
    })
    result = categorical_to_frequency(df, 'cat')
    expected = [2, 2, 1, 1, 3, 3, 3]
    assert np.allclose(expected, result)


def test_grouped_lagged_decay():
    df = pd.DataFrame({
        'horse_name': ['x', 'x', 'x', 'x'],
        'win_flag': [1, 0, 1, 0]
        })

    result = grouped_lagged_decay(df, 'horse_name', 'win_flag')
    expected = [
        0, 1,
        (0 + 1*np.e**-1),
        (1 + 0*np.e**-1 + 1*np.e**-2)
    ]
    assert np.allclose(expected, result)


def test_grouped_lagged_decay_with_nans():
    df = pd.DataFrame({
        'horse_name': ['x', 'x', 'x', 'x'],
        'win_flag': [1, np.nan, 0, np.nan]
    })

    result = grouped_lagged_decay(df, 'horse_name', 'win_flag')
    expected = [
        0, 1,
        (0 + 1*np.e**-1),
        (0 + 0*np.e**-1 + 1*np.e**-2)
    ]
    assert np.allclose(expected, result)


def test_datys_to_first_event():
    df = pd.DataFrame({
        'scheduled_time': np.array(
            ['2011-01-01', '2011-01-02', '2011-01-12']
        ).astype('datetime64[ns]'),
        'horse_name': ['a', 'a', 'a']
    })
    result = days_to_first_event(df, 'horse_name', 'scheduled_time')
    expected = [0, 1, 11]
    assert np.allclose(expected, result)


def test_grouped_days_since_result():
    df = pd.DataFrame({
        'runner_id': np.array([1, 1, 1, 1, 1, 2, 2, 2, 2, 2]),
        'win_flag': np.array([1, 0, 1, 0, 0,
                              1, 0, 1, 0, 0]),
        'scheduled_time': np.array(
            ['2004-01-01', '2004-01-02', '2004-01-03', '2004-01-04', '2004-01-06',
             '2004-01-01', '2004-01-02', '2004-01-03', '2004-01-04', '2004-01-06']
            ).astype('datetime64[ms]'),
        })
    expected = [-1, 1, 2, 1, 3,
                -1, 1, 2, 1, 3]
    # result = features.days_to_previous_result(df, col='win_flag', value=1)
    result = grouped_days_since_result(df, groupby='runner_id', col='win_flag')
    assert np.allclose(expected, result)


def test_ema():
    array = np.array([5, 10, 1, 0, 0])
    assert np.allclose(array, ema(array, alpha=1))
    expected = [5., 7.5, 4.25, 2.125, 1.0625]
    assert np.allclose(expected, ema(array, alpha=0.5))


def test_grouped_ema():
    df = pd.DataFrame({
        'group': np.array([1, 1, 1, 1, 1, 2, 2, 2, 2, 2]),
        'price': np.array([5, 10, 1, 0, 0]*2),
        })
    assert np.allclose(df['price'], grouped_ema(df, 'price', 1, 'group'))
    expected = [5., 7.5, 4.25, 2.125, 1.0625] * 2
    assert np.allclose(expected, grouped_ema(df, 'price', .5, 'group'))
