#### Edit ninfo file with given weight factors
#### Usage: $ python3 edit_ninfo.py [2,2,2,2,2,2] ninfo_file
#### list order of scale factors [bond, angle, dihedral_1, dihedral_2, basepair, basestack]
#### a new ninfo file is produced with a name ninfo_file_mod

import sys

if len( sys.argv ) < 3:
    print("provide scale_factor and ninfo_file")
    sys.exit()

scale_factors = sys.argv[1]
input_file = sys.argv[2]
out_file = input_file.split(".")[0]+"_mod.ninfo"
[bd_factor, ba_factor, dih_factor1, dih_factor2, bp_factor, bs_factor] = map(float, scale_factors.strip('[]').split(','))
# print(scale_factors)
with open(input_file,'r') as fin, open(out_file, 'w') as fout :
    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("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)
