import logging
import argparse
from Bio.SeqRecord import SeqRecord
from Bio import SeqIO
from Bio.PDB import PDBParser
from Bio.PDB import PDBIO
from Bio.PDB.Atom import PDBConstructionWarning


pdb=('1kxk.pdb')
chain_ids = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U',
             'V', 'W', 'X', 'Y', 'Z']
argument = '1-10,11-20,23-30'
range_set= str(argument)
range_set=range_set.split(',')
range_set_list = []

for i in range_set:
    i = i.replace('-',',')
    #range_set_list = [i]
    i=i.split(',')
    range_set_list.append(i)
    #print i

print range_set_list
#print argument
struc = PDBParser().get_structure('struc', pdb)

def change_chains_ids(struc,chain_ids,range_set_list):
    r_ids =[]
    ch_ids = []
    for s in struc:
        for c in s:
            #print c
 #           for r in c:
            for r,ntr, ids_ch in zip(c,range_set_list,chain_ids):
                    #print int(ntr[0])
                    #print ids_ch,ntr
                for r.id in range(int(ntr[0]),int(ntr[1])):
                    c.id = ids_ch
                    #r_ids.append(r.id)
                    #ch_ids.append(c.id)
                    #print r.id,c.id
        struc1=struc
        #print struc1.get_chains(), r_ids
    return ch_ids, struc1

                        #for i in c.id:

# main
if __name__ == '__main__':
    print type(change_chains_ids(struc, chain_ids, range_set_list))
    struc1 = change_chains_ids(struc,chain_ids,range_set_list)
    print struc1

    for chain in struc1.get_chains():
        io.set_structure(chain)
        io.save(pdb.get_id() + "_" + chain.get_id() + ".pdb")

    #nowa = change_chains_ids(struc, chain_ids, range_set_list)[1]
    #io = PDBIO()
    #io.set_structure(nowa)
    #io.save("out.pdb")

#new_struc = PDBParser().set_structure(change_chains_ids(struc, chain_ids, range_set_list).set_structure())

    #new_struc = PDBParser().get_structure("new_str",change_chains_ids(struc, chain_ids, range_set_list))