#!/usr/bin/python3

#=======================================================================
# TASK: select specified chains from the structure & save a cropped structure with
# the selected chains. Useful for dimer, etc modelling.

# link for saving each chain as a separate file
# https://stackoverflow.com/questions/11685716/how-to-extract-chains-from-a-structure-file
#=======================================================================
#%%
import os, sys
from Bio.PDB import PDBParser, PDBIO, Select
#%% homdir and curr dir and local imports
homedir = os.path.expanduser('~') 
# set working dir
os.getcwd()
os.chdir(homedir + '/git/Data/ethambutol/input/')
os.getcwd()
#%%
# Select() Method to return True for every chain in 'chains'
class ChainExtract(Select):
    def __init__(self, chain):
        self.chain = chain
        
    def accept_chain(self, chain):
        #print dir(chain)
        if chain.id in self.chain:
            return 1
        else:
            return 0
             
def main():
    p = PDBParser(PERMISSIVE=1)  
    structure = p.get_structure("3byw", "3byw.pdb")
    
    my_chains = ['G', 'H'] # specify selected chains
    c_names = ''.join(my_chains)
   
    pdb_chains_file = 'pdb_crop_' + c_names + '.pdb'
    
    io = PDBIO()
    io.set_structure(structure)
    #io.save(structure.get_id() + "_crop.structure", ChainExtract(my_chains))
    io.save(pdb_chains_file, ChainExtract(my_chains))
    
if __name__ == '__main__':
    main()
#%%
# test
#my_chains = ['G', 'H'] # specify selected chains
#foo = ''.join(my_chains) # to append to file name
#pdb_chains_file = '_{}.pdb'.format(my_chains)