#!/usr/bin/env python
# PDBStructureFile.py
#   Copyright (C) 2006 CCLRC, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# A handler for PDB files for macromolecular structures. This is to
# enable analysis programs.
#

import os

residue_letters = { }

residue_letters['ALA'] = 'A'
residue_letters['ARG'] = 'R'
residue_letters['ASN'] = 'N'
residue_letters['ASP'] = 'D'
residue_letters['CYS'] = 'C'
residue_letters['GLN'] = 'Q'
residue_letters['GLU'] = 'E'
residue_letters['GLY'] = 'G'
residue_letters['HIS'] = 'H'
residue_letters['ILE'] = 'I'
residue_letters['LEU'] = 'L'
residue_letters['LYS'] = 'K'
residue_letters['MET'] = 'M'
residue_letters['PHE'] = 'F'
residue_letters['PRO'] = 'P'
residue_letters['SER'] = 'S'
residue_letters['THR'] = 'T'
residue_letters['TRP'] = 'W'
residue_letters['TYR'] = 'Y'
residue_letters['VAL'] = 'V'

def parse_pdb_remark_200(pdb_file_records):
  '''Look for REMARK 200 data and transform into a dictionary.'''

  j = 0

  results = []

  while j < len(pdb_file_records):
    record = pdb_file_records[j]
    if not 'REMARK 200' in record:
      j += 1
      continue

    if ':' in record:
      pair = record.replace('REMARK 200', '').strip(
          ).split(':')
      results.append((pair[0].strip(), pair[1].strip()))

    j += 1

  return results

found = []

def compute(pdb_file):
  '''Compute some interesting stuff from the pdb file.'''
  data = open(pdb_file).readlines()
  remarks = parse_pdb_remark_200(data)

  dmin = 0.0
  multiplicity = 0.0
  completeness = 0.0
  r_merge = 0.0
  i_sigma = 0.0

  for remark in remarks:
    if 'HIGHEST RESOLUTION SHELL, RANGE HIGH (A)' in \
       remark[0] and not 'NULL' in remark[1]:
      dmin = float(remark[1])
    if 'COMPLETENESS FOR SHELL     (%)' in \
       remark[0] and not 'NULL' in remark[1]:
      completeness = float(remark[1])
    if 'DATA REDUNDANCY IN SHELL' in \
       remark[0] and not 'NULL' in remark[1]:
      multiplicity = float(remark[1])
    if '<I/SIGMA(I)> FOR SHELL' in \
       remark[0] and not 'NULL' in remark[1]:
      i_sigma = float(remark[1])
    if 'R MERGE FOR SHELL' in \
       remark[0] and not 'NULL' in remark[1]:
      r_merge = float(remark[1])
    if 'R SYM FOR SHELL' in \
       remark[0] and not 'NULL' in remark[1]:
      r_merge = float(remark[1])

  if dmin > 0.0 and completeness > 0.0 and multiplicity > 0.0 \
         and i_sigma > 0.0 and r_merge > 0.0 and i_sigma < 200:
    pdb = os.path.split(pdb_file)[-1][:4]
    datum = (dmin, completeness, multiplicity,
             i_sigma, r_merge)
    if not datum in found:
      found.append(datum)

    # print '%4s %4.2f %5.1f %4.1f %5.1f %6.4f' % \
    # (pdb, dmin, completeness, multiplicity,
    # i_sigma, r_merge)

def analyse():
  '''Analyse all of the data already collected.'''

  for f in found:
    dmin, completeness, multiplicity, i_sigma, r_merge = f

    predicted_r = 0.7979 / i_sigma
    print '%4.2f %5.1f %4.1f %5.1f %6.4f %6.4f %4.2f' % \
          (dmin, completeness, multiplicity,
           i_sigma, r_merge, predicted_r, r_merge /
           predicted_r)

def read_modifications(pdb_file):

  modifications = { }

  for record in open(pdb_file):
    if not 'MODRES' in record[:6]:
      continue

    modifications[record.split()[2]] = record.split()[5]

  return modifications

def read_sequence(pdb_file):
  '''Read the sequence from the SEQRES records and return as a dictionary.'''

  sequences = { }
  chains = []
  lengths = { }

  modifications = read_modifications(pdb_file)

  for record in open(pdb_file):
    if not 'SEQRES' in record[:6]:
      continue

    chain = record.split()[2]
    residues = record.split()[4:]

    if not chain in chains:
      chains.append(chain)

    if not chain in sequences:
      sequences[chain] = ''

    for r in residues:
      if r in modifications:
        r = modifications[r]
      sequences[chain] += residue_letters[r]

    if not chain in lengths:
      lengths[chain] = int(record.split()[3])

  for chain in chains:
    assert(len(sequences[chain]) == lengths[chain])

  return sequences

if __name__ == '__main__' and False:

  import sys

  if len(sys.argv) < 2:
    raise RuntimeError, '%s pdb_file ...' % sys.argv[0]

  for arg in sys.argv[1:]:
    compute(arg)

  analyse()

if __name__ == '__main__':

  import sys

  if len(sys.argv) < 2:
    raise RuntimeError, '%s pdb_file ...' % sys.argv[0]

  for arg in sys.argv[1:]:
    sequences = read_sequence(arg)
    for k in sorted(sequences):
      sequence = sequences[k]
      print k
      print sequence