#! /usr/bin/env python
#
# test_random_parameter.py
#
# This file is part of NEST.
#
# Copyright (C) 2004 The NEST Initiative
#
# NEST is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# NEST is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with NEST.  If not, see <http://www.gnu.org/licenses/>.
"""
Tests for random topology parameter distributions. This is an implementation
of the Kolmogorov-Smirnov test [1] to check that the distribution of weights
fits the expected distribution to level alpha=0.05 when using random
parameters. Also serves as a regression test for ticket #687.

[1] http://www.itl.nist.gov/div898/handbook/eda/section3/eda35g.htm
"""

import unittest
from math import sqrt
try:
    import numpy
    have_numpy = True
except ImportError:
    have_numpy = False
try:
    from math import erf
    have_erf = True
except ImportError:
    have_erf = False
import nest
import nest.topology as topo
from nest.tests.decorators import _skipIf

@_skipIf(not have_numpy, 'Python numpy package not installed', 'testcase')
class RandomParameterTestCase(unittest.TestCase):

    def ks_test(self,weight_dict,expected_cdf_func):
        """Create connections with given distribution of weights and test that it fits the given expected cumulative distribution using K-S."""

        # n = rows * cols * Nconn
        rows = 10
        cols = 10
        Nconn = 100

        nest.ResetKernel()

        # Create layer and connect with given weight distribution
        layer = topo.CreateLayer({'rows':rows,'columns':cols,'elements':'iaf_neuron'})
        topo.ConnectLayers(layer, layer, {'connection_type': 'convergent', 'number_of_connections':Nconn,'weights':weight_dict})

        # Get connection weights and sort
        connectome = nest.GetConnections()
        weights = numpy.array(nest.GetStatus(connectome,'weight'))
        weights.sort()
        n = len(weights)

        # The observed (empirical) cdf is simply i/n for weights[i]
        observed_cdf = numpy.arange(n+1,dtype=float)/n
        expected_cdf = expected_cdf_func(weights)

        D = max( numpy.abs(expected_cdf-observed_cdf[:-1]).max(), numpy.abs(expected_cdf-observed_cdf[1:]).max() )

        # Code to find Kalpha corresponding to level alpha:
        # alpha = 0.05
        # import scipy.optimize,scipy.stats
        # Kalpha = scipy.optimize.fmin( lambda x: abs(alpha-scipy.stats.ksprob(x)), 1)[0]
        Kalpha = 1.3581054687500012

        self.assertTrue( sqrt(n)*D < Kalpha )


    def test_uniform(self):
        """Test uniform distribution of weights."""

        w_min = -1.3
        w_max = 2.7
        
        weight_dict = {'uniform': {'min':w_min, 'max':w_max}}

        uniform_cdf_func = lambda w:(w>w_min)*(w<w_max)*((w-w_min)/(w_max-w_min)) + (w>=w_max)*1.0

        self.ks_test(weight_dict, uniform_cdf_func)


    @_skipIf(not have_erf, 'Python function math.erf not available')
    def test_normal(self):
        """Test normal distribution of weights."""

        mean = 2.3
        sigma = 1.7

        weight_dict = {'normal': {'mean':mean, 'sigma':sigma}}

        numpy_erf = numpy.vectorize(erf)

        normal_cdf_func = lambda w:0.5*(1.0 + numpy_erf( (w-mean)/sqrt(2)/sigma ))

        self.ks_test(weight_dict, normal_cdf_func)


    @_skipIf(not have_erf, 'Python function math.erf not available')
    def test_lognormal(self):
        """Test lognormal distribution of weights."""

        mu = 1.7
        sigma = 0.9

        weight_dict = {'lognormal': {'mu':mu, 'sigma':sigma}}

        numpy_erf = numpy.vectorize(erf)

        lognormal_cdf_func = lambda w:0.5*(1.0 + numpy_erf( (numpy.log(w)-mu)/sqrt(2)/sigma ))

        self.ks_test(weight_dict, lognormal_cdf_func)


def suite():

    suite = unittest.makeSuite(RandomParameterTestCase,'test')
    return suite


if __name__ == "__main__":

    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite())