#!/usr/bin/python3

#=======================================================================
# TASK: extract chain from pdb and save each chain as a separate file

# link for saving each chain as a separate file
# https://stackoverflow.com/questions/11685716/how-to-extract-chains-from-a-pdb-file
# command line args
# https://stackoverflow.com/questions/15753701/how-can-i-pass-a-list-as-a-command-line-argument-with-argparse
#=======================================================================   

#%%
import os, sys
from Bio.PDB import Select, PDBIO
from Bio.PDB.PDBParser import PDBParser
#%% homdir and curr dir and local imports
homedir = os.path.expanduser('~') 
# set working dir
os.getcwd()
os.chdir(homedir + '/git/LSHTM_analysis/scripts')
os.getcwd()
#%%
class ChainSelect(Select):
    def __init__(self, chain):
        self.chain = chain

    def accept_chain(self, chain):
        if chain.get_id() == self.chain:
            return 1
        else:          
            return 0

def main(): 
    chains = ['A','B','C','F']
    p = PDBParser(PERMISSIVE=1)       
    #structure = p.get_structure(pdb_file, pdb_file)
    structure = p.get_structure('/home/tanu/git/Data/ethambutol/input/3byw', '/home/tanu/git/Data/ethambutol/input/3byw.pdb')
    #print('STRUCTURE:', structure.get_id())
#    pdb_filename = print()
    for chain in chains:
        pdb_chain_file = 'pdb_file_chain_{}.pdb'.format(chain)         
                        
        io = PDBIO()               
        io.set_structure(structure)
        io.save('{}'.format(pdb_chain_file), ChainSelect(chain))

# If run from command line...
if __name__ == "__main__":
    main()