# -*- coding: utf-8 -*-

import re
import os
import tempfile
import unittest
import numpy
import math

from distutils.version import StrictVersion


from aiida.backends.sqlalchemy.tests.base import SqlAlchemyTests
from aiida.common.datastructures import calc_states
from aiida.common.exceptions import ModificationNotAllowed, ValidationError

from aiida.orm import JobCalculation, Node, Data
from aiida.orm.data.singlefile import SinglefileData
from aiida.orm.data.cif import (CifData, has_pycifrw,
                                symop_string_from_symop_matrix_tr,
                                parse_formula, refine_inline)
from aiida.orm.data.structure import (has_ase, has_pymatgen, has_pyspglib,
                                      get_pymatgen_version, Kind, StructureData,
                                      _atomic_masses, Site, get_formula,
                                      ase_refine_cell)

from aiida.orm.data.array.kpoints import KpointsData
from aiida.orm.data.structure import StructureData
from aiida.orm.data.array.trajectory import TrajectoryData
from aiida.orm.data.array import ArrayData
from aiida.common.constants import elements






__copyright__ = u"Copyright (c), This file is part of the AiiDA platform. For further information please visit http://www.aiida.net/. All rights reserved."
__license__ = "MIT license, see LICENSE.txt file."
__authors__ = "The AiiDA team."
__version__ = "0.7.0.1"

if has_pycifrw():
    from aiida.orm.data.cif import pycifrw_from_cif

if has_ase():
    import ase
    from ase.lattice.spacegroup import crystal

if has_pymatgen():
    from pymatgen.io.smartio import read_mol, read_structure


class TestCalcStatus(SqlAlchemyTests):
    """
    Test the functionality of calculation states.
    """

    def test_state(self):

        # Use the SqlAlchemyTests test computer

        c = JobCalculation(computer=self.computer,
                           resources={
                               'num_machines': 1,
                               'num_mpiprocs_per_machine': 1}
                           )
        # Should be in the NEW state before storing
        self.assertEquals(c.get_state(), calc_states.NEW)

        with self.assertRaises(ModificationNotAllowed):
            c._set_state(calc_states.TOSUBMIT)

        c.store()
        # Should be in the NEW state right after storing
        self.assertEquals(c.get_state(), calc_states.NEW)

        # Set a different state and check
        c._set_state(calc_states.TOSUBMIT)
        self.assertEquals(c.get_state(), calc_states.TOSUBMIT)

        # Set a different state and check
        c._set_state(calc_states.SUBMITTING)
        self.assertEquals(c.get_state(), calc_states.SUBMITTING)

        # Try to reset a state and check that the proper exception is raised
        with self.assertRaises(ModificationNotAllowed):
            c._set_state(calc_states.SUBMITTING)

        # Set a different state and check
        c._set_state(calc_states.FINISHED)
        self.assertEquals(c.get_state(), calc_states.FINISHED)

        # Try to set a previous state (that is, going backward from
        # FINISHED to WITHSCHEDULER, for instance)
        # and check that this is not allowed
        with self.assertRaises(ModificationNotAllowed):
            c._set_state(calc_states.WITHSCHEDULER)

class TestSinglefileData(SqlAlchemyTests):
    """
    Test the SinglefileData class.
    """

    def test_reload_singlefiledata(self):

        file_content = 'some text ABCDE'
        with tempfile.NamedTemporaryFile() as f:
            filename = f.name
            basename = os.path.split(filename)[1]
            f.write(file_content)
            f.flush()
            a = SinglefileData(file=filename)

        the_uuid = a.uuid

        self.assertEquals(a.get_folder_list(), [basename])

        with open(a.get_abs_path(basename)) as f:
            self.assertEquals(f.read(), file_content)

        a.store()

        with open(a.get_abs_path(basename)) as f:
            self.assertEquals(f.read(), file_content)
        self.assertEquals(a.get_folder_list(), [basename])

        b = Node.get_subclass_from_uuid(the_uuid)

        # I check the retrieved object
        self.assertTrue(isinstance(b, SinglefileData))
        self.assertEquals(b.get_folder_list(), [basename])
        with open(b.get_abs_path(basename)) as f:
            self.assertEquals(f.read(), file_content)


