import os
import sys

# PYTHONPATH shenanigans - to allow importing fastats
current_dir = os.path.abspath(os.curdir)
sys.path.append(os.path.abspath(current_dir + "/../../.."))

import fastats

from fastats import fs
from fastats.core.windowed_pass import windowed_pass, windowed_stateful_pass

import numpy as np
import pandas as pd

input_arr = np.arange(10e6, dtype='float')


@fs
def plain_sum(x):
    return np.sum(x)

pd_series = pd.Series(input_arr)

%timeit pd_series.rolling(10).sum()

# this will re-compile on each run
%timeit windowed_pass(input_arr, 10, value=plain_sum)

"""
Explanation: Quick overview of windowed_stateful_pass - this probably should go to docs once API is stabilized!
End of explanation
"""


# this compiles the rolling fuction only once
callable_pass = windowed_pass(value=plain_sum, return_callable=True)
callable_pass(input_arr, 10)

%timeit callable_pass(input_arr, 10)

"""
Explanation: Even with the compile overhead, it still outperforms pandas. Let's see how it will perform in pre-compiled mode (return_callable).
End of explanation
"""


@fs
def rolling_sum(x, val_in, val_out, state):
    if state.size == 0:   # state is empty
        state = np.array([np.sum(x)], dtype=x.dtype)
    else:
        state[0] += val_in - val_out
    return state[0], state

%timeit windowed_stateful_pass(input_arr, 10, value=rolling_sum)

callable_stateful_pass = windowed_stateful_pass(value=rolling_sum, return_callable=True)
callable_stateful_pass(input_arr, 10)

%timeit callable_stateful_pass(input_arr, 10)

"""
Explanation: Now, let's see how fast we can get it with windowed_stateful_pass.
End of explanation
"""


print("pandas")
%timeit pd_series.rolling(100).sum()

print("\nfastats windowed_pass")
%timeit callable_pass(input_arr, 100)

print("\nfastats windowed_stateful_pass")
%timeit callable_stateful_pass(input_arr, 100)

"""
Explanation: Quite impressive! The "stateful" version avoids re-calculating the full sum in each iteration for a significant speedup.
End of explanation
"""


# the same run time for window size 10, 100 etc
%timeit callable_stateful_pass(input_arr, 1000)

"""
Explanation: We can see that the re-summing overhead for large window values is significant ($O(n * win)$ complexity). windowed_pass is therefore slower than pandas here. However, pandas and windowed_stateful_pass maintain the same run time, the window size doesn't matter.
End of explanation
"""


from numba import jit


@jit(nopython=True)
def windowed_stateful_pass_numba(x, win):
    result = np.full_like(x, np.nan)
    state = np.sum(x[0:win])
    
    result[win-1] = state

    for i in range(win+1, x.shape[0]+1):
        start = i - win
        in_idx = i - 1

        val_in = x[in_idx]
        val_out = x[start-1]
        state += val_in - val_out
        result[in_idx] = state

    return result

from numpy.testing import assert_allclose


assert_allclose(windowed_stateful_pass_numba(input_arr, 100),
                callable_stateful_pass(input_arr, 100))

print("Pure numba")
%timeit windowed_stateful_pass_numba(input_arr, 100)

print("\nfastats")
%timeit callable_stateful_pass(input_arr, 100)

"""
Explanation: Lastly, let's compare fastats against pure numba rolling sum function.
End of explanation
"""


