print('''
      
QM:MM中文引导程序ver.1.0
---------------------------------说明区---------------------------------
Generated by zxt 
程序虽好，切莫贪杯
原理不懂，bug一堆
请仔细阅读说明，完成qmmm基础学习后再使用该程序，并不对结果负责
---------------------------------------------
按键说明：  
      s：进行下一步
      b：回退到上一个选项
      q：退出当前选项支
      0：程序初始化
      qq:退出程序
--------------------------初始化选择-------------------------
1：非标准残基参数化
2：使用Amber程序进行经典动力学模拟
3：QM：MM计算
4：含金属元素的非标准残基参数

''')

import sys
import os

def re(n):
    print('---------------------------reloading------------------------\n\n\n\n\n')
    python = sys.executable
    os.execl(python, python, *sys.argv)

def g1(path, fn, pn, nn):
    with open(f'{path}/{fn}', 'w+') as wr:
        wr.write(f'''
source leaprc.protein.ff14SB
source leaprc.water.tip3p
{nn} = loadmol2 {nn}.mol2
sys = loadpdb {pn}-addH.pdb
savepdb sys {pn}-addH.pdb
quit\n''')

def g2(path, fn, pn, nn, tmps2):
    with open(f'{path}/{fn}', 'w+') as wr:
        wr.write(f'''
source leaprc.protein.ff14SB
source leaprc.water.tip3p
{nn} = loadmol2 {nn}.mol2
loadamberparams {nn}.frcmod
sys = loadpdb {pn}-addH.pdb
bond sys.{tmps2-1}.C sys.{tmps2}.N
bond sys.{tmps2}.C sys.{tmps2+1}.N
solvatebox sys TIP3PBOX 15.0
addions sys Na+ 0 
addions sys Cl- 0
savepdb sys {pn}-solv.pdb
saveamberparm sys {pn}-solv.prmtop {pn}-solv.inpcrd
quit\n''')

log0 = sys.argv

path = input('请输入临时文件输入地址，程序运行中的文件将会输出在这里，默认位置为本程序所在文件夹\n') or log0
ck1 = eval(input(
    '''请输入计算内容
            1：非标准残基参数化
            2：使用Amber程序进行经典动力学模拟
            3：QM：MM计算
            4：含金属元素的非标准残基参数\n

           '''))
n = 1
while True:
    if n == 0:
        re(n)
    else:
        #这里是算法区，按键基本功能在这里
        if n != -1:
            cz1 = input('''
      s/Enter：进行下一步
      b：回退到上一个选项
      q：退出当前选项支
      0：程序初始化
      qq:退出程序

''') or "s"
            if cz1 == '0':
                n = 0
                continue
            if cz1 == "q":
                  n=1
                  ck1=eval(input(
'''请输入计算内容
      1：非标准残基参数化
      2：使用Amber程序进行经典动力学模拟
      3：QM：MM计算
      4：含金属元素的非标准残基参数\n''') )
            if cz1 == "b":
                if n == 1:
                    print('这是第一步')
                    continue
                if n > 1:
                    n -= 1
            if cz1 == "s":
                n += 1
                print(ck1,n)
            if cz1 == "qq":
                exit()
            if cz1 not in ["b", "s", "qq", "q", "0"]:
                print('按键无效，请重新输入')
            #这里是程序输出区
            if n == 2 and ck1 == 1:
                print('-----------------------非标准残基参数化step1/4-----------------------')
                print('''
说明：
      1、这部分程序只适用于单个蛋白质的单个非标准残基，如有多个非标准残基，请多次循环使用该部分
      2、若该非标准残基存在两个以上需要封端的地方，请谨慎使用接下来使用的所有程序，或者在该步骤结束后按qq一劳永逸
                      ''')
                print('请检查你已经获得了一个pdb文件，做了删除与提取')
                porname=input("请输入蛋白质名\n") or None
                nsrname=input("请输入非标准残基名\n") or None
                print('随后请检查是否有{}.pdb和{}.pdb文件'.format(porname, nsrname))
                print('使用gview打开{}.pdb文件，补氢,修改成键方式（结构暂时先忽略），保存为{}-bond.mol2（注意文件名不能错）\n要求：不考虑共轭，每个原子键小于4即可'.format(nsrname, nsrname))
                print('-----------------------------step1 Finish---------------------------')
            if n == 3 and ck1 == 1:
                print('-----------------------非标准残基参数化step2/4-----------------------')
                print('打开{}.pdb文件，正确加氢后保留残基，初步调整结构'.format(nsrname))
                print('完成后进行封端，将封端后结果保留为{}-large.pdb和{}-large.gjf（opt freq）'.format(nsrname, nsrname))
                print('完成后输入下列指令：\nantechamber -i {}-large.pdb -fi pdb -o {}-large.mol2 -fo mol2 -pf y -at amber'.format(nsrname, nsrname))
                print('parmchk2 -i {}-large.mol2 -f mol2 -o {}.frcmod -a y'.format(nsrname, nsrname))
                print('-----------------------------step2 Finish---------------------------')
            if n == 4 and ck1 == 1:
                print('-----------------------非标准残基参数化step3/4-----------------------')
                print('将{}-large.gjf结果{}-large.log投入计算'.format(nsrname, nsrname))
                tmps1 = eval(input('含封端集团在内有多少个原子\n')) or 12
                print('请输入：\necho "{},{},0 " > cons.txt\necho -e "7\\n18\\n6\\ncons.txt\\n2\\ny\\n0\\nq" | Multiwfn {}-large.fchk'.format(tmps1-11, tmps1, nsrname))
                print('检查当前文件下有无{}-large.chg'.format(nsrname))
                print('-----------------------------step3 Finish---------------------------')
            if n == 5 and ck1 == 1:
                print('-----------------------非标准残基参数化step4/4-----------------------')
                print('检查当前文件夹下是否存在{}-large.mol，{}-large.chg，{}-bond.mol2'.format(nsrname, nsrname, nsrname))
                print('请输入：\npython ionsrp.py {}'.format(nsrname))
                print('-----------------------------step4 Finish---------------------------')
            if n==2 and ck1==2: 
                print('-----------------------使用Amber程序进行经典动力学模拟step1/3-----------------------')
                print('''
说明：本部分同样仅支持一个蛋白质与一个非标准残基''')
                if 'orname' and 'nsrname' in dir():
                    ck2=input('是否使用已有蛋白质与非标准残基？\ny/n')
                    if ck2=='y':
                        print(f'正在使用蛋白质名{porname}和非标准残基名{nsrname}')
                    if ck2=='n':
                        porname=input("请输入蛋白质名\n") or None
                        nsrname=input("请输入非标准残基名\n") or None
                else:
                    porname=input("请输入蛋白质名\n") or None
                    nsrname=input("请输入非标准残基名\n") or None
                print(f'请打开pka网站https://www.ddl.unimi.it/vegaol/propka.htm\n目录下存在{porname}.pdb,{porname}.pka的情况下输入：\npython protonation.py {porname}')
                print('-----------------------------step1 Finish---------------------------')
            if n==3 and ck1==2:
                print('-----------------------使用Amber程序进行经典动力学模拟step2/3-----------------------')
                print('正在生成leap-addH.in')
                print(f'请检查{nsrname}.mol2和{porname}.pdb中命名是否统一')
                g1(path,'leap_addH.in',porname,nsrname)
                print('请输入：\ntleap -s -f leap-addH.in > leap-addH.out\n完成后检查是否需要将TER删掉')
                print('-----------------------------step2 Finish---------------------------')
            if n==4 and ck1==2:
                print('-----------------------使用Amber程序进行经典动力学模拟step3/3-----------------------')
                tmps2=eval(input('请输入非标准残基编号'))
                g2(path,'leap-params.in',porname,nsrname,tmps2)
                print('警告：请手动修改成键原子名\n修改完后请输入\n tleap -s -f leap-params.in > leap-params.out')
                print('-----------------------------step3 Finish---------------------------')
