#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Mar 17 22:59:55 2017

Calculates center of mass iteratively for a bunch of frames in a gromacs .gro file.
You have to modify some values(described below) according to your system.

@author: Sajeewa Walimuni Dewage (a.k.a.   Sajeewa Pemasinghe :) )

"""

import numpy as np

def is_number(s): # handy function to make sure whether a given string corresponds to a number
    try:
        float(s)
        return True
    except ValueError:
        return False

def assign_mass(s):
    if (s=="C"):
        return 12.0107
    elif (s=="H"):
        return 1.00794
    elif (s=="N"):
        return 14.0067
    elif (s=="O"):
        return 15.9994
    elif (s=="P"):
        return 30.973762
    elif (s=="S"):
        return 32.065

def calc_com(masses,xcoords,ycoords,zcoords):
    xproduct=0.0
    yproduct=0.0
    zproduct=0.0
    for i in range(len(masses)):
        xproduct=xproduct+masses[i]*xcoords[i]
        yproduct=yproduct+masses[i]*ycoords[i]
        zproduct=zproduct+masses[i]*zcoords[i]
    return (xproduct/sum(masses),yproduct/sum(masses),zproduct/sum(masses))


atoms_in_file = 1960 #change this number according to your system
frames_in_file = 101 #change this number according to your system       
position_x = list()
position_y = list()
position_z = list()
masslist = list()
com = np.ones((frames_in_file,3))
counter = 0

read_location = 'lysozyme_1ns.gro' #give the .gro file location
write_file = open('COMlist.txt','w') #give the output file location

file1 = open(read_location)

#reading coordinates
for i, line in enumerate (file1.readlines()):
    if(  (is_number(line[20:28].strip())) and (is_number(line[36:44].strip()))  ): # this if condition makes sure that I am reading the lines with coordinates 
        element_temp = line[10:15].strip()
        masslist.append(assign_mass(element_temp[0]))
        position_x.append(float(line[20:28].strip()))
        position_y.append(float(line[28:36].strip()))
        position_z.append(float(line[36:44].strip()))
       
file1.close()

#calculating COM for each frame
for i in range(frames_in_file):
    com[i]=calc_com(masslist[counter:counter+atoms_in_file],position_x[counter:counter+atoms_in_file],position_y[counter:counter+atoms_in_file],position_z[counter:counter+atoms_in_file])
    counter=counter+atoms_in_file

#writing COM values to file
for i in range(len(com)):
     write_file.write(str(round(com[i][0],3))+'\t'+str(round(com[i][1],3))+'\t'+str(round(com[i][2],3))+'\n')
     
write_file.close()