"""
Tests for the stats.mstats module (support for maskd arrays)
"""


import numpy as np
from numpy import nan
import numpy.ma as ma
from numpy.ma import masked, nomask

import scipy.stats.mstats as mstats
from numpy.testing import *
from numpy.ma.testutils import assert_equal, assert_almost_equal, \
    assert_array_almost_equal



class TestMquantiles(TestCase):
    """Regression tests for mstats module.""" 
    def test_mquantiles_limit_keyword(self): 
        """Ticket #867""" 
        data = np.array([[   6.,    7.,    1.], 
                         [  47.,   15.,    2.], 
                         [  49.,   36.,    3.], 
                         [  15.,   39.,    4.], 
                         [  42.,   40., -999.], 
                         [  41.,   41., -999.], 
                         [   7., -999., -999.], 
                         [  39., -999., -999.], 
                         [  43., -999., -999.], 
                         [  40., -999., -999.], 
                         [  36., -999., -999.]])
        desired = [[19.2, 14.6, 1.45], 
                   [40.0, 37.5, 2.5 ], 
                   [42.8, 40.05, 3.55]] 
        quants = mstats.mquantiles(data, axis=0, limit=(0, 50)) 
        assert_almost_equal(quants, desired) 



class TestGMean(TestCase):
    def test_1D(self):
        a = (1,2,3,4)
        actual= mstats.gmean(a)
        desired = np.power(1*2*3*4,1./4.)
        assert_almost_equal(actual, desired,decimal=14)

        desired1 = mstats.gmean(a,axis=-1)
        assert_almost_equal(actual, desired1, decimal=14)
        assert not isinstance(desired1, ma.MaskedArray)
        #
        a = ma.array((1,2,3,4),mask=(0,0,0,1))
        actual= mstats.gmean(a)
        desired = np.power(1*2*3,1./3.)
        assert_almost_equal(actual, desired,decimal=14)

        desired1 = mstats.gmean(a,axis=-1)
        assert_almost_equal(actual, desired1, decimal=14)
    #
    def test_2D(self):
        a = ma.array(((1,2,3,4),(1,2,3,4),(1,2,3,4)),
                     mask=((0,0,0,0),(1,0,0,1),(0,1,1,0)))
        actual= mstats.gmean(a)
        desired = np.array((1,2,3,4))
        assert_array_almost_equal(actual, desired, decimal=14)
        #
        desired1 = mstats.gmean(a,axis=0)
        assert_array_almost_equal(actual, desired1, decimal=14)
        #
        actual= mstats.gmean(a, -1)
        desired = ma.array((np.power(1*2*3*4,1./4.),
                            np.power(2*3,1./2.),
                            np.power(1*4,1./2.)))
        assert_array_almost_equal(actual, desired, decimal=14)

class TestHMean(TestCase):
    def test_1D(self):
        a = (1,2,3,4)
        actual= mstats.hmean(a)
        desired =  4. / (1./1 + 1./2 + 1./3 + 1./4)
        assert_almost_equal(actual, desired, decimal=14)
        desired1 = mstats.hmean(ma.array(a),axis=-1)
        assert_almost_equal(actual, desired1, decimal=14)
        #
        a = ma.array((1,2,3,4),mask=(0,0,0,1))
        actual= mstats.hmean(a)
        desired = 3. / (1./1 + 1./2 + 1./3)
        assert_almost_equal(actual, desired,decimal=14)
        desired1 = mstats.hmean(a,axis=-1)
        assert_almost_equal(actual, desired1, decimal=14)

    def test_2D(self):
        a = ma.array(((1,2,3,4),(1,2,3,4),(1,2,3,4)),
                     mask=((0,0,0,0),(1,0,0,1),(0,1,1,0)))
        actual= mstats.hmean(a)
        desired = ma.array((1,2,3,4))
        assert_array_almost_equal(actual, desired, decimal=14)
        #
        actual1 = mstats.hmean(a,axis=-1)
        desired = (4./(1/1.+1/2.+1/3.+1/4.),
                   2./(1/2.+1/3.),
                   2./(1/1.+1/4.)
                   )
        assert_array_almost_equal(actual1, desired, decimal=14)


