#!/usr/bin/python

# See file COPYING distributed with fsutils for copyright and license.

import sys
import os
import argparse
import datetime
import dateutil.parser

version = '0.3.1'

class FSStatusError(Exception):

    """base class for exceptions"""

    def __init__(self, msg):
        self.msg = msg
        return

    def __str__(self):
        return self.msg

class SubjectError(FSStatusError):

    """error reading subject status"""

class NoRunError(FSStatusError):

    """no run found"""

    def __init__(self):
        return

    def __str__(self):
        return 'no run found'

class LogError(FSStatusError):

    """error in recon-all.log"""

class Subject:

    def __init__(self, spec, subjects_dir=None):
        self.subjects_dir = subjects_dir
        self.spec = spec
        if not self.subjects_dir or '/' in spec:
            self._init_from_path(self.spec)
        else:
            log = os.path.join(self.subjects_dir, 
                               spec, 
                               'scripts', 
                               'recon-all.log')
            self._init_from_path(log)
        return

    def __getitem__(self, key):
        if isinstance(key, int):
            return self.runs[key]
        raise KeyError(key)

    def __len__(self):
        return len(self.runs)

    def __iter__(self):
        for run in self.runs:
            yield run
        return

    def _init_from_path(self, path):
        if os.path.isdir(path):
            log = os.path.join(path, 'scripts', 'recon-all.log')
            self._init_from_path(log)
            return
        self._read_log(path)
        return

    def _read_log(self, path):
        if not os.path.exists(path):
            raise SubjectError('%s does not exist' % path)
        self.runs = []
        with open(path) as fo:
            while True:
                try:
                    run = Run(self, len(self.runs), fo)
                except NoRunError:
                    break
                except Exception:
                    import traceback
                    traceback.print_exc()
                    raise LogError('error parsing log file')
                else:
                    self.runs.append(run)
        if not self.runs:
            raise SubjectError('no runs found')
        for run in self.runs:
            run.order_vars['NR'] = len(self.runs)
            run.print_vars['NR'] = len(self.runs)
        return

