#! /usr/bin/env python
"""Show the status of Freesurfer recon-all processing."""
import os
import sys
import argparse
import os.path as op
from glob import glob
from datetime import datetime
import numpy as np


def main(arglist):

    # Parse commandline
    args = parse_cmdline(arglist)

    # Figure out the subjects directory
    if args.subjects_dir is not None:
        subjects_dir = args.subjects_dir
    else:
        subjects_dir = os.environ["SUBJECTS_DIR"]

    if args.subjects:
        # Possibly look up a specific set of subjects
        subjects = args.subjects
        subj_logs = [op.join(subjects_dir, s,
                             "scripts/recon-all-status.log")
                     for s in args.subjects]
    else:
        # Otherwise look at everyone with a log
        subj_logs = glob(op.join(subjects_dir, "*",
                         "scripts/recon-all-status.log"))
        subj_logs = [s for s in subj_logs if "fsaverage" not in s]
        subjects = [p.split("/")[-3] for p in subj_logs]
        sorter = np.argsort(subjects)
        subjects = np.array(subjects)[sorter]
        subj_logs = np.array(subj_logs)[sorter]

    # Get the current time
    now = datetime.now()

    # Iterate through the subjects and print status messages
    for subj, log in zip(subjects, subj_logs):
        try:
            log_file = open(log)
        except IOError:
            print "%s: Could not open status log" % log
            continue
        log_entries = log_file.readlines()
        log_file.close()
        # Prune off empty lines
        while not log_entries[-1]:
            log_entries.pop()
        status = log_entries[-1]
        timestr = " ".join(status.split()[-6:])
        timestamp = datetime.strptime(timestr, "%a %b %d %H:%M:%S %Z %Y")

        if status.startswith("#@#"):
            # This string means it's still going
            uptime = str(now - timestamp)[:-7]
            stage = " ".join(status.split()[1:-6])
            print "%s: Stage - %s  Uptime - %s" % (subj, stage, uptime)
        else:
            # Otherwise it's done
            donefor = (now - timestamp).days
            if not args.inprogress and not (args.recent and donefor > 5):
                msg = status.replace("recon-all -s %s " % subj, "")[:-9]
                msg = msg.replace("finished", "Finished")
                msg = msg.replace("exited", "Exited")
                msg = msg.replace("without error ", "")
                print "%s: %s" % (subj, msg)


def parse_cmdline(arglist):

    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument("-dir", dest="subjects_dir",
                        help="override SUBJECTS_DIR variable")
    parser.add_argument("-subjects", nargs="*", metavar="SUBJ",
                        help="show status for specific subject(s)")
    parser.add_argument("-inprogress", action="store_true",
                        help="skip if recon is finished")
    parser.add_argument("-recent", action="store_true",
                        help="skip if finished more than 5 days ago")
    return parser.parse_args(arglist)


if __name__ == "__main__":
    main(sys.argv[1:])
