#!/opt/hall-lab/python-3.7.0/bin/python -u

import os, sys, time, datetime
import psutil

# add eagle to PATH
os.environ['PATH'] = os.pathsep.join(['/opt/hall-lab/eagle-2.4.1/bin', os.environ['PATH']])

def log(msg):
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %T")
    print('---> [{}] {}'.format(timestamp, msg), file=sys.stderr)
    sys.stderr.flush()

def create_directory(dirpath):
    if os.path.isdir(dirpath):
        return
    log(f"Creating output directory: {dirpath}")
    os.makedirs(dirpath)

def collect_eagle_args():
    in_vcf = os.environ['INPUT_VCF']
    out_vcf = os.environ['OUTPUT_VCF']
    out_dir = os.path.dirname(out_vcf)
    create_directory(out_dir)
    num_threads = os.environ['INPUT_THREADS']
    input_gmap = os.environ['INPUT_GMAP']
    kpbwt = 200000
    out_prefix = os.path.join(out_dir, os.environ['INPUT_PREFIX'])
    args = [
        '--numThreads', f"{num_threads}",
        '--geneticMapFile', input_gmap,
        '--vcf', in_vcf,
        '--Kpbwt', f"{kpbwt}",
        '--outPrefix', out_prefix
    ]
    return args

def exec_eagle():
    pid = os.fork()
    if pid > 0:
        return pid
    else:
        args = collect_eagle_args()
        log(f"Eagle Command: eagle {' '.join(args)}")
        os.execlp("eagle", "eagle", *args)

def monitor(pid, poll_interval=60):
    process = psutil.Process(pid)
    (pid, exit_code) = os.waitpid(pid, os.WNOHANG)
    while(pid == 0):
        mem_info = process.memory_full_info()
        uss = mem_info.uss
        (uss_bytes, uss_kb, uss_mb, uss_gb) = \
                (uss, uss/1024.0, uss/1024.0/1024.0, uss/1024.0/1024.0/1024.0)
        mem_pct = process.memory_percent(memtype="uss")
        log(f"Memory Usage (uss): {mem_pct:.2f}% -- {uss_bytes} bytes / {uss_kb:.2f} kb / {uss_mb:.2f} mb / {uss_gb:.2f} gb")
        time.sleep(poll_interval)
        (pid, exit_code) = os.waitpid(pid, os.WNOHANG)
    return exit_code

def main(input_vcf):
    pid = exec_eagle()
    log(f"Invoked eagle as process id: {pid}")
    exit_code = monitor(pid, poll_interval=10)
    print_exit_code = exit_code//256
    log(f"Eagle has terminated -- exit code: {print_exit_code}")
    log("All Done!")

if __name__ == "__main__":
    main()