class TestRanking(TestCase):
    #
    def __init__(self, *args, **kwargs):
        TestCase.__init__(self, *args, **kwargs)
    #
    def test_ranking(self):
        x = ma.array([0,1,1,1,2,3,4,5,5,6,])
        assert_almost_equal(mstats.rankdata(x),[1,3,3,3,5,6,7,8.5,8.5,10])
        x[[3,4]] = masked
        assert_almost_equal(mstats.rankdata(x),[1,2.5,2.5,0,0,4,5,6.5,6.5,8])
        assert_almost_equal(mstats.rankdata(x,use_missing=True),
                            [1,2.5,2.5,4.5,4.5,4,5,6.5,6.5,8])
        x = ma.array([0,1,5,1,2,4,3,5,1,6,])
        assert_almost_equal(mstats.rankdata(x),[1,3,8.5,3,5,7,6,8.5,3,10])
        x = ma.array([[0,1,1,1,2], [3,4,5,5,6,]])
        assert_almost_equal(mstats.rankdata(x),[[1,3,3,3,5],[6,7,8.5,8.5,10]])
        assert_almost_equal(mstats.rankdata(x,axis=1),[[1,3,3,3,5],[1,2,3.5,3.5,5]])
        assert_almost_equal(mstats.rankdata(x,axis=0),[[1,1,1,1,1],[2,2,2,2,2,]])


class TestCorr(TestCase):
    #
    def test_pearsonr(self):
        "Tests some computations of Pearson's r"
        x = ma.arange(10)
        assert_almost_equal(mstats.pearsonr(x,x)[0], 1.0)
        assert_almost_equal(mstats.pearsonr(x,x[::-1])[0], -1.0)
        #
        x = ma.array(x, mask=True)
        pr = mstats.pearsonr(x,x)
        assert(pr[0] is masked)
        assert(pr[1] is masked)
    #
    def test_spearmanr(self):
        "Tests some computations of Spearman's rho"
        (x, y) = ([5.05,6.75,3.21,2.66],[1.65,2.64,2.64,6.95])
        assert_almost_equal(mstats.spearmanr(x,y)[0], -0.6324555)
        (x, y) = ([5.05,6.75,3.21,2.66,np.nan],[1.65,2.64,2.64,6.95,np.nan])
        (x, y) = (ma.fix_invalid(x), ma.fix_invalid(y))
        assert_almost_equal(mstats.spearmanr(x,y)[0], -0.6324555)
        #
        x = [ 2.0, 47.4, 42.0, 10.8, 60.1,  1.7, 64.0, 63.1,
              1.0,  1.4,  7.9,  0.3,  3.9,  0.3,  6.7]
        y = [22.6, 08.3, 44.4, 11.9, 24.6,  0.6,  5.7, 41.6,
              0.0,  0.6,  6.7,  3.8,  1.0,  1.2,  1.4]
        assert_almost_equal(mstats.spearmanr(x,y)[0], 0.6887299)
        x = [ 2.0, 47.4, 42.0, 10.8, 60.1,  1.7, 64.0, 63.1,
              1.0,  1.4,  7.9,  0.3,  3.9,  0.3,  6.7, np.nan]
        y = [22.6, 08.3, 44.4, 11.9, 24.6,  0.6,  5.7, 41.6,
              0.0,  0.6,  6.7,  3.8,  1.0,  1.2,  1.4, np.nan]
        (x, y) = (ma.fix_invalid(x), ma.fix_invalid(y))
        assert_almost_equal(mstats.spearmanr(x,y)[0], 0.6887299)
    #
    def test_kendalltau(self):
        "Tests some computations of Kendall's tau"
        x = ma.fix_invalid([5.05, 6.75, 3.21, 2.66,np.nan])
        y = ma.fix_invalid([1.65, 26.5, -5.93, 7.96, np.nan])
        z = ma.fix_invalid([1.65, 2.64, 2.64, 6.95, np.nan])
        assert_almost_equal(np.asarray(mstats.kendalltau(x,y)),
                            [+0.3333333,0.4969059])
        assert_almost_equal(np.asarray(mstats.kendalltau(x,z)),
                            [-0.5477226,0.2785987])
        #
        x = ma.fix_invalid([ 0, 0, 0, 0,20,20, 0,60, 0,20,
                            10,10, 0,40, 0,20, 0, 0, 0, 0, 0, np.nan])