<reponame>NatureGeorge/trRosetta2
import os
import sys
import numpy as np
import tensorflow as tf
import logging
tf.get_logger().setLevel(logging.ERROR)
from model import TrRefiner

n_chunk_default = 5
n_filter_default = 128
def parse_input_args():
    import argparse
    opt = argparse.ArgumentParser\
            (description='''AngPred: Backbone torsion angle prediction using ResNet''')
    #
    # Training config
    opt.add_argument('-train_param', dest='train_param', metavar='TRAIN_PARAM', default='param_train.json',\
                         help="json file containing all the training related parameters (e.g. learning rate, etc)")
    #
    # Model config
    opt.add_argument('-num_chunks', '--num_chunks', dest='num_chunks', metavar='num_chunks', type=int, default=n_chunk_default, \
                         help='The number of ResNet chunks [%d]'%n_chunk_default)
    opt.add_argument('-n2d_filters', '--n2d_filters', dest='n2d_filters', metavar='n2d_filters', type=int, default=n_filter_default, \
                         help='The number of 2D filters [%d]'%n_filter_default)
    #
    # Flag for training or not
    opt.add_argument('-train', '--train', dest='train', default=False, action='store_true',\
                         help='for training [False]')
    opt.add_argument('-att', '--att', dest='att', default=False, action='store_true',\
                         help='use pixel-attention')
    #
    # Inputs for prediction mode
    opt.add_argument('-a3m_fn', dest='a3m_fn', default=None, \
                        help="MSA file in a3m format")
    opt.add_argument('-pdb_fn_s', dest='pdb_fn_s', default=[None], nargs='+',\
                        help="Input structures in pdb format")
    opt.add_argument('-npz_fn_s', dest='npz_fn_s', default=[None], nargs='+',\
                        help="Previous 6D coords distribution prediction results in npz format. It must contain dist, omega, theta, phi")
    opt.add_argument('-SS_fn_s', dest='SS_fn_s', default=[None], nargs='+',\
                        help="predicted SS in npz file from Emin")
    opt.add_argument('-acc_fn_s', dest='acc_fn_s', default=[None], nargs='+',\
                        help="predicted acc in npz file for all input pdbs")
    opt.add_argument('-out_fn', dest='out_fn', default=None,\
                        help="Output file name")
    opt.add_argument('-model_dir', dest='model_dir', default="%s/model"%(os.path.dirname(__file__)),\
                        help="Directory for pre-trained ML model checkpoint files")
    
    if len(sys.argv) == 1:
        opt.print_help()
        sys.exit()
    #
    FLAGS = opt.parse_args()

    return FLAGS

def main():
    args = parse_input_args()
    #
    if not args.train: # prediction mode
        error = False
        if args.a3m_fn == None:
            sys.stderr.write("ERROR: Please provide MSA in a3m format\n")
            error = True
        if args.out_fn == None:
            sys.stderr.write("ERROR: Please provide output file name\n")
            error = True
        if args.pdb_fn_s[0] == None:
            sys.stderr.write("ERROR: Please provide structure model generated by trRosetta\n")
            error = True
        if args.npz_fn_s[0] == None:
            sys.stderr.write("ERROR: please provide 6d coords prediction results in npz format\n")
            error = True
        if not os.path.exists("%s_rep1"%args.model_dir):
            sys.stderr.write("ERROR: please check directory path for pre-trained checkpoint files, %s\n"%args.model_dir)
            error = True
        if error:
            return

    run_config = tf.ConfigProto(
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.92)
    )
    with tf.Session(config=run_config) as sess:

        if args.train:
            model = TrRefiner(sess, num_chunks=args.num_chunks, n2d_filters=args.n2d_filters, pdb_in=False, train=True, config_file=args.train_param, attention=args.att)
            model.train(args.train_param)
        else:
            model = TrRefiner(sess, num_chunks=args.num_chunks, n2d_filters=args.n2d_filters, pdb_in=True, attention=args.att)
            model.predict(args.model_dir, args.a3m_fn, args.npz_fn_s, args.pdb_fn_s, args.acc_fn_s, args.SS_fn_s, args.out_fn)

main()
