#### Edit ninfo file with given weight factors
#### Usage: $ python3 edit_ninfo.py [2,2,2,2,2,2,2] monomer.ninfo restraints.txt cg_pdb 
#### list order of scale factors [bond, angle, dihedral_1, dihedral_2,long_range, basepair, basestack]
#### Value of each scale factor range (0:9)
#### a new ninfo file is produced with a name ninfo_file_mod

import sys
from configparser import ConfigParser

def scale_Coef(input_file, out_file, scale_factors, contact_info):
  [bd_factor, ba_factor, dih_factor1, dih_factor2,long_factor, bp_factor, bs_factor] = map(float, scale_factors.strip('[]').split(','))
  total_contact = get_totalContact(input_file) ## Total number of tertiary contacts in ninfo file
  with open(input_file,'r') as fin, open(out_file, 'w') as fout :
    count_contact = 0
    for line in fin:
        if line.startswith("bond"):
            item1_list = [line.split()[0]]
            item2_list = [item.rjust(7) for item in line.split()[1:8]]
            item3_list = [item.rjust(13) for item in line.split()[8:11]]
            item4 = '%.4f'% (float(line.split()[11])*bd_factor)
            item4_list = [item4.rjust(13)]
            item5_list = [line.split()[12].rjust(3)]
            sum_list = item1_list+item2_list+item3_list+item4_list+item5_list
            sum_string = "".join(sum_list)
            fout.write(sum_string+"\n")
        elif line.startswith("angl"):
            item1_list = [line.split()[0]]
            item2_list = [item.rjust(7) for item in line.split()[1:10]]
            item3_list = [item.rjust(13) for item in line.split()[10:13]]
            item4 = '%.4f'% (float(line.split()[13])*ba_factor)
            item4_list = [item4.rjust(13)]
            item5_list = [line.split()[14].rjust(4)]
            sum_list = item1_list+item2_list+item3_list+item4_list+item5_list
            sum_string = "".join(sum_list)
            fout.write(sum_string+"\n")
        elif line.startswith("dihd"):
            item1_list = [line.split()[0]]
            item2_list = [item.rjust(7) for item in line.split()[1:12]]
            item3_list = [item.rjust(13) for item in line.split()[12:15]]
            item4 = '%.4f'% (float(line.split()[15])*dih_factor1)
            item5 = '%.4f'% (float(line.split()[16])*dih_factor2)
            item4_list = [item4.rjust(13)]
            item5_list = [item5.rjust(13)]
            item6_list = [line.split()[17].rjust(5)]
            sum_list = item1_list+item2_list+item3_list+item4_list+item5_list+item6_list
            sum_string = "".join(sum_list)
            fout.write(sum_string+"\n")
        elif line.startswith("contact"):
            item1_list = [line.split()[0]]
            item2_list = [item.rjust(7) for item in line.split()[1:8]]
            item3_list = [item.rjust(13) for item in line.split()[8:10]]
            item4_list = [item.rjust(7) for item in line.split()[10]]
            item5 = '%.4f'% (float(line.split()[11])*long_factor)
            item5_list = [item5.rjust(13)]
            item6_list = [line.split()[12].rjust(4)]
            sum_list = item1_list+item2_list+item3_list+item4_list+item5_list+item6_list
            sum_string = "".join(sum_list)
            fout.write(sum_string+"\n")
            count_contact += 1
            if (count_contact == total_contact):
                [res_i, res_j, dist, weight, type_i, type_j] = contact_info
                fmt_contact = '%-7s%7d%7d%7d%7d%7d%7d%7d%13.4f%13.4f%7d%13.4f%4s'
                for rank in range(len(res_i)):
                    count_contact += 1
                    fout.write(fmt_contact  % ("contact",count_contact,1,1,res_i[rank], res_j[rank],
                                               res_i[rank],res_j[rank], float(dist[rank]), 1.0,1,
                                               float(weight[rank]),type_i[rank]+"-"+type_j[rank]))
                    fout.write("\n")


        elif line.startswith("basepair"):
            item1_list = [line.split()[0]]
            item2_list = [item.rjust(7) for item in line.split()[1:8]]
            item3_list = [item.rjust(13) for item in line.split()[8:10]]
            item4_list = [line.split()[10].rjust(7)]
            item5 = '%.4f'% (float(line.split()[11])*bp_factor)
            item5_list = [item5.rjust(13)]
            item6_list = [line.split()[12].rjust(5)]
            item7_list = [line.split()[13].rjust(2)]
            sum_list = item1_list+item2_list+item3_list+item4_list+item5_list+item6_list+item7_list
            sum_string = "".join(sum_list)
            fout.write(sum_string+"\n")
        elif line.startswith("basestack"):
            item1_list = [line.split()[0]]
            item2_list = [item.rjust(7) for item in line.split()[1:8]]
            item3_list = [item.rjust(13) for item in line.split()[8:10]]
            item4_list = [line.split()[10].rjust(7)]
            item5 = '%.4f'% (float(line.split()[11])*bs_factor)
            item5_list = [item5.rjust(13)]
            item6_list = [line.split()[12].rjust(5)]
            sum_list = item1_list+item2_list+item3_list+item4_list+item5_list+item6_list
            sum_string = "".join(sum_list)
            fout.write(sum_string+"\n")

        else:
            fout.write(line)
  return