class TestCifData(SqlAlchemyTests):
    """
    Tests for CifData class.
    """

    @unittest.skipIf(not has_pycifrw(), "Unable to import PyCifRW")
    def test_reload_cifdata(self):
        file_content = "data_test _cell_length_a 10(1)"
        with tempfile.NamedTemporaryFile() as f:
            filename = f.name
            basename = os.path.split(filename)[1]
            f.write(file_content)
            f.flush()
            a = CifData(file=filename,
                        source={'version': '1234',
                                'db_name': 'COD',
                                'id': '0000001'})

        # Key 'db_kind' is not allowed in source description:
        with self.assertRaises(KeyError):
            a.source = {'db_kind': 'small molecule'}

        the_uuid = a.uuid

        self.assertEquals(a.get_folder_list(), [basename])

        with open(a.get_abs_path(basename)) as f:
            self.assertEquals(f.read(), file_content)

        a.store()

        self.assertEquals(a.source, {
            'db_name': 'COD',
            'id': '0000001',
            'version': '1234',
        })

        with open(a.get_abs_path(basename)) as f:
            self.assertEquals(f.read(), file_content)
        self.assertEquals(a.get_folder_list(), [basename])

        b = Node.get_subclass_from_uuid(the_uuid)

        # I check the retrieved object
        self.assertTrue(isinstance(b, CifData))
        self.assertEquals(b.get_folder_list(), [basename])
        with open(b.get_abs_path(basename)) as f:
            self.assertEquals(f.read(), file_content)

        # Checking the get_or_create() method:
        with tempfile.NamedTemporaryFile() as f:
            f.write(file_content)
            f.flush()
            c, created = CifData.get_or_create(f.name, store_cif=False)

        self.assertTrue(isinstance(c, CifData))
        self.assertTrue(not created)

        with open(c.get_file_abs_path()) as f:
            self.assertEquals(f.read(), file_content)

        other_content = "data_test _cell_length_b 10(1)"
        with tempfile.NamedTemporaryFile() as f:
            f.write(other_content)
            f.flush()
            c, created = CifData.get_or_create(f.name, store_cif=False)

        self.assertTrue(isinstance(c, CifData))
        self.assertTrue(created)

        with open(c.get_file_abs_path()) as f:
            self.assertEquals(f.read(), other_content)

    @unittest.skipIf(not has_pycifrw(), "Unable to import PyCifRW")
    def test_parse_cifdata(self):

        file_content = "data_test _cell_length_a 10(1)"
        with tempfile.NamedTemporaryFile() as f:
            f.write(file_content)
            f.flush()
            a = CifData(file=f.name)

        self.assertEquals(a.values.keys(), ['test'])

    @unittest.skipIf(not has_pycifrw(), "Unable to import PyCifRW")
    def test_change_cifdata_file(self):
        file_content_1 = "data_test _cell_length_a 10(1)"
        file_content_2 = "data_test _cell_length_a 11(1)"
        with tempfile.NamedTemporaryFile() as f:
            f.write(file_content_1)
            f.flush()
            a = CifData(file=f.name)

        self.assertEquals(a.values['test']['_cell_length_a'], '10(1)')

        with tempfile.NamedTemporaryFile() as f:
            f.write(file_content_2)
            f.flush()
            a.set_file(f.name)

        self.assertEquals(a.values['test']['_cell_length_a'], '11(1)')

    @unittest.skipIf(not has_ase() or not has_pycifrw(),
                     "Unable to import ase or pycifrw")
    def test_get_aiida_structure(self):
        with tempfile.NamedTemporaryFile() as f:
            f.write('''
data_test
_cell_length_a    10
_cell_length_b    10
_cell_length_c    10
_cell_angle_alpha 90
_cell_angle_beta  90
_cell_angle_gamma 90
loop_
_atom_site_label
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
C 0 0 0
O 0.5 0.5 0.5
            ''')
            f.flush()
            a = CifData(file=f.name)

        with self.assertRaises(ValueError):
            a._get_aiida_structure(converter='none')

        c = a._get_aiida_structure()

        self.assertEquals(c.get_kind_names(), ['C', 'O'])

    @unittest.skipIf(not has_ase() or not has_pycifrw(),
                     "Unable to import ase or pycifrw")
    def test_ase_primitive_and_conventional_cells_ase(self):
        """
        Checking the number of atoms per primitive/conventional cell
        returned by ASE ase.io.cif.read_cif() method. Test input is
        adapted from http://www.crystallography.net/cod/9012064.cif@120115
        """
        with tempfile.NamedTemporaryFile() as f:
            f.write('''
                data_9012064
                _space_group_IT_number           166
                _symmetry_space_group_name_H-M   'R -3 m :H'
                _cell_angle_alpha                90
                _cell_angle_beta                 90
                _cell_angle_gamma                120
                _cell_length_a                   4.395
                _cell_length_b                   4.395
                _cell_length_c                   30.440
                _cod_database_code               9012064
                loop_
                _atom_site_label
                _atom_site_fract_x
                _atom_site_fract_y
                _atom_site_fract_z
                _atom_site_U_iso_or_equiv
                Bi 0.00000 0.00000 0.40046 0.02330
                Te1 0.00000 0.00000 0.00000 0.01748
                Te2 0.00000 0.00000 0.79030 0.01912
            ''')
            f.flush()
            c = CifData(file=f.name)

        ase = c._get_aiida_structure(converter='ase',
                                     primitive_cell=False).get_ase()
        self.assertEquals(ase.get_number_of_atoms(), 15)

        ase = c._get_aiida_structure(converter='ase').get_ase()
        self.assertEquals(ase.get_number_of_atoms(), 15)

        ase = c._get_aiida_structure(converter='ase',
                                     primitive_cell=True).get_ase()
        self.assertEquals(ase.get_number_of_atoms(), 5)

    @unittest.skipIf(not has_ase() or not has_pycifrw() or
                     not has_pymatgen() or
                     StrictVersion(get_pymatgen_version()) <
                     StrictVersion('3.0.13'), "Unable to import pymatgen"
                     "Unable to import ase, pycifrw or pymatgen")
    def test_ase_primitive_and_conventional_cells_pymatgen(self):
        """
        Checking the number of atoms per primitive/conventional cell
        returned by ASE ase.io.cif.read_cif() method. Test input is
        adapted from http://www.crystallography.net/cod/9012064.cif@120115
        """
        with tempfile.NamedTemporaryFile() as f:
            f.write('''
data_9012064
_space_group_IT_number           166
_symmetry_space_group_name_H-M   'R -3 m :H'
_cell_angle_alpha                90
_cell_angle_beta                 90
_cell_angle_gamma                120
_cell_length_a                   4.395
_cell_length_b                   4.395
_cell_length_c                   30.440
_cod_database_code               9012064
loop_
_atom_site_label
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
_atom_site_U_iso_or_equiv
Bi 0.00000 0.00000 0.40046 0.02330
Te1 0.00000 0.00000 0.00000 0.01748
Te2 0.00000 0.00000 0.79030 0.01912
            ''')
            f.flush()
            c = CifData(file=f.name)

        ase = c._get_aiida_structure(converter='pymatgen',
                                     primitive_cell=False).get_ase()
        self.assertEquals(ase.get_number_of_atoms(), 15)

        ase = c._get_aiida_structure(converter='pymatgen').get_ase()
        self.assertEquals(ase.get_number_of_atoms(), 15)

        ase = c._get_aiida_structure(converter='pymatgen',
                                     primitive_cell=True).get_ase()
        self.assertEquals(ase.get_number_of_atoms(), 5)

    @unittest.skipIf(not has_pycifrw(), "Unable to import PyCifRW")
    def test_pycifrw_from_datablocks(self):
        """
        Tests CifData.pycifrw_from_cif()
        """
        datablocks = [
            {
                '_atom_site_label': ['A', 'B', 'C'],
                '_atom_site_occupancy': [1.0, 0.5, 0.5],
                '_publ_section_title': 'Test CIF'
            }
        ]
        lines = pycifrw_from_cif(datablocks).WriteOut().split('\n')
        non_comments = []
        for line in lines:
            if not re.search('^#', line):
                # XXX we change single space line into an empty char because
                # some editor are removing trailing whitespace on save.
                if line == ' ':
                    line = ''
                non_comments.append(line)
        s = '''
data_0
loop_
  _atom_site_label
   A
   B
   C

loop_
  _atom_site_occupancy
   1.0
   0.5
   0.5

_publ_section_title                     'Test CIF'
'''
        self.assertEquals("\n".join(non_comments), s)

        loops = {'_atom_site': ['_atom_site_label', '_atom_site_occupancy']}
        lines = pycifrw_from_cif(datablocks, loops).WriteOut().split('\n')
        non_comments = []
        for line in lines:
            if not re.search('^#', line):
                if line == ' ':
                    line = ''
                non_comments.append(line)
        self.assertEquals("\n".join(non_comments),
                          '''
data_0
loop_
  _atom_site_label
  _atom_site_occupancy
   A  1.0
   B  0.5
   C  0.5

_publ_section_title                     'Test CIF'
''')

    @unittest.skipIf(not has_ase() or not has_pycifrw(),
                     "Unable to import ase or pycifrw")
    def test_cif_roundtrip(self):
        with tempfile.NamedTemporaryFile() as f:
            f.write('''
                data_test
                _cell_length_a    10
                _cell_length_b    10
                _cell_length_c    10
                _cell_angle_alpha 90
                _cell_angle_beta  90
                _cell_angle_gamma 90
                loop_
                _atom_site_label
                _atom_site_fract_x
                _atom_site_fract_y
                _atom_site_fract_z
                C 0 0 0
                O 0.5 0.5 0.5
                _cod_database_code 0000001
                _[local]_flags     ''
            ''')
            f.flush()
            a = CifData(file=f.name)

        b = CifData(values=a.values)
        c = CifData(values=b.values)
        self.assertEquals(b._prepare_cif(), c._prepare_cif())

        b = CifData(ase=a.ase)
        c = CifData(ase=b.ase)
        self.assertEquals(b._prepare_cif(), c._prepare_cif())

    def test_symop_string_from_symop_matrix_tr(self):
        self.assertEquals(
            symop_string_from_symop_matrix_tr(
                [[1, 0, 0], [0, 1, 0], [0, 0, 1]]),"x,y,z")

        self.assertEquals(
            symop_string_from_symop_matrix_tr(
                [[1, 0, 0], [0, -1, 0], [0, 1, 1]]),"x,-y,y+z")

        self.assertEquals(
            symop_string_from_symop_matrix_tr(
                [[-1, 0, 0], [0, 1, 0], [0, 0, 1]],[1,-1,0]),"-x+1,y-1,z")

    def test_parse_formula(self):
        self.assertEqual(parse_formula("C H"),
                         {'C': 1, 'H': 1})

        self.assertEqual(parse_formula("C5 H1"),
                         {'C': 5, 'H': 1})

        self.assertEqual(parse_formula("Ca5 Ho"),
                         {'Ca': 5, 'Ho': 1})

        self.assertEqual(parse_formula("H0.5 O"),
                         {'H': 0.5, 'O': 1})

        self.assertEqual(parse_formula("C0 O0"),
                         {'C': 0, 'O': 0})

        # Invalid literal for float()
        with self.assertRaises(ValueError):
            parse_formula("H0.5.2 O")

    @unittest.skipIf(not has_ase() or not has_pycifrw(),
                     "Unable to import ase or pycifrw")
    def test_attached_hydrogens(self):
        with tempfile.NamedTemporaryFile() as f:
            f.write('''
                data_test
                _cell_length_a    10
                _cell_length_b    10
                _cell_length_c    10
                _cell_angle_alpha 90
                _cell_angle_beta  90
                _cell_angle_gamma 90
                loop_
                _atom_site_label
                _atom_site_fract_x
                _atom_site_fract_y
                _atom_site_fract_z
                _atom_site_attached_hydrogens
                C 0 0 0 ?
                O 0.5 0.5 0.5 .
                H 0.75 0.75 0.75 0
            ''')
            f.flush()
            a = CifData(file=f.name)

        self.assertEqual(a.has_attached_hydrogens(), False)

        with tempfile.NamedTemporaryFile() as f:
            f.write('''
                data_test
                _cell_length_a    10
                _cell_length_b    10
                _cell_length_c    10
                _cell_angle_alpha 90
                _cell_angle_beta  90
                _cell_angle_gamma 90
                loop_
                _atom_site_label
                _atom_site_fract_x
                _atom_site_fract_y
                _atom_site_fract_z
                _atom_site_attached_hydrogens
                C 0 0 0 ?
                O 0.5 0.5 0.5 1
                H 0.75 0.75 0.75 0
            ''')
            f.flush()
            a = CifData(file=f.name)

        self.assertEqual(a.has_attached_hydrogens(), True)


    @unittest.skipIf(not has_ase() or not has_pycifrw() or
                     not has_pyspglib(),
                     "Unable to import ase, pycifrw or pyspglib")
    def test_refine(self):
        """
        Test case for refinement (space group determination) for a
        CifData object.
        """
        with tempfile.NamedTemporaryFile() as f:
            f.write('''
                data_test
                _cell_length_a    10
                _cell_length_b    10
                _cell_length_c    10
                _cell_angle_alpha 90
                _cell_angle_beta  90
                _cell_angle_gamma 90
                loop_
                _atom_site_label
                _atom_site_fract_x
                _atom_site_fract_y
                _atom_site_fract_z
                C 0.5 0.5 0.5
                O 0.25 0.5 0.5
                O 0.75 0.5 0.5
            ''')
            f.flush()
            a = CifData(file=f.name)

        ret_dict = refine_inline(a)
        b = ret_dict['cif']
        self.assertEqual(b.values.keys(),['test'])
        self.assertEqual(b.values['test']['_chemical_formula_sum'],'C O2')
        self.assertEqual(b.values['test']['_symmetry_equiv_pos_as_xyz'],[
            'x,y,z',
            '-x,-y,-z',
            '-y,x,z',
            'y,-x,-z',
            '-x,-y,z',
            'x,y,-z',
            'y,-x,z',
            '-y,x,-z',
            'x,-y,-z',
            '-x,y,z',
            '-y,-x,-z',
            'y,x,z',
            '-x,y,-z',
            'x,-y,z',
            'y,x,-z',
            '-y,-x,z'])

        with tempfile.NamedTemporaryFile() as f:
            f.write('''
                data_a
                data_b
            ''')
            f.flush()
            c = CifData(file=f.name)

        with self.assertRaises(ValueError):
            ret_dict = refine_inline(c)

    def test_parse_formula(self):
        self.assertEqual(parse_formula("C H"),
                         {'C': 1, 'H': 1})

        self.assertEqual(parse_formula("C5 H1"),
                         {'C': 5, 'H': 1})

        self.assertEqual(parse_formula("Ca5 Ho"),
                         {'Ca': 5, 'Ho': 1})

        self.assertEqual(parse_formula("H0.5 O"),
                         {'H': 0.5, 'O': 1})

        # Invalid literal for float()
        with self.assertRaises(ValueError):
            parse_formula("H0.5.2 O")

class TestKindValidSymbols(SqlAlchemyTests):
    """
    Tests the symbol validation of the
    aiida.orm.data.structure.Kind class.
    """

    def test_bad_symbol(self):
        """
        Should not accept a non-existing symbol.
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols='Hxx')

    def test_empty_list_symbols(self):
        """
        Should not accept an empty list
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols=[])

    def test_valid_list(self):
        """
        Should not raise any error.
        """
        _ = Kind(symbols=['H', 'He'], weights=[0.5, 0.5])


