import numpy
import time

from amuse.test import amusetest
from amuse.units import units
from amuse.datamodel.memory_storage import InMemoryGridAttributeStorage
from amuse.datamodel.memory_storage import InMemoryVectorQuantityAttribute


class TestInMemoryGridAttributeStorage(amusetest.TestCase):

    def test1(self):
        x = InMemoryGridAttributeStorage(5, 4, 3)
        i = (0, 1, 2, 3, 4)
        j = (1, 3, 1, 3, 1)
        k = (0, 2, 0, 2, 0)
        x.set_values_in_store(
            (i, j, k),
            ['a', 'b'],
            [2.0 | units.kg, 1.0 | units.m]
        )

        b, a = x.get_values_in_store(None, ['b', 'a'])
        print(b.shape, a.shape)
        self.assertEqual(b[0][1][0], 1.0 | units.m)
        self.assertEqual(b[0][0][0], 0.0 | units.m)
        self.assertEqual(a[0][1][0], 2.0 | units.kg)
        self.assertEqual(a[1][3][2], 2.0 | units.kg)
        self.assertEqual(a[1][2][2], 0.0 | units.kg)

        (b,) = x.get_values_in_store((numpy.s_[0:4], numpy.s_[1:4], numpy.s_[:]), ['a'])

        self.assertEqual(b[0][0][0], 2.0 | units.kg)
        self.assertEqual(b[0][0][2], 0.0 | units.kg)
        self.assertEqual(b[1][2][2], 2.0 | units.kg)
        self.assertEqual(b[2][0][0], 2.0 | units.kg)
        self.assertEqual(b[3][2][2], 2.0 | units.kg)

        self.assertEqual(b.sum(), 8.0 | units.kg)

        self.assertEqual(sorted(x.get_defined_attribute_names()), ["a", "b"])

    def test2(self):
        x = InMemoryGridAttributeStorage(5, 4, 3)
        i = (0, 1, 2, 3, 4)
        j = (1, 3, 1, 3, 1)
        k = (0, 2, 0, 2, 0)
        x.set_values_in_store(
            (i, j, k),
            ['a', 'b'],
            [2.0, 1.0]
        )

        b, a = x.get_values_in_store(None, ['b', 'a'])
        print(a)
        self.assertEqual(b[0][1][0], 1.0)
        self.assertEqual(b[0][0][0], 0.0)
        self.assertEqual(a[0][1][0], 2.0)
        self.assertEqual(a[1][3][2], 2.0)
        self.assertEqual(a[1][2][2], 0.0)

        (b,) = x.get_values_in_store((numpy.s_[0:4], numpy.s_[1:4], numpy.s_[:]), ['a'])

        self.assertEqual(b[0][0][0], 2.0)
        self.assertEqual(b[0][0][2], 0.0)
        self.assertEqual(b[1][2][2], 2.0)
        self.assertEqual(b[2][0][0], 2.0)
        self.assertEqual(b[3][2][2], 2.0)

        self.assertEqual(b.sum(), 8.0)

        self.assertEqual(sorted(x.get_defined_attribute_names()), ["a", "b"])

        y = x.copy()

        (b,) = y.get_values_in_store((numpy.s_[0:4], numpy.s_[1:4], numpy.s_[:]), ['a'])

        self.assertEqual(b[0][0][0], 2.0)
        self.assertEqual(b[0][0][2], 0.0)
        self.assertEqual(b[1][2][2], 2.0)

    def test3(self):
        x = InMemoryGridAttributeStorage(5, 4, 3)
        i = (0, 1, 2, 3, 4)
        j = (1, 3, 1, 3, 1)
        k = (0, 2, 0, 2, 0)
        x.set_values_in_store(
            (i, j, k),
            ['a', 'b'],
            [2.0 | units.kg, 1.0 | units.m]
        )

        b, a = x.get_values_in_store((0, 1, 0), ['b', 'a'])
        print(b, a)
        self.assertEqual(b, 1.0 | units.m)
        self.assertEqual(a, 2.0 | units.kg)
        b, a = x.get_values_in_store((0, 0, 0), ['b', 'a'])
        print(b, a)
        self.assertEqual(b, 0.0 | units.m)
        self.assertEqual(a, 0.0 | units.kg)


class TestInMemoryVectorQuantityAttribute(amusetest.TestCase):

    def test1(self):
        quantity = units.m.new_quantity(numpy.array([1.0, 2.0, 3.0]))
        attribute = InMemoryVectorQuantityAttribute('test', quantity.shape, quantity.unit)
        attribute.set_values(None, quantity)

        self.assertEqual(attribute.get_length(), 3)
        self.assertEqual(attribute.get_shape(), (3,))
        self.assertEqual(attribute.get_values([1, 2]), [2.0, 3.0] | units.m)

        attribute.increase_to_length(5)
        self.assertEqual(attribute.get_values(None), [1.0, 2.0, 3.0, 0.0, 0.0] | units.m)

    def test2(self):
        quantity = units.m.new_quantity(numpy.array([1.0, 2.0, 3.0]))
        attribute = InMemoryVectorQuantityAttribute('test', quantity.shape, quantity.unit)
        attribute.set_values(None, quantity)
        attribute.set_values([1, 2], [4.0, 5.0] | units.m)

        attribute.increase_to_length(5)
        self.assertEqual(attribute.get_values(None), [1.0, 4.0, 5.0, 0.0, 0.0] | units.m)

    def test3(self):
        quantity = units.m.new_quantity(numpy.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]))
        attribute = InMemoryVectorQuantityAttribute('test', quantity.shape, quantity.unit)
        attribute.set_values(None, quantity)
        self.assertEqual(attribute.get_values([1]),  [4.0, 5.0, 6.0] | units.m)
        self.assertEqual(attribute.get_shape(), (2, 3))
        attribute.increase_to_length(4)
        self.assertEqual(attribute.get_shape(), (4, 3))
        self.assertEqual(attribute.get_values(None), [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] | units.m)
