"""
@package pmsco.cluster
cluster tools for MSC and EDAC

the Cluster class is provided to facilitate the construction and import/export of clusters.
a cluster can be built by adding single atoms, layers, or a half-space bulk lattice.
the class can import from/export to EDAC, MSC, and XYZ cluster files.
XYZ allows for export to 3D visualizers, e.g. Avogadro.

@pre requires the periodictable package (https://pypi.python.org/pypi/periodictable)
@code{.sh}
pip install --user periodictable
@endcode

@author Matthias Muntwiler

@copyright (c) 2015-18 by Paul Scherrer Institut
"""

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

import math
import numpy as np
import periodictable as pt
import sys

## default file format identifier
FMT_DEFAULT = 0
## MSC file format identifier
FMT_MSC = 1
## EDAC file format identifier
FMT_EDAC = 2
## XYZ file format identifier
FMT_XYZ = 3

# python version dependent type of chemical symbol
if sys.version_info[0] >= 3:
    _SYMBOL_TYPE = 'U2'
else:
    _SYMBOL_TYPE = 'S2'

## numpy.array datatype of Cluster.data array
DTYPE_CLUSTER_INTERNAL = [('i', 'i4'), ('t', 'i4'), ('s', _SYMBOL_TYPE), ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
                          ('e', 'u1')]
## file format of internal Cluster.data array
FMT_CLUSTER_INTERNAL = ["%5u", "%2u", "%s", "%7.3f", "%7.3f", "%7.3f", "%1u"]
## field (column) names of internal Cluster.data array
FIELDS_CLUSTER_INTERNAL = ['i', 't', 's', 'x', 'y', 'z', 'e']