class TestSiteValidWeights(SqlAlchemyTests):
    """
    Tests valid weight lists.
    """

    def test_isnot_list(self):
        """
        Should not accept a non-list, non-number weight
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols='Ba', weights='aaa')

    def test_empty_list_weights(self):
        """
        Should not accept an empty list
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols='Ba', weights=[])

    def test_symbol_weight_mismatch(self):
        """
        Should not accept a size mismatch of the symbols and weights
        list.
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols=['Ba', 'C'], weights=[1.])

        with self.assertRaises(ValueError):
            _ = Kind(symbols=['Ba'], weights=[0.1, 0.2])

    def test_negative_value(self):
        """
        Should not accept a negative weight
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols=['Ba', 'C'], weights=[-0.1, 0.3])

    def test_sum_greater_one(self):
        """
        Should not accept a sum of weights larger than one
        """
        with self.assertRaises(ValueError):
            _ = Kind(symbols=['Ba', 'C'],
                     weights=[0.5, 0.6])

    def test_sum_one_weights(self):
        """
        Should accept a sum equal to one
        """
        _ = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 2. / 3.])

    def test_sum_less_one_weights(self):
        """
        Should accept a sum equal less than one
        """
        _ = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 1. / 3.])

    def test_none(self):
        """
        Should accept None.
        """
        _ = Kind(symbols='Ba', weights=None)


class TestKindTestGeneral(SqlAlchemyTests):
    """
    Tests the creation of Kind objects and their methods.
    """

    def test_sum_one_general(self):
        """
        Should accept a sum equal to one
        """
        a = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 2. / 3.])
        self.assertTrue(a.is_alloy())
        self.assertFalse(a.has_vacancies())

    def test_sum_less_one_general(self):
        """
        Should accept a sum equal less than one
        """

        a = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 1. / 3.])
        self.assertTrue(a.is_alloy())
        self.assertTrue(a.has_vacancies())

    def test_no_position(self):
        """
        Should not accept a 'positions' parameter
        """

        with self.assertRaises(ValueError):
            _ = Kind(position=[0., 0., 0.], symbols=['Ba'],
                     weights=[1.])

    def test_simple(self):
        """
        Should recognize a simple element.
        """

        a = Kind(symbols='Ba')
        self.assertFalse(a.is_alloy())
        self.assertFalse(a.has_vacancies())

        b = Kind(symbols='Ba', weights=1.)
        self.assertFalse(b.is_alloy())
        self.assertFalse(b.has_vacancies())

        c = Kind(symbols='Ba', weights=None)
        self.assertFalse(c.is_alloy())
        self.assertFalse(c.has_vacancies())


    def test_automatic_name(self):
        """
        Check the automatic name generator.
        """

        a = Kind(symbols='Ba')
        self.assertEqual(a.name, 'Ba')

        a = Kind(symbols=('Si', 'Ge'), weights=(1. / 3., 2. / 3.))
        self.assertEqual(a.name, 'GeSi')

        a = Kind(symbols=('Si', 'Ge'), weights=(0.4, 0.5))
        self.assertEqual(a.name, 'GeSiX')

        # Manually setting the name of the species
        a.name = 'newstring'
        self.assertEqual(a.name, 'newstring')


class TestKindTestMasses(SqlAlchemyTests):
    """
    Tests the management of masses during the creation of Kind objects.
    """

    def test_auto_mass_one(self):
        """
        mass for elements with sum one
        """

        a = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 2. / 3.])
        self.assertAlmostEqual(a.mass,
                               (_atomic_masses['Ba'] +
                                2. * _atomic_masses['C']) / 3.)

    def test_sum_less_one_masses(self):
        """
        mass for elements with sum less than one
        """

        a = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 1. / 3.])
        self.assertAlmostEqual(a.mass,
                               (_atomic_masses['Ba'] +
                                _atomic_masses['C']) / 2.)

    def test_sum_less_one_singleelem(self):
        """
        mass for a single element
        """

        a = Kind(symbols=['Ba'])
        self.assertAlmostEqual(a.mass,
                               _atomic_masses['Ba'])

    def test_manual_mass(self):
        """
        mass set manually
        """

        a = Kind(symbols=['Ba', 'C'],
                 weights=[1. / 3., 1. / 3.],
                 mass=1000.)
        self.assertAlmostEqual(a.mass, 1000.)


class TestStructureDataInit(SqlAlchemyTests):
    """
    Tests the creation of StructureData objects (cell and pbc).
    """

    def test_cell_wrong_size_1(self):
        """
        Wrong cell size (not 3x3)
        """
        with self.assertRaises(ValueError):
            _ = StructureData(cell=((1., 2., 3.),))

    def test_cell_wrong_size_2(self):
        """
        Wrong cell size (not 3x3)
        """
        with self.assertRaises(ValueError):
            _ = StructureData(cell=((1., 0., 0.), (0., 0., 3.), (0., 3.)))

    def test_cell_zero_vector(self):
        """
        Wrong cell (one vector has zero length)
        """
        with self.assertRaises(ValueError):
            _ = StructureData(cell=((0., 0., 0.), (0., 1., 0.), (0., 0., 1.)))

    def test_cell_zero_volume(self):
        """
        Wrong cell (volume is zero)
        """
        with self.assertRaises(ValueError):
            _ = StructureData(cell=((1., 0., 0.), (0., 1., 0.), (1., 1., 0.)))

    def test_cell_ok_init(self):
        """
        Correct cell
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell)
        out_cell = a.cell

        for i in range(3):
            for j in range(3):
                self.assertAlmostEqual(cell[i][j], out_cell[i][j])

    def test_volume(self):
        """
        Check the volume calculation
        """
        a = StructureData(cell=((1., 0., 0.), (0., 2., 0.), (0., 0., 3.)))
        self.assertAlmostEqual(a.get_cell_volume(), 6.)

    def test_wrong_pbc_1(self):
        """
        Wrong pbc parameter (not bool or iterable)
        """
        with self.assertRaises(ValueError):
            cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
            _ = StructureData(cell=cell, pbc=1)

    def test_wrong_pbc_2(self):
        """
        Wrong pbc parameter (iterable but with wrong len)
        """
        with self.assertRaises(ValueError):
            cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
            _ = StructureData(cell=cell, pbc=[True, True])

    def test_wrong_pbc_3(self):
        """
        Wrong pbc parameter (iterable but with wrong len)
        """
        with self.assertRaises(ValueError):
            cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
            _ = StructureData(cell=cell, pbc=[])

    def test_ok_pbc_1(self):
        """
        Single pbc value
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell, pbc=True)
        self.assertEquals(a.pbc, tuple([True, True, True]))

        a = StructureData(cell=cell, pbc=False)
        self.assertEquals(a.pbc, tuple([False, False, False]))

    def test_ok_pbc_2(self):
        """
        One-element list
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell, pbc=[True])
        self.assertEqual(a.pbc, tuple([True, True, True]))

        a = StructureData(cell=cell, pbc=[False])
        self.assertEqual(a.pbc, tuple([False, False, False]))

    def test_ok_pbc_3(self):
        """
        Three-element list
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell, pbc=[True, False, True])
        self.assertEqual(a.pbc, tuple([True, False, True]))



