#!/usr/bin/env python3
import os
import sys
import time
from config import para_json
from utils import job_queue, submitjob, command_header

docstring='''
Run_DeepMSA2.py 
    Run DeepMSA2 for all component chains.

Usage: Run_DeepMSA2.py [option]

    required options:
    -o=/home/simth/test (/home/simith/test should contain the fasta sequence file seq.fasta)

    optional options:
    -run_type=local (default) or sbatch

'''

if __name__=="__main__":

    datadir = ""
    run_type_arg = "local"

    argv = []
    for arg in sys.argv[1:]:
        if arg.startswith("-o="):
            datadir = os.path.abspath(arg[len("-o="):])
        elif arg.startswith("-run_type="):
            run_type_arg = arg[len("-run_type="):]
        elif arg.startswith('-'):
            sys.stderr.write("ERROR! No such option %s\n" % arg)
            exit()
        else:
            argv.append(arg)

    if len(argv) != 0:
        sys.stderr.write(docstring)
        exit()

    if not os.path.exists(os.path.join(datadir,"seq.fasta")):
        print("There is no sequence file seq.fasta in the datadir %s"%datadir)
        sys.stderr.write(docstring)
        exit(1)

    if run_type_arg != "":
        run_type_arg = run_type_arg
    else:
        run_type_arg = para_json['run_type']

    if run_type_arg == "sbatch":
        run_type = 2
    else:
        run_type = 1 ### local

    chains=[]
    chain_file=open(os.path.join(datadir,'chainnr.list'),'r')
    lines=chain_file.readlines()
    chain_file.close()
    for line in lines:
        chains.append(line.strip('\n'))

    #print(run_type)
    #print(chains)
    #bsub=submitjob("test.sh",1)
    #print(bsub)
    PKGdir=para_json['programrootpath']
    python3=para_json['python_DeepPotential']

    ####### DeepMSA noIMG
    while 1:
        noimg_complete=0
        for i in range(0,len(chains)):
            chain=chains[i]
            chain_datadir=os.path.join(datadir,chain)
            chain_recdir=os.path.join(chain_datadir,"record")
            chain_seq=os.path.join(chain_datadir,"seq.fasta")
            cmdcontent = f"#!/bin/sh \n" \
                     f"%s %s/DeepMSA2_noIMG.py -o=%s -i=%s -run_type=%s \n"%(python3,PKGdir,chain_datadir,chain_seq,run_type_arg) 
            #print(cmdcontent)
            os.system("mkdir -p %s"%chain_recdir)
            jobname=os.path.join(chain_recdir,chain+"-deepmsa2noimg.sh")
            with open(jobname, "w") as fp:
                fp.write(cmdcontent)
            os.system(f"chmod a+x {jobname}")
            bsub = submitjob(jobname, 1)
            print(bsub)
            if bsub.__contains__("DeepMSA2_noIMG has been complete for %s"%chain):
                noimg_complete+=1

        if noimg_complete>=len(chains):
            print("DeepMSA2 non-IMG for all component chains has been complete!")
            break    
        time.sleep(600)

    ######## DeepMSA IMG
    while 1:
        img_complete=0
        for i in range(0,len(chains)):
            chain=chains[i]
            chain_datadir=os.path.join(datadir,chain)
            chain_recdir=os.path.join(chain_datadir,"record")
            chain_seq=os.path.join(chain_datadir,"seq.fasta")
            cmdcontent = f"#!/bin/sh \n" \
                     f"%s %s/DeepMSA2_IMG.py -o=%s -i=%s -run_type=%s \n"%(python3,PKGdir,chain_datadir,chain_seq,run_type_arg) 
            #print(cmdcontent)
            os.system("mkdir -p %s"%chain_recdir)
            jobname=os.path.join(chain_recdir,chain+"-deepmsa2img.sh")
            with open(jobname, "w") as fp:
                fp.write(cmdcontent)
            os.system(f"chmod a+x {jobname}")
            bsub = submitjob(jobname, 1)
            print(bsub)
            if bsub.__contains__("DeepMSA2_IMG for %s is finished"%chain) or bsub.__contains__("%s does not need additional JGI search"%chain):
                img_complete+=1

        if img_complete>=len(chains):
            print("DeepMSA2 IMG for all component chains has been complete!")
            break    
        time.sleep(600)

    ######## MSA ranking
    while 1:
        msaranking_complete=0
        for i in range(0,len(chains)):
            chain=chains[i]
            chain_datadir=os.path.join(datadir,chain)
            chain_recdir=os.path.join(chain_datadir,"record")
            chain_seq=os.path.join(chain_datadir,"seq.fasta")
            cmdcontent = f"#!/bin/sh \n" \
                     f"%s %s/MSA_selection.py -o=%s -i=%s -run_type=%s -method=alphafold2 \n"%(python3,PKGdir,chain_datadir,chain_seq,run_type_arg) 
            #print(cmdcontent)
            os.system("mkdir -p %s"%chain_recdir)
            jobname=os.path.join(chain_recdir,chain+"-msaranking.sh")
            with open(jobname, "w") as fp:
                fp.write(cmdcontent)
            os.system(f"chmod a+x {jobname}")
            bsub = submitjob(jobname, 1)
            print(bsub)
            if bsub.__contains__("MSA ranking for %s ended normally!"%chain):
                msaranking_complete+=1

        if msaranking_complete>=len(chains):
            print("DeepMSA2 MSA ranking for all component chains has been complete!")
            break    
        time.sleep(600)