class Run:

    def __init__(self, subject, run_number, fo):
        self.subject = subject
        self.run_number = run_number
        self.t_end = None
        self.steps = []
        end_state = self.read_log(fo)
        if end_state == 'start':
            raise NoRunError()
        self.print_vars = {}
        self.order_vars = {}
        self.print_vars['BUILD'] = self.build_stamp
        self.order_vars['BUILD'] = self.build_stamp
        self.print_vars['UNAME'] = self.uname
        self.order_vars['UNAME'] = self.uname
        self.print_vars['ARGS'] = self.args
        self.order_vars['ARGS'] = self.args
        # NR is set by the subject object
        # RN is printed as one-based while run_number is zero-based
        self.order_vars['RN'] = self.run_number
        self.print_vars['RN'] = self.run_number+1
        self.order_vars['TSTART'] = self.t_start
        self.print_vars['TSTART'] = format_time(self.t_start)
        if not self.t_end:
            self.order_vars['TFINISH'] = None
            if args.dash:
                self.print_vars['TFINISH'] = '-          -'
            else:
                self.print_vars['TFINISH'] = ''
            self.order_vars['TRUN'] = now-self.t_start
            self.print_vars['TRUN'] = format_delta(now-self.t_start)
            self.order_vars['R'] = True
            self.print_vars['R'] = 'R'
            self.order_vars['E'] = None
            if args.dash:
                self.print_vars['E'] = '-'
            else:
                self.print_vars['E'] = ''
            if self.steps:
                self.order_vars['SN'] = len(self.steps)
                self.print_vars['SN'] = str(len(self.steps))
                self.order_vars['STEPNAME'] = self.steps[-1]['name']
                self.print_vars['STEPNAME'] = self.steps[-1]['name']
                tstart = self.steps[-1]['tstart']
                self.order_vars['STEPTSTART'] = tstart
                self.print_vars['STEPTSTART'] = format_time(tstart)
                self.order_vars['STEPTRUN'] = now-tstart
                self.print_vars['STEPTRUN'] = format_delta(now-tstart)
            else:
                self.order_vars['SN'] = None
                self.order_vars['STEPNAME'] = None
                self.order_vars['STEPTSTART'] = None
                self.order_vars['STEPTRUN'] = None
                if args.dash:
                    self.print_vars['SN'] = '-'
                    self.print_vars['STEPNAME'] = '-'
                    self.print_vars['STEPTSTART'] = '-'
                    self.print_vars['STEPTRUN'] = '-'
                else:
                    self.print_vars['SN'] = ''
                    self.print_vars['STEPNAME'] = ''
                    self.print_vars['STEPTSTART'] = ''
                    self.print_vars['STEPTRUN'] = ''
        else:
            self.order_vars['TFINISH'] = self.t_end
            self.print_vars['TFINISH'] = format_time(self.t_end)
            self.order_vars['TRUN'] = self.t_end-self.t_start
            self.print_vars['TRUN'] = format_delta(self.t_end-self.t_start)
            self.order_vars['R'] = False
            if args.dash:
                self.print_vars['R'] = '-'
            else:
                self.print_vars['R'] = ''
            if self.error:
                self.order_vars['E'] = True
                self.print_vars['E'] = 'E'
            else:
                self.order_vars['E'] = False
                if args.dash:
                    self.print_vars['E'] = '-'
                else:
                    self.print_vars['E'] = ''
            self.order_vars['SN'] = None
            self.order_vars['STEPNAME'] = None
            self.order_vars['STEPTSTART'] = None
            self.order_vars['STEPTRUN'] = None
            if args.dash:
                self.print_vars['SN'] = '-'
                self.print_vars['STEPNAME'] = '-'
                self.print_vars['STEPTSTART'] = '-'
                self.print_vars['STEPTRUN'] = '-'
            else:
                self.print_vars['SN'] = ''
                self.print_vars['STEPNAME'] = ''
                self.print_vars['STEPTSTART'] = ''
                self.print_vars['STEPTRUN'] = ''
        self.order_vars['SPEC'] = self.subject.spec
        self.print_vars['SPEC'] = self.subject.spec
        return

    def read_log(self, fo):
        """read a recon-all.log

        states are:

            start
            header block
            post-header
            memory block
            post-memory
            verions block
            step
            end

        lines in the header block:

            date
            pwd
            $0
            args
            subjid
            SUBJECTS_DIR
            FREESURFER_HOME
            actual FREESURFER_HOME
            build stamp
            uname -a
        """
        state = 'start'
        headers = []
        for line in fo:
            line = line.strip()
            if line.startswith('To report a problem, see'):
                # this appears sometimes after a run (after we are done with 
                # this processing), so it will show up at the start of our 
                # processing of the next run
                continue
            if line.startswith('New invocation'):
                # appears between runs; ignore
                continue
            if line == r'\n\n':
                # sometimes appears between runs; ignore
                continue
            if state == 'start':
                if line:
                    state = 'header block'
                    headers.append(line)
                continue
            if state == 'header block':
                if line:
                    headers.append(line)
                else:
                    state = 'post-header'
                    self.t_start = dateutil.parser.parse(headers[0])
                    self.pwd = headers[1]
                    self.script_name = headers[2]
                    self.args = headers[3]
                    self.subjid = headers[4]
                    self.subjects_dir = headers[5]
                    self.freesurfer_home = headers[6]
                    self.actual_freesurfer_home = headers[7]
                    self.build_stamp = headers[8]
                    if self.build_stamp.startswith('build-stamp.txt: '):
                        self.build_stamp = self.build_stamp[17:]
                    self.uname = headers[9]
                continue
            if state == 'post-header':
                if line:
                    state = 'memory block'
                continue
            if state == 'memory block':
                if not line:
                    state = 'post-memory'
                continue
            if state == 'post-memory':
                if line:
                    state = 'versions block'
                continue
            if line.startswith('#@# '):
                step = line[4:-28].strip()
                t = dateutil.parser.parse(line[-28:])
                self.steps.append({'name': step, 'tstart': t})
                state = 'step'
                continue
            # state == 'step'
            if 'exited with ERRORS' in line:
                self.error = True
                self.t_end = dateutil.parser.parse(line[-28:])
                state = 'end'
                break
            if 'finished without error' in line:
                self.error = False
                self.t_end = dateutil.parser.parse(line[-28:])
                state = 'end'
                break
        return state

