# -*- coding: utf-8 -*-
"""
Created on Sun Sep 25 21:23:38 2011

Author: Josef Perktold and Scipy developers
"""
from __future__ import print_function
from statsmodels.compat.python import range
import numpy as np
from scipy import stats

from numpy import exp

def anderson_statistic(x, dist='norm', fit=True, params=(), axis=0):
'''calculate anderson-darling A2 statistic

Parameters
----------
x : array_like
data
dist : 'norm' or callable
null distribution for the test statistic
fit : bool
If True, then the distribution parameters are estimated.
Currently only for 1d data x, except in case dist='norm'
params : tuple
optional distribution parameters if fit is False
axis : integer
If dist is 'norm' or fit is False, then data can be an n-dimensional
and axis specifies the axis of a variable

Returns
-------
Anderson-Darling statistic

'''
x = np.asarray(x)
y = np.sort(x, axis=axis)
N = y.shape[axis]
if fit:
if dist == 'norm':
xbar = np.expand_dims(np.mean(x, axis=axis), axis)
s = np.expand_dims(np.std(x, ddof=1, axis=axis), axis)
w = (y-xbar)/s
z = stats.norm.cdf(w)
#print z
elif hasattr(dist, '__call__'):
params = dist.fit(x)
#print params
z = dist.cdf(y, *params)
print(z)
else:
if hasattr(dist, '__call__'):
z = dist.cdf(y, *params)
else:
raise ValueError('if fit is false, then dist needs to be callable')

i = np.arange(1,N+1)
sl1 = [None]*x.ndim
sl1[axis] = slice(None)
sl2 = [slice(None)]*x.ndim
sl2[axis] = slice(None,None,-1)
S = np.sum((2*i[sl1]-1.0)/N*(np.log(z)+np.log(1-z[sl2])), axis=axis)
A2 = -N-S
return A2