'''
Lowess testing suite.

Expected outcomes are generated by R's lowess function given the same
arguments. The R script test_lowess_r_outputs.R can be used to
generate the expected outcomes.

The delta tests utilize Silverman's motorcycle collision data,
available in R's MASS package.
'''

import os
import numpy as np
from numpy.testing import (assert_almost_equal, assert_, assert_raises,
                           assert_equal)
#import statsmodels.api as sm
from statsmodels.nonparametric.smoothers_lowess import lowess

# Number of decimals to test equality with.
# The default is 7.
testdec = 7
curdir = os.path.dirname(os.path.abspath(__file__))
rpath = os.path.join(curdir, 'results')


class  TestLowess(object):

    def test_import(self):
        #this doesn't work
        #from statsmodels.api.nonparametric import lowess as lowess1
        import statsmodels.api as sm
        lowess1 = sm.nonparametric.lowess
        assert_(lowess is lowess1)


    def test_simple(self):
        rfile = os.path.join(rpath, 'test_lowess_simple.csv')
        test_data = np.genfromtxt(open(rfile, 'rb'),
                                  delimiter = ',', names = True)
        expected_lowess = np.array([test_data['x'], test_data['out']]).T

        actual_lowess = lowess(test_data['y'], test_data['x'])

        assert_almost_equal(expected_lowess, actual_lowess, decimal = testdec)


    def test_iter(self):
        rfile = os.path.join(rpath, 'test_lowess_iter.csv')
        test_data = np.genfromtxt(open(rfile, 'rb'),
                                  delimiter = ',', names = True)

        expected_lowess_no_iter = np.array([test_data['x'], test_data['out_0']]).T
        expected_lowess_3_iter = np.array([test_data['x'], test_data['out_3']]).T

        actual_lowess_no_iter = lowess(test_data['y'], test_data['x'], it = 0)
        actual_lowess_3_iter = lowess(test_data['y'], test_data['x'], it = 3)

        assert_almost_equal(expected_lowess_no_iter, actual_lowess_no_iter, decimal = testdec)
        assert_almost_equal(expected_lowess_3_iter, actual_lowess_3_iter, decimal = testdec)


    def test_frac(self):
        rfile = os.path.join(rpath, 'test_lowess_frac.csv')
        test_data = np.genfromtxt(open(rfile, 'rb'),
                                  delimiter = ',', names = True)

        expected_lowess_23 = np.array([test_data['x'], test_data['out_2_3']]).T
        expected_lowess_15 = np.array([test_data['x'], test_data['out_1_5']]).T

        actual_lowess_23 = lowess(test_data['y'], test_data['x'] ,frac = 2./3)
        actual_lowess_15 = lowess(test_data['y'], test_data['x'] ,frac = 1./5)

        assert_almost_equal(expected_lowess_23, actual_lowess_23, decimal = testdec-1)
        assert_almost_equal(expected_lowess_15, actual_lowess_15, decimal = testdec)


    def test_delta(self):
        rfile = os.path.join(rpath, 'test_lowess_delta.csv')
        test_data = np.genfromtxt(open(rfile, 'rb'),
                                  delimiter = ',', names = True)

        expected_lowess_del0 = np.array([test_data['x'], test_data['out_0']]).T
        expected_lowess_delRdef = np.array([test_data['x'], test_data['out_Rdef']]).T
        expected_lowess_del1 = np.array([test_data['x'], test_data['out_1']]).T

        actual_lowess_del0    = lowess(test_data['y'], test_data['x'], frac=0.1)
        actual_lowess_delRdef = lowess(test_data['y'], test_data['x'], frac=0.1,
                       delta = 0.01 * np.ptp(test_data['x']))
        actual_lowess_del1    = lowess(test_data['y'], test_data['x'], frac = 0.1, delta = 1.0 + 1e-10)

        assert_almost_equal(expected_lowess_del0, actual_lowess_del0, decimal = testdec)
        assert_almost_equal(expected_lowess_delRdef, actual_lowess_delRdef, decimal = testdec)
        assert_almost_equal(expected_lowess_del1, actual_lowess_del1, decimal = 10) #testdec)

    def test_options(self):
        rfile = os.path.join(rpath, 'test_lowess_simple.csv')
        test_data = np.genfromtxt(open(rfile, 'rb'),
                                  delimiter = ',', names = True)
        y, x = test_data['y'], test_data['x']
        res1_fitted = test_data['out']
        expected_lowess = np.array([test_data['x'], test_data['out']]).T

        # check skip sorting
        actual_lowess1 = lowess(y, x, is_sorted=True)
        assert_almost_equal(actual_lowess1, expected_lowess, decimal=13)

        # check skip missing
        actual_lowess = lowess(y, x, is_sorted=True, missing='none')
        assert_almost_equal(actual_lowess, actual_lowess1, decimal=13)

        # check order/index, returns yfitted only
        actual_lowess = lowess(y[::-1], x[::-1], return_sorted=False)
        assert_almost_equal(actual_lowess, actual_lowess1[::-1, 1], decimal=13)

        # check returns yfitted only
        actual_lowess = lowess(y, x, return_sorted=False, missing='none',
                               is_sorted=True)
        assert_almost_equal(actual_lowess, actual_lowess1[:, 1], decimal=13)

        # check integer input
        actual_lowess = lowess(np.round(y).astype(int), x, is_sorted=True)
        actual_lowess1 = lowess(np.round(y), x, is_sorted=True)
        assert_almost_equal(actual_lowess, actual_lowess1, decimal=13)
        assert_(actual_lowess.dtype is np.dtype(float))
        # this will also have duplicate x
        actual_lowess = lowess(y, np.round(x).astype(int), is_sorted=True)
        actual_lowess1 = lowess(y, np.round(x), is_sorted=True)
        assert_almost_equal(actual_lowess, actual_lowess1, decimal=13)
        assert_(actual_lowess.dtype is np.dtype(float))

        # check with nans,  this changes the arrays
        y[[5, 6]] = np.nan
        x[3] = np.nan
        mask_valid = np.isfinite(x) & np.isfinite(y)
        #actual_lowess1[[3, 5, 6], 1] = np.nan
        actual_lowess = lowess(y, x, is_sorted=True)
        actual_lowess1 = lowess(y[mask_valid], x[mask_valid], is_sorted=True)
        assert_almost_equal(actual_lowess, actual_lowess1, decimal=13)
        assert_raises(ValueError, lowess, y, x, missing='raise')

        perm_idx = np.arange(len(x))
        np.random.shuffle(perm_idx)
        yperm = y[perm_idx]
        xperm = x[perm_idx]
        actual_lowess2 = lowess(yperm, xperm, is_sorted=False)
        assert_almost_equal(actual_lowess, actual_lowess2, decimal=13)

        actual_lowess3 = lowess(yperm, xperm, is_sorted=False,
                                return_sorted=False)
        mask_valid = np.isfinite(xperm) & np.isfinite(yperm)
        assert_equal(np.isnan(actual_lowess3), ~mask_valid)
        # get valid sorted smoothed y from actual_lowess3
        sort_idx = np.argsort(xperm)
        yhat = actual_lowess3[sort_idx]
        yhat = yhat[np.isfinite(yhat)]
        assert_almost_equal(yhat, actual_lowess2[:,1], decimal=13)




if __name__ == "__main__":
    import nose
    nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb'], exit=False)
