from amuse.support.testing import amusetest
from io import BytesIO
from collections import namedtuple

import os.path
import math
from amuse import io
from amuse.io import gadget
from amuse.io import nemobin
from amuse.units import nbody_system
from amuse.datamodel import Particles


class GadgetFileFormatProcessorTests(amusetest.TestCase):
    header_parts = (
        b'\x00\x01\x00\x00 N\x00\x00 \xa1\x07\x00 N\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\xf6`Q\xc6#\xcc\x9c>\x8d',
        b'\xed\xb5\xa0\xf7\xc6\x90>\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00 N\x00\x00 ',
        b'\xa1\x07\x00 N\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00',
        b'\x00\x00\x00\x00\x00\x00\x01\x00\x00',
    )

    def test1(self):
        header = b''.join(self.header_parts)
        x = gadget.GadgetFileFormatProcessor()
        file = BytesIO(header)
        x.load_header(file)
        print(x.header_struct)
        self.assertEqual(x.header_struct.Npart[0], 20000)
        self.assertEqual(x.header_struct.Npart[1], 500000)
        self.assertEqual(x.header_struct.Npart[2], 20000)
        self.assertEqual(x.header_struct.Npart[3], 0)
        self.assertEqual(x.header_struct.Npart[4], 0)
        self.assertEqual(x.header_struct.Npart[5], 0)

        self.assertEqual(x.header_struct.Massarr[0], 0.0)
        self.assertAlmostRelativeEqual(x.header_struct.Massarr[1], 4.2911501e-07, 8)
        print(x.header_struct.Massarr[2])
        self.assertAlmostRelativeEqual(x.header_struct.Massarr[2], 2.5000000e-07, 8)

        self.assertEqual(x.header_struct.FlagSfr, 0)
        self.assertEqual(x.header_struct.FlagFeedback, 0)
        self.assertEqual(x.header_struct.FlagAge, 0)
        self.assertEqual(x.header_struct.HubbleParam, 0)

    def test2(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'gadget_snapshot')
        x = gadget.GadgetFileFormatProcessor()
        file = open(filename, 'rb')
        result = x.load_file(file)
        file.close()
        self.assertEqual(len(result[0]), 1000)
        self.assertEqual(len(result[1]), 10000)

    def test3(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'gadget_snapshot')
        x = gadget.GadgetFileFormatProcessor()
        result = io.read_set_from_file(filename, format='gadget')
        self.assertEqual(len(result[0]), 1000)
        self.assertEqual(len(result[1]), 10000)

    def test4(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
        x = gadget.GadgetFileFormatProcessor()
        result = io.read_set_from_file(filename, format='gadget')
        self.assertEqual(len(result[0]), 1472)
        self.assertEqual(len(result[1]), 0)

    def test5(self):
        options = io.get_options_for_format('gadget')
        found_has_acceleration = False
        for name, description, defaultval in options:
            if name == 'has_acceleration':
                found_has_acceleration = True

        self.assertTrue(found_has_acceleration)

    def test6(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
        x = gadget.GadgetFileFormatProcessor()
        gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(filename, format='gadget')
        self.assertEqual(len(gas), 1472)
        self.assertEqual(len(halo), 0)
        self.assertEqual(gas[0].key, 1)
        self.assertEqual(gas[1].key, 2)
        self.assertEqual(gas[2].key, 3)
        self.assertEqual(gas[1471].key, 1472)

    def test7(self):
        """test returned ids from gadget file
        for ticket #245.
        All the 'uneven' particles have key "1", and identical velocities/positions. This is incorrect
        upon further inspection, the test file is incorrect
        """
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'ticket245.dat')
        gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(filename, format='gadget')

        self.assertEqual(len(gas), 0)
        self.assertEqual(len(halo), 1324)
        self.assertEqual(len(disk), 0)
        self.assertEqual(len(bulge), 0)
        self.assertEqual(len(stars), 0)
        self.assertEqual(len(bndry), 0)
        self.assertEqual(halo[0].key, 544418538)
        self.assertEqual(halo[1].key, 544511335)
        self.assertEqual(halo[2].key, 544511457)
        self.assertAlmostRelativeEquals(halo[0].velocity[0], -24.785614 | nbody_system.speed, 7)
        print(halo[1].velocity)
        self.assertAlmostRelativeEquals(halo[1].velocity[0], -25.346375 | nbody_system.speed, 7)
        self.assertAlmostRelativeEquals(halo[2].velocity[0], -25.394440 | nbody_system.speed, 7)

    def test8(self):
        """test returned ids from gadget file
        for ticket #245.
        added option to not use the ids as a key, should fix the problem
        for incorrect id's
        """
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'ticket245.dat')
        gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(filename, format='gadget', ids_are_keys=False)

        self.assertEqual(len(gas), 0)
        self.assertEqual(len(halo), 1324)
        self.assertEqual(len(disk), 0)
        self.assertEqual(len(bulge), 0)
        self.assertEqual(len(stars), 0)
        self.assertEqual(len(bndry), 0)
        self.assertEqual(halo[0].id, 544418538)
        self.assertEqual(halo[1].id, 544511335)
        self.assertEqual(halo[2].id, 544511457)
        self.assertAlmostRelativeEquals(halo[0].velocity[0], -24.785614 | nbody_system.speed, 7)
        self.assertAlmostRelativeEquals(halo[1].velocity[0], -25.346375 | nbody_system.speed, 7)
        self.assertAlmostRelativeEquals(halo[2].velocity[0], -25.394440 | nbody_system.speed, 7)

    def test9(self):
        class FakeList(object):
            def __init__(self, _len):
                self._len = _len

            def __len__(self):
                return self._len

        set = (FakeList(20000), FakeList(500000), FakeList(20000), (), (), ())
        x = gadget.GadgetFileFormatProcessor(set=set)
        x.equal_mass_array = (0.0, 4.291150104743886e-07, 2.5e-07, 0.0, 0.0, 0.0) | nbody_system.mass
        file = BytesIO()
        x.store_header(file)
        print(x.header_struct)
        self.assertEqual(x.header_struct.Npart[0], 20000)
        self.assertEqual(x.header_struct.Npart[1], 500000)
        self.assertEqual(x.header_struct.Npart[2], 20000)
        self.assertEqual(x.header_struct.Npart[3], 0)
        self.assertEqual(x.header_struct.Npart[4], 0)
        self.assertEqual(x.header_struct.Npart[5], 0)

        print(repr(file.getvalue()))
        print(repr(b''.join(self.header_parts)))
        self.assertEqual(repr(file.getvalue()[0:30]), repr(b''.join(self.header_parts)[0:30]))

    def test10(self):
        p = Particles(2)
        p[0].position = [1.0, 2.0, 3.0] | nbody_system.length
        p[1].position = [4.0, 5.0, 6.0] | nbody_system.length
        p[0].velocity = [7.0, 8.0, 10.0] | nbody_system.length / nbody_system.time
        p[1].velocity = [11.0, 12.0, 13.0] | nbody_system.length / nbody_system.time
        p.u = [3, 4] | nbody_system.potential
        p.rho = [5, 6] | nbody_system.density
        p.mass = [5, 6] | nbody_system.mass
        x = gadget.GadgetFileFormatProcessor(set=p)
        file = BytesIO()
        x.store_body(file)
        input = BytesIO(file.getvalue())
        positions = x.read_fortran_block_float_vectors(input)
        self.assertEqual(positions[0], [1.0, 2.0, 3.0])
        self.assertEqual(positions[1], [4.0, 5.0, 6.0])
        velocities = x.read_fortran_block_float_vectors(input)
        self.assertEqual(velocities[0], [7.0, 8.0, 10.0])
        self.assertEqual(velocities[1], [11.0, 12.0, 13.0])
        ids = x.read_fortran_block_ulongs(input)
        self.assertEqual(ids[0], p[0].key)
        self.assertEqual(ids[1], p[1].key)
        masses = x.read_fortran_block_floats(input)
        self.assertEqual(masses[0], 5)
        self.assertEqual(masses[1], 6)
        u = x.read_fortran_block_floats(input)
        self.assertEqual(u[0], 3)
        self.assertEqual(u[1], 4)

    def test11(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
        gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(filename, format='gadget')
        self.assertEqual(len(gas), 1472)
        self.assertEqual(len(halo), 0)
        self.assertEqual(gas[0].key, 1)
        self.assertEqual(gas[1].key, 2)
        self.assertEqual(gas[2].key, 3)
        self.assertEqual(gas[1471].key, 1472)
        self.assertAlmostRelativeEquals(gas[0:5].x, [-0.0713372901082, 0.0713372901082, -0.21178227663, -0.0698266476393, 0.0698266476393] | nbody_system.length, 7)
        self.assertAlmostRelativeEquals(gas[0:5].u, [0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451] | (nbody_system.length / nbody_system.time)**2, 7)

        outputfilename = 'gadgettest.output'
        try:
            io.write_set_to_file((gas, halo, disk, bulge, stars, bndry), outputfilename, format='gadget', ids_are_long=False)

            gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(outputfilename, format='gadget')
            self.assertEqual(len(gas), 1472)
            self.assertEqual(len(halo), 0)
            self.assertEqual(gas[0].key, 1)
            self.assertEqual(gas[1].key, 2)
            self.assertEqual(gas[2].key, 3)
            self.assertEqual(gas[1471].key, 1472)
        finally:
            if os.path.exists(outputfilename):
                os.remove(outputfilename)

    def test12(self):
        print("Test return_header for Gadget read_set_from_file")
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
        data = io.read_set_from_file(filename, format='gadget', return_header=False)  # (default)
        self.assertTrue(isinstance(data, tuple))
        self.assertEqual(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry)")

        data = io.read_set_from_file(filename, format='gadget', return_header=True)
        self.assertTrue(isinstance(data, tuple))
        self.assertEqual(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry, "
            "Npart, Massarr, Time, Redshift, FlagSfr, FlagFeedback, Nall, FlagCooling, "
            "NumFiles, BoxSize, Omega0, OmegaLambda, HubbleParam, FlagAge, FlagMetals, "
            "NallHW, flag_entr_ics)")

        self.assertEqual(len(data.gas), 1472)
        self.assertEqual(len(data.halo), 0)
        self.assertEqual(data.gas[0].key, 1)
        self.assertEqual(data.gas[1].key, 2)
        self.assertEqual(data.gas[2].key, 3)
        self.assertEqual(data.gas[1471].key, 1472)
        self.assertAlmostRelativeEquals(data.gas[0:5].x, [-0.0713372901082, 0.0713372901082, -0.21178227663, -0.0698266476393, 0.0698266476393] | nbody_system.length, 7)
        self.assertAlmostRelativeEquals(data.gas[0:5].u, [0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451] | (nbody_system.length / nbody_system.time)**2, 7)

        self.assertEqual(data.Npart, (1472, 0, 0, 0, 0, 0))
        self.assertEqual(data.Time, 0.0)
        self.assertEqual(data.Redshift, 0.0)

    def test13(self):
        print("Test convert_gadget_w_to_velocity and return_header for Gadget read_set_from_file")
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'tiny_lcdm_data_littleendian.dat')
        data = io.read_set_from_file(filename, format='gadget', return_header=False, convert_gadget_w_to_velocity=False)  # (default)
        self.assertTrue(isinstance(data, tuple))
        self.assertEqual(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry)")
        self.assertEqual(len(data.gas), 32)
        self.assertEqual(len(data.halo), 32)
        self.assertEqual(data.gas[0].key, 1)
        self.assertEqual(data.halo[0].key, 32**3 + 1)
        self.assertAlmostRelativeEquals(data.gas[:3].position, [[395.23443604, 395.75210571, 1244.31152344],
            [310.17266846, 440.21728516, 2817.06396484], [191.95669556, 465.57223511, 4430.20068359]] | nbody_system.length, 7)

        data_converted = io.read_set_from_file(filename, format='gadget', return_header=True, convert_gadget_w_to_velocity=True)
        self.assertTrue(isinstance(data_converted, tuple))
        self.assertEqual(data_converted.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry, "
            "Npart, Massarr, Time, Redshift, FlagSfr, FlagFeedback, Nall, FlagCooling, "
            "NumFiles, BoxSize, Omega0, OmegaLambda, HubbleParam, FlagAge, FlagMetals, "
            "NallHW, flag_entr_ics)")

        self.assertEqual(len(data_converted.gas), 32)
        self.assertEqual(len(data_converted.halo), 32)
        self.assertEqual(data_converted.gas[0].key, 1)
        self.assertEqual(data_converted.halo[0].key, 32**3 + 1)
        self.assertEqual(data_converted.Npart, (32, 32, 0, 0, 0, 0))
        self.assertEqual(data_converted.Time, 1/11.0)
        self.assertEqual(data_converted.Redshift, 10.0)
        self.assertEqual(data.gas.position, data_converted.gas.position)
        self.assertAlmostRelativeEquals(data.gas.velocity, math.sqrt(data_converted.Time) * data_converted.gas.velocity, 7)


