"""
@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-17 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
"""

import unittest
import math
import numpy as np
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()
        xy2 = clu.data[['x', 'y']].copy()
        xy3 = xy2.view((xy2.dtype[0], len(xy2.dtype.names)))
        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()
        xy2 = clu.data[['x', 'y']].copy()
        xy3 = xy2.view((xy2.dtype[0], len(xy2.dtype.names)))
        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()
        expect = [(0., 0., 0., 1), (1., 0., 1., 10)]
        self.assertItemsEqual(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_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_add_layer(self):
        clu = mc.Cluster()
        # from hbncu project
        b_surf = 2.50
        clu.set_rmax(4.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))
        a1 = -10.0 * b1 - 10.0 * b2
        emitter = np.array((0.0, 0.0, 0.0))
        clu.add_layer(7, a1, b1, b2)
        pos = clu.find_positions(pos=emitter)
        self.assertEqual(len(pos), 1)

    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)
        clu2.add_atom(4, np.asarray([0, -0.2, 0]), 0)
        clu2.add_atom(5, np.asarray([0, 0.05, -1]), 0)
        clu2.add_atom(5, np.asarray([0, 0, -1.01]), 0)
        clu2.add_atom(6, np.asarray([0, 0, -1.99]), 0)

        clu1.set_rmax(1.5)
        clu1.add_cluster(clu2, check_rmax=True, check_unique=True, tol=0.1)
        self.assertEqual(clu1.get_atom_count(), 5+2)

    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(len(pos), 1)
        self.assertEqual(pos[0], 206)

    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(len(idx), 2)
        self.assertEqual(clu.get_atomtype(idx[0]), 8)
        self.assertEqual(clu.get_atomtype(idx[1]), 20)
        idx = clu.find_index_cylinder(pos, rxy, rz, 8)
        self.assertEqual(len(idx), 1)

    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)
        clu.trim_cylinder(2.3, 4.2)
        self.assertEqual(clu.data.dtype, clu.dtype)
        self.assertEqual(clu.data.shape[0], 21 * 5)
        self.assertEqual(clu.data[1]['i'], 2)
        self.assertEqual(clu.data[1]['s'], 'N')
        self.assertEqual(clu.data[1]['t'], 7)
        self.assertEqual(clu.get_emitter_count(), 1)

    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)
        clu.trim_sphere(2.3)
        self.assertEqual(clu.data.dtype, clu.dtype)
        self.assertEqual(clu.data.shape[0], 39)
        self.assertEqual(clu.data[1]['i'], 2)
        self.assertEqual(clu.data[1]['s'], 'N')
        self.assertEqual(clu.data[1]['t'], 7)
        self.assertEqual(clu.get_emitter_count(), 1)

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