extract_segments

#! /usr/bin/python3

import re
import math
import os
from sys import argv

my_dir = os.getcwd()
my_file = "%s/4GBX.pdb" % my_dir

my_chains = ["A","B","D","E"]

#Get number of chains
no_chains = len(my_chains)

#Instantiate lists
in_first_res = []
in_last_res = []
segments = []

#Instantiate lists for first_res and last_res of all segments
all_first_res = []
all_last_res = []
all_centre = []

flex_res = [["45","57","77"],["51","64","91"],["91"],["92"]]

continue_var = 'y'

cter = 0

for j in my_chains:

    f = open(my_file)
    i = 0
    prev_res = ""

    for line in f:
        if "ATOM" in line and line[21] in j and line[23:26] != prev_res:
            
            current_res = line[23:26].strip()

            if i == 0:
                in_first_res.append(int(line[23:26]))

            else:
                pass

            i += 1
            prev_res = line[23:26]

        else:
            pass
    
    in_last_res.append(int(in_first_res[cter]) + i-1)

    f.close()
    cter += 1


# A function to calculate centres between flexible residues i.e. segments
def calc_centres(start,end,flex):

    start = int(start)-1
    end = int(end)+1
    
    seg_start = []
    seg_end= []
    
    # Add start and end residues to flexible residues list for calculation of segment lengths
    flex.insert(0,start)
    flex.append(end)

    # loop through flexible residues list to calculate segment lengths and centres
    for i in range (0,len(flex)-1):

        # make lists of last and first res
        seg_start.append(int(flex[i]) + 1)
        seg_end.append(int(flex[i + 1]) - 1)

    return(seg_start,seg_end)

# Store ouput from calc_centres in separate lists
for i in range(int(no_chains)):

    segments.append(calc_centres(in_first_res[i],in_last_res[i],flex_res[i][:]))

    all_first_res.append(segments[i][0])
    all_last_res.append(segments[i][1])

def make_regions(all_first_res,all_last_res,seg_comb):

    seg_comb_list = seg_comb.split()
    seg_com_list_split = [i.split(':') for i in seg_comb_list]

    custom_start = []
    custom_end = []
    custom_centre = []
    
    no_segments = int()

    for i in range(len(seg_com_list_split)):

        chain = int(seg_com_list_split[i][0])
        segment = int(seg_com_list_split[i][1])
    
        custom_start.append(my_chains[chain]+str(all_first_res[chain][segment]))
        custom_end.append(my_chains[chain]+str(all_last_res[chain][segment]))
        
    no_segments = len(custom_start)

    return(custom_start,custom_end,no_segments,chain)


seg_comb = ["0:1","0:2","0:3","1:1","1:2","1:3","0:1 0:2","0:2 0:3","1:1 1:2","1:2 1:3","0:3 1:3","2:0 3:0","2:1","3:1","2:1 3:1"]

def write_segments(first,last,seg_no):

    f = open('segment_'+str(seg_no)+'.pdb','w')
    #f2 = open(my_file)
    
    print(first)

    for i in range(len(first)):

        f2 = open(my_file)

        for line in f2:

            if "ATOM" in line and line[21] in first[i][0] and int(line[23:26]) >= int(first[i][1:]) and int(line[23:26]) <= int(last[i][1:]):

                f.write(line)

            else:
                pass

        f2.close()

    f.close()

    return
                    

for i in range(len(seg_comb)):

    first,last,no_segments,chain = make_regions(all_first_res,all_last_res,seg_comb[i])
    write_segments(first,last,i)