#!/usr/bin/env python

"""Calculation of RMSD of ADPs between 2 molecular structures in PDB format
and their average ADP.

Usage: python3 adp_rmsd.py first.pdb second.pdb

Both PDB files must contain the same number of atoms.

ADP_RMSD_raw is calculated using formula sqrt{{\sum_{i=0}^{n} [B_1i - B_2i]^2}/{N}}
ADP_RMSD_central is calculated using formula sqrt{{\sum_{i=0}^{n} [B_1i - mean{B_1} - B_2i + mean{B_2}]^2}/{N}}
"""

__author__ = "Martin Maly, Petr Kolenko"
__license__ = "Creative Commons Attribution 4.0 International License"
__email__ = "malymar9@fjfi.cvut.cz"

import math
import sys
import statistics

array = [[], []]
for i in range(2):
    with open(sys.argv[i + 1], 'r') as pdb:
        lines = pdb.readlines()
    after_CRYST1 = False
    for j in range(len(lines)):
        if 'CRYST1' in lines[j]:
            after_CRYST1 = True
        if after_CRYST1:
            if ('ATOM' in lines[j] or 'HETATM' in lines[j]):
                array[i].append(float(lines[j][60:66]))

assert len(array[0]) == len(array[1])
adp_avrg0 = statistics.mean(array[0])
adp_avrg1 = statistics.mean(array[1])
lenght = len(array[0])
adp_diff_sum_raw = 0
adp_diff_sum_central = 0
for j in range(lenght):
    adp_diff_sum_raw += ((array[0][j] - adp_avrg0 - array[1][j] + adp_avrg1) ** 2)
    adp_diff_sum_central += ((array[0][j] - array[1][j]) ** 2)
adp_diff_avrg_raw = math.sqrt(float(adp_diff_sum_raw/lenght))
adp_diff_avrg_central = math.sqrt(float(adp_diff_sum_central/lenght))

print("ADP_RMSD_raw: " + str('{:.3f}'.format(round(adp_diff_avrg_raw, 3))) +
      "\tADP-factor_RMSD_central: " + str('{:.3f}'.format(round(adp_diff_avrg_central, 3))) +
      "\tmean_ADP_PDB1: " + str('{:.3f}'.format(round(adp_avrg0, 3))) +
      "\tmean_ADP_PDB2: " + str('{:.3f}'.format(round(adp_avrg1, 3))))