class TestStructureData(SqlAlchemyTests):
    """
    Tests the creation of StructureData objects (cell and pbc).
    """
    def test_cell_ok_and_atoms(self):
        """
        Test the creation of a cell and the appending of atoms
        """
        cell = ((2., 0., 0.), (0., 2., 0.), (0., 0., 2.))

        a = StructureData(cell=cell)
        out_cell = a.cell
        self.assertAlmostEquals(cell, out_cell)

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'])
        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])
        a.append_atom(position=(1.2, 1.4, 1.6), symbols=['Ti'])
        self.assertFalse(a.is_alloy())
        self.assertFalse(a.has_vacancies())
        # There should be only two kinds! (two atoms of kind Ti should
        # belong to the same kind)
        self.assertEquals(len(a.kinds), 2)

        a.append_atom(position=(0.5, 1., 1.5), symbols=['O', 'C'],
                      weights=[0.5, 0.5])
        self.assertTrue(a.is_alloy())
        self.assertFalse(a.has_vacancies())

        a.append_atom(position=(0.5, 1., 1.5), symbols=['O'], weights=[0.5])
        self.assertTrue(a.is_alloy())
        self.assertTrue(a.has_vacancies())

        a.clear_kinds()
        a.append_atom(position=(0.5, 1., 1.5), symbols=['O'], weights=[0.5])
        self.assertFalse(a.is_alloy())
        self.assertTrue(a.has_vacancies())

    def test_kind_1(self):
        """
        Test the management of kinds (automatic detection of kind of
        simple atoms).
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'])
        a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba'])
        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])

        self.assertEqual(len(a.kinds), 2)  # I should only have two types
        # I check for the default names of kinds
        self.assertEqual(set(k.name for k in a.kinds),
                         set(('Ba', 'Ti')))

    def test_kind_2(self):
        """
        Test the management of kinds (manual specification of kind name).
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'], name='Ba1')
        a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba'], name='Ba2')
        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])

        kind_list = a.kinds
        self.assertEqual(len(kind_list), 3)  # I should have now three kinds
        self.assertEqual(set(k.name for k in kind_list),
                         set(('Ba1', 'Ba2', 'Ti')))

    def test_kind_3(self):
        """
        Test the management of kinds (adding an atom with different mass).
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'], mass=100.)
        with self.assertRaises(ValueError):
            # Shouldn't allow, I am adding two sites with the same name 'Ba'
            a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba'],
                          mass=101., name='Ba')

            # now it should work because I am using a different kind name
        a.append_atom(position=(0.5, 0.5, 0.5),
                      symbols=['Ba'], mass=101., name='Ba2')

        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])

        self.assertEqual(len(a.kinds), 3)  # I should have now three types
        self.assertEqual(len(a.sites), 3)  # and 3 sites
        self.assertEqual(set(k.name for k in a.kinds), set(('Ba', 'Ba2', 'Ti')))

    def test_kind_4(self):
        """
        Test the management of kind (adding an atom with different symbols
        or weights).
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols=['Ba', 'Ti'],
                      weights=(1., 0.), name='mytype')

        with self.assertRaises(ValueError):
            # Shouldn't allow, different weights
            a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba', 'Ti'],
                          weights=(0.9, 0.1), name='mytype')

        with self.assertRaises(ValueError):
            # Shouldn't allow, different weights (with vacancy)
            a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba', 'Ti'],
                          weights=(0.8, 0.1), name='mytype')

        with self.assertRaises(ValueError):
            # Shouldn't allow, different symbols list
            a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba'],
                          name='mytype')

        with self.assertRaises(ValueError):
            # Shouldn't allow, different symbols list
            a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Si', 'Ti'],
                          weights=(1., 0.), name='mytype')

            # should allow because every property is identical
        a.append_atom(position=(0., 0., 0.), symbols=['Ba', 'Ti'],
                      weights=(1., 0.), name='mytype')

        self.assertEquals(len(a.kinds), 1)

    def test_kind_5(self):
        """
        Test the management of kinds (automatic creation of new kind
        if name is not specified and properties are different).
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols='Ba', mass=100.)
        a.append_atom(position=(0., 0., 0.), symbols='Ti')
        # The name does not exist
        a.append_atom(position=(0., 0., 0.), symbols='Ti', name='Ti2')
        # The name already exists, but the properties are identical => OK
        a.append_atom(position=(1., 1., 1.), symbols='Ti', name='Ti2')
        # The name already exists, but the properties are different!
        with self.assertRaises(ValueError):
            a.append_atom(position=(1., 1., 1.), symbols='Ti', mass=100., name='Ti2')
        # Should not complain, should create a new type
        a.append_atom(position=(0., 0., 0.), symbols='Ba', mass=150.)

        # There should be 4 kinds, the automatic name for the last one
        # should be Ba1
        self.assertEquals([k.name for k in a.kinds],
                          ['Ba', 'Ti', 'Ti2', 'Ba1'])
        self.assertEquals(len(a.sites), 5)

    def test_kind_5_bis(self):
        """
        Test the management of kinds (automatic creation of new kind
        if name is not specified and properties are different).
        This test was failing in, e.g., commit f6a8f4b.
        """
        s = StructureData(cell=((6., 0., 0.), (0., 6., 0.), (0., 0., 6.)))

        s.append_atom(symbols='Fe', position=[0, 0, 0], mass=12)
        s.append_atom(symbols='Fe', position=[1, 0, 0], mass=12)
        s.append_atom(symbols='Fe', position=[2, 0, 0], mass=12)
        s.append_atom(symbols='Fe', position=[2, 0, 0])
        s.append_atom(symbols='Fe', position=[4, 0, 0])

        # I expect only two species, the first one with name 'Fe', mass 12,
        # and referencing the first three atoms; the second with name
        # 'Fe1', mass = elements[26]['mass'], and referencing the last two atoms
        self.assertEquals(
            set([(k.name, k.mass) for k in s.kinds]),
            set([('Fe', 12.0), ('Fe1', elements[26]['mass'])]))

        kind_of_each_site = [site.kind_name for site in s.sites]
        self.assertEquals(kind_of_each_site,
                          ['Fe', 'Fe', 'Fe', 'Fe1', 'Fe1'])

    @unittest.skipIf(not has_ase(), "Unable to import ase")
    def test_kind_5_bis_ase(self):
        """
        Same test as test_kind_5_bis, but using ase
        """
        asecell = ase.Atoms('Fe5', cell=((6., 0., 0.), (0., 6., 0.), (0., 0., 6.)))
        asecell.set_positions([
            [0, 0, 0],
            [1, 0, 0],
            [2, 0, 0],
            [3, 0, 0],
            [4, 0, 0],
        ])

        asecell[0].mass = 12.
        asecell[1].mass = 12.
        asecell[2].mass = 12.

        s = StructureData(ase=asecell)

        # I expect only two species, the first one with name 'Fe', mass 12,
        # and referencing the first three atoms; the second with name
        # 'Fe1', mass = elements[26]['mass'], and referencing the last two atoms
        self.assertEquals(
            set([(k.name, k.mass) for k in s.kinds]),
            set([('Fe', 12.0), ('Fe1', asecell[3].mass)]))

        kind_of_each_site = [site.kind_name for site in s.sites]
        self.assertEquals(kind_of_each_site,
                          ['Fe', 'Fe', 'Fe', 'Fe1', 'Fe1'])

    def test_kind_6(self):
        """
        Test the returning of kinds from the string name (most of the code
        copied from :py:meth:`.test_kind_5`).
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols='Ba', mass=100.)
        a.append_atom(position=(0., 0., 0.), symbols='Ti')
        # The name does not exist
        a.append_atom(position=(0., 0., 0.), symbols='Ti', name='Ti2')
        # The name already exists, but the properties are identical => OK
        a.append_atom(position=(1., 1., 1.), symbols='Ti', name='Ti2')
        # Should not complain, should create a new type
        a.append_atom(position=(0., 0., 0.), symbols='Ba', mass=150.)
        # There should be 4 kinds, the automatic name for the last one
        # should be Ba1 (same check of test_kind_5
        self.assertEquals([k.name for k in a.kinds],
                          ['Ba', 'Ti', 'Ti2', 'Ba1'])
        #############################
        # Here I start the real tests
        # No such kind
        with self.assertRaises(ValueError):
            a.get_kind('Ti3')
        k = a.get_kind('Ba1')
        self.assertEqual(k.symbols, ('Ba',))
        self.assertAlmostEqual(k.mass, 150.)

    def test_kind_7(self):
        """
        Test the functions returning the list of kinds, symbols, ...
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols='Ba', mass=100.)
        a.append_atom(position=(0., 0., 0.), symbols='Ti')
        # The name does not exist
        a.append_atom(position=(0., 0., 0.), symbols='Ti', name='Ti2')
        # The name already exists, but the properties are identical => OK
        a.append_atom(position=(0., 0., 0.), symbols=['O', 'H'], weights=[0.9, 0.1], mass=15.)

        self.assertEquals(a.get_symbols_set(), set(['Ba', 'Ti', 'O', 'H']))

    @unittest.skipIf(not has_ase() or not has_pyspglib(),
                     "Unable to import ase or pyspglib")
    def test_kind_8(self):
        """
        Test the ase_refine_cell() function
        """
        a = ase.Atoms(cell=[10,10,10])
        a.append(ase.Atom('C',[0,0,0]))
        a.append(ase.Atom('C',[5,0,0]))
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.get_chemical_symbols(),['C'])
        self.assertEquals(b.cell.tolist(),[[10,0,0],[0,10,0],[0,0,5]])
        self.assertEquals(sym,{'hall': '-P 4 2', 'hm': 'P4/mmm', 'tables': 123})

        a = ase.Atoms(cell=[10,2*math.sqrt(75),10])
        a.append(ase.Atom('C',[0,0,0]))
        a.append(ase.Atom('C',[5,math.sqrt(75),0]))
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.get_chemical_symbols(),['C'])
        self.assertEquals(numpy.round(b.cell,2).tolist(),
                          [[10,0,0],[-5,8.66,0],[0,0,10]])
        self.assertEquals(sym,{'hall': '-P 6 2', 'hm': 'P6/mmm', 'tables': 191})

        a = ase.Atoms(cell=[[10,0,0],[-10,10,0],[0,0,10]])
        a.append(ase.Atom('C',[5,5,5]))
        a.append(ase.Atom('F',[0,0,0]))
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.get_chemical_symbols(),['C','F'])
        self.assertEquals(b.cell.tolist(),[[10,0,0],[0,10,0],[0,0,10]])
        self.assertEquals(b.get_scaled_positions().tolist(),
                          [[0.5,0.5,0.5],[0,0,0]])
        self.assertEquals(sym,{'hall': '-P 4 2 3', 'hm': 'Pm-3m', 'tables': 221})

        a = ase.Atoms(cell=[[10,0,0],[-10,10,0],[0,0,10]])
        a.append(ase.Atom('C',[0,0,0]))
        a.append(ase.Atom('F',[5,5,5]))
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.get_chemical_symbols(),['C','F'])
        self.assertEquals(b.cell.tolist(),[[10,0,0],[0,10,0],[0,0,10]])
        self.assertEquals(b.get_scaled_positions().tolist(),
                          [[0,0,0],[0.5,0.5,0.5]])
        self.assertEquals(sym,{'hall': '-P 4 2 3', 'hm': 'Pm-3m', 'tables': 221})

        a = ase.Atoms(cell=[[12.132,0,0],[0,6.0606,0],[0,0,8.0956]])
        a.append(ase.Atom('Ba',[1.5334848,1.3999986,2.00042276]))
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.cell.tolist(),[[12.132,0,0],[0,6.0606,0],[0,0,8.0956]])
        self.assertEquals(b.get_scaled_positions().tolist(),
                          [[0,0,0]])

        a = ase.Atoms(cell=[10,10,10])
        a.append(ase.Atom('C',[5,5,5]))
        a.append(ase.Atom('O',[2.5,5,5]))
        a.append(ase.Atom('O',[7.5,5,5]))
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.get_chemical_symbols(),['C','O'])
        self.assertEquals(sym,{'hall': '-P 4 2', 'hm': 'P4/mmm', 'tables': 123})

        # Generated from COD entry 1507756
        # (http://www.crystallography.net/cod/1507756.cif@87343)
        a = crystal(['Ba','Ti','O','O'],
                    [
                      [0,0,0],
                      [0.5,0.5,0.482],
                      [0.5,0.5,0.016],
                      [0.5,0,0.515]
                    ],
                    cell=[3.9999,3.9999,4.0170],
                    spacegroup=99)
        b,sym = ase_refine_cell(a)
        sym.pop('rotations')
        sym.pop('translations')
        self.assertEquals(b.get_chemical_symbols(),['Ba','Ti','O','O'])
        self.assertEquals(sym,{'hall': 'P 4 -2', 'hm': 'P4mm', 'tables': 99})

    def test_get_formula(self):
        """
        Tests the generation of formula
        """
        self.assertEquals(get_formula(['Ba', 'Ti'] + ['O'] * 3),
                          'BaO3Ti')
        self.assertEquals(get_formula(['Ba', 'Ti', 'C'] + ['O'] * 3,
                                      separator=" "),
                          'C Ba O3 Ti')
        self.assertEquals(get_formula(['H'] * 6 + ['C'] * 6),
                          'C6H6')
        self.assertEquals(get_formula(['H'] * 6 + ['C'] * 6,
                                      mode="hill_compact"),
                          'CH')
        self.assertEquals(get_formula((['Ba', 'Ti'] + ['O'] * 3) * 2 + \
                                      ['Ba'] + ['Ti'] * 2 + ['O'] * 3,
                                      mode="group"),
                          '(BaTiO3)2BaTi2O3')
        self.assertEquals(get_formula((['Ba', 'Ti'] + ['O'] * 3) * 2 + \
                                      ['Ba'] + ['Ti'] * 2 + ['O'] * 3,
                                      mode="group", separator=" "),
                          '(Ba Ti O3)2 Ba Ti2 O3')
        self.assertEquals(get_formula((['Ba', 'Ti'] + ['O'] * 3) * 2 + \
                                      ['Ba'] + ['Ti'] * 2 + ['O'] * 3,
                                      mode="reduce"),
                          'BaTiO3BaTiO3BaTi2O3')
        self.assertEquals(get_formula((['Ba', 'Ti'] + ['O'] * 3) * 2 + \
                                      ['Ba'] + ['Ti'] * 2 + ['O'] * 3,
                                      mode="reduce", separator=", "),
                          'Ba, Ti, O3, Ba, Ti, O3, Ba, Ti2, O3')
        self.assertEquals(get_formula((['Ba', 'Ti'] + ['O'] * 3) * 2,
                                      mode="count"),
                          'Ba2Ti2O6')
        self.assertEquals(get_formula((['Ba', 'Ti'] + ['O'] * 3) * 2,
                                      mode="count_compact"),
                          'BaTiO3')

    def test_get_cif(self):
        """
        Tests the conversion to CifData
        """
        a = StructureData(cell=((2., 0., 0.), (0., 2., 0.), (0., 0., 2.)))

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'])
        a.append_atom(position=(0.5, 0.5, 0.5), symbols=['Ba'])
        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])

        try:
            c = a._get_cif()
        # Exception thrown if ase can't be found
        except ImportError:
            return
        # XXX remove trailing space
        cif_prepared = c._prepare_cif()
        cif_prepared = "\n".join(
            map(lambda _: re.sub(r"( )+$", "", _), cif_prepared.split("\n"))
        )
        s = """#\#CIF1.1