def run_cmp(a, b):
    """cmp() on runs, ordering by the fields in order_fields"""
    for field in order_fields:
        c = cmp(a.order_vars[field], b.order_vars[field])
        if c != 0:
            return c
    return 0

# available output fields
fields = ('BUILD', 'UNAME', 'ARGS', 'NR', 'RN', 'TSTART', 'TFINISH', 'TRUN', 
          'E', 'R', 'SPEC', 'SN', 'STEPNAME', 'STEPTSTART', 'STEPTRUN')

progname = os.path.basename(sys.argv[0])

description = """

Report on the status of FreeSurfer runs.

By default, {progname} reports on the latest run for each subject in 
$SUBJECTS_DIR.

An alternate subjects directory can be specified using the -S flag.

Individual subjects can be selected at the command line.  If these 
specifiers include a forward slash or if no subjects directory is given, 
they are taken as paths to subject directories or recon-all.log files.

Examples:

    {progname} -- report on the last run for all subjects in $SUBJECTS_DIR

    {progname} -S /data/study_2 S001 S002 S003 -- report on the last run 
        for subjects S001, S002, and S003 in SUBJECTS_DIR=/data/study_2

    {progname} ./S001 /tmp/recon-all.log -- report on the last run for 
        the subject in S001 and the subject described by /tmp/recon-all.log

Output can be customized by a comma-separated list of fields:

    BUILD -- the FreeSurfer build stamp

    UNAME -- the uname -a report

    ARGS -- the run arguments

    NR -- number of runs

    RN -- the run number being reported on

    TSTART -- the run start time

    TFINISH -- the run end time

    TRUN -- the run time

    SN -- the number of the running step

    STEPNAME -- the name of the running step

    STEPTSTART -- the start time of the running step

    STEPTRUN -- the run time of the running step

    E -- 'E' if the run exited with errors

    R -- 'R' if the case is still running

    SPEC -- the subject specifier

Reporting order can be customized by a comma-separated list of the same 
fields.

The run number to be reported on can be specified using the -r option and 
defaults to the last run.  Use "-r all" to report on all runs.

""".format(progname=progname)

def format_time(t):
    return t.strftime('%Y-%m-%d %H:%M:%S')

def format_delta(d):
    s = d.seconds % 60
    minutes = d.seconds / 60
    m = minutes % 60
    hours = minutes / 60
    h = 24 * d.days + hours
    return '%02d:%02d:%02d' % (h, m, s)

