#!/usr/bin/env python3
#
# parsemarks version 2.0
# NIMH/MEG Core staff

import os
import sys
import re
import argparse

class MarkerFile():
    def __init__(self, dataset, num_markers):
        self.dataset = dataset
        self.num_markers = num_markers
        self.markers = []

    def __str__(self):
        return 'MarkerFile[dataset={0.dataset},num_markers={0.num_markers}]'.format(self)

class Marker():
    def __init__(self, name, classgroupid, classid, num_samples):
        self.name = name
        self.classgroupid = classgroupid
        self.classid = classid  # 1,2,3, ... # order of marker set in dataset
        self.classgroup = self.getGroup()
        self.num_samples = num_samples
        self.trials = []  # list of tuples

    def __str__(self):
        return 'Marker[name={0.name},classgroup={0.classgroup},classid={0.classid},num_samples={0.num_samples}]'.format(self)

    def getGroup(self):
        group = {
            0:'TriggerGroup', 
            1:'TemplateGroup', 
            2:'ConditionGroup', 
            3:'ManualGroup',
        }
        return group.get(self.classgroupid, None)

    def get_trials(self, lines, start):
        trials = []
        for i in range(self.num_samples):
            trial_num, time_from_sync = lines[start+i].split()
            trial_num = int(trial_num)
            time_from_sync = float(time_from_sync)
            trials.append((self.name, trial_num, time_from_sync))
        self.trials = trials

def process_markerfile(filename):
    lines = []
    with open(filename, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            lines.append(line)

    for i, line in enumerate(lines):
        if line.startswith('PATH OF DATASET:'):
            dataset = lines[i+1]
        elif line.startswith('NUMBER OF MARKERS:'):
            num_markers = int(lines[i+1])
            markerfile = MarkerFile(dataset, num_markers)
        elif line.startswith('CLASSGROUPID:'):
            classgroupid = int(lines[i+1])
        elif line.startswith('NAME:'):
            marker_name = lines[i+1]
        elif line.startswith('CLASSID:'):
            classid = int(lines[i+1])
        elif line.startswith('NUMBER OF SAMPLES:'):
            num_samples = int(lines[i+1])
            marker = Marker(marker_name, classgroupid, classid, num_samples)
        elif line.startswith('LIST OF SAMPLES:'):         
            marker.get_trials(lines, i+2)
            markerfile.markers.append(marker)
    return markerfile

def write_markers(mf, labels):
    all_trials = []
    width = 0
    for marker in mf.markers:
        for trial in marker.trials:
            if len(trial[0]) > width:
                width = len(trial[0])
            all_trials.append(trial)

    for trial in sorted(all_trials, key=lambda trial: (trial[1], trial[2])):
        label, trial_num, time = trial
        if labels:
            print('{:{}} {:<2d} {:.3f}'.format(label, width, trial_num, time))
        else:
            print('{:<2d} {:.3f}'.format(trial_num, time))
            
if __name__ == '__main__':
    parser = argparse.ArgumentParser(usage="parsemarks.py [-h] [-l] dataset",
        description="Extract the marks from the marker file associated with dataset and print them in a useful format.")
    parser.add_argument('-l', action='store_true', help="the marks are labeled in the output. Useful for debugging.")
    parser.add_argument('dataset', help="path to CTF dataset or MarkerFile.mrk (required)")

    if len(sys.argv) == 1:
        sys.exit(parser.print_help())
    args = parser.parse_args()

    dsname = args.dataset
    if dsname.endswith('MarkerFile.mrk'):
        filename = dsname
    elif os.path.isdir(dsname):
        filename = os.path.join(dsname, 'MarkerFile.mrk')
        if not os.path.isfile(filename):
            sys.exit('Error opening marker file: {}'.format(filename))
    else:
        sys.exit('Error opening marker file: {}'.format(dsname))

    mf = process_markerfile(filename)
    write_markers(mf, args.l)