"""
@package tests.test_cluster
unit tests for pmsco.cluster

the purpose of these tests is to check whether the code runs as expected in a particular environment.

to run the tests, change to the directory which contains the tests directory, and execute =nosetests=.

@pre nose must be installed (python-nose package on Debian).

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015-19 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from io import BytesIO
import math
import numpy as np
import unittest
import pmsco.cluster as mc


class TestClusterFunctions(unittest.TestCase):
    def setUp(self):
        # before each test method
        pass

    def tearDown(self):
        # after each test method
        pass
 
    @classmethod
    def setup_class(cls):
        # before any methods in this class
        pass
 
    @classmethod
    def teardown_class(cls):
        # teardown_class() after any methods in this class
        pass

    @staticmethod
    def create_cube():
        """
        create a cluster object with atoms on the corners, faces and body center of the unit cube.

        the atom types are unique in an arbitrary sequence.
        the emitter is at the origin, atom type 1.

        @return: cluster.Cluster object.
        """
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([0, 0, 0]), 1)
        clu.add_atom(2, np.asarray([1, 0, 0]), 0)
        clu.add_atom(3, np.asarray([0, 1, 0]), 0)
        clu.add_atom(4, np.asarray([0, 0, 1]), 0)
        clu.add_atom(5, np.asarray([-1, 0, 0]), 0)
        clu.add_atom(6, np.asarray([0, -1, 0]), 0)
        clu.add_atom(7, np.asarray([0, 0, -1]), 0)
        clu.add_atom(8, np.asarray([1, 1, 0]), 0)
        clu.add_atom(9, np.asarray([0, 1, 1]), 0)
        clu.add_atom(10, np.asarray([1, 0, 1]), 0)
        clu.add_atom(11, np.asarray([-1, 1, 0]), 0)
        clu.add_atom(12, np.asarray([0, -1, 1]), 0)
        clu.add_atom(13, np.asarray([1, 0, -1]), 0)
        clu.add_atom(14, np.asarray([-1, -1, 0]), 0)
        clu.add_atom(15, np.asarray([0, -1, -1]), 0)
        clu.add_atom(16, np.asarray([-1, 0, -1]), 0)
        clu.add_atom(17, np.asarray([1, -1, 0]), 0)
        clu.add_atom(18, np.asarray([0, 1, -1]), 0)
        clu.add_atom(19, np.asarray([-1, 0, 1]), 0)
        clu.add_atom(20, np.asarray([1, 1, 1]), 0)
        clu.add_atom(21, np.asarray([-1, 1, 1]), 0)
        clu.add_atom(22, np.asarray([1, -1, 1]), 0)
        clu.add_atom(23, np.asarray([1, 1, -1]), 0)
        clu.add_atom(24, np.asarray([-1, -1, -1]), 0)
        clu.add_atom(25, np.asarray([1, -1, -1]), 0)
        clu.add_atom(26, np.asarray([-1, 1, -1]), 0)
        clu.add_atom(27, np.asarray([-1, -1, 1]), 0)
        return clu

    def test_numpy_extract(self):
        """
        test array extraction code which should be compatible to numpy versions before and after 1.14.

        numpy 1.14 introduces changes to multi-column indexing of structured arrays such as data[['x','y']].
        first, it will return a view instead of a copy.
        second, it will assign fields by position rather than by name.

        the first change affects our cluster code in several places
        where we extract XYZ coordinates from the cluster data array.
        this test checks whether the new code works with a particular numpy version.

        @return: None
        """
        clu = self.create_cube()
        xy3 = np.empty((clu.data.shape[0], 2), np.float32)
        xy3[:, 0] = clu.data['x']
        xy3[:, 1] = clu.data['y']
        ctr = np.asarray((1.0, 0.0, 0.0))
        dist = np.linalg.norm(xy3 - ctr[0:2], axis=1)
        self.assertAlmostEqual(1.0, dist[0])
        self.assertAlmostEqual(0.0, dist[1])

        clu.clear()
        xy3 = np.empty((clu.data.shape[0], 2), np.float32)
        xy3[:, 0] = clu.data['x']
        xy3[:, 1] = clu.data['y']
        ctr = np.asarray((1.0, 0.0, 0.0))
        dist = np.linalg.norm(xy3 - ctr[0:2], axis=1)
        self.assertEqual(0, dist.shape[0])

    def test_get_positions(self):
        """
        check that we get an independent copy of the original data.

        @return: None
        """
        clu = self.create_cube()
        pos = clu.get_positions()
        self.assertEqual(clu.data.shape[0], pos.shape[0])
        self.assertEqual(3, pos.shape[1])
        self.assertEqual(np.float32, pos.dtype)
        self.assertEqual(1.0, pos[1, 0])
        self.assertEqual(0.0, pos[1, 1])
        self.assertEqual(0.0, pos[1, 2])
        pos[0, 0] = 15.0
        self.assertEqual(0.0, clu.data['x'][0])

        # empty cluster
        clu.clear()
        self.assertEqual(clu.data.shape[0], 0)
        self.assertEqual(3, pos.shape[1])
        self.assertEqual(np.float32, pos.dtype)

    def test_set_positions(self):
        clu = mc.Cluster()
        clu.data = np.zeros(2, dtype=clu.dtype)
        pos = np.array([[1., 2., 3.], [4., 5., 6.]])
        clu.set_positions(pos)
        self.assertEqual(1., clu.data['x'][0])
        self.assertEqual(2., clu.data['y'][0])
        self.assertEqual(3., clu.data['z'][0])
        self.assertEqual(4., clu.data['x'][1])
        self.assertEqual(5., clu.data['y'][1])
        self.assertEqual(6., clu.data['z'][1])

    def test_get_emitters(self):
        clu = self.create_cube()
        clu.set_emitter(idx=0)
        clu.set_emitter(idx=9)
        self.assertEqual(2, clu.get_emitter_count())
        result = clu.get_emitters(['x', 'y', 'z', 't'])
        expect = [(0., 0., 0., 1), (1., 0., 1., 10)]
        self.assertEqual(expect, result)

    def test_get_z_layers(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 0.1]), 0)
        clu.add_atom(2, np.asarray([0, 1, -0.3]), 0)
        clu.add_atom(1, np.asarray([0, 1, -0.2]), 0)
        clu.add_atom(1, np.asarray([1, 0, 0]), 1)
        clu.add_atom(1, np.asarray([0, 1, -0.2001]), 0)
        clu.add_atom(2, np.asarray([0, 1, -0.1]), 0)
        clu.add_atom(1, np.asarray([0, 1, -0.1999]), 0)
        layers = clu.get_z_layers(0.01)

        np.testing.assert_allclose(layers, np.asarray([-0.3, -0.2, -0.1, 0.0, +0.1]), atol=0.001)

    def test_get_center(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 0]), 0)
        clu.add_atom(2, np.asarray([0, 0, 1]), 0)
        clu.add_atom(1, np.asarray([0, 1, 0]), 0)
        clu.add_atom(2, np.asarray([-1, -1, -1]), 0)
        v0 = np.asarray([0, 0, 0])
        v1 = np.asarray([1/2, 1/2, 0])
        v2 = np.asarray([-1/2, -1/2, 0])
        v = clu.get_center()
        np.testing.assert_allclose(v, v0, atol=0.001)
        v = clu.get_center(element=1)
        np.testing.assert_allclose(v, v1, atol=0.001)
        v = clu.get_center(element="He")
        np.testing.assert_allclose(v, v2, atol=0.001)

    def test_relax(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 1]), 0)
        clu.add_atom(1, np.asarray([1, 0, 0]), 1)
        clu.add_atom(2, np.asarray([0, 1, -1]), 0)
        clu.add_atom(1, np.asarray([0, 1, -2]), 0)
        clu.add_atom(2, np.asarray([0, 1, -3]), 0)
        idx = clu.relax(-0.3, -0.1, 2)

        np.testing.assert_almost_equal(idx, np.asarray([2, 4]))
        np.testing.assert_allclose(clu.get_position(0), np.asarray([1, 0, 1]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(1), np.asarray([1, 0, 0]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(2), np.asarray([0, 1, -1.1]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(3), np.asarray([0, 1, -2.0]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(4), np.asarray([0, 1, -3.1]), atol=1e-6)

    def test_rotate_x(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 0]), 0)
        clu.add_atom(1, np.asarray([0, 1, 0]), 0)
        clu.add_atom(1, np.asarray([0, 0, 1]), 0)
        clu.rotate_x(90)

        np.testing.assert_allclose(clu.get_position(0), np.asarray([1, 0, 0]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(1), np.asarray([0, 0, 1]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(2), np.asarray([0, -1, 0]), atol=1e-6)

    def test_rotate_y(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 0]), 0)
        clu.add_atom(1, np.asarray([0, 1, 0]), 0)
        clu.add_atom(1, np.asarray([0, 0, 1]), 0)
        clu.rotate_y(90)

        np.testing.assert_allclose(clu.get_position(0), np.asarray([0, 0, -1]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(1), np.asarray([0, 1, 0]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(2), np.asarray([1, 0, 0]), atol=1e-6)

    def test_rotate_z(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 0]), 0)
        clu.add_atom(1, np.asarray([0, 1, 0]), 0)
        clu.add_atom(1, np.asarray([0, 0, 1]), 0)
        clu.rotate_z(90)
        
        np.testing.assert_allclose(clu.get_position(0), np.asarray([0, 1, 0]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(1), np.asarray([-1, 0, 0]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(2), np.asarray([0, 0, 1]), atol=1e-6)

    def test_translate(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([1, 0, 0]), 0)
        clu.add_atom(2, np.asarray([0, 1, 0]), 0)
        clu.add_atom(3, np.asarray([0, 0, 1]), 0)

        v = np.array((0.1, 0.2, 0.3))
        shift = clu.translate(v)
        np.testing.assert_allclose(clu.get_position(0), np.asarray([1.1, 0.2, 0.3]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(1), np.asarray([0.1, 1.2, 0.3]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(2), np.asarray([0.1, 0.2, 1.3]), atol=1e-6)
        np.testing.assert_allclose(shift, np.asarray([0, 1, 2]))

        shift = clu.translate(v, element=3)
        np.testing.assert_allclose(clu.get_position(0), np.asarray([1.1, 0.2, 0.3]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(1), np.asarray([0.1, 1.2, 0.3]), atol=1e-6)
        np.testing.assert_allclose(clu.get_position(2), np.asarray([0.2, 0.4, 1.6]), atol=1e-6)
        np.testing.assert_allclose(shift, np.asarray([2]))

    def test_add_layer(self):
        clu = mc.Cluster()
        clu.set_rmax(2.0)
        b1 = np.array((1.0, 0.0, 0.0))
        b2 = np.array((0.0, 1.0, 0.0))
        a1 = np.array((0.1, 0.0, -0.1))
        clu.add_layer(7, a1, b1, b2)

        exp_pos = [[-0.9, -1.0, -0.1], [0.1, -1.0, -0.1], [1.1, -1.0, -0.1],
                   [-1.9, 0.0, -0.1], [-0.9, 0.0, -0.1], [0.1, 0.0, -0.1], [1.1, 0.0, -0.1],
                   [-0.9, 1.0, -0.1], [0.1, 1.0, -0.1], [1.1, 1.0, -0.1]]

        nn = len(exp_pos)
        self.assertEqual(nn, clu.data.shape[0])
        self.assertEqual(nn, clu.get_atom_count())
        act_pos = np.sort(clu.get_positions(), axis=0)
        exp_pos = np.sort(np.array(exp_pos), axis=0)
        np.testing.assert_allclose(act_pos, exp_pos)
        act_idx = np.unique(clu.data['i'])
        self.assertEqual(nn, act_idx.shape[0])
        act_typ = (clu.data['t'] == 7).nonzero()[0]
        self.assertEqual(nn, act_typ.shape[0])

    def test_add_bulk(self):
        clu = mc.Cluster()
        clu.set_rmax(2.0)
        b1 = np.array((1.0, 0.0, 0.0))
        b2 = np.array((0.0, 1.0, 0.0))
        b3 = np.array((0.0, 0.0, 1.0))
        a1 = np.array((0.1, 0.0, -0.1))
        z_surf = 0.8
        clu.add_bulk(7, a1, b1, b2, b3, z_surf=z_surf)

        r_great = max(clu.rmax, np.linalg.norm(a1))
        n1 = max(int(r_great / np.linalg.norm(b1)) + 1, 4) * 3
        n2 = max(int(r_great / np.linalg.norm(b2)) + 1, 4) * 3
        n3 = max(int(r_great / np.linalg.norm(b3)) + 1, 4) * 3
        exp_pos = []
        nn = 0
        for i1 in range(-n1, n1 + 1):
            for i2 in range(-n2, n2 + 1):
                for i3 in range(-n3, n3 + 1):
                    v = a1 + b1 * i1 + b2 * i2 + b3 * i3
                    if np.linalg.norm(v) <= clu.rmax and v[2] <= z_surf:
                        exp_pos.append(v)
                        nn += 1

        self.assertEqual(nn, clu.data.shape[0])
        self.assertEqual(nn, clu.get_atom_count())
        act_pos = np.sort(clu.get_positions(), axis=0)
        exp_pos = np.sort(np.array(exp_pos), axis=0)
        np.testing.assert_allclose(act_pos, exp_pos)
        act_idx = np.unique(clu.data['i'])
        self.assertEqual(nn, act_idx.shape[0])
        act_typ = (clu.data['t'] == 7).nonzero()[0]
        self.assertEqual(nn, act_typ.shape[0])

    def test_add_cluster(self):
        clu1 = mc.Cluster()
        clu1.add_atom(1, np.asarray([0, 0, 0]), 1)
        clu1.add_atom(2, np.asarray([1, 0, 0]), 0)
        clu1.add_atom(3, np.asarray([0, 1, 0]), 0)
        clu1.add_atom(4, np.asarray([0, 0, -1]), 0)
        clu1.add_atom(5, np.asarray([0, 0, -2]), 0)

        clu2 = mc.Cluster()
        clu2.add_atom(3, np.asarray([-0.2, 0, 0]), 0)   # unique
        clu2.add_atom(4, np.asarray([0, -0.2, 0]), 0)   # unique
        clu2.add_atom(5, np.asarray([0, 0.05, -1]), 0)  # not unique
        clu2.add_atom(5, np.asarray([0, 0, -1.09]), 0)  # just within tolerance of uniqueness
        clu2.add_atom(6, np.asarray([0, 0, -1.99]), 0)  # not unique
        clu2.add_atom(7, np.asarray([0, 0, -1.10]), 0)  # just out of tolerance of uniqueness

        clu1.set_rmax(1.5)
        clu1.add_cluster(clu2, check_rmax=True, check_unique=True, tol=0.1)
        self.assertEqual(5+3, clu1.get_atom_count())
        self.assertEqual(7, clu1.data['t'][-1])
        self.assertEqual(6, clu2.data.shape[0])

    def test_find_positions(self):
        clu = mc.Cluster()
        # from hbncu project
        b_surf = 2.50
        clu.set_rmax(b_surf * 10.0)
        b1 = np.array((b_surf, 0.0, 0.0))
        b2 = np.array((b_surf / 2.0, b_surf * math.sqrt(3.0) / 2.0, 0.0))
        a_N = np.array((0.0, 0.0, 0.0))
        a_B = np.array(((b1[0] + b2[0]) / 3.0, (b1[1] + b2[1]) / 3.0, 0.0))

        emitter = a_N + b1 * 1 + b2 * 2

        clu.add_layer(7, a_N, b1, b2)
        clu.add_layer(5, a_B, b1, b2)
        pos = clu.find_positions(pos=emitter)
        self.assertEqual(1, len(pos))
        self.assertEqual(206, pos[0])

        # position in the format returned by get_emitters
        emitter = (emitter[0], emitter[1], emitter[2], 7)
        pos = clu.find_positions(pos=emitter)
        self.assertEqual(1, len(pos))
        self.assertEqual(206, pos[0])

    def test_find_index_cylinder(self):
        clu = self.create_cube()
        pos = np.array((0.8, 0.8, 0.8))
        rxy = 0.5
        rz = 1.0
        idx = clu.find_index_cylinder(pos, rxy, rz, None)
        self.assertEqual(2, len(idx))
        self.assertEqual(8, clu.get_atomtype(idx[0]))
        self.assertEqual(20, clu.get_atomtype(idx[1]))
        idx = clu.find_index_cylinder(pos, rxy, rz, 8)
        self.assertEqual(1, len(idx))

    def test_trim_cylinder(self):
        clu = mc.Cluster()
        clu.set_rmax(10.0)
        v_pos = np.asarray([0, 0, 0])
        v_lat1 = np.asarray([1, 0, 0])
        v_lat2 = np.asarray([0, 1, 0])
        v_lat3 = np.asarray([0, 0, 1])
        clu.add_bulk(7, v_pos, v_lat1, v_lat2, v_lat3)
        clu.set_emitter(pos=v_pos)
        r0 = 2.3
        z0 = 4.2
        clu.trim_cylinder(r0, z0)
        self.assertEqual(clu.dtype, clu.data.dtype)
        self.assertEqual(21 * 5, clu.data.shape[0])
        self.assertEqual(2, clu.data[1]['i'])
        self.assertEqual('N', clu.data[1]['s'])
        self.assertEqual(7, clu.data[1]['t'])
        self.assertEqual(1, clu.get_emitter_count())
        n_low = np.sum(clu.data['z'] < -z0)
        self.assertEqual(0, n_low)
        n_high = np.sum(clu.data['z'] > z0)
        self.assertEqual(0, n_high)
        n_out = np.sum(clu.data['x']**2 + clu.data['y']**2 > r0**2)
        self.assertEqual(0, n_out)

    def test_trim_sphere(self):
        clu = mc.Cluster()
        clu.set_rmax(10.0)
        v_pos = np.asarray([0, 0, 0])
        v_lat1 = np.asarray([1, 0, 0])
        v_lat2 = np.asarray([0, 1, 0])
        v_lat3 = np.asarray([0, 0, 1])
        clu.add_bulk(7, v_pos, v_lat1, v_lat2, v_lat3)
        clu.set_emitter(pos=v_pos)
        r0 = 2.3
        clu.trim_sphere(r0)
        self.assertEqual(clu.dtype, clu.data.dtype)
        self.assertEqual(39, clu.data.shape[0])
        self.assertEqual(2, clu.data[1]['i'])
        self.assertEqual('N', clu.data[1]['s'])
        self.assertEqual(7, clu.data[1]['t'])
        self.assertEqual(1, clu.get_emitter_count())
        n_out = np.sum(clu.data['x']**2 + clu.data['y']**2 + clu.data['z'] > r0**2)
        self.assertEqual(0, n_out)

    def test_trim_paraboloid(self):
        clu = mc.Cluster()
        clu.set_rmax(10.0)
        v_pos = np.asarray([0, 0, 0])
        v_lat1 = np.asarray([1, 0, 0])
        v_lat2 = np.asarray([0, 1, 0])
        v_lat3 = np.asarray([0, 0, 1])
        clu.add_bulk(7, v_pos, v_lat1, v_lat2, v_lat3)
        clu.set_emitter(pos=v_pos)
        r0 = 3.5
        z0 = -2.3
        clu.trim_paraboloid(r0, z0)
        self.assertEqual(clu.data.dtype, clu.dtype)
        self.assertEqual(63, clu.data.shape[0])
        self.assertEqual(2, clu.data[1]['i'])
        self.assertEqual('N', clu.data[1]['s'])
        self.assertEqual(7, clu.data[1]['t'])
        self.assertEqual(1, clu.get_emitter_count())
        n_low = np.sum(clu.data['z'] < z0)
        self.assertEqual(0, n_low)
        n_out = np.sum(clu.data['x']**2 + clu.data['y']**2 > r0**2)
        self.assertEqual(0, n_out)

    def test_trim_slab(self):
        clu = self.create_cube()
        clu.trim_slab('x', 0.5, 1.1)
        self.assertEqual(clu.dtype, clu.data.dtype)
        self.assertEqual(9 * 2, clu.data.shape[0])
        self.assertEqual(1, clu.get_emitter_count())

    def test_save_to_file(self):
        clu = self.create_cube()
        f = BytesIO()
        pos = np.asarray((-1, -1, 0))
        clu.set_emitter(pos=pos)
        clu.save_to_file(f, mc.FMT_PMSCO, "qwerty", emitters_only=True)
        f.seek(0)
        line = f.readline()
        self.assertEqual(line, b"# index element symbol class x y z emitter charge\n", b"line 1: " + line)
        line = f.readline()
        self.assertRegex(line, b"[0-9]+ +1 +H +[0-9]+ +[0.]+ +[0.]+ +[0.]+ +1 +[0.]", b"line 3: " + line)
        line = f.readline()
        self.assertRegex(line, b"[0-9]+ +14 +Si +[0-9]+ +[01.-]+ +[01.-]+ +[0.]+ +1 +[0.]", b"line 4: " + line)
        line = f.readline()
        self.assertEqual(b"", line, b"end of file")

        f = BytesIO()
        clu.save_to_file(f, mc.FMT_XYZ, "qwerty", emitters_only=True)
        f.seek(0)
        line = f.readline()
        self.assertEqual(b"2\n", line, b"line 1: " + line)
        line = f.readline()
        self.assertEqual(b"qwerty\n", line, b"line 2: " + line)
        line = f.readline()
        self.assertRegex(line, b"H +[0.]+ +[0.]+ +[0.]+", b"line 3: " + line)
        line = f.readline()
        self.assertRegex(line, b"Si +[01.-]+ +[01.-]+ +[0.]+", b"line 4: " + line)
        line = f.readline()
        self.assertEqual(b"", line, b"end of file")

    def test_load_from_file(self):
        f = BytesIO()
        f.write(b"2\n")
        f.write(b"qwerty\n")
        f.write(b"H 0.5 0.6 0.7\n")
        f.write(b"Si -1.5 -1.6 -1.7\n")
        f.seek(0)
        clu = mc.Cluster()
        clu.load_from_file(f, fmt=mc.FMT_XYZ)
        np.testing.assert_allclose(clu.data['t'], np.array([1, 14]))
        np.testing.assert_allclose(clu.data['x'], np.array([0.5, -1.5]))
        np.testing.assert_allclose(clu.data['y'], np.array([0.6, -1.6]))
        np.testing.assert_allclose(clu.data['z'], np.array([0.7, -1.7]))
        np.testing.assert_allclose(clu.data['e'], np.array([0, 0]))
        np.testing.assert_allclose(clu.data['q'], np.array([0, 0]))

        f = BytesIO()
        f.write(b"# index element symbol class x y z emitter charge\n")
        # ['i', 't', 's', 'c', 'x', 'y', 'z', 'e', 'q']
        f.write(b"1 6 C 1 0.5 0.6 0.7 0 -0.5\n")
        f.write(b"2 14 Si 2 -1.5 -1.6 -1.7 1 0.5\n")
        f.seek(0)
        clu = mc.Cluster()
        clu.load_from_file(f, fmt=mc.FMT_PMSCO)
        np.testing.assert_allclose(clu.data['t'], np.array([6, 14]))
        np.testing.assert_allclose(clu.data['x'], np.array([0.5, -1.5]))
        np.testing.assert_allclose(clu.data['y'], np.array([0.6, -1.6]))
        np.testing.assert_allclose(clu.data['z'], np.array([0.7, -1.7]))
        np.testing.assert_allclose(clu.data['e'], np.array([0, 1]))
        np.testing.assert_allclose(clu.data['q'], np.array([-0.5, 0.5]))

    def test_update_atoms(self):
        clu = mc.Cluster()
        clu.add_atom(1, np.asarray([0, 0, 0]), 1)
        clu.add_atom(3, np.asarray([0, 1, 0]), 0)
        clu.add_atom(5, np.asarray([-1, 0, 0]), 0)
        clu.add_atom(6, np.asarray([0, -1, 0]), 0)
        clu.add_atom(2, np.asarray([1, 0, 0]), 0)
        clu.add_atom(4, np.asarray([0, 0, 1]), 0)

        other = mc.Cluster()
        other.add_atom(1, np.asarray([0, 0, 0]), 1)
        other.add_atom(5, np.asarray([-1, 0, 0]), 0)
        other.add_atom(2, np.asarray([1, 0, 0]), 0)
        other.add_atom(6, np.asarray([0, -1, 0]), 0)
        other.add_atom(3, np.asarray([0, 1, 0]), 0)
        other.add_atom(4, np.asarray([0, 0, 1]), 0)

        other.data['c'] = np.asarray((1, 2, 2, 3, 3, 4))
        clu.update_atoms(other, {'c'})
        expected = np.asarray((1, 3, 2, 3, 2, 4))
        np.testing.assert_array_equal(expected, clu.data['c'])

    def test_calc_scattering_angles(self):
        clu = mc.Cluster()
        ref_em = np.asarray([0.1, -0.1, 0.5])
        ref_th = np.asarray([0., 15., 90., 100., 120.])
        ref_ph = np.asarray([0., 90., 180., 270., 360.])
        ref_di = np.asarray([0.5, 1.0, 1.5, 2.0, 2.5])
        exp_th = ref_th[0:4]
        exp_ph = ref_ph[0:4]
        exp_di = ref_di[0:4]
        sel_ph = exp_ph > 180.
        exp_ph[sel_ph] = exp_ph[sel_ph] - 360.

        idx_em = clu.add_atom(1, ref_em, 1)
        for i, r in enumerate(ref_di):
            v = np.asarray([
                r * math.cos(math.radians(ref_ph[i])) * math.sin(math.radians(ref_th[i])) + ref_em[0],
                r * math.sin(math.radians(ref_ph[i])) * math.sin(math.radians(ref_th[i])) + ref_em[1],
                r * math.cos(math.radians(ref_th[i])) + ref_em[2]])
            clu.add_atom(i, v, 0)

        result = clu.calc_scattering_angles(idx_em, 2.2)
        np.testing.assert_allclose(result['index'], np.arange(1, exp_di.shape[0] + 1))
        np.testing.assert_allclose(result['polar'], exp_th, atol=1e-3)
        np.testing.assert_allclose(result['azimuth'], exp_ph, atol=1e-3)
        np.testing.assert_allclose(result['dist'], exp_di, rtol=1e-5)
