
from chempy import Storage,Atom
from chempy.models import Indexed
import string

class PDB(Storage):
   
#---------------------------------------------------------------------------------
   def fromList(self,list):   # currently no handling of conect records 

      model = Indexed()
      
# read atoms
      cnt = 0
      at = None
      for rec in list:
         if (rec[0:4]=='ATOM') or (rec[0:6]=='HETATM'):
            at = Atom()
            if rec[0]=='A': at.hetatm=0  # default is 1
            at.index = cnt
            at.name = string.strip(rec[12:16])
            at.alt = string.strip(rec[16:17])
            at.resn = string.strip(rec[17:20])
            at.chain = string.strip(rec[21:22])
            at.resi = string.strip(rec[22:27]) # note: insertion is part of resi
            at.resi_number = int(rec[22:26])
            at.coord = [float(rec[30:38]), 
                        float(rec[38:46]),
                        float(rec[46:54])]
            try:
               at.q = float(rec[54:60])
            except ValueError:
               at.q = 1.0
            try:               
               at.b = float(rec[60:66])
            except ValueError:
               at.b = 0.0
            at.segi = string.strip(rec[72:76])
            at.symbol = string.strip(rec[76:78])
            if not len(at.symbol):
               at.symbol = at.name[0:1]
               if at.symbol in '012345678':
                  at.symbol = at.name[1:2]
            cnt = cnt + 1
            model.add_atom(at)
         elif (rec[0:3]=='TER'):
            if at:
               at.ter=1
      return(model)

#---------------------------------------------------------------------------------
   def toList(self,model):

      list = []

      cnt = 1
      for at in model.atom:
         if at.hetatm:
            het = 'HETATM'
         else:
            het = 'ATOM  '
         latn = len(at.name)
         if latn and (latn<4):
            if at.name[0] in '0123456789':
               name=at.name
            else:
               name=' '+at.name
         else:
            name=at.name
         lrsi = len(at.resi)
         if lrsi and (lrsi<5):
            if at.resi[lrsi-1] in '0123456789':
               resi=at.resi+' '
            else:
               resi=at.resi
         list.append(
            "%6s%5i %-4s%1s%3s %1s%5s   %8.3f%8.3f%8.3f%6.2f%6.2f      %-4s%2s\n" % 
            (het,cnt,name,at.alt,at.resn,at.chain,resi,
            at.coord[0],at.coord[1],at.coord[2],at.q,at.b,at.segi,at.symbol))
         if hasattr(at,'ter'):
            if at.ter:
               cnt = cnt + 1
               list.append(
                  "%6s%5i      %3s %1s%4s\n" %
                  ('TER   ',cnt,at.resn,at.chain,at.resi))
         cnt = cnt + 1
      list.append("END\n")
      return(list)

