<commit_before>import os
import sys
from subprocess import Popen, PIPE
import gzip
import errno
import pipes
import re
import threading
import Queue
import time
import cPickle as pickle
import shutil
if __name__ == "__main__" and __package__ is None:
    sys.path.insert(1,
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from dimple import utils
from dimple import mtz
from dimple import pdb
from dimple import coots


_jobindex_fmt = "%3d "
_jobname_fmt = "%-15s"
_elapsed_fmt = "%5.1fs  "

PICKLE_FILENAME = "workflow.pickle"


# heh, python from MSYS2 is handy, but needs some monkey-patching
if sys.platform == 'msys':
    old_opjoin = os.path.join
    def new_opjoin(*args):
        for n in range(len(args)-1, -1, -1):
            if n == 0 or os.path.isabs(args[n]):
                return old_opjoin(*args[n:])
    os.path.join = new_opjoin


class JobError(Exception):
    def __init__(self, msg, note=None): # pylint: disable=super-init-not-called
        self.msg = msg
        self.note = note


class Output:
    "storage for Job's stdout/stderr"
    def __init__(self, role):
        self.role = role  # "out" or "err"
        self.file_extension = {"out":"log", "err":"err"}[role]
        self.lines = []
        self.saved_to = None
        self.que = None

    def __nonzero__(self):
        return bool(self.lines or self.saved_to)

    def size_as_str(self):
        if self.lines:
            return "%d lines" % len(self.lines)
        else:
            return "-      "

    def read_line(self):
        while self.que is not None:
            try:
                line = self.que.get_nowait()
            except Queue.Empty:
                break
            self.lines.append(line)
            yield line

    def finish_que(self):
        if self.que:
            while not self.que.empty():
                self.lines.append(self.que.get_nowait())
            self.que = None

    def save_output(self, output_dir, basename, remove_long_list=True):
        filename = basename + '.' + self.file_extension
        if self.lines:
            with open(os.path.join(output_dir, filename), "w") as f:
                for line in self.lines:
                    f.write(line)
            self.saved_to = filename
            utils.log_value("std"+self.role, filename)
            if remove_long_list and len(self.lines) > 5:
                self.lines = []

    def summary(self):
        n = len(self.lines)
        if 0 < n <= 3:
            return "".join(self.lines)
        elif self.saved_to:
            return "-> %s" % self.saved_to
        elif n == 0:
            return ""
        else: # n > 3
            return "".join(self.lines[:3]) + ("%s more lines" % (n-3))


class Job:
    def __init__(self, workflow, prog):
        self.name = os.path.basename(prog) or prog  # only used to show info
        self.workflow = workflow
        self.args = [prog]
        self.std_input = ""
        self.stdin_file = None # if set, it overwrites std_input
        # the rest is set after the job is run
        self.exit_status = None
        self.out = Output("out")
        self.err = Output("err")
        self.started = None  # will be set to time.time() at start
        self.total_time = None  # will be set when job ends
        # possible values: None (stdout preview),
        #                  string starting with space (' ') that is just shown,
        #                  or name of global function that parses output
        self.parser = None
        # job-specific data from output parsing
        self.data = {}

    def __repr__(self):
        if self.started:
            t = time.strftime(" %Y-%m-%d %H:%M", time.localtime(self.started))
        else:
            t = ""
        return "<Job %s%s>" % (self.name, t)

    def args_as_str(self):
        s = " ".join(pipes.quote(a) for a in self.args)
        if self.stdin_file:
            s += " < " + self.stdin_file
        elif self.std_input:
            s += " << EOF\n%s\nEOF" % self.std_input
        return s

    def run(self, show_progress=True, new_line=True, may_fail=False):
        self.workflow.run_job(job=self,
                              show_progress=show_progress, new_line=new_line)
        # exit_status may be None if --from-step is used
        if not may_fail and self.exit_status:
            notes = [self.args_as_str(), ""]
            if self.out.saved_to:
                notes += ["stdout -> %s/%s" % (self.workflow.output_dir,
                                               self.out.saved_to)]
            if self.err:
                notes += ["stderr:", self.err.summary()]
            raise JobError("%s failed (exit status %d)" % (self.name,
                                                           self.exit_status),
                           note="\n".join(notes))
        return self

    def parse(self):
        if self.parser is None:  # preview mode
            # generic non-parser
            line = ""
            for line in self.out.read_line():
                pass
            if line:
                # we only remove \e[1m (bold), it can be generalized if needed
                trimmed = line.strip().replace('\033[1m', '')
                ret = "[%d] %-44.44s" % (len(self.out.lines), trimmed)
            else:
                ret = "stdout:%11s" % self.out.size_as_str()
                if self.err:
                    ret += " stderr: %s" % self.err.size_as_str()
                ret = ret.ljust(50)
            return ret

        elif self.parser == '' or self.parser[0] == ' ':
            return self.parser

        else:
            p = globals()[self.parser]
            return p(self)


def _format(fmt, arg):
    return (fmt % arg) if arg else ""

# parsers for various programs
def _find_blobs_parser(job):
    if "blobs" not in job.data:
        job.data["blobs"] = []
        job.data["scores"] = []
    for line in job.out.read_line():
        if line.startswith("#"):
            sp = line.split(None, 6)
            score = float(sp[5])
            if True: #score > 150: XXX: better scoring may be needed
                x, y, z = sp[-1].strip("() \t\r\n").split(",")
                job.data["blobs"].append((float(x), float(y), float(z)))
                job.data["scores"].append(score)
        elif line.startswith("Protein mass center:"):
            sp = line.split("(")[1].rstrip("\r\n )").split(",")
            ctr = tuple(float(x) for x in sp)
            job.data["center"] = ctr
        elif line.startswith("Density std.dev"):
            job.data["density_info"] = line.strip()
    scores = job.data["scores"]
    if scores:
        return "Blob scores: " + " ".join("%.0f" % sc for sc in scores)
    elif "density_info" in job.data:
        return 'searching with' + job.data["density_info"].split(',')[1]
    else:
        return ""

def _rwcontents_parser(job):
    d = job.data
    for line in job.out.read_line():
        if line.startswith(' Cell volume:'):
            d["volume"] = float(line.split(':')[-1])
        elif line.startswith(' Molecular Weight of protein:'):
            d["weight"] = float(line.split(':')[-1])
        if line.startswith(' The Matthews Coefficient is :'):
            Vm = float(line.split(':')[-1])
            d["Vm"] = Vm
            # 1.23 is used in phaser/src/Composition.cc
            d["solvent_percent"] = (1 - 1.23/Vm) * 100
    if 'volume' in d and 'weight' in d and 'Vm' in d and 'num_mol' not in d:
        d['num_mol'] = int(round(d['volume'] / d['weight'] / d['Vm']))
    return u"%d x %.0fkDa in %.fnm3  Vm=%.2f (%.0f%% of solvent)" % (
            d.get('num_mol', 0),
            d.get('weight', 0) / 1000, # Da -> kDa
            d.get('volume', 0) / 1000, # A^2 -> nm^3
            d.get('Vm', 0),
            d.get('solvent_percent', 0))

def _cad_parser(job):
    for line in job.out.read_line():
        if line.startswith(' Final Total of Unique records to HKLOUT ='):
            job.data["refl_out"] = int(line.split('=')[1])
    return "#refl -> %s" % job.data.get("refl_out", "")


class Ccp4LogTable(object):
    def __init__(self, title_line):
        assert '$TABLE:' in title_line
        self.title = title_line.split(':')[1]
        self.columns = []
        self.data = []
        self.section = 0  # 0=header, 1=columns, 2=data, 3=end

    def send_line(self, line):
        line = line.strip()
        if line.startswith('$$'):
            self.section += 1
            if len(line) > 2:
                self.send_line(line[2:])
            return self.section != 3

        if self.section == 1:
            self.columns += line.rstrip('$').split()
        elif self.section == 2:
            self.data.append(line.split())
        return True

    def column(self, name):
        try:
            idx = self.columns.index(name)
            return [float(a[idx]) for a in self.data]
        except (ValueError, IndexError):
            return


def _refmac_parser(job):
    if "cycle" not in job.data:
        # ini_free_r, free_r and iter_free_r are set optionally
        job.data["cycle"] = 0
        # iter_*_r values were added in dimple 1.5
        job.data["iter_overall_r"] = []
    sink = None
    for line in job.out.read_line():
        if sink:
            more = sink.send_line(line)
            if not more:
                for name in ['rmsBOND', 'rmsANGL', 'rmsCHIRAL']:
                    col_data = sink.column(name)
                    if col_data and any(x != 0 for x in col_data):
                        job.data[name] = col_data
                sink = None
        elif line.startswith("Free R factor"):
            job.data['free_r'] = float(line.split('=')[-1])
            if 'ini_free_r' not in job.data:
                job.data['ini_free_r'] = job.data['free_r']
                job.data['iter_free_r'] = []
            job.data['iter_free_r'].append(job.data['free_r'])
        elif line.startswith("Overall R factor"):
            job.data['overall_r'] = float(line.split('=')[-1])
            if 'ini_overall_r' not in job.data:
                job.data['ini_overall_r'] = job.data['overall_r']
            job.data['iter_overall_r'].append(job.data['overall_r'])
        elif (line.startswith("     Rigid body cycle =") or
              line.startswith("     CGMAT cycle number =")):
            job.data['cycle'] = int(line.split('=')[-1])
        elif line.startswith("$TABLE: Rfactor analysis, stats vs cycle"):
            sink = Ccp4LogTable(line)
    cycle_str = "%2d/%d" % (job.data["cycle"], job.data.get("ncyc", -1))
    if 'ini_overall_r' in job.data:
        if 'ini_free_r' in job.data:
            return "%s   R/Rfree  %.4f/%.4f  ->  %.4f/%.4f" % (
                    cycle_str,
                    job.data["ini_overall_r"], job.data["ini_free_r"],
                    job.data["overall_r"], job.data["free_r"])
        return "%s   R  %.4f  ->  %.4f" % (
                cycle_str, job.data["ini_overall_r"], job.data["overall_r"])
    return cycle_str

# example:
#     Alternative reindexing        Lklhd      CC     R(E^2)    Number Cell_deviation
#    [-k,-l,h+k+l]           0.079    0.029    0.506     61253      2.99
_POINTLESS_ALTREINDEX_MATCH = re.compile(r"^\s+\[[hkl+, -]+\][ \t0-9+.eE-]+$")

def _float_or_nan(s):
    try:
        x = float(s)
    except ValueError:
        x = float('nan')
    return x

def _pointless_parser(job):
    # the line with 'reflections copied' is the last we read
    # to avoid reading 'Alternative reindexing' duplicated in the summary
    if "refl_out" not in job.data:
        for line in job.out.read_line():
            if line.startswith("Maximum resolution used:"):
                job.data["resol"] = float(line.split(":")[1])
            elif line.startswith("Number of reflections:"):
                job.data["refl_ref"] = int(line.split(":")[1])
            elif _POINTLESS_ALTREINDEX_MATCH.match(line):
                s = line.split()
                job.data.setdefault('alt_reindex', []).append(
                        {'op': s[0],
                         'cc': _float_or_nan(s[2]),
                         'cell_deviat': _float_or_nan(s[-1])})
            elif line.startswith("   Cell:") and 'output_cell' not in job.data:
                s = line.split()[1:]
                job.data["output_cell"] = tuple(float(i) for i in s)
            elif "reflections copied to output file" in line:
                job.data["refl_out"] = int(line.split()[0])
                break
    resol_txt = _format("%.2f", job.data.get("resol"))
    refl_out = job.data.get("refl_out", "")
    #refl_ref = job.data.get("refl_ref", "")
    return "resol. %4s A   #refl: %5s" % (resol_txt, refl_out)


def _phaser_parser(job):
    d = job.data
    for line in job.out.read_line():
        if line.startswith('*** Phaser Module:'):
            d['status'] = '[%s]' % line[19:70].strip().lower()
        if line.startswith('   SOLU SET ') and 'LLG=' in line:
            d['status'] = line[12:].strip()
            if len(d['status']) > 52:
                d['status'] = d['status'][:50].rsplit(' ', 1)[0] + '...'
        if line.startswith('   Sorry - No solution'):
            d['status'] = line.strip()
            if 'No solution with all components' in line:
                d['partial_solution'] = 'yes'
        if line.startswith('   SOLU SPAC '):
            d['SG'] = line[13:].strip()
    return "%-48s" % d.get('status', '')

def _truncate_parser(job):
    for line in job.out.read_line():
        if line.startswith(' Least squares straight line gives:'):
            b_str = line.partition(':')[-1].strip(' \tB=').split()[0]
            job.data["B-factor"] = float(b_str)
    return "B=%4s" % _format("%.1f", job.data.get("B-factor"))

def _ctruncate_parser(job):
    d = job.data
    for line in job.out.read_line():
        if line.startswith("$GRAPHS: Wilson plot - estimated B factor ="):
            d["B-factor"] = float(line.partition('=')[-1].split()[0])
        elif line.startswith("Eigenvalue ratios:"):
            d["eigval-ratios"] = tuple(float(i) for i in line.split()[-3:])
        elif line.startswith("L statistic ="):
            d["L-test"] = float(line.partition('=')[-1].split()[0])
    return "B=%4s   aniso %14s   L-test:%s" % (
            _format("%.1f", d.get("B-factor")),
            _format("%.2f:%.2f:%.2f", d.get("eigval-ratios")),
            _format("%.2f", d.get("L-test")))


def ccp4_job(workflow, prog, logical=None, ki="", parser=None, add_end=True):
    """Handle traditional convention for arguments of CCP4 programs.
    logical - dictionary with where keys are so-called logical names.
    ki (string or list of lines) - Keyworded Input to be passed through stdin.
    add_end - adds "end" as the last line of stdin

    Note: "colin" and "labin" mean the same (column label),
    but different programs use different keywords.
    """
    job = Job(workflow, utils.cbin(prog))
    if logical:
        for a in ["hklin", "hklout", "hklref", "xyzin", "xyzout", "libin"]:
            if logical.get(a):
                job.args += [a.upper(), logical[a]]
    lines = ki.splitlines() if isinstance(ki, basestring) else ki
    stripped = [a.strip() for a in lines if a and not a.isspace()]
    if add_end and not (stripped and stripped[-1].lower() == "end"):
        stripped.append("end")
    job.std_input = "\n".join(stripped)
    job.parser = parser
    return job


def _print_progress(job, event):
    while not event.wait(0.5):
        p = job.parse()
        if p is not None:
            text = (_elapsed_fmt % (time.time() - job.started)) + p
            utils.put_temporarily(text)
            utils.reset_color()


def _start_enqueue_thread(file_obj):
    def enqueue_lines(f, q):
        for line in iter(f.readline, b''):
            q.put(line)
        f.close()
    que = Queue.Queue()
    thr = threading.Thread(target=enqueue_lines, args=(file_obj, que))
    thr.daemon = True
    thr.start()
    return thr, que

def _get_input_as_string(job):
    if job.stdin_file:
        path = job.workflow.path(job.stdin_file)
        try:
            return open(path, "rb").read()
        except IOError:
            raise JobError("cannot read input from: %s" % job.stdin_file)
    else:
        return job.std_input

def _run_and_parse(process, job):
    try:
        # job.*.que can be used by parsers (via Output.read_line() or directly)
        out_t, job.out.que = _start_enqueue_thread(process.stdout)
        err_t, job.err.que = _start_enqueue_thread(process.stderr)
        try:
            job_input = _get_input_as_string(job)
            process.stdin.write(job_input)
        except IOError as e:
            utils.put("\nWarning: passing input to %s failed.\n" % job.name)
            if e.errno not in (errno.EPIPE, e.errno != errno.EINVAL):
                raise
        process.stdin.close()
        out_t.join()
        err_t.join()
        process.wait()
        # nothing is written to the queues at this point
        # parse what's left in the queues
        job.parse()
    finally:
        # take care of what is left by the parser
        job.out.finish_que()
        job.err.finish_que()


class Workflow:
    def __init__(self, output_dir, from_job=0):
        self.output_dir = os.path.abspath(output_dir)
        self.jobs = []
        self.file_info = {}
        self.temporary_files = set()
        self.from_job = from_job  # skip jobs before from_job (for testing)
        if from_job >= 1:
            try:
                _pkl = self.load_pickle()
                self.repl_jobs = _pkl.jobs
                self.file_info = _pkl.file_info
            except:
                self.repl_jobs = None
        self.dry_run = False
        self.argv = sys.argv
        if not os.path.isdir(self.output_dir):
            try:
                os.makedirs(self.output_dir)
            except OSError as e:
                utils.put_error(e)
                sys.exit(1)
        # this can seriously affect Refmac compiled with GFortran
        bad_var = os.getenv('GFORTRAN_UNBUFFERED_ALL')
        if bad_var and bad_var[0] not in ('0', 'n', 'N'):
            utils.put_error(
                    '$GFORTRAN_UNBUFFERED_ALL may terribly slow down Refmac',
                    comment='It is unset internally in dimple.')
            del os.environ['GFORTRAN_UNBUFFERED_ALL']
        # avoid html-like crap in the output of CCP4 program
        os.environ['CCP_SUPPRESS_HTML'] = '1'


    def __str__(self):
        return "Workflow with %d jobs @ %s" % (len(self.jobs), self.output_dir)

    def path(self, rel_path):
        return os.path.join(self.output_dir, rel_path)

    def dump_pickle(self):
        with open(self.path(PICKLE_FILENAME), "wb") as f:
            pickle.dump(self, f, -1)

    def load_pickle(self):
        with open(self.path(PICKLE_FILENAME), "rb") as f:
            return pickle.load(f)

    def run_job(self, job, show_progress, new_line=True):
        if not hasattr(sys.stdout, 'isatty') or not sys.stdout.isatty():
            show_progress = False
        self.jobs.append(job)
        job_num = len(self.jobs)
        if new_line:
            utils.put("\n" + _jobindex_fmt % job_num)
            utils.put_green(_jobname_fmt % job.name)
        else:
            utils.put(" / %d" % job_num)
        sys.stdout.flush()
        utils.log_section(job.name)

        job_idx = len(self.jobs) - 1
        if job_idx < self.from_job - 1: # from_job is 1-based
            # unpickle or skip
            if self.repl_jobs and len(self.repl_jobs) > job_idx:
                old_job = self.repl_jobs[job_idx]
                if old_job.name == job.name:
                    job = old_job
                    utils.put("unpickled")
                    utils.log_value("not_run", "unpickled")
                    self.jobs[-1] = job
                else:
                    utils.put("skipped (mismatch)")
                    utils.log_value("not_run", "unpickled/mismatch")
            else:
                utils.put("skipped")
                utils.log_value("not_run", "skipped")
            return

        job.started = time.time()
        utils.log_time("start_time", job.started)
        if job.stdin_file:
            utils.log_value("stdin", job.stdin_file)
        elif job.std_input:
            utils.log_value("input", job.std_input)
        utils.log_value("prog", job.args[0])
        utils.log_value("args", " ".join(pipes.quote(a) for a in job.args[1:]))
        #job.args[0] = "true"  # for debugging
        try:
            process = Popen(job.args, stdin=PIPE, stdout=PIPE, stderr=PIPE,
                            cwd=self.output_dir)
        except OSError as e:
            if e.errno == errno.ENOENT:
                raise JobError("Program not found: %s\n" % job.args[0])
            else:
                raise

        if self.dry_run:
            return

        if show_progress:
            event = threading.Event()
            progress_thread = threading.Thread(target=_print_progress,
                                               args=(job, event))
            progress_thread.daemon = True
            progress_thread.start()

        try:
            if job.parser is not None or show_progress:
                _run_and_parse(process, job)
            else:
                job_input = _get_input_as_string(job)
                out, err = process.communicate(input=job_input)
                job.out.lines = out.splitlines(True)
                job.err.lines = err.splitlines(True)
        except KeyboardInterrupt:
            raise JobError("KeyboardInterrupt while running %s" % job.name,
                           note=job.args_as_str())
        finally:
            if show_progress:
                event.set()
                progress_thread.join()
            end_time = time.time()
            job.total_time = end_time - job.started
            utils.log_time("end_time", end_time)
            job.exit_status = process.poll()
            if new_line:
                utils.put(_elapsed_fmt % job.total_time)
            parse_output = job.parse()
            utils.put("%s" % (parse_output or ""))
            if parse_output:
                utils.log_value("info", parse_output)
            self._write_logs(job)
            for k, v in job.data.iteritems():
                utils.log_value(k, v)
        if job.exit_status != 0:
            utils.log_value("exit_status", job.exit_status)

    def _write_logs(self, job):
        log_basename = "%02d-%s" % (len(self.jobs), job.name.replace(" ","_"))
        for output in (job.out, job.err):
            output.save_output(self.output_dir, log_basename)

    def remove_hetatm(self, xyzin, xyzout, remove_all):
        with open(self.path(xyzout), "wb") as out:
            return pdb.remove_hetatm(self.path(xyzin), out, remove_all)

    def read_pdb_metadata(self, xyzin):
        if xyzin not in self.file_info:
            self.file_info[xyzin] = pdb.read_metadata(self.path(xyzin))
        return self.file_info[xyzin]

    def read_mtz_metadata(self, hklin):
        if hklin not in self.file_info:
            self.file_info[hklin] = mtz.read_metadata(self.path(hklin))
        return self.file_info[hklin]

    def count_mtz_missing(self, hklin, col):
        key = (hklin, 'missing', col)
        if key not in self.file_info:
            self.file_info[key] = mtz.get_num_missing(self.path(hklin), col)
        return self.file_info[key]

    def molrep(self, f, m, keys=""):
        job = Job(self, utils.cbin("molrep"))
        job.args += ["-f", f, "-m", m]
        if keys:
            job.args.append("-i")
            job.std_input = keys.strip() + "\nend"
        return job

    def phaser_auto(self, hklin, labin, model, root, solvent_percent,
                    sg_alt="NONE"):
        lines = ['MODE MR_AUTO',
                 'SEARCH METHOD FAST',
                 'SEARCH DEEP OFF',
                 'ENSEMBLE p PDBFILE "%(pdb)s" IDENTITY %(identity)g' % model,
                 'SEARCH ENSEMBLE p NUM %(num)d' % model,
                 'COMPOSITION BY SOLVENT',
                 'COMPOSITION PERCENTAGE %f' % solvent_percent,
                 # Since Phaser 2.5.6 template matched solutions are moved
                 # to the template solution origin. Which is better than
                 # getting a solution one cell away, so we set template here.
                 'SOLUTION TEMPLATE original_model',
                 'SOLUTION 6DIM ENSE p EULER 0 0 0 FRAC 0 0 0',
                 'HKLIN "%s"' % hklin,
                 'LABIN %s' % labin,
                 'SGALTERNATIVE SELECT %s' % sg_alt,
                 'KILL TIME 120',  # 2h is much more than we want
                 'ROOT %s' % root,
                 ]
        # tNCS: we go with what phaser does by default -- tNCS of order 2
        # are handled automatically. While we could specify tNCS for
        # pseudo-tripling/quadrupling of the cell (TNCS NMOL 3) I don't know
        # if it'd do more good or bad.<commit_msg>temporarily, set RESOLUTION HIGH 3.0 in phaser<commit_after>import os
import sys
from subprocess import Popen, PIPE
import gzip
import errno
import pipes
import re
import threading
import Queue
import time
import cPickle as pickle
import shutil
if __name__ == "__main__" and __package__ is None:
    sys.path.insert(1,
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from dimple import utils
from dimple import mtz
from dimple import pdb
from dimple import coots


_jobindex_fmt = "%3d "
_jobname_fmt = "%-15s"
_elapsed_fmt = "%5.1fs  "

PICKLE_FILENAME = "workflow.pickle"


# heh, python from MSYS2 is handy, but needs some monkey-patching
if sys.platform == 'msys':
    old_opjoin = os.path.join
    def new_opjoin(*args):
        for n in range(len(args)-1, -1, -1):
            if n == 0 or os.path.isabs(args[n]):
                return old_opjoin(*args[n:])
    os.path.join = new_opjoin


class JobError(Exception):
    def __init__(self, msg, note=None): # pylint: disable=super-init-not-called
        self.msg = msg
        self.note = note


class Output:
    "storage for Job's stdout/stderr"
    def __init__(self, role):
        self.role = role  # "out" or "err"
        self.file_extension = {"out":"log", "err":"err"}[role]
        self.lines = []
        self.saved_to = None
        self.que = None

    def __nonzero__(self):
        return bool(self.lines or self.saved_to)

    def size_as_str(self):
        if self.lines:
            return "%d lines" % len(self.lines)
        else:
            return "-      "

    def read_line(self):
        while self.que is not None:
            try:
                line = self.que.get_nowait()
            except Queue.Empty:
                break
            self.lines.append(line)
            yield line

    def finish_que(self):
        if self.que:
            while not self.que.empty():
                self.lines.append(self.que.get_nowait())
            self.que = None

    def save_output(self, output_dir, basename, remove_long_list=True):
        filename = basename + '.' + self.file_extension
        if self.lines:
            with open(os.path.join(output_dir, filename), "w") as f:
                for line in self.lines:
                    f.write(line)
            self.saved_to = filename
            utils.log_value("std"+self.role, filename)
            if remove_long_list and len(self.lines) > 5:
                self.lines = []

    def summary(self):
        n = len(self.lines)
        if 0 < n <= 3:
            return "".join(self.lines)
        elif self.saved_to:
            return "-> %s" % self.saved_to
        elif n == 0:
            return ""
        else: # n > 3
            return "".join(self.lines[:3]) + ("%s more lines" % (n-3))


class Job:
    def __init__(self, workflow, prog):
        self.name = os.path.basename(prog) or prog  # only used to show info
        self.workflow = workflow
        self.args = [prog]
        self.std_input = ""
        self.stdin_file = None # if set, it overwrites std_input
        # the rest is set after the job is run
        self.exit_status = None
        self.out = Output("out")
        self.err = Output("err")
        self.started = None  # will be set to time.time() at start
        self.total_time = None  # will be set when job ends
        # possible values: None (stdout preview),
        #                  string starting with space (' ') that is just shown,
        #                  or name of global function that parses output
        self.parser = None
        # job-specific data from output parsing
        self.data = {}

    def __repr__(self):
        if self.started:
            t = time.strftime(" %Y-%m-%d %H:%M", time.localtime(self.started))
        else:
            t = ""
        return "<Job %s%s>" % (self.name, t)

    def args_as_str(self):
        s = " ".join(pipes.quote(a) for a in self.args)
        if self.stdin_file:
            s += " < " + self.stdin_file
        elif self.std_input:
            s += " << EOF\n%s\nEOF" % self.std_input
        return s

    def run(self, show_progress=True, new_line=True, may_fail=False):
        self.workflow.run_job(job=self,
                              show_progress=show_progress, new_line=new_line)
        # exit_status may be None if --from-step is used
        if not may_fail and self.exit_status:
            notes = [self.args_as_str(), ""]
            if self.out.saved_to:
                notes += ["stdout -> %s/%s" % (self.workflow.output_dir,
                                               self.out.saved_to)]
            if self.err:
                notes += ["stderr:", self.err.summary()]
            raise JobError("%s failed (exit status %d)" % (self.name,
                                                           self.exit_status),
                           note="\n".join(notes))
        return self

    def parse(self):
        if self.parser is None:  # preview mode
            # generic non-parser
            line = ""
            for line in self.out.read_line():
                pass
            if line:
                # we only remove \e[1m (bold), it can be generalized if needed
                trimmed = line.strip().replace('\033[1m', '')
                ret = "[%d] %-44.44s" % (len(self.out.lines), trimmed)
            else:
                ret = "stdout:%11s" % self.out.size_as_str()
                if self.err:
                    ret += " stderr: %s" % self.err.size_as_str()
                ret = ret.ljust(50)
            return ret

        elif self.parser == '' or self.parser[0] == ' ':
            return self.parser

        else:
            p = globals()[self.parser]
            return p(self)


def _format(fmt, arg):
    return (fmt % arg) if arg else ""

# parsers for various programs
def _find_blobs_parser(job):
    if "blobs" not in job.data:
        job.data["blobs"] = []
        job.data["scores"] = []
    for line in job.out.read_line():
        if line.startswith("#"):
            sp = line.split(None, 6)
            score = float(sp[5])
            if True: #score > 150: XXX: better scoring may be needed
                x, y, z = sp[-1].strip("() \t\r\n").split(",")
                job.data["blobs"].append((float(x), float(y), float(z)))
                job.data["scores"].append(score)
        elif line.startswith("Protein mass center:"):
            sp = line.split("(")[1].rstrip("\r\n )").split(",")
            ctr = tuple(float(x) for x in sp)
            job.data["center"] = ctr
        elif line.startswith("Density std.dev"):
            job.data["density_info"] = line.strip()
    scores = job.data["scores"]
    if scores:
        return "Blob scores: " + " ".join("%.0f" % sc for sc in scores)
    elif "density_info" in job.data:
        return 'searching with' + job.data["density_info"].split(',')[1]
    else:
        return ""

def _rwcontents_parser(job):
    d = job.data
    for line in job.out.read_line():
        if line.startswith(' Cell volume:'):
            d["volume"] = float(line.split(':')[-1])
        elif line.startswith(' Molecular Weight of protein:'):
            d["weight"] = float(line.split(':')[-1])
        if line.startswith(' The Matthews Coefficient is :'):
            Vm = float(line.split(':')[-1])
            d["Vm"] = Vm
            # 1.23 is used in phaser/src/Composition.cc
            d["solvent_percent"] = (1 - 1.23/Vm) * 100
    if 'volume' in d and 'weight' in d and 'Vm' in d and 'num_mol' not in d:
        d['num_mol'] = int(round(d['volume'] / d['weight'] / d['Vm']))
    return u"%d x %.0fkDa in %.fnm3  Vm=%.2f (%.0f%% of solvent)" % (
            d.get('num_mol', 0),
            d.get('weight', 0) / 1000, # Da -> kDa
            d.get('volume', 0) / 1000, # A^2 -> nm^3
            d.get('Vm', 0),
            d.get('solvent_percent', 0))

def _cad_parser(job):
    for line in job.out.read_line():
        if line.startswith(' Final Total of Unique records to HKLOUT ='):
            job.data["refl_out"] = int(line.split('=')[1])
    return "#refl -> %s" % job.data.get("refl_out", "")


class Ccp4LogTable(object):
    def __init__(self, title_line):
        assert '$TABLE:' in title_line
        self.title = title_line.split(':')[1]
        self.columns = []
        self.data = []
        self.section = 0  # 0=header, 1=columns, 2=data, 3=end

    def send_line(self, line):
        line = line.strip()
        if line.startswith('$$'):
            self.section += 1
            if len(line) > 2:
                self.send_line(line[2:])
            return self.section != 3

        if self.section == 1:
            self.columns += line.rstrip('$').split()
        elif self.section == 2:
            self.data.append(line.split())
        return True

    def column(self, name):
        try:
            idx = self.columns.index(name)
            return [float(a[idx]) for a in self.data]
        except (ValueError, IndexError):
            return


def _refmac_parser(job):
    if "cycle" not in job.data:
        # ini_free_r, free_r and iter_free_r are set optionally
        job.data["cycle"] = 0
        # iter_*_r values were added in dimple 1.5
        job.data["iter_overall_r"] = []
    sink = None
    for line in job.out.read_line():
        if sink:
            more = sink.send_line(line)
            if not more:
                for name in ['rmsBOND', 'rmsANGL', 'rmsCHIRAL']:
                    col_data = sink.column(name)
                    if col_data and any(x != 0 for x in col_data):
                        job.data[name] = col_data
                sink = None
        elif line.startswith("Free R factor"):
            job.data['free_r'] = float(line.split('=')[-1])
            if 'ini_free_r' not in job.data:
                job.data['ini_free_r'] = job.data['free_r']
                job.data['iter_free_r'] = []
            job.data['iter_free_r'].append(job.data['free_r'])
        elif line.startswith("Overall R factor"):
            job.data['overall_r'] = float(line.split('=')[-1])
            if 'ini_overall_r' not in job.data:
                job.data['ini_overall_r'] = job.data['overall_r']
            job.data['iter_overall_r'].append(job.data['overall_r'])
        elif (line.startswith("     Rigid body cycle =") or
              line.startswith("     CGMAT cycle number =")):
            job.data['cycle'] = int(line.split('=')[-1])
        elif line.startswith("$TABLE: Rfactor analysis, stats vs cycle"):
            sink = Ccp4LogTable(line)
    cycle_str = "%2d/%d" % (job.data["cycle"], job.data.get("ncyc", -1))
    if 'ini_overall_r' in job.data:
        if 'ini_free_r' in job.data:
            return "%s   R/Rfree  %.4f/%.4f  ->  %.4f/%.4f" % (
                    cycle_str,
                    job.data["ini_overall_r"], job.data["ini_free_r"],
                    job.data["overall_r"], job.data["free_r"])
        return "%s   R  %.4f  ->  %.4f" % (
                cycle_str, job.data["ini_overall_r"], job.data["overall_r"])
    return cycle_str

# example:
#     Alternative reindexing        Lklhd      CC     R(E^2)    Number Cell_deviation
#    [-k,-l,h+k+l]           0.079    0.029    0.506     61253      2.99
_POINTLESS_ALTREINDEX_MATCH = re.compile(r"^\s+\[[hkl+, -]+\][ \t0-9+.eE-]+$")

def _float_or_nan(s):
    try:
        x = float(s)
    except ValueError:
        x = float('nan')
    return x

def _pointless_parser(job):
    # the line with 'reflections copied' is the last we read
    # to avoid reading 'Alternative reindexing' duplicated in the summary
    if "refl_out" not in job.data:
        for line in job.out.read_line():
            if line.startswith("Maximum resolution used:"):
                job.data["resol"] = float(line.split(":")[1])
            elif line.startswith("Number of reflections:"):
                job.data["refl_ref"] = int(line.split(":")[1])
            elif _POINTLESS_ALTREINDEX_MATCH.match(line):
                s = line.split()
                job.data.setdefault('alt_reindex', []).append(
                        {'op': s[0],
                         'cc': _float_or_nan(s[2]),
                         'cell_deviat': _float_or_nan(s[-1])})
            elif line.startswith("   Cell:") and 'output_cell' not in job.data:
                s = line.split()[1:]
                job.data["output_cell"] = tuple(float(i) for i in s)
            elif "reflections copied to output file" in line:
                job.data["refl_out"] = int(line.split()[0])
                break
    resol_txt = _format("%.2f", job.data.get("resol"))
    refl_out = job.data.get("refl_out", "")
    #refl_ref = job.data.get("refl_ref", "")
    return "resol. %4s A   #refl: %5s" % (resol_txt, refl_out)


def _phaser_parser(job):
    d = job.data
    for line in job.out.read_line():
        if line.startswith('*** Phaser Module:'):
            d['status'] = '[%s]' % line[19:70].strip().lower()
        if line.startswith('   SOLU SET ') and 'LLG=' in line:
            d['status'] = line[12:].strip()
            if len(d['status']) > 52:
                d['status'] = d['status'][:50].rsplit(' ', 1)[0] + '...'
        if line.startswith('   Sorry - No solution'):
            d['status'] = line.strip()
            if 'No solution with all components' in line:
                d['partial_solution'] = 'yes'
        if line.startswith('   SOLU SPAC '):
            d['SG'] = line[13:].strip()
    return "%-48s" % d.get('status', '')

def _truncate_parser(job):
    for line in job.out.read_line():
        if line.startswith(' Least squares straight line gives:'):
            b_str = line.partition(':')[-1].strip(' \tB=').split()[0]
            job.data["B-factor"] = float(b_str)
    return "B=%4s" % _format("%.1f", job.data.get("B-factor"))

def _ctruncate_parser(job):
    d = job.data
    for line in job.out.read_line():
        if line.startswith("$GRAPHS: Wilson plot - estimated B factor ="):
            d["B-factor"] = float(line.partition('=')[-1].split()[0])
        elif line.startswith("Eigenvalue ratios:"):
            d["eigval-ratios"] = tuple(float(i) for i in line.split()[-3:])
        elif line.startswith("L statistic ="):
            d["L-test"] = float(line.partition('=')[-1].split()[0])
    return "B=%4s   aniso %14s   L-test:%s" % (
            _format("%.1f", d.get("B-factor")),
            _format("%.2f:%.2f:%.2f", d.get("eigval-ratios")),
            _format("%.2f", d.get("L-test")))


def ccp4_job(workflow, prog, logical=None, ki="", parser=None, add_end=True):
    """Handle traditional convention for arguments of CCP4 programs.
    logical - dictionary with where keys are so-called logical names.
    ki (string or list of lines) - Keyworded Input to be passed through stdin.
    add_end - adds "end" as the last line of stdin

    Note: "colin" and "labin" mean the same (column label),
    but different programs use different keywords.
    """
    job = Job(workflow, utils.cbin(prog))
    if logical:
        for a in ["hklin", "hklout", "hklref", "xyzin", "xyzout", "libin"]:
            if logical.get(a):
                job.args += [a.upper(), logical[a]]
    lines = ki.splitlines() if isinstance(ki, basestring) else ki
    stripped = [a.strip() for a in lines if a and not a.isspace()]
    if add_end and not (stripped and stripped[-1].lower() == "end"):
        stripped.append("end")
    job.std_input = "\n".join(stripped)
    job.parser = parser
    return job


def _print_progress(job, event):
    while not event.wait(0.5):
        p = job.parse()
        if p is not None:
            text = (_elapsed_fmt % (time.time() - job.started)) + p
            utils.put_temporarily(text)
            utils.reset_color()


def _start_enqueue_thread(file_obj):
    def enqueue_lines(f, q):
        for line in iter(f.readline, b''):
            q.put(line)
        f.close()
    que = Queue.Queue()
    thr = threading.Thread(target=enqueue_lines, args=(file_obj, que))
    thr.daemon = True
    thr.start()
    return thr, que

def _get_input_as_string(job):
    if job.stdin_file:
        path = job.workflow.path(job.stdin_file)
        try:
            return open(path, "rb").read()
        except IOError:
            raise JobError("cannot read input from: %s" % job.stdin_file)
    else:
        return job.std_input

def _run_and_parse(process, job):
    try:
        # job.*.que can be used by parsers (via Output.read_line() or directly)
        out_t, job.out.que = _start_enqueue_thread(process.stdout)
        err_t, job.err.que = _start_enqueue_thread(process.stderr)
        try:
            job_input = _get_input_as_string(job)
            process.stdin.write(job_input)
        except IOError as e:
            utils.put("\nWarning: passing input to %s failed.\n" % job.name)
            if e.errno not in (errno.EPIPE, e.errno != errno.EINVAL):
                raise
        process.stdin.close()
        out_t.join()
        err_t.join()
        process.wait()
        # nothing is written to the queues at this point
        # parse what's left in the queues
        job.parse()
    finally:
        # take care of what is left by the parser
        job.out.finish_que()
        job.err.finish_que()


class Workflow:
    def __init__(self, output_dir, from_job=0):
        self.output_dir = os.path.abspath(output_dir)
        self.jobs = []
        self.file_info = {}
        self.temporary_files = set()
        self.from_job = from_job  # skip jobs before from_job (for testing)
        if from_job >= 1:
            try:
                _pkl = self.load_pickle()
                self.repl_jobs = _pkl.jobs
                self.file_info = _pkl.file_info
            except:
                self.repl_jobs = None
        self.dry_run = False
        self.argv = sys.argv
        if not os.path.isdir(self.output_dir):
            try:
                os.makedirs(self.output_dir)
            except OSError as e:
                utils.put_error(e)
                sys.exit(1)
        # this can seriously affect Refmac compiled with GFortran
        bad_var = os.getenv('GFORTRAN_UNBUFFERED_ALL')
        if bad_var and bad_var[0] not in ('0', 'n', 'N'):
            utils.put_error(
                    '$GFORTRAN_UNBUFFERED_ALL may terribly slow down Refmac',
                    comment='It is unset internally in dimple.')
            del os.environ['GFORTRAN_UNBUFFERED_ALL']
        # avoid html-like crap in the output of CCP4 program
        os.environ['CCP_SUPPRESS_HTML'] = '1'


    def __str__(self):
        return "Workflow with %d jobs @ %s" % (len(self.jobs), self.output_dir)

    def path(self, rel_path):
        return os.path.join(self.output_dir, rel_path)

    def dump_pickle(self):
        with open(self.path(PICKLE_FILENAME), "wb") as f:
            pickle.dump(self, f, -1)

    def load_pickle(self):
        with open(self.path(PICKLE_FILENAME), "rb") as f:
            return pickle.load(f)

    def run_job(self, job, show_progress, new_line=True):
        if not hasattr(sys.stdout, 'isatty') or not sys.stdout.isatty():
            show_progress = False
        self.jobs.append(job)
        job_num = len(self.jobs)
        if new_line:
            utils.put("\n" + _jobindex_fmt % job_num)
            utils.put_green(_jobname_fmt % job.name)
        else:
            utils.put(" / %d" % job_num)
        sys.stdout.flush()
        utils.log_section(job.name)

        job_idx = len(self.jobs) - 1
        if job_idx < self.from_job - 1: # from_job is 1-based
            # unpickle or skip
            if self.repl_jobs and len(self.repl_jobs) > job_idx:
                old_job = self.repl_jobs[job_idx]
                if old_job.name == job.name:
                    job = old_job
                    utils.put("unpickled")
                    utils.log_value("not_run", "unpickled")
                    self.jobs[-1] = job
                else:
                    utils.put("skipped (mismatch)")
                    utils.log_value("not_run", "unpickled/mismatch")
            else:
                utils.put("skipped")
                utils.log_value("not_run", "skipped")
            return

        job.started = time.time()
        utils.log_time("start_time", job.started)
        if job.stdin_file:
            utils.log_value("stdin", job.stdin_file)
        elif job.std_input:
            utils.log_value("input", job.std_input)
        utils.log_value("prog", job.args[0])
        utils.log_value("args", " ".join(pipes.quote(a) for a in job.args[1:]))
        #job.args[0] = "true"  # for debugging
        try:
            process = Popen(job.args, stdin=PIPE, stdout=PIPE, stderr=PIPE,
                            cwd=self.output_dir)
        except OSError as e:
            if e.errno == errno.ENOENT:
                raise JobError("Program not found: %s\n" % job.args[0])
            else:
                raise

        if self.dry_run:
            return

        if show_progress:
            event = threading.Event()
            progress_thread = threading.Thread(target=_print_progress,
                                               args=(job, event))
            progress_thread.daemon = True
            progress_thread.start()

        try:
            if job.parser is not None or show_progress:
                _run_and_parse(process, job)
            else:
                job_input = _get_input_as_string(job)
                out, err = process.communicate(input=job_input)
                job.out.lines = out.splitlines(True)
                job.err.lines = err.splitlines(True)
        except KeyboardInterrupt:
            raise JobError("KeyboardInterrupt while running %s" % job.name,
                           note=job.args_as_str())
        finally:
            if show_progress:
                event.set()
                progress_thread.join()
            end_time = time.time()
            job.total_time = end_time - job.started
            utils.log_time("end_time", end_time)
            job.exit_status = process.poll()
            if new_line:
                utils.put(_elapsed_fmt % job.total_time)
            parse_output = job.parse()
            utils.put("%s" % (parse_output or ""))
            if parse_output:
                utils.log_value("info", parse_output)
            self._write_logs(job)
            for k, v in job.data.iteritems():
                utils.log_value(k, v)
        if job.exit_status != 0:
            utils.log_value("exit_status", job.exit_status)

    def _write_logs(self, job):
        log_basename = "%02d-%s" % (len(self.jobs), job.name.replace(" ","_"))
        for output in (job.out, job.err):
            output.save_output(self.output_dir, log_basename)

    def remove_hetatm(self, xyzin, xyzout, remove_all):
        with open(self.path(xyzout), "wb") as out:
            return pdb.remove_hetatm(self.path(xyzin), out, remove_all)

    def read_pdb_metadata(self, xyzin):
        if xyzin not in self.file_info:
            self.file_info[xyzin] = pdb.read_metadata(self.path(xyzin))
        return self.file_info[xyzin]

    def read_mtz_metadata(self, hklin):
        if hklin not in self.file_info:
            self.file_info[hklin] = mtz.read_metadata(self.path(hklin))
        return self.file_info[hklin]

    def count_mtz_missing(self, hklin, col):
        key = (hklin, 'missing', col)
        if key not in self.file_info:
            self.file_info[key] = mtz.get_num_missing(self.path(hklin), col)
        return self.file_info[key]

    def molrep(self, f, m, keys=""):
        job = Job(self, utils.cbin("molrep"))
        job.args += ["-f", f, "-m", m]
        if keys:
            job.args.append("-i")
            job.std_input = keys.strip() + "\nend"
        return job

    def phaser_auto(self, hklin, labin, model, root, solvent_percent,
                    sg_alt="NONE"):
        lines = ['MODE MR_AUTO',
                 'SEARCH METHOD FAST',
                 'SEARCH DEEP OFF',
                 'ENSEMBLE p PDBFILE "%(pdb)s" IDENTITY %(identity)g' % model,
                 'SEARCH ENSEMBLE p NUM %(num)d' % model,
                 'COMPOSITION BY SOLVENT',
                 'COMPOSITION PERCENTAGE %f' % solvent_percent,
                 # Since Phaser 2.5.6 template matched solutions are moved
                 # to the template solution origin. Which is better than
                 # getting a solution one cell away, so we set template here.
                 'SOLUTION TEMPLATE original_model',
                 'SOLUTION 6DIM ENSE p EULER 0 0 0 FRAC 0 0 0',
                 #'PURGE ROT NUM 20',
                 #'PURGE TRA NUM 20',
                 'RESOLUTION HIGH 3.0',
                 #'MACANO PROTOCOL OFF',
                 'HKLIN "%s"' % hklin,
                 'LABIN %s' % labin,
                 'SGALTERNATIVE SELECT %s' % sg_alt,
                 'KILL TIME 120',  # 2h is much more than we want
                 'ROOT %s' % root,
                 ]
        # tNCS: we go with what phaser does by default -- tNCS of order 2
        # are handled automatically. While we could specify tNCS for
        # pseudo-tripling/quadrupling of the cell (TNCS NMOL 3) I don't know
        # if it'd do more good or bad.
{'fraction_non_alphanumeric': 0.13201162466658706, 'fraction_numerical': 0.008479636928221665, 'mean_word_length': 1.837560011296244, 'pattern_counts': {'":': 8, '<': 17, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 2, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1242393', 'n_tokens_mistral': 16535, 'n_tokens_neox': 15893, 'n_words': 3869}