parser = argparse.ArgumentParser(description=description, 
                                 formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument('-v', '--version', 
                    default=False, 
                    action='store_true', 
                    dest='version', 
                    help='show version and exit')
parser.add_argument('-H', '--suppress-header', 
                    default=False, 
                    action='store_true', 
                    dest='suppress_header', 
                    help='suppress header')
parser.add_argument('-o', '--output', 
                    dest='output_fields', 
                    default='NR,RN,TSTART,TFINISH,TRUN,E,SPEC', 
                    help='output fields')
parser.add_argument('-O', '--order', 
                    dest='order_fields', 
                    default='SPEC,RN', 
                    help='order fields')
parser.add_argument('--dash', '-d', 
                    default=False, 
                    action='store_true', 
                    help='print dashes for empty values')
parser.add_argument('--long', '-l', 
                    default=False, 
                    action='store_true', 
                    help='print a long listing (one field per line)')
parser.add_argument('--run', '-r', 
                    default=None, 
                    help='run number')
parser.add_argument('--subjects-dir', '-S')
parser.add_argument('subjects', 
                    metavar='subject', 
                    nargs='*')

args = parser.parse_args()

if args.version:
    print '%s %s' % (progname, version)
    sys.exit(0)

output_fields = []
for field in args.output_fields.split(','):
    field = field.upper().strip()
    if field not in fields:
        sys.stderr.write('%s: unknown output field "%s"\n' % (progname, field))
        sys.exit(1)
    output_fields.append(field)

order_fields = []
for field in args.order_fields.split(','):
    field = field.upper().strip()
    if field not in fields:
        sys.stderr.write('%s: unknown order field "%s"\n' % (progname, field))
        sys.exit(1)
    order_fields.append(field)

if args.run is None:
    run = 'last'
elif args.run == 'all':
    run = 'all'
else:
    try:
        run = int(args.run)
    except ValueError:
        sys.stderr.write('%s: bad value for run number\n' % progname)
        sys.exit(2)
    if run <= 0:
        sys.stderr.write('%s: bad value for run number\n' % progname)
        sys.exit(2)

subjects_dir = args.subjects_dir
if not subjects_dir and 'SUBJECTS_DIR' in os.environ:
    subjects_dir = os.environ['SUBJECTS_DIR']

subjects = []

now = datetime.datetime.now(dateutil.tz.tzlocal())

if not subjects_dir:
    if not args.subjects:
        parser.print_usage(sys.stderr)
        fmt = '%s: error: no subjects directory or subjects given\n'
        sys.stderr.write(fmt % progname)
        sys.exit(2)
    for spec in args.subjects:
        try:
            subjects.append(Subject(spec))
        except FSStatusError, data:
            sys.stderr.write('%s: %s: %s\n' % (progname, spec, str(data)))
else:
    if not os.path.isdir(subjects_dir):
        msg = '%s: %s is not a directory\n' % (progname, subjects_dir)
        sys.stderr.write(msg)
        sys.exit(1)
    if not args.subjects:
        for subject in os.listdir(subjects_dir):
            try:
                subjects.append(Subject(subject, subjects_dir=subjects_dir))
            except SubjectError:
                # subjects directory given but no specific subjects given 
                # means that we're asked to scan for subjects, so in this 
                # case we won't report on errors in subjects we try
                pass
            except LogError, data:
                msg = '%s: %s: %s\n' % (progname, subject, str(data))
                sys.stderr.write(msg)
    else:
        for subject in args.subjects:
            try:
                subjects.append(Subject(subject, subjects_dir=subjects_dir))
            except FSStatusError, data:
                msg = '%s: %s: %s\n' % (progname, subject, str(data))
                sys.stderr.write(msg)

if not subjects:
    sys.stderr.write('%s: no subjects found\n' % progname)
    sys.exit(1)

runs = []
for subject in subjects:
    if run == 'all':
        runs.extend(subject)
    elif run == 'last':
        runs.append(subject[-1])
    else:
        try:
            runs.append(subject[run-1])
        except IndexError:
            fmt = '%s: %s has no run %d\n'
            sys.stderr.write(fmt % (progname, subject.spec, run))

if not runs:
    sys.stderr.write('%s: no runs found\n' % progname)
    sys.exit(1)

runs.sort(run_cmp)

if args.long:

    for run in runs:
        for field in output_fields:
            print '%s: %s' % (field, run.print_vars[field])
        print

else:

    headers = []
    fmts = []
    for field in output_fields:
        widths = [ len(field) ]
        for run in runs:
            widths.append(len('{0}'.format(run.print_vars[field])))
        width = max(widths)
        headers.append(field.ljust(width))
        fmts.append('{%s:%d}' % (field, width))

    header = '  '.join(headers).strip()
    fmt = '  '.join(fmts).strip()

    if not args.suppress_header:
        print header
    for run in runs:
        print fmt.format(**run.print_vars).rstrip()

sys.exit(0)

# eof