##########################################################################
#               Crystallographic Information Format file
#               Produced by PyCifRW module
#
#  This is a CIF file.  CIF has been adopted by the International
#  Union of Crystallography as the standard for data archiving and
#  transmission.
#
#  For information on this file format, follow the CIF links at
#  http://www.iucr.org
##########################################################################

data_0
loop_
  _atom_site_label
  _atom_site_fract_x
  _atom_site_fract_y
  _atom_site_fract_z
  _atom_site_type_symbol
   Ba1  0.0  0.0  0.0  Ba
   Ba2  0.25  0.25  0.25  Ba
   Ti1  0.5  0.5  0.5  Ti

_cell_angle_alpha                       90.0
_cell_angle_beta                        90.0
_cell_angle_gamma                       90.0
_cell_length_a                          2.0
_cell_length_b                          2.0
_cell_length_c                          2.0
loop_
  _symmetry_equiv_pos_as_xyz
   'x, y, z'

_symmetry_Int_Tables_number             1
_symmetry_space_group_name_H-M          'P 1'
_symmetry_space_group_name_Hall         'P 1'
_cell_formula_units_Z                   1
_chemical_formula_sum                   'Ba2 Ti'
"""
        self.assertEquals(cif_prepared, s)


class TestStructureDataLock(SqlAlchemyTests):
    """
    Tests that the structure is locked after storage
    """

    def test_lock(self):
        """
        Start from a StructureData object, convert to raw and then back
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell)

        a.pbc = [False, True, True]

        k = Kind(symbols='Ba', name='Ba')
        s = Site(position=(0., 0., 0.), kind_name='Ba')
        a.append_kind(k)
        a.append_site(s)

        a.append_atom(symbols='Ti', position=[0., 0., 0.])

        a.store()

        k2 = Kind(symbols='Ba', name='Ba')
        # Nothing should be changed after store()
        with self.assertRaises(ModificationNotAllowed):
            a.append_kind(k2)
        with self.assertRaises(ModificationNotAllowed):
            a.append_site(s)
        with self.assertRaises(ModificationNotAllowed):
            a.clear_sites()
        with self.assertRaises(ModificationNotAllowed):
            a.clear_kinds()
        with self.assertRaises(ModificationNotAllowed):
            a.cell = cell
        with self.assertRaises(ModificationNotAllowed):
            a.pbc = [True, True, True]

        _ = a.get_cell_volume()
        _ = a.is_alloy()
        _ = a.has_vacancies()

        b = a.copy()
        # I check that I can edit after copy
        b.append_site(s)
        b.clear_sites()
        # I check that the original did not change
        self.assertNotEquals(len(a.sites), 0)
        b.cell = cell
        b.pbc = [True, True, True]


