# Date: August 2017
# Author: Kutay B. Sezginel
"""
Mean squared displacement calculation for Lammps trajectory.
"""
import numpy as np
import periodictable


def center_of_mass(atoms, coordinates):
    """ Calculate center of mass for given coordinates and atom names of a single frame """
    masses = np.array([periodictable.elements.symbol(atom).mass for atom in atoms])
    total_mass = masses.sum()
    x_cm = (masses * np.array([i[0] for i in coordinates])).sum() / total_mass
    y_cm = (masses * np.array([i[1] for i in coordinates])).sum() / total_mass
    z_cm = (masses * np.array([i[2] for i in coordinates])).sum() / total_mass
    return [x_cm, y_cm, z_cm]


def mean_squared_displacement(pos_data, dt=1):
    """ Calculate MSD for single dimension """
    msd_sum = 0
    n_points = len(pos_data) - dt
    for i in range(n_points):
        msd_sum += (pos_data[i + dt] - pos_data[i]) ** 2
    return (msd_sum / n_points)


def get_msd(trajectory, dt=1):
    """ Calculate MSD for given trajectory """
    msd = []
    for direction in range(3):
        pos_data = [i[direction] for i in trajectory['com']]
        msd.append(mean_squared_displacement(pos_data, dt))
    return msd


def msd_distance(coordinates, frames='all', atom=0):
    """ Calculate MSD for single atom using a reference frame and given frames
        - coordinates: list of coordinates for each frame
        - frames: frames to use for calculation ('all' or tuple -> (start, end))
        - atom: atom index to use for calculation
    """
    if frames == 'all':
        frames = (0, len(coordinates))
    d_sum = 0
    n_frames = len(coordinates[frames[0]:frames[1]])
    for frame in range(1, n_frames):
        coor_ref = coordinates[frame - 1][atom]         # Atom position at t -> r(t)
        coor = coordinates[frame][atom]                 # Atom position at t + 1 => r(t + 1)
        d_sum += np.sqrt(((coor[0] - coor_ref[0]) ** 2 +
                          (coor[1] - coor_ref[1]) ** 2 +
                          (coor[2] - coor_ref[2]) ** 2))
    return (d_sum / n_frames)

@@ -279,16 +279,222 @@
 le frame
+%0A%0A    Args:%0A        - atoms (list): List of element names%0A        - coordinates (list): List of coordinates (2D list)%0A%0A    Returns:%0A        - list: Center of mass coordinate for list of atom coordinates%0A   
  %22%22%22%0A   

Add docstring for center_of_mass function.