import unittest
import numpy as np
from pynibs.optimization import constraint_min_current, get_current_init


class TestConstraintMinCurrent(unittest.TestCase):

    def test_constraint_met(self):
        x = np.array([0.8, 0.8, 0.8, 0.8, 0.8, 0.8])
        n_stim = 2
        result = constraint_min_current(x, n_stim, channels=(0, 1), min_current=0.7, frac=1, factor=100., verbose=False)
        self.assertEqual(result, 0)

    def test_constraint_not_met(self):
        x = np.array([0.6, 0.6, 0.6, 0.6, 0.6, 0.6])
        n_stim = 2
        result = constraint_min_current(x, n_stim, channels=(0, 1), min_current=0.7, frac=0.95, factor=100.,
                                        verbose=False)
        self.assertLess(result, 0)

    def test_verbose_output(self):
        x = np.array([0.8, 0.8, 0.8, 0.8, 0.8, 0.8])
        n_stim = 2
        result = constraint_min_current(x, n_stim, channels=(0, 1), min_current=0.7, frac=0.95, factor=100.,
                                        verbose=True)
        self.assertGreaterEqual(result, 0)

    def test_different_channels(self):
        x = np.array([0.8, 0.8, 0.8, 0.8, 0.8, 0.8])
        n_stim = 2
        result = constraint_min_current(x, n_stim, channels=(1, 2), min_current=0.7, frac=0.95, factor=100.,
                                        verbose=False)
        self.assertGreaterEqual(result, 0)

    def test_different_min_current(self):
        x = np.array([0.8, 0.8, 0.8, 0.8, 0.8, 0.8])
        n_stim = 2
        result = constraint_min_current(x, n_stim, channels=(0, 1), min_current=0.9, frac=0.95, factor=100.,
                                        verbose=False)
        self.assertLess(result, 0)


class TestGetCurrentInit(unittest.TestCase):

    def test_default_parameters(self):
        n_channel = 5
        n_stim = 10
        x0 = get_current_init(n_channel, n_stim)
        self.assertEqual(x0.shape, (n_channel * n_stim,))
        self.assertTrue(np.all(x0 >= -1) and np.all(x0 <= 1))

    def test_with_bounds(self):
        n_channel = 5
        n_stim = 10
        bounds = [(-0.5, 0.5) for _ in range(n_channel)]
        x0 = get_current_init(n_channel, n_stim, bounds=bounds)
        self.assertEqual(x0.shape, (n_channel * n_stim,))
        for i in range(n_channel):
            self.assertTrue(
                np.all(x0[n_stim * i: n_stim * (i + 1)] >= -0.5) and np.all(x0[n_stim * i: n_stim * (i + 1)] <= 0.5))

    def test_only_positive(self):
        n_channel = 5
        n_stim = 10
        bounds = [(0, 1) for _ in range(n_channel)]
        only_pos = [True for _ in range(n_channel)]
        x0 = get_current_init(n_channel, n_stim, bounds=bounds, only_pos=only_pos)
        self.assertEqual(x0.shape, (n_channel * n_stim,))
        self.assertTrue(np.all(x0 >= 0) and np.all(x0 <= 1))

    def test_mixed_bounds_and_only_positive(self):
        n_channel = 5
        n_stim = 10
        bounds = [(0, 1), (-1, 1), (0.5, 1), (-0.5, 0.5), (0, 0.5)]
        only_pos = [True, False, True, False, True]
        x0 = get_current_init(n_channel, n_stim, bounds=bounds, only_pos=only_pos)
        self.assertEqual(x0.shape, (n_channel * n_stim,))
        for i in range(n_channel):
            self.assertTrue(np.all(x0[n_stim * i: n_stim * (i + 1)] >= bounds[i][0]) and np.all(
                x0[n_stim * i: n_stim * (i + 1)] <= bounds[i][1]))
            if only_pos[i]:
                self.assertTrue(np.all(x0[n_stim * i: n_stim * (i + 1)] >= 0))


if __name__ == '__main__':
    unittest.main()
