#Sebastian Raschka 2014
# Script that calculates the weighted center of mass of PDB files.
# Usage python3 center_of_mass.py file.pdb

# Atomic masses
# list of atomic weights from http://en.wikipedia.org/wiki/List_of_elements


ATOMIC_WEIGHTS = {'H': 1.008, 'HE': 4.002602, 'LI': 6.94, 'BE': 9.012182,
                  'B': 10.81, 'C': 12.011, 'N': 14.007, 'O': 15.999, 'F': 18.9984032,
                  'NE': 20.1797, 'NA': 22.98976928, 'MG': 24.305, 'AL': 26.9815386,
                  'SI': 28.085, 'P': 30.973762, 'S': 32.06, 'CL': 35.45, 'AR': 39.948,
                  'K': 39.0983, 'CA': 40.078, 'SC': 44.955912, 'TI': 47.867, 'V': 50.9415,
                  'CR': 51.9961, 'MN': 54.938045, 'FE': 55.845, 'CO': 58.933195,
                  'NI': 58.6934, 'CU': 63.546, 'ZN': 65.38, 'GA': 69.723, 'GE': 72.630,
                  'AS': 74.92160, 'SE': 78.96, 'BR': 79.904, 'RB': 85.4678, 'SR': 87.62,
                  'Y': 88.90585, 'ZR': 91.224, 'NB': 92.90638, 'MO': 95.96, 'TC': 98,
                  'RU': 101.07, 'RH': 102.90550, 'PD': 106.42, 'AG': 107.8682, 'CD': 112.411,
                  'IN': 114.818, 'SN': 118.710, 'SB': 121.760, 'TE': 127.60, 'I': 126.90447,
                  'XE': 131.293, 'CS': 132.9054519, 'BA': 137.327, 'LA': 138.90547,
                  'CE': 140.116, 'PR': 140.90765, 'ND': 144.242, 'PM': 145, 'SM': 150.36,
                  'EU': 151.964, 'GD': 157.25, 'TB': 158.92535, 'DY': 162.500, 'HO': 164.93032,
                  'ER': 167.259, 'TM': 168.93421, 'YB': 173.054, 'LU': 174.9668, 'HF': 178.49,
                  'TA': 180.94788, 'W': 183.84, 'RE': 186.207, 'OS': 190.23, 'IR': 192.217,
                  'PT': 195.084, 'AU': 196.966569, 'HG': 200.592, 'TL': 204.38, 'PB': 207.2,
                  'BI': 208.98040, 'PO': 209, 'AT': 210, 'RN': 222, 'FR': 223, 'RA': 226,
                  'AC': 227, 'TH': 232.03806, 'PA': 231.03588, 'U': 238.02891, 'NP': 237,
                  'PU': 244, 'AM': 243, 'CM': 247, 'BK': 247, 'CF': 251, 'ES': 252, 'FM': 257,
                  'MD': 258, 'NO': 259, 'LR': 262, 'RF': 267, 'DB': 268, 'SG': 269, 'BH': 270,
                  'HS': 269, 'MT': 278, 'DS': 281, 'RG': 281, 'CN': 285, 'UUT': 286, 'FL': 289,
                  'UUP': 288, 'LV': 293, 'UUS': 294}


def center_of_mass(pdbfile, include='ATOM,HETATM'):
  """
  Calculates center of mass of a protein and/or ligand structure.
  Returns:
      center (list): List of float coordinates [x,y,z] that represent the
      center of mass (precision 3).
  """

  center = [None, None, None]
  include = tuple(include.split(','))

  with open(pdbfile, 'r') as pdb:

    # extract coordinates [ [x1,y1,z1], [x2,y2,z2], ... ]
    coordinates = []
    masses = []
    for line in pdb:
      if line.startswith(include):
        coordinates.append([float(line[30:38]),    # x_coord
                            float(line[38:46]),    # y_coord
                            float(line[46:54])     # z_coord
                            ])
        element_name = line[76:].strip()
        if element_name not in ATOMIC_WEIGHTS:
          element_name = line.split()[2].strip()[0]
        masses.append(ATOMIC_WEIGHTS[element_name])

    assert len(coordinates) == len(masses)

    # calculate relative weight of every atomic mass
    total_mass = sum(masses)
    weights = [float(atom_mass / total_mass) for atom_mass in masses]

    # calculate center of mass
    center = [sum([coordinates[i][j] * weights[i]
                   for i in range(len(weights))]) for j in range(3)]
    center_rounded = [round(center[i], 3) for i in range(3)]
    return center_rounded


if __name__ == '__main__':

  import argparse

  parser = argparse.ArgumentParser(
      description='Calculates the weighted center of mass for structures in a PDB file.\n'
                  'By default, all atoms in the PDB file are included in the calculation.',
      epilog='Example:\n'
      'center_of_mass.py ~/Desktop/protein.pdb \n'
      '[-8.125, 20.461, -10.438]'
      '\n\nNote that for the center of mass calculation, the relative\natomic'
      ' weights are taken into account (atomic mass unit [u]).\n\n'
      'A list of the atomic weights can be found, e.g., at\n'
      'http://en.wikipedia.org/wiki/List_of_elements',
      formatter_class=argparse.RawTextHelpFormatter
  )

  parser.add_argument('PDBfile')
  parser.add_argument('-i', '--include', type=str,
                      default='ATOM,HETATM',
                      metavar='coordinate-ID',
                      help='Coordinate lines to include (default: "ATOM,HETATM")')

  parser.add_argument('-v', '--version', action='version', version='center_of_mass v. 1.0')

  args = parser.parse_args()

  center_of_mass = center_of_mass(args.PDBfile, include=args.include)

  for component in center_of_mass:
    print(component)
