"""
Fractional differentiation is a technique to make a time series stationary but also
retain as much memory as possible.  This is done by differencing by a positive real
number. Fractionally differenced series can be used as a feature in machine learning
process.
Updated to use Numba for performance-critical loops by Vadim Surin
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from statsmodels.tsa.stattools import adfuller
from numba import njit


class FractionalDifferentiation:
    """
    FractionalDifferentiation class encapsulates the functions that can
    be used to compute fractionally differentiated series.

    Optimizations:
    - Vectorized weight computations
    - Use Numba for performance-critical loops
    - Vectorized convolution approach for frac_diff when possible
    """

    @staticmethod
    @njit
    def _get_weights_numba(diff_amt: float, size: int) -> np.ndarray:
        """
        Numba-accelerated version of weight computation.

        :param diff_amt: (float) Differencing amount
        :param size: (int) Length of the series
        :return: (np.ndarray) Weight vector
        """
        weights = np.zeros(size, dtype=np.float64)
        weights[0] = 1.0
        for k in range(1, size):
            weights[k] = -weights[k - 1] * (diff_amt - k + 1) / k
        return weights[::-1]  # reverse the array

    @staticmethod
    def get_weights(diff_amt: float, size: int) -> np.ndarray:
        """
        Compute the vector of weights for fractional differencing.
        Advances in Financial Machine Learning, Chapter 5, section 5.4.2, page 79.

        The helper function generates weights that are used to compute fractionally
        differentiated series. It computes the weights that get used in the computation
        of  fractionally differentiated series. This generates a non-terminating series
        that approaches zero asymptotically. The side effect of this function is that
        it leads to negative drift "caused by an expanding window's added weights"
        (see page 83 AFML)

        When diff_amt is real (non-integer) positive number then it preserves memory.

        The book does not discuss what should be expected if d is a negative real
        number. Conceptually (from set theory) negative d leads to set of negative
        number of elements. And that translates into a set whose elements can be
        selected more than once or as many times as one chooses (multisets with
        unbounded multiplicity) - see http://faculty.uml.edu/jpropp/msri-up12.pdf.

        :param diff_amt : float Differencing amount. When positive real (non-integer), it preserves memory.
        :param size : int Length of the weight vector to generate.
        :return: (np.ndarray) Weight vector for fractional differencing.
        """
        return FractionalDifferentiation._get_weights_numba(diff_amt, size).reshape(-1, 1)


    @staticmethod
    @njit
    def _get_weights_ffd_numba(diff_amt: float, thresh: float, lim: int) -> np.ndarray:
        """
        Numba-accelerated version for fixed-fractional differencing weights.
        :param diff_amt: (float) Differencing amount
        :param thresh: (float) Threshold for minimum weight
        :param lim: (int) Maximum length of the weight vector
        :return: (np.ndarray) Weight vector
        """
        weights = [1.0]
        k = 1
        # The algorithm below executes the iterativetive estimation (section 5.4.2, page 78)
        # The output weights array is of the indicated length (specified by lim)
        ctr = 0
        while True:
            # compute the next weight
            w_ = -weights[-1] * (diff_amt - k + 1) / k
            if np.abs(w_) < thresh:
                break
            weights.append(w_)
            k += 1
            ctr += 1
            if ctr == lim - 1:  # if we have reached the size limit, exit the loop
                break
        # Now, reverse the list, convert into a numpy column vector
        arr = np.array(weights[::-1], dtype=np.float64).reshape(-1, 1)
        return arr
    
    @staticmethod
    def get_weights_ffd(diff_amt, thresh, lim):
        """
        Advances in Financial Machine Learning, Chapter 5, section 5.4.2, page 83.

        The helper function generates weights that are used to compute fractionally
        differentiate dseries. It computes the weights that get used in the computation
        of fractionally differentiated series. The series is of fixed width and same
        weights (generated by this function) can be used when creating fractional
        differentiated series.
        This makes the process more efficient. But the side-effect is that the
        fractionally differentiated series is skewed and has excess kurtosis. In
        other words, it is not Gaussian any more.

        The discussion of positive and negative d is similar to that in get_weights
        (see the function get_weights)

        :param diff_amt: (float) Differencing amount
        :param thresh: (float) Threshold for minimum weight
        :param lim: (int) Maximum length of the weight vector
        :return: (np.ndarray) Weight vector
        """

        return FractionalDifferentiation._get_weights_ffd_numba(diff_amt, thresh, lim)


    @staticmethod
    def frac_diff(series: pd.DataFrame, diff_amt: float, thresh: float = 0.01) -> pd.DataFrame:
        """
        Expanding window fractional differentiation.
        Uses convolution to speed up computation

        Advances in Financial Machine Learning, Chapter 5, section 5.5, page 82.

        References:
        https://www.wiley.com/en-us/Advances+in+Financial+Machine+Learning-p-9781119482086
        https://wwwf.imperial.ac.uk/~ejm/M3S8/Problems/hosking81.pdf
        https://en.wikipedia.org/wiki/Fractional_calculus

        The steps are as follows:
        - Compute weights (this is a one-time exercise)
        - Iteratively apply the weights to the price series and generate output points

        This is the expanding window variant of the fracDiff algorithm
        Note 1: For thresh-1, nothing is skipped
        Note 2: diff_amt can be any positive fractional, not necessarility bounded [0, 1]

        :param series: (pd.DataFrame) A time series that needs to be differenced
        :param diff_amt: (float) Differencing amount
        :param thresh: (float) Threshold or epsilon
        :return: (pd.DataFrame) Differenced series
        """

        # 1. Compute weights for the longest series
        size = series.shape[0]
        weights = FractionalDifferentiation.get_weights(diff_amt, size)[:, 0] # 1D array

        # 2. Determine initial calculations to be skipped based on weight-loss threshold
        cum_w = np.cumsum(abs(weights))
        cum_w /= cum_w[-1]
        skip = np.sum(cum_w < thresh)

        # We'll use convolution: For an expanding window, the value at time i is:
        # (weights[0]*series[i] + weights[1]*series[i-1] + ...).
        # This is essentially a reversed convolution of the series with weights.
        # Since weights are reversed, we can directly apply convolution and then align indices.

        # 3. Apply weights to values
        output_df = {}
        for col in series.columns:
            col_values = series[col].ffill().to_numpy(np.float64)
            # Convolve series with weights (full mode)
            conv_res = np.convolve(col_values, weights, mode='full')[:len(col_values)]
            # conv_res[i] gives the dot product up to i-th element.
            # We only start using values after 'skip', and align with the original index:
            # Original series length: N
            # conv_res length: N + N - 1 = 2N - 1
            # The valid fractional differences start from index (skip-1), basically.
            # But we must ensure we properly slice to get a length == series length.
            # The expanding window result should have length = original series length.
            # conv_res[:N] corresponds to partial sums up to each index.
            # Actually, for an expanding window, the i-th output corresponds to conv_res[i].
            # Start from skip:
            result = np.full_like(col_values, np.nan)
            if skip < len(result):
                result[skip:] = conv_res[skip:skip+len(result)]
            output_df[col] = result

        return pd.DataFrame(output_df, index=series.index, columns=series.columns)



    @staticmethod
    def frac_diff_ffd(series: pd.DataFrame, diff_amt: float, thresh: float = 1e-5) -> pd.DataFrame:
        """
        Fixed-window fractional differencing.
        We can use a rolling window convolution here as well.

        Advances in Financial Machine Learning, Chapter 5, section 5.5, page 83.

        References:

        * https://www.wiley.com/en-us/Advances+in+Financial+Machine+Learning-p-9781119482086
        * https://wwwf.imperial.ac.uk/~ejm/M3S8/Problems/hosking81.pdf
        * https://en.wikipedia.org/wiki/Fractional_calculus

        The steps are as follows:

        - Compute weights (this is a one-time exercise)
        - Iteratively apply the weights to the price series and generate output points

        Constant width window (new solution)
        Note 1: thresh determines the cut-off weight for the window
        Note 2: diff_amt can be any positive fractional, not necessarity bounded [0, 1].

        :param series: (pd.DataFrame) A time series that needs to be differenced
        :param diff_amt: (float) Differencing amount
        :param thresh: (float) Threshold for minimum weight
        :return: (pd.DataFrame) A data frame of differenced series
        """

        # 1) Compute weights for the longest series
        size = series.shape[0]
        weights = FractionalDifferentiation.get_weights_ffd(diff_amt, thresh, size)[:, 0] # 1D array
        width = len(weights)

        # 2) Apply weights to values
        # 2.1) Start by creating a dictionary to hold all the fractionally differenced series
        output_df = {}
        for col in series.columns:
            col_values = series[col].ffill().to_numpy(np.float64)
            # Convolution with a fixed window:
            # The weights are reversed, so we use 'valid' mode to get the correct length
            conv_res = np.convolve(col_values, weights[::-1], mode='valid')
            # 'valid' mode returns length = N - width + 1
            # Align this with the original index:
            result = np.full_like(col_values, np.nan)
            # The first valid fractional difference appears only after width-1 steps:   
            result[width-1:] = conv_res
            output_df[col] = result

        # transform the dictionary into a data frame
        output_df = pd.DataFrame(output_df, index=series.index, columns=series.columns)
        return output_df


def get_weights(diff_amt: float, size: int) -> np.ndarray:
    """ This is a pass-through function """
    return FractionalDifferentiation.get_weights(diff_amt, size)

def frac_diff(series: pd.DataFrame, diff_amt: float, thresh: float = 0.01) -> pd.DataFrame:
    """ This is a pass-through function """
    return FractionalDifferentiation.frac_diff(series, diff_amt, thresh)


def get_weights_ffd(diff_amt: float, thresh: float, lim: int) -> np.ndarray:
    """ This is a pass-through function """
    return FractionalDifferentiation.get_weights_ffd(diff_amt, thresh, lim)


def frac_diff_ffd(series: pd.DataFrame, diff_amt: float, thresh: float = 1e-5) -> pd.DataFrame:
    """
    Advances in Financial Machine Learning, Chapter 5, section 5.5, page 83.

    References:

    * https://www.wiley.com/en-us/Advances+in+Financial+Machine+Learning-p-9781119482086
    * https://wwwf.imperial.ac.uk/~ejm/M3S8/Problems/hosking81.pdf
    * https://en.wikipedia.org/wiki/Fractional_calculus

    The steps are as follows:

    - Compute weights (this is a one-time exercise)
    - Iteratively apply the weights to the price series and generate output points

    Constant width window (new solution)
    Note 1: thresh determines the cut-off weight for the window
    Note 2: diff_amt can be any positive fractional, not necessarity bounded [0, 1].

    :param series: (pd.Series) A time series that needs to be differenced
    :param diff_amt: (float) Differencing amount
    :param thresh: (float) Threshold for minimum weight
    :return: (pd.DataFrame) A data frame of differenced series
    """
    return FractionalDifferentiation.frac_diff_ffd(series, diff_amt, thresh)

def plot_min_ffd(series: pd.DataFrame) -> plt.Axes:
    """
    Advances in Financial Machine Learning, Chapter 5, section 5.6, page 85.

    References:

    * https://www.wiley.com/en-us/Advances+in+Financial+Machine+Learning-p-9781119482086

    This function plots the graph to find the minimum D value that passes the ADF test.

    It allows to determine d - the amount of memory that needs to be removed to achieve
    stationarity. This function covers the case of 0 < d << 1, when the original series is
    "mildly non-stationary."

    The right y-axis on the plot is the ADF statistic computed on the input series downsampled
    to a daily frequency.

    The x-axis displays the d value used to generate the series on which the ADF statistic is computed.

    The left y-axis plots the correlation between the original series (d=0) and the differentiated
    series at various d values.

    Examples on how to interpret the results of this function are available in the corresponding part
    in the book Advances in Financial Machine Learning.

    :param series: (pd.DataFrame) Dataframe that contains a 'close' column with prices to use.
    :return: (plt.AxesSubplot) A plot that can be displayed or used to obtain resulting data.
    """

    results = pd.DataFrame(columns=['adfStat', 'pVal', 'lags', 'nObs', '95% conf', 'corr'])

    # Iterate through d values with 0.1 step
    for d_value in np.linspace(0, 1, 11):
        close_prices = np.log(series[['close']]).resample('1D').last()  # Downcast to daily obs
        close_prices.dropna(inplace=True)

        # Applying fractional differentiation
        differenced_series = frac_diff_ffd(close_prices, diff_amt=d_value, thresh=0.01).dropna()

        # Correlation between the original and the differentiated series
        corr = np.corrcoef(close_prices.loc[differenced_series.index, 'close'],
                           differenced_series['close'])[0, 1]
        # Applying ADF
        differenced_series = adfuller(differenced_series['close'], maxlag=1, regression='c', autolag=None)

        # Results to dataframe
        results.loc[d_value] = list(differenced_series[:4]) + [differenced_series[4]['5%']] + [corr]  # With critical value

    # Plotting
    plot = results[['adfStat', 'corr']].plot(secondary_y='adfStat', figsize=(10, 8))
    mean_conf = results['95% conf'].mean()
    plt.axhline(mean_conf, linewidth=1, color='r', linestyle='dotted')
    plt.axhline(0, linewidth=1, color='g', linestyle='dotted')

    return plot
