import os
import iotbx.pdb
from scitbx.array_family import flex

def combine_pdbs(pdb_file_names):
  print len(pdb_file_names)
  r = iotbx.pdb.hierarchy.root()
  for i_seq, f in enumerate(pdb_file_names):
    pdb_inp = iotbx.pdb.input(file_name = f)
    h = pdb_inp.construct_hierarchy()
    models = h.models()
    if(len(models)==1):
      mdc = models[0].detached_copy()
      mdc.id = str(i_seq)
      r.append_model(mdc)
    #if i_seq == 100: break
  r.atoms().set_occ(flex.double(len(list(r.atoms())), 0))
  r.write_pdb_file(file_name = "all.pdb")
  #sel = r.atom_selection_cache().selection(
  #  string = "name N or name C or name CA")
  #r = r.select(sel)
  #r.write_pdb_file(file_name = "all_overfit_bb.pdb")
  #r.write_pdb_file(file_name = "all_fem.pdb")
  #r.write_pdb_file(file_name = "all_fem.pdb")
      
def run():
  pdb_file_names = []
  for file_name in os.listdir("."):
    if(file_name.endswith(".pdb")):
      pdb_file_names.append(file_name)
  combine_pdbs(pdb_file_names = pdb_file_names)
  
if (__name__ == "__main__"):
  run()