import numpy as np 
import scipy.stats as st 
import matplotlib.pyplot as plt 
from getCriticalValuesKS import *

np.random.seed(12345) # Fix seed to allow rerpoducibility

def LCG(X0 = 1, n = 100, a = 3, b = 0, m = 31):
    """
    Linear Congruential Generator
    """
    X = np.zeros(n+1)
    X[0] = X0
    for i in range(1,n+1):
            X[i] = (a * X[i-1] + b) % m

    U = X / m
    return U

def KSTest(X, alpha = 0.1):
    """
    Kolmogorov Smirnov Test for data X and significance alpha
    """
    xgrid = np.linspace(0, 1, 10001)
    Fhat = [(X<=x0).sum()/float(n) for x0 in xgrid]
    Dn = np.max(np.abs(Fhat - F(xgrid)))
    #[Dn1, p] = st.kstest(X, F) # Comparison with the scipy.stats' builtin function
    print('KS Test statistic: ' + str(Dn) )

    val = getCriticalValuesKS(n, alpha)
    true = int( (Dn > val) ) # Convert boolean to int
    rej = ['cannot be', 'is']
    stat = {'Statistic': Dn, 'Quantile': val, 'Significance': alpha}
    message = 'KS test: the null hypothesis H0 ' + rej[true] + ' rejected at level alpha = ' + str(alpha) 
    return message, stat

def ChiSquareTest(X, K = 10, alpha = 0.1):
    """
    Chi Squared Test for data X and significance alpha with K degrees of freedom
    """
    n = len(X)
    p = np.ones(K) / K
    N = np.array([np.sum((float(i)/K < X) & (X <= (i+1.)/K)) for i in range(K)])
    QK = np.sum( (N-n*p)**2. / (n*p)) # test statistic
    critval = st.chi2.ppf(1-alpha, K-1)
    true_chi = 1 * (QK > critval)
    rej = ['cannot be', 'is']
    stat = {'Statistic': QK, 'Quantile': critval, 'Significance': alpha}
    message = 'Chi2 test: the null hypothesis H0 ' + rej[true_chi] + ' rejected at level alpha = ' + str(alpha)
    return message, stat


def SerialTest(X, d = 2, alpha = 0.1):
    """
    Serial statistical test 
    """
    assert X.shape[0] % 2 == 0, 'Random sample length should be even.'
    Y = X.reshape(2, int(X.shape[0]/2)).T
    nY = Y.shape[0]
    m = 10
    K = m**d
    Nm = np.zeros(K)
    p = np.ones(K) / float(K) # True probabilities (uniform partition)
    for k in range(K):
            xl = (k % m) / float(m)
            xu = xl + 1./m
            yl = np.max([0., np.floor(k/float(m))]) / m
            yu = yl + 1./m
            Nm[k] = np.sum( ((xl< Y[:,0]) & (Y[:,0] <= xu)) * ( (yl<Y[:,1]) & (Y[:,1]<= yu)))

    rej = ['cannot be', 'is']
    Qm = np.sum( (Nm - nY*p)**2 / (nY*p))
    serval = st.chi2.ppf(1-alpha, K-1)
    true_ser = 1 * (Qm > serval)
    stat = {'Statistic' : Qm, 'Quantile': serval, 'Significance': alpha}
    message = 'Serial test: the null hypothesis H0 ' + rej[true_ser] + ' rejected at level alpha = ' + str(alpha)
    return message, stat

def GapTest(X, alpha = 0.1, aa=0., bb=0.5, r=5):
    """
    Gap test
    """
    idx = list(set(np.where(X > aa)[0]) & set(np.where(X < bb)[0]))
    idx = np.hstack([0, np.array(idx) + 1])
    Z = idx[1:] - idx[:-1] - 1
    prob = bb - aa
    p = prob * (1 - prob) ** np.array([i for i in range(r)])
    p = np.hstack([p, (1 - prob)**r])

    Nr = np.zeros(r+1)
    for j in range(r):
            Nr[j] = np.sum( Z == j)

    nZ = len(Z)
    Nr[-1] = nZ - np.sum(Nr)
    Qr = np.sum( (Nr - p*nZ)**2. / (p*nZ))
    critval = st.chi2.ppf(1 - alpha, r)
    rej = ['cannot be', 'is']
    true = 1 * (Qr > critval)
    stat = {'Statistic': Qr, 'Quantile': critval, 'Significance' : alpha}
    message = 'Gap test: the null hypothesis H0 ' + rej[true] + ' rejected at level alpha = ' + str(alpha)
    return message, stat

def cdf(X):
    """
    Empirical CDF
    """
    Xs = np.sort(X) # Sorted version of data
    n = len(X)
    x = np.array(range(1, n+1)) / (n+1.)
    Nx = np.array([(x< x_i).sum() for x_i in x]) / float(n) # Computes empirical CDF
    xx =  np.hstack([Xs.reshape(n,1), Xs.reshape(n,1)]).flatten()
    Nxx = np.hstack([ np.hstack([Nx.reshape(n,1), Nx.reshape(n,1)]).flatten()[1:2*n], 1])
    return xx, Nxx

F = lambda u: u * (u>=0) * (u<=1) + (u>1) # CDF
Finv = lambda u: u * (u>=0) * (u<=1) + (u>1) # Inverse CDF

if __name__=='__main__':
    # Generate Data
    n = 1000 # number of samples 
    Y = st.uniform.rvs(size = n) # Using Scipy's built-in function
    X = LCG(n = n)[1:] # Using LCG function above 

    Xs = np.sort(X) # Sorted version of data

    x = np.array(range(1, n+1)) / (n+1.)
    [xx, Nxx] = cdf(X)

    fig = plt.figure(figsize = (8,4))
    ax_qq = fig.add_subplot(121)
    ax_qq.plot(x, x, 'k--', linewidth = 1.)
    ax_qq.plot(Finv(x), Xs, '-.', linewidth = 1.)
    ax_qq.set_title('Q-Q')
    ax_cdf = fig.add_subplot(122)
    ax_cdf.plot(x, F(x), 'k--', linewidth = 1.)
    ax_cdf.plot(xx, Nxx, '-', linewidth = 1.)
    ax_cdf.set_title('CDF')

    #plt.savefig('../figures/qq_cdf_LCG_n1000.png')
    plt.show()

    [message_KS, stat_KS] = KSTest(X, 0.01)
    print(message_KS)

    [message_chi2, stat_chi] = ChiSquareTest(X, alpha=0.01)
    print(message_chi2)

    [message_ser, stat] = SerialTest(X)
    print(message_ser)

    [message_gap, stat_gap] = GapTest(X)
    print(message_gap)

    fig1 = plt.figure(figsize = (8, 4))
    ax1 = fig1.add_subplot(121)
    ax1.plot(Y[:-1], Y[1:], '-', linewidth = 1.)
    ax2 = fig1.add_subplot(122)
    ax2.plot(X[:-1], X[1:], '-')
    plt.show()