class NemoBinaryFileFormatProcessorTests(amusetest.TestCase):

    def test1(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        nemofile = nemobin.NemoBinaryFile(file)
        tagcharacter, tagstring, dim, mustswap = nemofile.get_item_header()
        self.assertEqual(tagcharacter, 'c')
        self.assertEqual(tagstring, 'Headline')
        self.assertEqual(len(dim), 1)
        self.assertEqual(dim[0], 28)
        file.close()

    def test2(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        nemofile = nemobin.NemoBinaryFile(file)
        item = nemofile.read_item()

        self.assertEqual(item.data, "init_xrandom: seed used 123")
        file.close()

    def test3(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        nemofile = nemobin.NemoBinaryFile(file)
        data = nemofile.read()
        file.close()
        self.assertEqual(len(data), 3)
        tags = list(data.keys())
        self.assertEqual(tags[0], 'Headline')
        self.assertEqual(tags[1], 'History')
        self.assertEqual(tags[2], 'SnapShot')
        self.assertEqual(data['History'][0].data, 'mkplummer out=plummer128.nemo nbody=128 seed=123 VERSION=2.8b')

        self.assertEqual(len(data['SnapShot'][0].data), 2)
        tags = list(data['SnapShot'][0].data.keys())
        self.assertEqual(tags[0], 'Parameters')
        self.assertEqual(tags[1], 'Particles')

    def test4(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        x = nemobin.NemoBinaryFileFormatProcessor()

        set = x.load_file(file).previous_state()
        file.close()
        self.assertEqual(len(set), 128)
        self.assertEqual(set.get_timestamp(), 0.0 | nbody_system.time)
        self.assertAlmostRelativeEquals(set.kinetic_energy(), 0.230214395174 | nbody_system.energy, 8)
        self.assertAlmostRelativeEquals(set.potential_energy(G=nbody_system.G), -0.473503040144 | nbody_system.energy, 8)

    def test5(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        nemofile = nemobin.NemoBinaryFile(file)
        data = nemofile.read()
        file.close()

        outputfile = BytesIO()
        nemooutputfile = nemobin.NemoBinaryFile(outputfile)
        nemooutputfile.write(data)
        string = outputfile.getvalue()
        outputfile.close()
        inputfile = BytesIO(string)
        nemoinputfile = nemobin.NemoBinaryFile(inputfile)
        tagcharacter, tagstring, dim, mustswap = nemoinputfile.get_item_header()
        self.assertEqual(tagcharacter, 'c')
        self.assertEqual(tagstring, 'Headline')
        self.assertEqual(len(dim), 1)
        self.assertEqual(dim[0], 28)
        inputfile.close()

    def test6(self):
        inputfile = BytesIO()
        nemoinputfile = nemobin.NemoBinaryFile(inputfile)
        data = nemoinputfile.read()
        self.assertEqual(len(data), 0)

    def test7(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        nemofile = nemobin.NemoBinaryFile(file)
        data = nemofile.read()
        file.close()

        outputfile = BytesIO()
        nemooutputfile = nemobin.NemoBinaryFile(outputfile)
        nemooutputfile.write(data)
        string = outputfile.getvalue()
        outputfile.close()

        file = open(filename, 'rb')
        original = file.read()
        file.close()

        self.assertEqual(len(original), len(string))
        self.assertEqual(original, string)

    def test8(self):
        directory_name = os.path.dirname(__file__)
        filename = os.path.join(directory_name, 'plummer128.nemo')
        file = open(filename, 'rb')
        x = nemobin.NemoBinaryFileFormatProcessor()

        set = x.load_file(file)
        file.close()

        outputfile = BytesIO()

        y = nemobin.NemoBinaryFileFormatProcessor()
        y.set = set
        y.store_file(outputfile)
        string = outputfile.getvalue()
        outputfile.close()

        inputfile = BytesIO(string)
        x = nemobin.NemoBinaryFileFormatProcessor()
        set = x.load_file(inputfile)
        inputfile.close()
        self.assertEqual(len(set), 128)
        self.assertAlmostRelativeEquals(set.kinetic_energy(), 0.230214395174 | nbody_system.energy, 8)
        self.assertAlmostRelativeEquals(set.potential_energy(G=nbody_system.G), -0.473503040144 | nbody_system.energy, 8)

    def test9(self):
        filename = os.path.join(os.path.dirname(__file__), 'plummer128.nemo')
        particles = io.read_set_from_file(filename, format="nemobin")
        self.assertEqual(len(particles), 128)
        self.assertAlmostEqual(particles.total_mass(), 1.0 | nbody_system.mass)
        self.assertAlmostEqual(particles.center_of_mass(), 0.0 | nbody_system.length)
        self.assertAlmostEqual(particles.center_of_mass_velocity(), 0.0 | nbody_system.speed)
        self.assertAlmostEqual(particles.kinetic_energy(), 0.230214395174 | nbody_system.energy)