## numpy.array datatype of cluster for MSC cluster file input/output
DTYPE_CLUSTER_MSC = [('i', 'i4'), ('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('t', 'i4')]
## file format of MSC cluster file
FMT_CLUSTER_MSC = ["%5u", "%7.3f", "%7.3f", "%7.3f", "%2u"]
## field (column) names of MSC cluster file
FIELDS_CLUSTER_MSC = ['i', 'x', 'y', 'z', 't']

## numpy.array datatype of cluster for EDAC cluster file input/output
DTYPE_CLUSTER_EDAC= [('i', 'i4'), ('t', 'i4'), ('x', 'f4'), ('y', 'f4'), ('z', 'f4')]
## file format of EDAC cluster file
FMT_CLUSTER_EDAC = ["%5u", "%2u", "%7.3f", "%7.3f", "%7.3f"]
## field (column) names of EDAC cluster file
FIELDS_CLUSTER_EDAC = ['i', 't', 'x', 'y', 'z']

## numpy.array datatype of cluster for XYZ file input/output
DTYPE_CLUSTER_XYZ= [('s', _SYMBOL_TYPE), ('x', 'f4'), ('y', 'f4'), ('z', 'f4')]
## file format of XYZ cluster file
FMT_CLUSTER_XYZ = ["%s", "%10.5f", "%10.5f", "%10.5f"]
## field (column) names of XYZ cluster file
FIELDS_CLUSTER_XYZ = ['s', 'x', 'y', 'z']


class Cluster(object):
    """
    Represents a cluster of atoms by their coordinates and chemical element.
    
    the object stores the following information per atom in the @ref data array:
    
    - sequential atom index (1-based)
    - atom type (chemical element number)
    - chemical element symbol
    - x coordinate of the atom position
    - t coordinate of the atom position
    - z coordinate of the atom position
    - emitter flag
    
    the class also defines methods that add or manipulate atoms of the cluster.
    see most importantly the set_rmax, add_atom, add_layer and add_bulk functions.
    emitters can be flagged by the set_emitter method.
    
    you may also manipulate the data array directly.
    in this case, be sure to keep the data array consistent.
    the update methods can help to recreate the index, atom type or symbol columns.
    
    the class can also load and save files in some simple formats.
    """

    ## @var rmax
    #       maximum distance of atoms from the origin.
    #
    #       float, default = 0
    #
    #       this parameter restricts the addition of new atoms.
    #       changing the parameter does not affect existing atoms.
    #       the default is 0 (no atom will be added!).
    #       you must set this parameter explicitly!

    ##    @var dtype
    #       data type of the internal numpy.ndarray.
    
    ##    @var file_format
    #       default file format.
    #
    #       must be one of the FMT_MSC, FMT_EDAC, FMT_XYZ constants.
    #       the initial value is FMT_XYZ.
    
    ##    @var data
    #       structured numpy array holding the atom positions.
    #
    #       the columns of the array are:
    #       @arg @c 'i' (int) atom index (1-based)
    #       @arg @c 't' (int) atom type (chemical element number)
    #       @arg @c 's' (string) chemical element symbol
    #       @arg @c 'x' (float32) x coordinate of the atom position
    #       @arg @c 'y' (float32) t coordinate of the atom position
    #       @arg @c 'z' (float32) z coordinate of the atom position
    #       @arg @c 'e' (uint8)   1 = emitter, 0 = regular atom

    ##  @var comment (str)
    #   one-line comment that can be included in some cluster files

    def __init__(self):
        self.data = None
        self.rmax = 0.0
        self.dtype = DTYPE_CLUSTER_INTERNAL
        self.file_format = FMT_XYZ
        self.comment = ""
        self.clear()

    def clear(self):
        """
        Remove all atoms from the cluster.
        """
        n_atoms = 0
        self.data = np.zeros(n_atoms, dtype=self.dtype)

    def copy_from(self, cluster):
        """
        Copy the data from another cluster.

        @param cluster: (Cluster) other Cluster object.
        """
        self.data = cluster.data.copy()

    def set_rmax(self, r):
        """
        set rmax, the maximum distance of atoms from the origin.
        
        atoms with norm greater than rmax will not be added to the cluster
        by the add_layer() and add_bulk() methods.
        existing atoms are not affected when changing rmax.

        you must set this parameter explicitly, as the default value is 0
        (no atom will be added)!
        """
        self.rmax = r

    def build_element(self, index, element_number, x, y, z, emitter):
        """
        build a tuple in the format of the internal data array.
        
        @param index: (int) index
        
        @param element_number: (int) chemical element number
        
        @param x, y, z: (float) atom coordinates in the cluster
        
        @param emitter: (int or bool) True = emitter, False = scatterer
        """
        symbol = pt.elements[element_number].symbol
        element = (index, element_number, symbol, x, y, z, int(emitter))
        return element

    def add_atom(self, atomtype, v_pos, is_emitter):
        """
        add a single atom to the cluster.
        
        @param atomtype: (int) chemical element number
        
        @param v_pos: (numpy.ndarray, shape = (3)) position vector
        
        @param is_emitter: (int or bool) True = emitter, False = scatterer
        """
        n0 = self.data.shape[0] + 1
        element = self.build_element(n0, atomtype, v_pos[0], v_pos[1], v_pos[2], int(is_emitter))
        self.data = np.append(self.data, np.array(element,
            dtype=self.data.dtype))

    def add_layer(self, atomtype, v_pos, v_lat1, v_lat2):
        """
        add a layer of atoms to the cluster.
        
        the layer is expanded up to the limit given by
        self.rmax (maximum distance from the origin).
        all atoms are non-emitters.
        
        @param atomtype: (int) chemical element number
        
        @param v_pos: (numpy.ndarray, shape = (3))
            position vector of the first atom (basis vector)
        
        @param v_lat1, v_lat2: (numpy.ndarray, shape = (3))
            lattice vectors.
        """
        r_great = max(self.rmax, np.linalg.norm(v_pos))
        n0 = self.data.shape[0] + 1
        n1 = max(int(r_great / np.linalg.norm(v_lat1)) + 1, 3) * 2
        n2 = max(int(r_great / np.linalg.norm(v_lat2)) + 1, 3) * 2
        nn = 0
        buf = np.empty((2 * n1 + 1) * (2 * n2 + 1), dtype=self.dtype)
        for i1 in range(-n1, n1 + 1):
            for i2 in range(-n2, n2 + 1):
                v = v_pos + v_lat1 * i1 + v_lat2 * i2
                if np.linalg.norm(v) <= self.rmax:
                    buf[nn] = self.build_element(nn + n0, atomtype, v[0], v[1], v[2], 0)
                    nn += 1
        buf = np.resize(buf, nn)
        self.data = np.append(self.data, buf)

    def add_bulk(self, atomtype, v_pos, v_lat1, v_lat2, v_lat3, z_surf=0.0):
        """
        add bulk atoms to the cluster.
        
        the lattice is expanded up to the limits given by 
        self.rmax (maximum distance from the origin)
        and z_surf (position of the surface).
        all atoms are non-emitters.

        @param atomtype: (int) chemical element number
        
        @param v_pos: (numpy.ndarray, shape = (3))
            position vector of the first atom (basis vector)
        
        @param v_lat1, v_lat2, v_lat3: (numpy.ndarray, shape = (3))
            lattice vectors.

        @param z_surf: (float) position of surface.
            atoms with z > z_surf are not added.
        """
        r_great = max(self.rmax, np.linalg.norm(v_pos))
        n0 = self.data.shape[0] + 1
        n1 = max(int(r_great / np.linalg.norm(v_lat1)) + 1, 4) * 3
        n2 = max(int(r_great / np.linalg.norm(v_lat2)) + 1, 4) * 3
        n3 = max(int(r_great / np.linalg.norm(v_lat3)) + 1, 4) * 3
        nn = 0
        buf = np.empty((2 * n1 + 1) * (2 * n2 + 1) * (n3 + 1), dtype=self.dtype)
        for i1 in range(-n1, n1 + 1):
            for i2 in range(-n2, n2 + 1):
                for i3 in range(-n3, n3 + 1):
                    v = v_pos + v_lat1 * i1 + v_lat2 * i2 + v_lat3 * i3
                    if np.linalg.norm(v) <= self.rmax and v[2] <= z_surf:
                        buf[nn] = self.build_element(nn + n0, atomtype, v[0], v[1], v[2], 0)
                        nn += 1
        buf = np.resize(buf, nn)
        self.data = np.append(self.data, buf)

    def add_cluster(self, cluster, check_rmax=False, check_unique=False, tol=0.001):
        """
        add atoms from another cluster object.
        
        @note the order of atoms in the internal data array may change during this operation.
            the atom index is updated.
            
        @param cluster: Cluster object to be added.
         
        @param check_rmax: if True, atoms outside self.rmax are not added.
            if False (default), all atoms of the other cluster are added.
        
        @param check_unique: if True, atoms occupying the same position as an existing atom will not be added.
            if False (default), all atoms are added even if they occupy the same position.
            
        @param tol: tolerance for checking uniqueness.
            positions of two atoms are considered equal if all coordinates lie within the tolerance interval.

        @return: None 
        """
        assert isinstance(cluster, Cluster)
        data = self.data.copy()
        source = cluster.data.copy()

        if check_rmax and source.shape[0] > 0:
            source_xyz = source[['x', 'y', 'z']].copy()
            source_xyz = source_xyz.view((source_xyz.dtype[0], len(source_xyz.dtype.names)))
            b_rmax = np.linalg.norm(source_xyz, axis=1) <= self.rmax
            idx = np.where(b_rmax)
            source = source[idx]
        data = np.append(data, source)

        if check_unique and data.shape[0] > 0:
            data_xyz = data[['x', 'y', 'z']].copy()
            data_xyz = data_xyz.view((data_xyz.dtype[0], len(data_xyz.dtype.names)))
            tol_xyz = np.round(data_xyz / tol)
            uni_xyz = tol_xyz.view(tol_xyz.dtype.descr * 3)
            _, idx = np.unique(uni_xyz, return_index=True)
            data = data[np.sort(idx)]

        self.data = data
        self.update_index()

    def get_z_layers(self, tol=0.001):
        """
        return the z-coordinates of atomic layers.
        the layers are stacked in the z-direction.

        the function gathers unique z-coordinates.
        coordinates which are within the given tolerance are assigned to the same layer.

        @param tol: tolerance
        @return: (numpy.ndarray) z-coordinates of the layers.
            the coordinates are numerically ordered, the top layer appears last.
            the returned coordinates may not be identical to any atom coordinate of a layer
            but deviate up to the given tolerance.
        """
        self_z = self.data['z'].view(np.float32).reshape(self.data.shape)
        z2 = np.round(self_z.copy() / tol)
        layers = np.unique(z2) * tol
        return layers

    def relax(self, z_cut, z_shift, element=0):
        """
        shift atoms below a certain z coordinate by a fixed distance in the z direction.

        @param z_cut: atoms below this z coordinate are shifted.
        @param z_shift: amount of shift in z direction
            (positive to move towards the surface, negative to move into the bulk).
        @param element: (int) chemical element number if atoms of a specific element should be affected.
            by default (element = 0), all atoms are moved.
        @return: (numpy.ndarray) indices of the atoms that have been shifted.
        """
        self_z = self.data['z'].view(np.float32).reshape(self.data.shape)
        b_z = self_z <= z_cut
        b_all = b_z

        if element:
            try:
                b_el = self.data['t'] == int(element)
            except ValueError:
                b_el = self.data['s'] == element
            b_all = np.all([b_z, b_el], axis=0)

        idx = np.where(b_all)
        self.data['z'][idx] += z_shift

        return idx

    def translate(self, vector, element=0):
        """
        translate the cluster or all atoms of a specified element.

        @param vector: (numpy.ndarray) 3-dimensional displacement vector.
        @param element: (int) chemical element number if atoms of a specific element should be affected.
            by default (element = 0), all atoms are moved.
        @return: (numpy.ndarray) indices of the atoms that have been shifted.
        """
        if element:
            try:
                sel = self.data['t'] == int(element)
            except ValueError:
                sel = self.data['s'] == element
        else:
            sel = np.ones_like(self.data['t'])
        idx = np.where(sel)
        self.data['x'][idx] += vector[0]
        self.data['y'][idx] += vector[1]
        self.data['z'][idx] += vector[2]

        return idx

    def matrix_transform(self, matrix):
        """
        apply a transformation matrix to each atom of the cluster.
        
        the transformed atom positions are calculated as v = R * transpose(v)
        
        @param matrix: transformation matrix
        
        @return: None 
        """
        for atom in self.data:
            v = np.matrix([atom['x'], atom['y'], atom['z']])
            w = matrix * v.transpose()
            atom['x'] = float(w[0])
            atom['y'] = float(w[1])
            atom['z'] = float(w[2])

    def rotate_x(self, angle):
        """
        rotate cluster about the surface normal axis

        @param angle (float) in degrees
        """
        angle = math.radians(angle)
        s = math.sin(angle)
        c = math.cos(angle)
        matrix = np.matrix([[1, 0, 0], [0, c, -s], [0, s, c]])
        self.matrix_transform(matrix)

    def rotate_y(self, angle):
        """
        rotate cluster about the surface normal axis

        @param angle (float) in degrees
        """
        angle = math.radians(angle)
        s = math.sin(angle)
        c = math.cos(angle)
        matrix = np.matrix([[c, 0, s], [0, 1, 0], [-s, 0, c]])
        self.matrix_transform(matrix)

    def rotate_z(self, angle):
        """
        rotate cluster about the surface normal axis

        @param angle (float) in degrees
        """
        angle = math.radians(angle)
        s = math.sin(angle)
        c = math.cos(angle)
        matrix = np.matrix([[c, -s, 0], [s, c, 0], [0, 0, 1]])
        self.matrix_transform(matrix)

    def find_positions(self, pos, tol=0.001):
        """
        find all atoms which occupy a given position.

        @param pos: (numpy.array, shape = (3)) position vector.
        
        @param tol: (float) matching tolerance per coordinate.

        @return numpy.array of indices which match v_pos.
        """
        b2 = np.abs(pos - self.get_positions()) < tol
        b1 = np.all(b2, axis=1)
        idx = np.where(b1)
        return idx[0]

    def find_index_cylinder(self, pos, r_xy, r_z, element):
        """
        find atoms of a given element within a cylindrical volume and return their indices.

        @param pos: (numpy.array, shape = (3)) center position of the cylinder.

        @param r_xy: (float) radius of the cylinder.
            returned atoms must match |atom(x,y) - pos(x,y)| <= r_xy.

        @param r_z: (float) half height of the cylinder.
            returned atoms must match |atom(z) - pos(z)| <= r_z.

        @param element: (str or int) element symbol or atomic number.
            if None, the element is not checked.

        @return numpy.array of indices which match v_pos.
        """
        pos_xy = pos[0:2]
        self_xy = self.data[['x', 'y']].copy()
        self_xy = self_xy.view((self_xy.dtype[0], len(self_xy.dtype.names)))
        b_xy = np.linalg.norm(self_xy - pos_xy, axis=1) <= r_xy

        pos_z = pos[2]
        self_z = self.data['z']
        b_z = np.abs(self_z - pos_z) <= r_z

        if element is not None:
            try:
                b_el = self.data['t'] == int(element)
            except ValueError:
                b_el = self.data['s'] == element
            b_all = np.all([b_xy, b_z, b_el], axis=0)
        else:
            b_all = np.all([b_xy, b_z], axis=0)

        idx = np.where(b_all)
        return idx[0]

    def trim_cylinder(self, r_xy, r_z):
        """
        remove atoms outside a given cylinder.

        the cylinder is centered at the origin.

        @param r_xy: (float) radius of the cylinder.
            atoms to keep must match |atom(x,y)| <= r_xy.

        @param r_z: (float) half height of the cylinder.
            atoms to keep must match |atom(z)| <= r_z.

        @return: None
        """
        self_xy = self.data[['x', 'y']].copy()
        self_xy = self_xy.view((self_xy.dtype[0], len(self_xy.dtype.names)))
        b_xy = np.linalg.norm(self_xy, axis=1) <= r_xy

        self_z = self.data['z']
        b_z = np.abs(self_z) <= r_z

        b_all = np.all([b_xy, b_z], axis=0)
        idx = np.where(b_all)
        self.data = self.data[idx]
        self.update_index()

    def trim_paraboloid(self, rxy, z0):
        """
        remove atoms outside a given paraboloid.

        the paraboloid is defined by z(r) = z0 * (1 - r**2 / rxy**2),
        where r**2 = x**2 + y**2.
        its apex is at (0, 0, z0).
        z(r) = 0 for x**2 + y**2 = rxy**2.

        the coordinates (x,y,z) of atoms to keep must match z >= z(r).

        @param rxy: (float) radius of the paraboloid at z = 0.

        @param z0: (float) vertical coordinate of the apex (at x = y = 0).
            in the usual cluster layout where the surface is near z = 0 and the atoms at lower z,
            this value is negative.

        @return: None
        """
        rsq = self.data['x']**2 + self.data['y']**2
        pz = z0 * (1. - rsq / rxy**2)
        idx = np.where(self.data['z'] >= pz)
        self.data = self.data[idx]
        self.update_index()

    def trim_sphere(self, radius):
        """
        remove atoms outside a given sphere.

        the sphere is centered at the origin.

        @param radius: (float) radius of the sphere.
            atoms to keep must match |atom(x,y,z)| <= radius.

        @return: None
        """
        self_xyz = self.data[['x', 'y', 'z']].copy()
        self_xyz = self_xyz.view((self_xyz.dtype[0], len(self_xyz.dtype.names)))
        b_xyz = np.linalg.norm(self_xyz, axis=1) <= radius
        idx = np.where(b_xyz)
        self.data = self.data[idx]
        self.update_index()

    def trim_slab(self, axis, center, depth):
        """
        remove atoms outside a slab that is parallel to one of the coordinate planes.

        @param axis: axis to trim: 'x', 'y' or 'z'.
        @param center: center position of the slab.
        @param depth: thickness of the slab.

        @return: None
        """
        coord = self.data[axis].view(np.float32).reshape(self.data.shape)
        sel = np.abs(coord - center) <= depth / 2
        idx = np.where(sel)
        self.data = self.data[idx]
        self.update_index()

    def set_emitter(self, pos=None, idx=-1, tol=0.001):
        """
        select an atom as emitter.
        
        the emitter atom can be specified by position or index.
        either one of the pos or idx arguments must be specified.

        @param idx: (int) array index of the atom.
        
        @param pos: (numpy.array, shape = (3)) position vector.
        
        @param tol: (float) matching tolerance per component if pos argument is used.

        @raise IndexError if the position cannot be found
        """
        if pos is not None:
            ares = self.find_positions(pos, tol)
            idx = ares[0]
        item = self.data[idx]
        item['e'] = 1
        
    def move_to_first(self, pos=None, idx=0, tol=0.001):
        """
        move an atom to the first position.
        
        the emitter atom can be specified by position or index.
        either one of the pos or idx arguments must be specified.

        @param idx: (int) array index of the atom.
            must be greater than 1 to have an effect.
        
        @param pos: (numpy.array, shape = (3)) position vector.
        
        @param tol: (float) matching tolerance per component if pos argument is used.

        @raise IndexError if the position cannot be found
        """

        if pos is not None:
            ares = self.find_positions(pos, tol)
            idx = ares[0]
        if idx:
            em = self.data[idx]
            self.data = np.delete(self.data, idx)
            self.data = np.insert(self.data, 0, em)
            self.update_index()

    def get_positions(self):
        """
        get an array of the atom coordinates.

        the returned array is an independent copy of the original data.
        changes will not affect the original cluster.
        
        @return numpy.ndarray, shape = (N,3)
        """
        pos = self.data[['x', 'y', 'z']].copy()
        pos = pos.view((pos.dtype[0], len(pos.dtype.names)))
        return pos

    def set_positions(self, positions):
        """
        set atom coordinates from an array of shape (N,3).
        
        this method can be used on a modified array obtained from get_positions.
        N must be the number of atoms defined in the cluster.
        
        @param positions: numpy.ndarray of shape (N,3) where N is the number of atoms in this cluster.
         
        @return: None
        
        @raise AssertionError if the array sizes do not match.
        """
        assert isinstance(positions, np.ndarray)
        assert positions.shape == (self.data.shape[0], 3)
        self.data['x'] = positions[:, 0]
        self.data['y'] = positions[:, 1]
        self.data['z'] = positions[:, 2]

    def get_position(self, index):
        """
        get the position of a single atom.
        
        @param index: (int) index of the atom.
        
        @return numpy.array, shape = (3): position vector.
            the array instance is independent from the original array.
        """
        rec = self.data[index]
        return np.array((rec['x'], rec['y'], rec['z']))

    def get_atom_count(self):
        """
        get the number of atoms (positions) in the cluster.

        @return the number of atoms in the cluster.
        """
        return self.data.shape[0]

    def get_atomtype(self, index):
        """
        get the chemical element number of an atom.
        
        @param index: (int) index of the atom.
        
        @return int: chemical element number.
        """
        rec = self.data[index]
        return rec['t']

    def get_symbol(self, index):
        """
        get the chemical element symbol of an atom.
        
        @param index: (int) index of the atom.
        
        @return string: chemical element symbol.
        """
        rec = self.data[index]
        return rec['s']

    def get_emitters(self):
        """
        get a list of all emitters.
        
        @return list of tuples (x, y, z, atomtype)
        """
        idx = self.data['e'] != 0
        ems = self.data[['x', 'y', 'z', 't']][idx]
        return [tuple(em) for em in ems]

    def get_emitter_count(self):
        """
        get the number of emitters in the cluster.

        @return the number of atoms marked as emitter.
        """
        idx = self.data['e'] != 0
        return np.sum(idx)

    def load_from_file(self, f, fmt=FMT_DEFAULT):
        """
        load a cluster from a file created by the scattering program.
        
        @param f (string/handle): path name or open file handle of the cluster file.
        
        @param fmt (int): file format.
            must be one of the FMT_ constants.
            if FMT_DEFAULT, self.file_format is used.
        
        @remark if the filename ends in .gz, the file is loaded from compressed gzip format
        """
        if fmt == FMT_DEFAULT:
            fmt = self.file_format

        if fmt == FMT_MSC:
            dtype = DTYPE_CLUSTER_MSC
            fields = FIELDS_CLUSTER_MSC
            sh = 0
        elif fmt == FMT_EDAC:
            dtype = DTYPE_CLUSTER_EDAC
            fields = FIELDS_CLUSTER_EDAC
            sh = 1
        elif fmt == FMT_XYZ:
            dtype = DTYPE_CLUSTER_XYZ
            fields = FIELDS_CLUSTER_XYZ
            sh = 2
        else:
            dtype = DTYPE_CLUSTER_XYZ
            fields = FIELDS_CLUSTER_XYZ
            sh = 2

        data = np.genfromtxt(f, dtype=dtype, skip_header=sh)
        self.data = np.empty(data.shape, dtype=self.dtype)
        self.data['x'] = data['x']
        self.data['y'] = data['y']
        self.data['z'] = data['z']
        if 'i' in fields:
            self.data['i'] = data['i']
        else:
            self.update_index()
        if 't' in fields:
            self.data['t'] = data['t']
        if 's' in fields:
            self.data['s'] = data['s']
        else:
            self.update_symbols()
        if 't' not in fields:
            self.update_atomtypes()
        if 'e' in fields:
            self.data['e'] = data['e']
        else:
            self.data['e'] = 0

        pos = self.get_positions()
        # note: np.linalg.norm does not accept axis argument in version 1.7
        # (check np.version.version)
        norm = np.sqrt(np.sum(pos**2, axis=1))
        self.rmax = np.max(norm)

    def update_symbols(self):
        """
        update element symbols from element numbers.
        
        if you have modified the element numbers in the self.data array directly,
        this method updates the symbol column to make the data consistent.
        """
        for atom in self.data:
            atom['s'] = pt.elements[atom['t']].symbol

    def update_atomtypes(self):
        """
        update element numbers from element symbols.
        
        if you have modified the element symbols in the self.data array directly,
        this method updates the atom type column to make the data consistent.
        """
        for atom in self.data:
            atom['t'] = pt.elements.symbol(atom['s'].strip()).number

    def update_index(self):
        """
        update the index column.
        
        if you have modified the order or number of elements in the self.data array directly,
        you may need to re-index the atoms if your code uses functions that rely on the index. 
        
        @return: None 
        """
        self.data['i'] = np.arange(1, self.data.shape[0] + 1)

    def save_to_file(self, f, fmt=FMT_DEFAULT, comment="", emitters_only=False):
        """
        save the cluster to a file which can be read by the scattering program.

        the method updates the atom index because some file formats require an index column.
        
        @param f: (string/handle) path name or open file handle of the cluster file.
            if the filename ends in .gz, the file is saved in compressed gzip format

        @param fmt: (int) file format.
            must be one of the FMT_ constants.
            if FMT_DEFAULT, self.file_format is used.

        @param comment: (str) comment line (second line) in XYZ file.
            not used in other file formats.
            by default, self.comment is used.

        @param emitters_only: (bool) if True, only atoms marked as emitters are saved.
            by default, all atoms are saved.

        @return None
        """
        if fmt == FMT_DEFAULT:
            fmt = self.file_format

        if not comment:
            comment = self.comment

        self.update_index()
        if emitters_only:
            idx = self.data['e'] != 0
            data = self.data[idx]
        else:
            data = self.data

        if fmt == FMT_MSC:
            file_format = FMT_CLUSTER_MSC
            fields = FIELDS_CLUSTER_MSC
            header = ""
        elif fmt == FMT_EDAC:
            file_format = FMT_CLUSTER_EDAC
            fields = FIELDS_CLUSTER_EDAC
            header = "{nat} l(A)".format(nat=data.shape[0])
        elif fmt == FMT_XYZ:
            file_format = FMT_CLUSTER_XYZ
            fields = FIELDS_CLUSTER_XYZ
            header = "{nat}\n{com}".format(nat=data.shape[0], com=comment)
        else:
            file_format = FMT_CLUSTER_XYZ
            fields = FIELDS_CLUSTER_XYZ
            header = "{nat}\n{com}".format(nat=data.shape[0], com=comment)

        data = data[fields]
        np.savetxt(f, data, fmt=file_format, header=header, comments="")


class ClusterGenerator(object):
    """
    cluster generator class.

    this class bundles the cluster methods in one place
    so that it's easier to exchange them for different kinds of clusters.

    the project must override at least the create_cluster method.
    if emitters should be run in parallel tasks, the count_emitters method must be implemented as well.
    """

    def __init__(self, project):
        """
        initialize the cluster generator.

        @param project: reference to the project object.
            cluster generators may need to look up project parameters.
        """
        self.project = project

    def count_emitters(self, model, index):
        """
        return the number of emitter configurations for a particular model, scan and symmetry.

        the number of emitter configurations may depend on the model parameters, scan index and symmetry index.
        by default, the method returns 1, which means that there is only one emitter configuration.

        emitter configurations are mainly a way to distribute the calculations to multiple processes
        since the resulting diffraction patterns add up incoherently.
        for this to work, the create_cluster() method must pay attention to the emitter index
        and generate either a full cluster with all emitters (single process)
        or a cluster with only a subset of the emitters according to the emitter index (multiple processes).
        whether all emitters are calculated in one or multiple processes is decided at run-time
        based on the available resources.

        note that this function returns the number of _configurations_ not _atoms_.
        an emitter configuration (declared in a Cluster) may include more than one atom.
        it is up to the project, what is included in a particular configuration.

        to enable multiple emitter configurations, the derived project class must override this method
        and return a number greater than 1.

        @note in some cases it may be most efficient to call create_cluster and
            return Cluster.get_emitter_count() of the generated cluster.
            this is possible because the method is called with emitter index -1.
            model and index can be passed unchanged to create_cluster.

        @param model (dictionary) model parameters to be used in the calculation.

        @param index (named tuple CalcID) calculation index.
            the method should consider only the following attributes:
            @arg @c scan   scan index (index into Project.scans)
            @arg @c sym    symmetry index (index into Project.symmetries)
            @arg @c emit   emitter index must be -1.

        @return number of emitter configurations.
            this implementation returns the default value of 1.
        """
        return 1

    def create_cluster(self, model, index):
        """
        create a Cluster object given the model parameters and calculation index.

        the generated cluster will typically depend on the model parameters.
        depending on the project, it may also depend on the scan index, symmetry index and emitter index.

        the scan index can be used to generate a different cluster for different scan geometry,
        e.g., if some atoms can be excluded due to a longer mean free path.
        if this is not the case for the specific project, the scan index can be ignored.

        the symmetry index may select a particular domain that has a different atomic arrangement.
        in this case, depending on the value of index.sym, the function must generate a cluster corresponding
        to the particular domain/symmetry.
        the method can ignore the symmetry index if the project defines only one symmetry,
        or if the symmetry does not correspond to a different atomic structure.

        the emitter index selects a particular emitter configuration.
        depending on the value of the emitter index, the method must react differently:

        1. if the value is -1, return the full cluster and mark all inequivalent emitter atoms.
           emitters which are reproduced by a symmetry expansion in combine_emitters() should not be marked.
           the full diffraction scan will be calculated in one calculation.

        2. if the value is greater or equal to zero, generate the cluster with the emitter configuration
           selected by the emitter index.
           the index is in the range between 0 and the return value of count_emitters() minus 1.
           the results of the individual emitter calculations are summed up in combine_emitters().

        the code should ideally be written such that either case yields the same diffraction result.
        if count_emitters() always returns 1 (default), the second case does not have to be implemented,
        and the method can ignore the emitter index.

        the method must ignore the region index.

        if the creation of a cluster fails due to invalid model parameters,
        rather than raising an exception or constraining values,
        the method should send an error message to the logger and return an empty cluster or None.

        @param model (dictionary) model parameters to be used in the calculation.

        @param index (named tuple CalcID) calculation index.
            the method should consider only the following attributes:
            @arg @c scan   scan index (index into Project.scans)
            @arg @c sym    symmetry index (index into Project.symmetries)
            @arg @c emit   emitter index.
                            if -1, generate the full cluster and mark all emitters.
                            if greater or equal to zero, the value is a zero-based index of the emitter configuration.

        @return None.
            sub-classes must return a valid Cluster object or None in case of failure.
        """
        return None


class LegacyClusterGenerator(ClusterGenerator):
    """
    cluster generator class for projects that don't declare a generator.

    in previous versions, the create_cluster and count_emitters methods were implemented by the project class.
    this class redirects generator calls to the project methods
    providing compatibility to older project code.
    """

    def __init__(self, project):
        super(LegacyClusterGenerator, self).__init__(project)

    def count_emitters(self, model, index):
        """
        redirect the call to the corresponding project method if implemented.
        """
        try:
            return self.project.count_emitters(model, index)
        except AttributeError:
            return 1

    def create_cluster(self, model, index):
        """
        redirect the call to the corresponding project method.
        """
        return self.project.create_cluster(model, index)
