#!/home/conda/feedstock_root/build_artifacts/bld/rattler-build_gstlal-ugly_1767612080/host_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_pl/bin/python
#
# Copyright (C) 2017 Duncan Meacher, Alexander Pace
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

from optparse import OptionParser
import sched, time, os, sys, subprocess, re, pwd
import numpy as np
import multiprocessing
import itertools

class color:
   PURPLE    = '\033[95m'
   CYAN      = '\033[96m'
   DARKCYAN  = '\033[36m'
   BLUE      = '\033[94m'
   GREEN     = '\033[92m'
   YELLOW    = '\033[93m'
   RED       = '\033[91m'
   BOLD      = '\033[1m'
   UNDERLINE = '\033[4m'
   END       = '\033[0m'

def parse_command_line():
        parser = OptionParser(description = __doc__)
        parser.add_option("-u", "--username", help = "User name.")
        parser.add_option("-s", "--sampler", default = "top", help = "Choose sampling program to use (top or ps), default = top")
        parser.add_option("-c", "--cpu-limit", metavar = "value", type = "float", default = 1.0, help = "cpu% to recheck jobs below")

        options, filenames = parser.parse_args()

        if options.username is None:
                options.username = pwd.getpwuid(os.getuid()).pw_name

        if options.sampler is None:
                options.sampler = 'top'
        
        if options.sampler not in ['top','ps']:
                print("Invalid sampler. Please specify either top or ps.")
                exit()

        return options, filenames

options, filenames = parse_command_line()

# Run condor_q -dag -nobatch username
command = 'condor_q -dag -nobatch -wide:100 ' + options.username
cq_out  = subprocess.Popen(command.split(), stdout=subprocess.PIPE)
cq_out  = cq_out.communicate()[0].splitlines()

# Define the prefixes to search for:
#prefixes=['gstlal_reference_psd_','gstlal_inspiral_','gstlal_inspiral_inj_']
prefixes = ['gstlal_reference_psd_','gstlal_svd_bank_','gstlal_inspiral_injection_snr_','gstlal_inspiral_','gstlal_inspiral_inj_']

# Get column width for pretty printing. Add four for the hex suffix and two 
# for buffer.
col_width = max(len(pref) for pref in prefixes) + 4 + 2 

# Initialize the results list to zero:
results = []
queue = {}

# If an item in the output from condor_q contains our prefix and the hex
# regular expression, then add it to the results list. 
for pre in prefixes:
        results.extend([line for line in cq_out if re.search(pre+'[0-9A-F]',line) is not None])
for line in results:
        split_list = line.split()
        job_name, condor_id, job_status = split_list[1][2:], split_list[0], split_list[5]
        queue[condor_id] = [job_name, condor_id, job_status, None, None, None]

# Check to see is any matches
if not results:
        print(color.YELLOW + "No matches found in condor_q" + color.END)
        exit()
else:
        print(color.BOLD + "".join("job_name".ljust(col_width))+ \
              "".join("condor_id".ljust(18))+"pid\tcpu%\tmem%" + color.END)

def process_result(input_tuple, options = options):
        job_name, condor_id, job_status = input_tuple
        if job_status == "R":
                # Put together the command string for condor_ssh_to job. Note that the 
                # quotes and slashes are key for the correct command to get evaluated by
                # the interpreter:

                if options.sampler == 'top':
                        sample_command = """ "top b -n 1 -p \\${_CONDOR_JOB_PIDS}" """

                if options.sampler == 'ps':
                        sample_command = ' "ps u -p \\${_CONDOR_JOB_PIDS}"'

                condor_ssh_command = 'condor_ssh_to_job ' + \
                                     condor_id + \
                                     sample_command

                cstj_err = None
                try:
                        cstj_output = subprocess.check_output(condor_ssh_command, shell = True)
                except subprocess.CalledProcessError as e:
                        print(e)
                        return (job_name, condor_id, "U", "--",  "--", "--")

                # Collect output. The syntax should be apparent from looking at which 
                # line is being collected and which column is being pulled after the
                # split().

                if options.sampler == 'top':
                        try:
                                cstj_output = cstj_output.splitlines()[-1].split()
                        except IndexError:
                                return (job_name, condor_id, "U", "--", "--", "--")
                        pid = cstj_output[0]
                        cpu = cstj_output[8]
                        mem = cstj_output[9]
                if options.sampler == 'ps':
                        cstj_output = cstj_output.splitlines()[1].split()
                        pid = cstj_output[1]
                        cpu = cstj_output[2]
                        mem = cstj_output[3]

                return (job_name, condor_id, job_status, pid, cpu, mem)

        elif job_status == "H":
                return (job_name, condor_id, "H", "--", "--", "--")
        elif job_status == "I":
                return (job_name, condor_id, "I", "--", "--", "--")
        else:
                return (job_name, condor_id, job_status, "--", "--", "--")
                #raise ValueError(job_status)

if __name__ == '__main__':

        # Loop over all jobs
        pool = multiprocessing.Pool(100)
        recheck = {}
        for r in sorted(pool.map(process_result, [(v[0], v[1], v[2]) for v in queue.values()]), key = lambda x: x[3] if x else 0):
                job_name, condor_id, job_status, pid, cpu, mem = r
                queue[condor_id] = [job_name, condor_id, job_status, pid, cpu, mem]
                try:
                        if float(cpu) < options.cpu_limit:
                                recheck[condor_id] = [job_name, condor_id, job_status, pid, cpu, mem]
                except ValueError:
                        pass

        # check one more time for jobs with cpu less than the limit
        for r in sorted(map(process_result, [(v[0], v[1], v[2]) for v in recheck.values()]), key = lambda x: float(x[3]) if x and x[3] != "--" else 0):
                job_name, condor_id, job_status, pid, cpu, mem = r
                queue[condor_id] = [job_name, condor_id, job_status, pid, cpu, mem]

        for r in sorted(queue.values(), key = lambda x: x[4] if x else 0, reverse=True):
                job_name, condor_id, job_status, pid, cpu, mem = r
                try:
                        if float(cpu) < options.cpu_limit:
                                print(color.RED + "%s%s\t%s\t%s\t%s" % (job_name.ljust(col_width), condor_id.ljust(18), job_status, cpu, mem) + color.END)
                        else:
                                print(color.GREEN + "%s%s\t%s\t%s\t%s" % (job_name.ljust(col_width), condor_id.ljust(18), job_status, cpu, mem) + color.END)
                except ValueError:
                        print("%s%s\t%s\t%s\t%s" % (job_name.ljust(col_width), condor_id.ljust(18), job_status, cpu, mem))
