import numpy as np
import pandas as pd
from numba import njit

from qfi.util.returns import log_return


@njit
def _calc_cusum(return_values: np.ndarray, return_index: np.ndarray, threshold: float) -> np.ndarray:
    s_pos = 0
    s_neg = 0
    ts_events = np.empty(len(return_values), dtype=return_index.dtype)
    event_count = 0
    for log_ret, ts in zip(return_values, return_index):
        s_pos = max(0, s_pos + log_ret)
        s_neg = min(0, s_neg + log_ret)
        if s_pos > threshold:
            s_pos = 0
            ts_events[event_count] = ts
            event_count += 1
        elif s_neg < -threshold:
            s_neg = 0
            ts_events[event_count] = ts
            event_count += 1
    return ts_events[:event_count]


def cusum(close: pd.Series, threshold: float) -> pd.DatetimeIndex:
    """
    CUSUM filter to identify events in a time series.

    Parameters
    ----------
    close : pd.Series
        Time series data (e.g., closing prices).
    threshold : float
        Threshold for identifying events.

    Returns
    -------
    pd.DatetimeIndex
        Timestamps where events are identified.
    """
    ret = log_return(close)
    return_values = ret.to_numpy()
    return_index = ret.index.to_numpy()
    ts_events = _calc_cusum(return_values, return_index, threshold)
    return pd.DatetimeIndex(ts_events)