class TestStructureDataReload(SqlAlchemyTests):
    """
    Tests the creation of StructureData, converting it to a raw format and
    converting it back.
    """

    def test_reload(self):
        """
        Start from a StructureData object, convert to raw and then back
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell)

        a.pbc = [False, True, True]

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'])
        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])

        a.store()

        b = StructureData(dbnode=a.dbnode)

        for i in range(3):
            for j in range(3):
                self.assertAlmostEqual(cell[i][j], b.cell[i][j])

        self.assertEqual(b.pbc, (False, True, True))
        self.assertEqual(len(b.sites), 2)
        self.assertEqual(b.kinds[0].symbols[0], 'Ba')
        self.assertEqual(b.kinds[1].symbols[0], 'Ti')
        for i in range(3):
            self.assertAlmostEqual(b.sites[0].position[i], 0.)
        for i in range(3):
            self.assertAlmostEqual(b.sites[1].position[i], 1.)

        # Fully reload from UUID
        b = StructureData.get_subclass_from_uuid(a.uuid)

        for i in range(3):
            for j in range(3):
                self.assertAlmostEqual(cell[i][j], b.cell[i][j])

        self.assertEqual(b.pbc, (False, True, True))
        self.assertEqual(len(b.sites), 2)
        self.assertEqual(b.kinds[0].symbols[0], 'Ba')
        self.assertEqual(b.kinds[1].symbols[0], 'Ti')
        for i in range(3):
            self.assertAlmostEqual(b.sites[0].position[i], 0.)
        for i in range(3):
            self.assertAlmostEqual(b.sites[1].position[i], 1.)


    def test_copy(self):
        """
        Start from a StructureData object, copy it and see if it is preserved
        """
        cell = ((1., 0., 0.), (0., 2., 0.), (0., 0., 3.))
        a = StructureData(cell=cell)

        a.pbc = [False, True, True]

        a.append_atom(position=(0., 0., 0.), symbols=['Ba'])
        a.append_atom(position=(1., 1., 1.), symbols=['Ti'])

        b = a.copy()

        for i in range(3):
            for j in range(3):
                self.assertAlmostEqual(cell[i][j], b.cell[i][j])

        self.assertEqual(b.pbc, (False, True, True))
        self.assertEqual(len(b.kinds), 2)
        self.assertEqual(len(b.sites), 2)
        self.assertEqual(b.kinds[0].symbols[0], 'Ba')
        self.assertEqual(b.kinds[1].symbols[0], 'Ti')
        for i in range(3):
            self.assertAlmostEqual(b.sites[0].position[i], 0.)
        for i in range(3):
            self.assertAlmostEqual(b.sites[1].position[i], 1.)

        a.store()

        # Copy after store()
        c = a.copy()
        for i in range(3):
            for j in range(3):
                self.assertAlmostEqual(cell[i][j], c.cell[i][j])

        self.assertEqual(c.pbc, (False, True, True))
        self.assertEqual(len(c.kinds), 2)
        self.assertEqual(len(c.sites), 2)
        self.assertEqual(c.kinds[0].symbols[0], 'Ba')
        self.assertEqual(c.kinds[1].symbols[0], 'Ti')
        for i in range(3):
            self.assertAlmostEqual(c.sites[0].position[i], 0.)
        for i in range(3):
            self.assertAlmostEqual(c.sites[1].position[i], 1.)


class TestStructureDataFromAse(SqlAlchemyTests):
    """
    Tests the creation of Sites from/to a ASE object.
    """

    @unittest.skipIf(not has_ase(), "Unable to import ase")
    def test_ase(self):
        a = ase.Atoms('SiGe', cell=(1., 2., 3.), pbc=(True, False, False))
        a.set_positions(
            ((0., 0., 0.),
             (0.5, 0.7, 0.9),)
        )
        a[1].mass = 110.2

        b = StructureData(ase=a)
        c = b.get_ase()

        self.assertEqual(a[0].symbol, c[0].symbol)
        self.assertEqual(a[1].symbol, c[1].symbol)
        for i in range(3):
            self.assertAlmostEqual(a[0].position[i], c[0].position[i])
        for i in range(3):
            for j in range(3):
                self.assertAlmostEqual(a.cell[i][j], c.cell[i][j])

        self.assertAlmostEqual(c[1].mass, 110.2)

    @unittest.skipIf(not has_ase(), "Unable to import ase")
    def test_conversion_of_types_1(self):
        a = ase.Atoms('Si4Ge4', cell=(1., 2., 3.), pbc=(True, False, False))
        a.set_positions(
            ((0.0, 0.0, 0.0),
             (0.1, 0.1, 0.1),
             (0.2, 0.2, 0.2),
             (0.3, 0.3, 0.3),
             (0.4, 0.4, 0.4),
             (0.5, 0.5, 0.5),
             (0.6, 0.6, 0.6),
             (0.7, 0.7, 0.7),
            )
        )

        a.set_tags((0, 1, 2, 3, 4, 5, 6, 7))

        b = StructureData(ase=a)
        self.assertEquals([k.name for k in b.kinds],
                          ["Si", "Si1", "Si2", "Si3",
                           "Ge4", "Ge5", "Ge6", "Ge7"])
        c = b.get_ase()

        a_tags = list(a.get_tags())
        c_tags = list(c.get_tags())
        self.assertEqual(a_tags, c_tags)

    @unittest.skipIf(not has_ase(), "Unable to import ase")
    def test_conversion_of_types_2(self):
        a = ase.Atoms('Si4', cell=(1., 2., 3.), pbc=(True, False, False))
        a.set_positions(
            ((0.0, 0.0, 0.0),
             (0.1, 0.1, 0.1),
             (0.2, 0.2, 0.2),
             (0.3, 0.3, 0.3),
            )
        )

        a.set_tags((0, 1, 0, 1))
        a[2].mass = 100.
        a[3].mass = 300.

        b = StructureData(ase=a)
        # This will give funny names to the kinds, because I am using
        # both tags and different properties (mass). I just check to have
        # 4 kinds
        self.assertEquals(len(b.kinds), 4)

        # Do I get the same tags after one full iteration back and forth?
        c = b.get_ase()
        d = StructureData(ase=c)
        e = d.get_ase()
        c_tags = list(c.get_tags())
        e_tags = list(e.get_tags())
        self.assertEqual(c_tags, e_tags)


    @unittest.skipIf(not has_ase(), "Unable to import ase")
    def test_conversion_of_types_3(self):
        a = StructureData()
        a.append_atom(position=(0., 0., 0.), symbols='Ba', name='Ba')
        a.append_atom(position=(0., 0., 0.), symbols='Ba', name='Ba1')
        a.append_atom(position=(0., 0., 0.), symbols='Cu', name='Cu')
        # continues with a number
        a.append_atom(position=(0., 0., 0.), symbols='Cu', name='Cu2')
        # does not continue with a number
        a.append_atom(position=(0., 0., 0.), symbols='Cu', name='Cu_my')
        # random string
        a.append_atom(position=(0., 0., 0.), symbols='Cu', name='a_name')
        # a name of another chemical symbol
        a.append_atom(position=(0., 0., 0.), symbols='Cu', name='Fe')
        # lowercase! as if it were a random string
        a.append_atom(position=(0., 0., 0.), symbols='Cu', name='cu1')

        # Just to be sure that the species were saved with the correct name
        # in the first place
        self.assertEquals([k.name for k in a.kinds],
                          ['Ba', 'Ba1', 'Cu', 'Cu2', 'Cu_my',
                           'a_name', 'Fe', 'cu1'])

        b = a.get_ase()
        self.assertEquals(b.get_chemical_symbols(), ['Ba', 'Ba', 'Cu',
                                                     'Cu', 'Cu', 'Cu',
                                                     'Cu', 'Cu'])
        self.assertEquals(list(b.get_tags()), [0, 1, 0, 2, 3, 4, 5, 6])


class TestStructureDataFromPymatgen(SqlAlchemyTests):
    """
    Tests the creation of StructureData from a pymatgen Structure and
    Molecule objects.
    """
    @unittest.skipIf(not has_pymatgen() or
                     StrictVersion(get_pymatgen_version()) <
                     StrictVersion('3.0.13'), "Unable to import pymatgen")
    def test_1(self):
        """
        Test's imput is derived from COD entry 9011963, processed with
        cif_mark_disorder (from cod-tools) and abbreviated.
        """
        with tempfile.NamedTemporaryFile(suffix=".cif") as f:
            f.write("""data_9011963
                _space_group_IT_number           166
                _symmetry_space_group_name_Hall  '-R 3 2"'
                _symmetry_space_group_name_H-M   'R -3 m :H'
                _cell_angle_alpha                90
                _cell_angle_beta                 90
                _cell_angle_gamma                120
                _cell_length_a                   4.298
                _cell_length_b                   4.298
                _cell_length_c                   29.774
                loop_
                _atom_site_label
                _atom_site_fract_x
                _atom_site_fract_y
                _atom_site_fract_z
                _atom_site_occupancy
                _atom_site_U_iso_or_equiv
                _atom_site_disorder_assembly
                _atom_site_disorder_group
                Bi 0.00000 0.00000 0.39580 1.00000 0.02343 . .
                Te1 0.00000 0.00000 0.00000 0.66667 0.02343 A 1
                Se1 0.00000 0.00000 0.00000 0.33333 0.02343 A 2
                Te2 0.00000 0.00000 0.21180 0.66667 0.02343 B 1
                Se2 0.00000 0.00000 0.21180 0.33333 0.02343 B 2
                """)
            f.flush()
            pymatgen_struct = read_structure(f.name)

        structs_to_test = [StructureData(pymatgen=pymatgen_struct),
                           StructureData(pymatgen_structure=pymatgen_struct)]

        for struct in structs_to_test:
            self.assertEquals(struct.get_site_kindnames(),
                              ['Bi', 'Bi', 'SeTe', 'SeTe', 'SeTe'])
            self.assertEquals([x.symbols for x in struct.kinds],
                              [('Bi',), ('Se', 'Te')])
            self.assertEquals([x.weights for x in struct.kinds],
                              [(1.0,), (0.33333, 0.66667)])

        struct = StructureData(pymatgen_structure=pymatgen_struct)

        # Testing pymatgen Structure -> StructureData -> pymatgen Structure
        # roundtrip.

        pymatgen_struct_roundtrip = struct.get_pymatgen_structure()
        dict1 = pymatgen_struct.as_dict()
        dict2 = pymatgen_struct_roundtrip.as_dict()

        for i in dict1['sites']:
            i['abc'] = [round(j, 2) for j in i['abc']]
        for i in dict2['sites']:
            i['abc'] = [round(j, 2) for j in i['abc']]

        self.assertEquals(dict1, dict2)

    @unittest.skipIf(not has_pymatgen() or
                     StrictVersion(get_pymatgen_version()) <
                     StrictVersion('3.0.13'), "Unable to import pymatgen")
    def test_2(self):
        """
        Input source: http://pymatgen.org/_static/Molecule.html
        """
        with tempfile.NamedTemporaryFile(suffix=".xyz") as f:
            f.write("""5
                H4 C1
                C 0.000000 0.000000 0.000000
                H 0.000000 0.000000 1.089000
                H 1.026719 0.000000 -0.363000
                H -0.513360 -0.889165 -0.363000
                H -0.513360 0.889165 -0.363000""")
            f.flush()
            pymatgen_mol = read_mol(f.name)

        for struct in [StructureData(pymatgen=pymatgen_mol),
                       StructureData(pymatgen_molecule=pymatgen_mol)]:
            self.assertEquals(struct.get_site_kindnames(),
                              ['H', 'H', 'H', 'H', 'C'])
            self.assertEquals(struct.pbc, (False, False, False))
            self.assertEquals([round(x, 2) for x in list(struct.sites[0].position)],
                              [5.77, 5.89, 6.81])
            self.assertEquals([round(x, 2) for x in list(struct.sites[1].position)],
                              [6.8, 5.89, 5.36])
            self.assertEquals([round(x, 2) for x in list(struct.sites[2].position)],
                              [5.26, 5.0, 5.36])
            self.assertEquals([round(x, 2) for x in list(struct.sites[3].position)],
                              [5.26, 6.78, 5.36])
            self.assertEquals([round(x, 2) for x in list(struct.sites[4].position)],
                              [5.77, 5.89, 5.73])


class TestPymatgenFromStructureData(SqlAlchemyTests):
    """
    Tests the creation of pymatgen Structure and Molecule objects from
    StructureData.
    """
    @unittest.skipIf(not has_pymatgen() or
                     StrictVersion(get_pymatgen_version()) <
                     StrictVersion('3.0.13'), "Unable to import pymatgen")
    def test_1(self):
        """
        Test the check of periodic boundary conditions.
        """
        struct = StructureData()

        struct.pbc = [True, True, True]
        pmg_struct = struct.get_pymatgen_structure()

        struct.pbc = [True, True, False]
        with self.assertRaises(ValueError):
            pmg_struct = struct.get_pymatgen_structure()

    @unittest.skipIf(not has_ase() or
                     not has_pymatgen() or
                     StrictVersion(get_pymatgen_version()) <
                     StrictVersion('3.0.13'),
                     "Unable to import ase or pymatgen")
    def test_2(self):
        aseatoms = ase.Atoms('Si4', cell=(1., 2., 3.),
                             pbc=(True, True, True))
        aseatoms.set_scaled_positions(
            ((0.0, 0.0, 0.0),
             (0.1, 0.1, 0.1),
             (0.2, 0.2, 0.2),
             (0.3, 0.3, 0.3),
            )
        )

        a_struct = StructureData(ase=aseatoms)
        p_struct = a_struct.get_pymatgen_structure()

        coord_array = [x['abc'] for x in p_struct.as_dict()['sites']]
        for i in range(len(coord_array)):
            coord_array[i] = [round(x, 2) for x in coord_array[i]]

        self.assertEquals(coord_array,
                          [[0.0, 0.0, 0.0],
                           [0.1, 0.1, 0.1],
                           [0.2, 0.2, 0.2],
                           [0.3, 0.3, 0.3]])

    @unittest.skipIf(not has_ase() or
                     not has_pymatgen() or
                     StrictVersion(get_pymatgen_version()) <
                     StrictVersion('3.0.13'),
                     "Unable to import ase or pymatgen")
    def test_3(self):
        """
        Test the conversion of StructureData to pymatgen's Molecule.
        """
        aseatoms = ase.Atoms('Si4', cell=(10, 10, 10),
                             pbc=(True, True, True))
        aseatoms.set_scaled_positions(
            ((0.0, 0.0, 0.0),
             (0.1, 0.1, 0.1),
             (0.2, 0.2, 0.2),
             (0.3, 0.3, 0.3),
            )
        )

        a_struct = StructureData(ase=aseatoms)
        p_mol = a_struct.get_pymatgen_molecule()

        self.assertEquals([x['xyz'] for x in p_mol.as_dict()['sites']],
                          [[0.0, 0.0, 0.0],
                           [1.0, 1.0, 1.0],
                           [2.0, 2.0, 2.0],
                           [3.0, 3.0, 3.0]])


class TestArrayData(SqlAlchemyTests):
    """
    Tests the ArrayData objects.
    """

    def test_creation(self):
        """
        Check the methods to add, remove, modify, and get arrays and
        array shapes.
        """
        # Create a node with two arrays
        n = ArrayData()
        first = numpy.random.rand(2, 3, 4)
        n.set_array('first', first)

        second = numpy.arange(10)
        n.set_array('second', second)

        third = numpy.random.rand(6, 6)
        n.set_array('third', third)


        # Check if the arrays are there
        self.assertEquals(set(['first', 'second', 'third']), set(n.arraynames()))
        self.assertAlmostEquals(abs(first - n.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n.get_array('second')).max(), 0.)
        self.assertAlmostEquals(abs(third - n.get_array('third')).max(), 0.)
        self.assertEquals(first.shape, n.get_shape('first'))
        self.assertEquals(second.shape, n.get_shape('second'))
        self.assertEquals(third.shape, n.get_shape('third'))

        with self.assertRaises(KeyError):
            n.get_array('nonexistent_array')

        # Delete an array, and try to delete a non-existing one
        n.delete_array('third')
        with self.assertRaises(KeyError):
            n.delete_array('nonexistent_array')

        # Overwrite an array
        first = numpy.random.rand(4, 5, 6)
        n.set_array('first', first)

        # Check if the arrays are there, and if I am getting the new one
        self.assertEquals(set(['first', 'second']), set(n.arraynames()))
        self.assertAlmostEquals(abs(first - n.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n.get_array('second')).max(), 0.)
        self.assertEquals(first.shape, n.get_shape('first'))
        self.assertEquals(second.shape, n.get_shape('second'))

        n.store()

        # Same checks, after storing
        self.assertEquals(set(['first', 'second']), set(n.arraynames()))
        self.assertAlmostEquals(abs(first - n.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n.get_array('second')).max(), 0.)
        self.assertEquals(first.shape, n.get_shape('first'))
        self.assertEquals(second.shape, n.get_shape('second'))

        # Same checks, again (this is checking the caching features)
        self.assertEquals(set(['first', 'second']), set(n.arraynames()))
        self.assertAlmostEquals(abs(first - n.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n.get_array('second')).max(), 0.)
        self.assertEquals(first.shape, n.get_shape('first'))
        self.assertEquals(second.shape, n.get_shape('second'))


        # Same checks, after reloading
        n2 = ArrayData(dbnode=n.dbnode)
        self.assertEquals(set(['first', 'second']), set(n2.arraynames()))
        self.assertAlmostEquals(abs(first - n2.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n2.get_array('second')).max(), 0.)
        self.assertEquals(first.shape, n2.get_shape('first'))
        self.assertEquals(second.shape, n2.get_shape('second'))

        # Same checks, after reloading with UUID
        n2 = ArrayData.get_subclass_from_uuid(n.uuid)
        self.assertEquals(set(['first', 'second']), set(n2.arraynames()))
        self.assertAlmostEquals(abs(first - n2.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n2.get_array('second')).max(), 0.)
        self.assertEquals(first.shape, n2.get_shape('first'))
        self.assertEquals(second.shape, n2.get_shape('second'))


        # Check that I cannot modify the node after storing
        with self.assertRaises(ModificationNotAllowed):
            n.delete_array('first')
        with self.assertRaises(ModificationNotAllowed):
            n.set_array('second', first)

        # Again same checks, to verify that the attempts to delete/overwrite
        # arrays did not damage the node content
        self.assertEquals(set(['first', 'second']), set(n.arraynames()))
        self.assertAlmostEquals(abs(first - n.get_array('first')).max(), 0.)
        self.assertAlmostEquals(abs(second - n.get_array('second')).max(), 0.)
        self.assertEquals(first.shape, n.get_shape('first'))
        self.assertEquals(second.shape, n.get_shape('second'))

    def test_iteration(self):
        """
        Check the functionality of the iterarrays() iterator
        """
        # Create a node with two arrays
        n = ArrayData()
        first = numpy.random.rand(2, 3, 4)
        n.set_array('first', first)

        second = numpy.arange(10)
        n.set_array('second', second)

        third = numpy.random.rand(6, 6)
        n.set_array('third', third)

        for name, array in n.iterarrays():
            if name == 'first':
                self.assertAlmostEquals(abs(first - array).max(), 0.)
            if name == 'second':
                self.assertAlmostEquals(abs(second - array).max(), 0.)
            if name == 'third':
                self.assertAlmostEquals(abs(third - array).max(), 0.)


class TestTrajectoryData(SqlAlchemyTests):
    """
    Tests the TrajectoryData objects.
    """

    def test_creation(self):
        """
        Check the methods to set and retrieve a trajectory.
        """
        # Create a node with two arrays
        n = TrajectoryData()

        # I create sample data
        stepids = numpy.array([60, 70])
        times = stepids * 0.01
        cells = numpy.array([
            [[2., 0., 0., ],
             [0., 2., 0., ],
             [0., 0., 2., ]],
            [[3., 0., 0., ],
             [0., 3., 0., ],
             [0., 0., 3., ]]])
        symbols = numpy.array(['H', 'O', 'C'])
        positions = numpy.array([
            [[0., 0., 0.],
             [0.5, 0.5, 0.5],
             [1.5, 1.5, 1.5]],
            [[0., 0., 0.],
             [0.5, 0.5, 0.5],
             [1.5, 1.5, 1.5]]])
        velocities = numpy.array([
            [[0., 0., 0.],
             [0., 0., 0.],
             [0., 0., 0.]],
            [[0.5, 0.5, 0.5],
             [0.5, 0.5, 0.5],
             [-0.5, -0.5, -0.5]]])

        # I set the node
        n.set_trajectory(stepids=stepids, cells=cells, symbols=symbols,
                         positions=positions, times=times, velocities=velocities)

        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertAlmostEqual(abs(times - n.get_times()).sum(), 0.)
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertAlmostEqual(abs(velocities - n.get_velocities()).sum(), 0.)

        # get_step_data function check
        data = n.get_step_data(1)
        self.assertEqual(data[0], stepids[1])
        self.assertAlmostEqual(data[1], times[1])
        self.assertAlmostEqual(abs(cells[1] - data[2]).sum(), 0.)
        self.assertEqual(symbols.tolist(), data[3].tolist())
        self.assertAlmostEqual(abs(data[4] - positions[1]).sum(), 0.)
        self.assertAlmostEqual(abs(data[5] - velocities[1]).sum(), 0.)

        # Step 70 has index 1
        self.assertEqual(1, n.get_index_from_stepid(70))
        with self.assertRaises(ValueError):
            # Step 66 does not exist
            n.get_index_from_stepid(66)

        ########################################################
        # I set the node, this time without times or velocities (the same node)
        n.set_trajectory(stepids=stepids, cells=cells, symbols=symbols,
                         positions=positions)
        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertIsNone(n.get_times())
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertIsNone(n.get_velocities())

        # Same thing, but for a new node
        n = TrajectoryData()
        n.set_trajectory(stepids=stepids, cells=cells, symbols=symbols,
                         positions=positions)
        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertIsNone(n.get_times())
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertIsNone(n.get_velocities())

        ########################################################
        # I set the node, this time without velocities (the same node)
        n.set_trajectory(stepids=stepids, cells=cells, symbols=symbols,
                         positions=positions, times=times)
        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertAlmostEqual(abs(times - n.get_times()).sum(), 0.)
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertIsNone(n.get_velocities())

        # Same thing, but for a new node
        n = TrajectoryData()
        n.set_trajectory(stepids=stepids, cells=cells, symbols=symbols,
                         positions=positions, times=times)
        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertAlmostEqual(abs(times - n.get_times()).sum(), 0.)
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertIsNone(n.get_velocities())

        n.store()

        # Again same checks, but after storing
        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertAlmostEqual(abs(times - n.get_times()).sum(), 0.)
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertIsNone(n.get_velocities())

        # get_step_data function check
        data = n.get_step_data(1)
        self.assertEqual(data[0], stepids[1])
        self.assertAlmostEqual(data[1], times[1])
        self.assertAlmostEqual(abs(cells[1] - data[2]).sum(), 0.)
        self.assertEqual(symbols.tolist(), data[3].tolist())
        self.assertAlmostEqual(abs(data[4] - positions[1]).sum(), 0.)
        self.assertIsNone(data[5])

        # Step 70 has index 1
        self.assertEqual(1, n.get_index_from_stepid(70))
        with self.assertRaises(ValueError):
            # Step 66 does not exist
            n.get_index_from_stepid(66)

        ##############################################################
        # Again, but after reloading from uuid
        n = TrajectoryData.get_subclass_from_uuid(n.uuid)
        # Generic checks
        self.assertEqual(n.numsites, 3)
        self.assertEqual(n.numsteps, 2)
        self.assertAlmostEqual(abs(stepids - n.get_stepids()).sum(), 0.)
        self.assertAlmostEqual(abs(times - n.get_times()).sum(), 0.)
        self.assertAlmostEqual(abs(cells - n.get_cells()).sum(), 0.)
        self.assertEqual(symbols.tolist(), n.get_symbols().tolist())
        self.assertAlmostEqual(abs(positions - n.get_positions()).sum(), 0.)
        self.assertIsNone(n.get_velocities())

        # get_step_data function check
        data = n.get_step_data(1)
        self.assertEqual(data[0], stepids[1])
        self.assertAlmostEqual(data[1], times[1])
        self.assertAlmostEqual(abs(cells[1] - data[2]).sum(), 0.)
        self.assertEqual(symbols.tolist(), data[3].tolist())
        self.assertAlmostEqual(abs(data[4] - positions[1]).sum(), 0.)
        self.assertIsNone(data[5])

        # Step 70 has index 1
        self.assertEqual(1, n.get_index_from_stepid(70))
        with self.assertRaises(ValueError):
            # Step 66 does not exist
            n.get_index_from_stepid(66)

    def test_conversion_to_structure(self):
        """
        Check the methods to export a given time step to a StructureData node.
        """
        # Create a node with two arrays
        n = TrajectoryData()

        # I create sample data
        stepids = numpy.array([60, 70])
        times = stepids * 0.01
        cells = numpy.array([
            [[2., 0., 0., ],
             [0., 2., 0., ],
             [0., 0., 2., ]],
            [[3., 0., 0., ],
             [0., 3., 0., ],
             [0., 0., 3., ]]])
        symbols = numpy.array(['H', 'O', 'C'])
        positions = numpy.array([
            [[0., 0., 0.],
             [0.5, 0.5, 0.5],
             [1.5, 1.5, 1.5]],
            [[0., 0., 0.],
             [0.5, 0.5, 0.5],
             [1.5, 1.5, 1.5]]])
        velocities = numpy.array([
            [[0., 0., 0.],
             [0., 0., 0.],
             [0., 0., 0.]],
            [[0.5, 0.5, 0.5],
             [0.5, 0.5, 0.5],
             [-0.5, -0.5, -0.5]]])

        # I set the node
        n.set_trajectory(stepids=stepids, cells=cells, symbols=symbols,
                         positions=positions, times=times, velocities=velocities)

        from_step = n.get_step_structure(1)
        from_get_aiida_structure = n._get_aiida_structure(index=1)

        for struc in [from_step, from_get_aiida_structure]:
            self.assertEqual(len(struc.sites), 3)  # 3 sites
            self.assertAlmostEqual(abs(numpy.array(struc.cell) - cells[1]).sum(), 0)
            newpos = numpy.array([s.position for s in struc.sites])
            self.assertAlmostEqual(abs(newpos - positions[1]).sum(), 0)
            newkinds = [s.kind_name for s in struc.sites]
            self.assertEqual(newkinds, symbols.tolist())

            # Weird assignments (nobody should ever do this, but it is possible in
            # principle and we want to check
            k1 = Kind(name='C', symbols='Cu')
            k2 = Kind(name='H', symbols='He')
            k3 = Kind(name='O', symbols='Os', mass=100.)
            k4 = Kind(name='Ge', symbols='Ge')

            with self.assertRaises(ValueError):
                # Not enough kinds
                struc = n.get_step_structure(1, custom_kinds=[k1, k2])

            with self.assertRaises(ValueError):
                # Too many kinds
                struc = n.get_step_structure(1, custom_kinds=[k1, k2, k3, k4])

            with self.assertRaises(ValueError):
                # Wrong kinds
                struc = n.get_step_structure(1, custom_kinds=[k1, k2, k4])

            with self.assertRaises(ValueError):
                # Two kinds with the same name
                struc = n.get_step_structure(1, custom_kinds=[k1, k2, k3, k3])

            # Correct kinds
            struc = n.get_step_structure(1, custom_kinds=[k1, k2, k3])

            # Checks
            self.assertEqual(len(struc.sites), 3)  # 3 sites
            self.assertAlmostEqual(abs(numpy.array(struc.cell) - cells[1]).sum(), 0)
            newpos = numpy.array([s.position for s in struc.sites])
            self.assertAlmostEqual(abs(newpos - positions[1]).sum(), 0)
            newkinds = [s.kind_name for s in struc.sites]
            # Kinds are in the same order as given in the custm_kinds list
            self.assertEqual(newkinds, symbols.tolist())
            newatomtypes = [struc.get_kind(s.kind_name).symbols[0] for s in struc.sites]
            # Atoms remain in the same order as given in the positions list
            self.assertEqual(newatomtypes, ['He', 'Os', 'Cu'])
            # Check the mass of the kind of the second atom ('O' _> symbol Os, mass 100)
            self.assertAlmostEqual(struc.get_kind(struc.sites[1].kind_name).mass, 100.)

    def test_conversion_from_structurelist(self):
        """
        Check the method to create a TrajectoryData from list of AiiDA
        structures.
        """
        cells = [
            [[2., 0., 0., ],
             [0., 2., 0., ],
             [0., 0., 2., ]],
            [[3., 0., 0., ],
             [0., 3., 0., ],
             [0., 0., 3., ]]
        ]
        symbols = [['H', 'O', 'C'], ['H', 'O', 'C']]
        positions = [
            [[0., 0., 0.],
             [0.5, 0.5, 0.5],
             [1.5, 1.5, 1.5]],
            [[0., 0., 0.],
             [0.75, 0.75, 0.75],
             [1.25, 1.25, 1.25]]
        ]
        structurelist = []
        for i in range(0, 2):
            struct = StructureData(cell=cells[i])
            for j, symbol in enumerate(symbols[i]):
                struct.append_atom(symbols=symbol, position=positions[i][j])
            structurelist.append(struct)

        td = TrajectoryData(structurelist=structurelist)
        self.assertEqual(td.get_cells().tolist(), cells)
        self.assertEqual(td.get_symbols().tolist(), symbols[0])
        self.assertEqual(td.get_positions().tolist(), positions)

        symbols = [['H', 'O', 'C'], ['H', 'O', 'P']]
        structurelist = []
        for i in range(0, 2):
            struct = StructureData(cell=cells[i])
            for j, symbol in enumerate(symbols[i]):
                struct.append_atom(symbols=symbol, position=positions[i][j])
            structurelist.append(struct)

        with self.assertRaises(ValueError):
            td = TrajectoryData(structurelist=structurelist)


class TestKpointsData(SqlAlchemyTests):
    """
    Tests the TrajectoryData objects.
    """
    # TODO: I should find a way to check the special points, case by case

    def test_mesh(self):
        """
        Check the methods to set and retrieve a mesh.
        """
        # Create a node with two arrays
        k = KpointsData()

        # check whether the mesh can be set properly
        input_mesh = [4, 4, 4]
        k.set_kpoints_mesh(input_mesh)
        mesh, offset = k.get_kpoints_mesh()
        self.assertEqual(mesh, tuple(input_mesh))
        self.assertEqual(offset, (0., 0., 0.))  # must be a tuple of three 0 by default

        # a too long list should fail
        with self.assertRaises(ValueError):
            k.set_kpoints_mesh([4, 4, 4, 4])

        # now try to put explicitely an offset
        input_offset = [0.5, 0.5, 0.5]
        k.set_kpoints_mesh(input_mesh, input_offset)
        mesh, offset = k.get_kpoints_mesh()
        self.assertEqual(mesh, tuple(input_mesh))
        self.assertEqual(offset, tuple(input_offset))

        # verify the same but after storing
        k.store()
        self.assertEqual(mesh, tuple(input_mesh))
        self.assertEqual(offset, tuple(input_offset))

        # cannot modify it after storage
        with self.assertRaises(ModificationNotAllowed):
            k.set_kpoints_mesh(input_mesh)

    def test_list(self):
        """
        Test the method to set and retrieve a kpoint list.
        """
        k = KpointsData()

        input_klist = numpy.array([(0.0, 0.0, 0.0),
                                   (0.2, 0.0, 0.0),
                                   (0.0, 0.2, 0.0),
                                   (0.0, 0.0, 0.2),
        ])

        # set kpoints list
        k.set_kpoints(input_klist)
        klist = k.get_kpoints()

        # try to get the same
        self.assertTrue(numpy.array_equal(input_klist, klist))

        # if no cell is set, cannot convert into cartesian
        with self.assertRaises(AttributeError):
            _ = k.get_kpoints(cartesian=True)

        # try to set also weights
        # should fail if the weights length do not match kpoints
        input_weights = numpy.ones(6)
        with self.assertRaises(ValueError):
            k.set_kpoints(input_klist, weights=input_weights)

        # try a right one
        input_weights = numpy.ones(4)
        k.set_kpoints(input_klist, weights=input_weights)
        klist, weights = k.get_kpoints(also_weights=True)
        self.assertTrue(numpy.array_equal(weights, input_weights))
        self.assertTrue(numpy.array_equal(klist, input_klist))

        # verify the same, but after storing
        k.store()
        klist, weights = k.get_kpoints(also_weights=True)
        self.assertTrue(numpy.array_equal(weights, input_weights))
        self.assertTrue(numpy.array_equal(klist, input_klist))

        # cannot modify it after storage
        with self.assertRaises(ModificationNotAllowed):
            k.set_kpoints(input_klist)

    def test_kpoints_to_cartesian(self):
        """
        Test how the list of kpoints is converted to cartesian coordinates
        """
        k = KpointsData()

        input_klist = numpy.array([(0.0, 0.0, 0.0),
                                   (0.2, 0.0, 0.0),
                                   (0.0, 0.2, 0.0),
                                   (0.0, 0.0, 0.2),
        ])

        # define a cell
        alat = 4.
        cell = numpy.array([[alat, 0., 0.],
                            [0., alat, 0.],
                            [0., 0., alat],
        ])

        k.set_cell(cell)

        # set kpoints list
        k.set_kpoints(input_klist)

        # verify that it is not the same of the input
        # (at least I check that there something has been done)
        klist = k.get_kpoints(cartesian=True)
        self.assertFalse(numpy.array_equal(klist, input_klist))

        # put the kpoints in cartesian and get them back, they should be equal
        # internally it is doing two matrix transforms
        k.set_kpoints(input_klist, cartesian=True)
        klist = k.get_kpoints(cartesian=True)
        self.assertTrue(numpy.allclose(klist, input_klist, atol=1e-16))


class TestData(SqlAlchemyTests):
    """
    Tests generic Data class.
    """

    @unittest.skipIf(True, "Validation for source is disabled due to "
                           "Issue #9 (https://bitbucket.org/epfl_theos/aiida_epfl/issues/9)")
    def test_license_validation(self):
        """
        Test the validation of source licenses.
        """
        data = Data(source={'license': 'CC0'})
        data.store()

        data = Data(source={'license': 'CC-BY-SA'})
        # CC-BY* type licenses must be accompanied by attribution, given
        # in 'description' key of source dictionary. This is enforced in
        # Data._validate(), thus the ValidationError.
        with self.assertRaises(ValidationError):
            data.store()
