#!/usr/bin/env python
from __future__ import division

#-----------------------------------------------------------------------------
# Copyright (c) 2013--, scikit-bio development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from unittest import TestCase, main

import numpy as np

from skbio.maths.subsample import subsample


class SubsampleTests(TestCase):

    def test_subsample_nonrandom(self):
        """Should function correctly for nonrandom cases."""
        a = np.array([0, 5, 0])

        # Input has too few counts.
        np.testing.assert_equal(subsample(a, 6), a)
        np.testing.assert_equal(subsample(a, 5), a)

        # Can only choose from one bin.
        np.testing.assert_equal(subsample(a, 2), np.array([0, 2, 0]))

    def test_subsample_without_replacement(self):
        """Should return a random subsample (without replacement)."""
        # Selecting 2 counts from the vector 1000 times yields each of the two
        # possible results at least once each.
        a = np.array([2, 0, 1])
        actual = {}
        for i in range(1000):
            obs = subsample(a, 2)
            actual[tuple(obs)] = None
        self.assertEqual(actual, {(1, 0, 1): None, (2, 0, 0): None})

        obs = subsample(a, 2)
        self.assertTrue(np.array_equal(obs, np.array([1, 0, 1])) or
                        np.array_equal(obs, np.array([2, 0, 0])))

    def test_subsample_with_replacement(self):
        """Should return a random subsample (with replacement)."""
        # Can choose from all in first bin, all in last bin (since we're
        # sampling with replacement), or split across bins.
        a = np.array([2, 0, 1])
        actual = {}
        for i in range(1000):
            obs = subsample(a, 2, replace=True)
            actual[tuple(obs)] = None
        self.assertEqual(actual, {(1, 0, 1): None, (2, 0, 0): None,
                                  (0, 0, 2): None})

        # Selecting 35 counts from a 36-count vector 1000 times yields more
        # than 10 different subsamples. If we were subsampling *without*
        # replacement, there would be only 10 possible subsamples, but there
        # are more possibilities when sampling *with* replacement.
        a = np.array([2, 0, 1, 2, 1, 8, 6, 0, 3, 3, 5, 0, 0, 0, 5])
        actual = {}
        for i in range(1000):
            obs = subsample(a, 35, replace=True)
            self.assertEqual(obs.sum(), 35)
            actual[tuple(obs)] = None
        self.assertTrue(len(actual) > 10)

    def test_subsample_invalid_input(self):
        """Should raise an error on invalid input."""
        # Wrong number of dimensions.
        with self.assertRaises(ValueError):
            _ = subsample([[1, 2, 3], [4, 5, 6]], 2)

        # Floats.
        with self.assertRaises(TypeError):
            _ = subsample([1, 2.3, 3], 2)


if __name__ == '__main__':
    main()

@@ -2097,33 +2097,43 @@
 ne%7D)%0A%0A        # 
-S
+Test that s
 electing 35 coun
@@ -2160,32 +2160,42 @@
 ector 1000 times
+%0A        #
  yields more%0A   
@@ -2190,26 +2190,16 @@
 lds more
-%0A        #
  than 10
@@ -2247,18 +2247,8 @@
 ling
- *without*
 %0A   
@@ -2253,16 +2253,26 @@
        #
+ *without*
  replace
@@ -2323,38 +2323,106 @@
 ples
-, but there%0A        # are more
+%0A        # because there are 10 nonzero bins in array a. However, there are more%0A        # than 10
  pos

update test code comment for clarity