# -*- coding: utf-8 -*-
"""
Created on Mon Apr 10 17:03:16 2023

@author: FeiMao
"""
from contextlib import redirect_stdout
###############################################################################
#specify input file names 
###############################################################################
ncacx_filename = "ncacx4nextrd_Talos.txt" #pls update the ncacx data file name
num_lines = 230 #ncocx_filename = "ncocx_b.txt" #pls update the ncacx data file name
protein_seq = "RSV_seq.txt"
###############################################################################
asgn_filename = "TalosReference4RSVNoAmbi.tab"
################################################################################
###############################################################################
#define some functions frequently used
###############################################################################
#function to extract only the letter part of a residue type asgn, not suitable 
#if you want to extract a single letter from an ambiguous residue type asgm
def extract_letter(rtyp):
    #import re
    kresname=''
    for i in rtyp:
        if i.isalpha():
            kresname = "".join([kresname, i])
    return kresname
#end of extract letter function definition

#function to extract only the number part of a residue type asgn
def extract_number(rtyp):
    #import re
    kresnum=''
    for i in rtyp:
        if i.isdigit():
            kresnum = "".join([kresnum, i])
    return kresnum
#end of extract number function

#end of frequent used function definition
###############################################################################
###############################################################################
#read in protein sequence
###############################################################################
print('Read in the protein sequence from ',protein_seq,'.')
import re
protein=[]
protein_aa=''
with open(protein_seq, "r") as file1:  
    count = 0    
    lines=file1.readlines()
   
    
    for line in lines:
        for letter in line:
            if letter.isalpha():
                protein.append(letter) 
                protein_aa = protein_aa + letter
    protein.insert(0,'')
    nres=len(protein)-1
    print('The number of residue in protein is ',nres)
    print(protein)
#Read in ncacx
###############################################################################
res_id = [0 for i in range(num_lines+1)]
c_cs = [0 for i in range(num_lines+1)]
ca_cs = [0 for i in range(num_lines+1)]
cb_cs = [0 for i in range(num_lines+1)]
n15_cs = [0 for i in range(num_lines+1)]
print('Read ncacx input from ',ncacx_filename,'.')
with open(ncacx_filename, "r") as file1:  
    count = 0
    print(file1)
    for line in file1: 
        
        #Read in the first row, how many residue type asgn in ncacx and
        #how many fre columns in each residue type asgn
        if count == 2:
            kline2=line.split()
        if count <= num_lines and count> 0:
            kline_old=line.split()
            kline=kline_old#kline_old[0].split()
            if len(extract_number(kline[9])) > 0 :
                temp_resid = int(extract_number(kline[9]))
            else:
                temp_resid = 0 
            res_id[temp_resid]=temp_resid
            
            if float(kline[2]) < 200.0:
                c_cs[temp_resid] = float(kline[2])
            if float(kline[0]) < 200.0: 
                n15_cs[temp_resid] = float(kline[0])
            if float(kline[1]) < 200.0:
                ca_cs[temp_resid] = float(kline[1])
            if float(kline[3]) < 200.0: 
                cb_cs[temp_resid] = float(kline[3])
        count += 1
        
            
print('ncacx completes.')     
print('\n')       
#end of read in ncacx
###############################################################################
with open(asgn_filename, mode='w') as f: 
    with redirect_stdout(f):
        #print('istep is ',istep)
        for i in range(0,num_lines):
            print(res_id[i],'  ',protein[i],'  C   ',c_cs[i])
            print(res_id[i],'  ',protein[i],'  CA  ',ca_cs[i])
            print(res_id[i],'  ',protein[i],'  CB  ',cb_cs[i])
            print(res_id[i],'  ',protein[i],'  N  ',n15_cs[i])
        