import unittest
from acceleration import *
import numpy as np
from numpy.testing import assert_almost_equal

accelerationFunctions = [bruteForce, bruteForceNumba, bruteForceNumbaOptimized, bruteForceCPP]

class MyTest(unittest.TestCase):
    # Massive particles
    def test1(self):
        r_vec = np.array([[0.,0,0],[1,0,0]])
        m_vec = np.array([1., 1])
        result = np.array([[1., 0, 0],[-1, 0, 0]])
        for fun in accelerationFunctions:
            assert_almost_equal(fun(r_vec, m_vec, 0), result)

    # Massless particles
    def test2(self):
        r_vec = np.array([[0.,0,0],[1,0,0]])
        m_vec = np.array([1., 0])
        result = np.array([[0, 0, 0],[-1, 0, 0]])
        for fun in accelerationFunctions:
            assert_almost_equal(fun(r_vec, m_vec, 0), result)

    # Softening
    def test3(self):
        r_vec = np.array([[0.,0,0],[1,0,0]])
        m_vec = np.array([1., 1])
        result = np.array([[1/(1+.1)**2., 0, 0],[-1/(1+.1)**2, 0, 0]])
        for fun in accelerationFunctions:
            assert_almost_equal(fun(r_vec, m_vec, 0.1), result)

    # 3 dimensions
    def test4(self):
        r_vec = np.array([[0.,0,0],[1,1,1]])
        m_vec = np.array([1., 0])
        result = np.array([[0, 0, 0],[-1, -1, -1]])/np.sqrt(3)**3
        for fun in accelerationFunctions:
            assert_almost_equal(fun(r_vec, m_vec, 0.), result)

    # More particles
    def test5(self):
        r_vec = np.array([[1.,0,0],[0,0,0],[2,0,0],[3,0,0]])
        m_vec = np.array([1., 0, 0, 0])
        result = np.array([[0, 0, 0],[1, 0, 0],[-1, 0, 0],[-1/4, 0, 0]])
        for fun in accelerationFunctions:
            assert_almost_equal(fun(r_vec, m_vec, 0.), result)

    # Many particles: self consistency
    def test5(self):
        r_vec = np.random.rand(100, 3)
        m_vec = np.random.rand(100)
        result = bruteForce(r_vec, m_vec, 0.1)
        for fun in [bruteForceNumba, bruteForceNumbaOptimized, bruteForceCPP]:
            assert_almost_equal(fun(r_vec, m_vec, 0.1), result)

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