#!/usr/bin/python3

#
# title:   pdb_bfac.py
# summary: Alter B-factor column of PDB to values in a table
# author:  Nicholas Fitzkee (nfitzkee at chemistry.msstate.edu)
# date:    February 25, 2021
#
# Can be run with:
#
#   python pdb_bfac.py 2oed.pdb 2oed_bfac.txt new.pdb A
#
# Show colors in PyMOL with:
#
#   spectrum b, blue_white_red, minimum=0.0, maximum=100.0   

from Bio import PDB, SeqUtils
import os, sys

# Default B-factor for residues not present in the table
DEFAULT_BFAC = -1.0

def load_table(fname):
    with open(fname) as tab:
        l = tab.readline()

        result = {}
        
        while(l):
            l = l.strip()

            if not l or l[0] == '#':
                l = tab.readline()
                continue

            l = l.split('#')[0]

            idx, bfac = l.split()

            idx  = int(idx)
            bfac = float(bfac)

            if idx in result:
                print('warning: %i appears more than once in %s' %
                      (idx, fname))

            result[idx] = bfac

            l = tab.readline()

    return result
            
def residue_bfac(pdb, bftab, out=sys.stdout, chain=None):
    bfdict = load_table(bftab)
    
    s = PDB.PDBParser().get_structure('target', pdb)
    m = s[0]             # First "model" 

    my_chain = None
    
    if chain is None:
        my_chain = m.child_list[0]
    else:
        if chain == '_': chain = ' '
      
        for c in m.child_list:
            if c.id == chain:
                my_chain = m[c.id]

    if not my_chain:
        print ('Chain "%s" not found in PDB file.' % chain)
        sys.exit(1)

    c = my_chain # First "chain"

    # Iterate through all residues in chain
    for res in c.child_list:
        het, id, icode = res.id
        new_bfac = DEFAULT_BFAC
        
        if id in bfdict:
            new_bfac = bfdict[id]
            
        for atom in res.child_list:
            atom.set_bfactor(new_bfac)

    # Save modified PDB structure; give the option to write it
    # to standard output (screen) or a file
    close_out = False
    if type(out) == type('string'):
        out = open(out, 'w')
        close_out = True

    io = PDB.PDBIO()
    io.set_structure(s)
    io.save(out)

    if close_out: out.close()

if __name__ == '__main__':
    try:
        pdb   = sys.argv[1]
        bftab = sys.argv[2]
        out   = sys.argv[3]
    except:
        print('usage: %s <pdb> <table> <new_pdb> [<chain>]' % os.path.split(sys.argv[0])[1])
        print('\n'
              '  <table>    is a file containing new b-factors; two\n'
              '             columns, the first is the residue number,\n'
              '             and the second is the b-factor. # can be\n'
              '             used to add comments.\n'
              '  <chain>    is the (optional) one character chain ID;\n'
              '             if none specified, the first chain is used.')
        
        sys.exit(1)

    print(pdb, out)
    try:
        chain = sys.argv[4]
    except:
        chain = None
    
    residue_bfac(pdb, bftab, out, chain)