import numpy as np
from scipy.stats import chi2
from statsmodels.tsa.api import stattools


class QlbChi2Test(object):
    """
    通过Q_lb统计量来检验序列是否为白噪音
    其中Q_lb统计量服从自由度为m的卡方分布
    """
    def __init__(self, arr):
        self.n = len(arr)
        # 延迟k阶自相关系数
        self.correlation = stattools.acf(arr)

    def get_q_lb(self, m):
        """
        计算延迟m期的Q_lb统计量： lb = n * (n + 2) * sum(rho_k^2 / n - k)
        :param m:
        :return:
        """
        rho_arr = np.power(self.correlation, 2)
        lb = 0
        for i in range(1, m + 1):
            lb += rho_arr[i] / (self.n - i)
        return self.n * (self.n + 2) * lb

    def chi2test(self, m, alpha=0.05):
        """
        对Q_lb进行卡方检验，检验其p值是否与检验水平的关系
        :param m:
        :param alpha:
        :return:
        """
        q_lb = self.get_q_lb(m)
        p = chi2.sf(x=q_lb, df=m)
        print('m: {}; q_lb: {}; p: {}'.format(m, q_lb, p))
        print('原序列为白噪音') if p > alpha else print('原序列不是白噪音')