def get_restraints(res_file, ninfo_file, pdb_name):
    """
    Parse the restraints file and format the ninfo
    """
    parser = ConfigParser()
    parser.read(res_file)
    for name, value in parser.items("Interaction Types"):
        if name == "tertiary":
            list_items = value.lstrip().splitlines()
            ## Initialize a list of lists
            features = [[] for i in range(len(list_items[0].split()))]
            for lines in list_items[1:len(list_items)]:
                for rank in range(0,len(features)):
                    features[rank].append(lines.split()[rank])
            ## assign name to elemens in the features list
            [res_i, res_j, dist, weight, type_i, type_j] = [i for i  in features]
            #print(res_i)
    ## Get the residue number to name map
    res_map = map_resid2name(pdb_name)
    for rank in range(len(res_i)):
        if ((int(res_i[rank]) == 1 and type_i[rank] == "P") or
        (int(res_j[rank]) == 1 and type_j[rank] == "P")):
            print("No phosphate in first residue, please review restraints file")
            sys.exit()
        if (type_i[rank] == "P"):
            res_i[rank] = (int(res_i[rank])-1)*3
        if (type_j[rank] == "P"):
            res_j[rank] = (int(res_j[rank])-1)*3
        if (type_i[rank] == "S"):
            res_i[rank] = (int(res_i[rank])-1)*3+1
        if (type_j[rank] == "S"):
            res_j[rank] = (int(res_j[rank])-1)*3+1
        if (type_i[rank] == "B"):
            type_i[rank] = list(res_map[int(res_i[rank])])[1]
            res_i[rank] = (int(res_i[rank])-1)*3+2
        if (type_j[rank] == "B"):
            type_j[rank] = list(res_map[int(res_j[rank])])[1]
            res_j[rank] = (int(res_j[rank])-1)*3+2
    return [res_i, res_j, dist, weight, type_i, type_j]

def get_totalContact(ninfo_file):
    """Count the number of lines in ninfo file starting with term 'contact' """
    with open(ninfo_file,'r') as fin:
        count_contact = 0
        line_count = 0
        for line in fin:
            line_count += 1
            if line.startswith("contact"):
                count_contact +=1
    return count_contact


def map_resid2name(pdb_name):
    """
read xyz from pdb
generate a resid to resname dictionary
    """
    xyz = []
    pdb_file = open(pdb_name, 'r')
    for line in pdb_file:
        n = line[12:16].strip()
        if (line.startswith("ATOM")):
            # extract n (atom name), rn (residue name), r (residue no.),x, y, z coordinates
            rn = line[17:20].strip()
            r = int(line[22:26].strip())
            x = float(line[30:38].strip())
            y = float(line[38:46].strip())
            z = float(line[46:54].strip())
            #if line[12:16].strip() == "CA":
            xyz.append([n, rn, r, x, y, z])
    pdb_file.close()
    # define a residue dictionary
    res_dis = {}
    first_res = xyz[0][2]
    res_dis[first_res] = xyz[0][1]
    #print(first_res)
    for item in xyz:
        if (item[2] != first_res):
            res_dis[item[2]] = item[1]
        first_res = item[2]
    #print(res_dis)
    return res_dis


if __name__ == '__main__':
  if len( sys.argv ) < 5:
      print("provide scale_factor, ninfo_file, restraint_file, and cg_pdbfile")
      sys.exit()

  scale_factors = sys.argv[1]
  input_file = sys.argv[2]  ## this is the ninfo file
  out_file = input_file.split(".")[0]+"_mod.ninfo"
  res_file = sys.argv[3]
  pdb_name = sys.argv[4]
  ## Scale the coefficient weights
  contact_info = get_restraints(res_file, input_file, pdb_name)
  #print(contact_info)
  scale_Coef(input_file, out_file, scale_factors, contact_info)

  #print(xyz[1][1])
