#!/usr/bin/env python3
#
# parsemarks_report version 1.0
# NIMH/MEG Core staff

import os
import sys
import re
import argparse
import xlsxwriter

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 study
        self.classgroup = self.getGroup()
        self.num_samples = num_samples
        self.comment = None
        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):
        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)

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('COMMENT:'):
            comment = 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)
            marker.comment = comment
        elif line.startswith('LIST OF SAMPLES:'):         
            marker.get_trials(lines, i+2)
            markerfile.markers.append(marker)
    return markerfile

def get_datasets(studydir):
    datasets = []
    for name in os.listdir(studydir):
        namedir = os.path.join(studydir, name)
        if name.endswith('.ds') and os.path.isdir(namedir):
            dsdir = namedir
            markerfile = os.path.join(dsdir, 'MarkerFile.mrk')
            if os.path.isfile(markerfile):
                datasets.append(markerfile)
    datasets.sort()
    return datasets

def get_allmarkers(datasets):
    recs = []
    for ds in datasets:
        mf = process_markerfile(ds)
        recs.append(mf)
    return recs

def mkExcel(xlsname, xlsdir, recs):
    try:
        if not os.path.isdir(xlsdir):
            os.mkdir(xlsdir, mode=0o775)
    except OSError as e:
        sys.exit('Error making directory: {}'.format(xlsdir))
    output = os.path.join(xlsdir, xlsname)
    print('MS Excel file: {}'.format(output))

    workbook = xlsxwriter.Workbook(output)
    worksheet1 = workbook.add_worksheet('all_markers')
    worksheet2 = workbook.add_worksheet('sheet2')
    worksheet3 = workbook.add_worksheet('sheet3')

    format1 = workbook.add_format({'font_name':'Calibri', 'bold':1, 'align':'left', 'valign':'vcenter', 'fg_color':'#D7E4BC', 'border':1, 'size':10})
    format2 = workbook.add_format({'font_name':'Calibri', 'align':'left', 'size':10})
    format3 = workbook.add_format({'font_name':'Calibri', 'num_format':'yyyy-mm-dd', 'align':'left', 'size':10})

    # write column headers
    worksheet1.write(0,0,'dataset',format1)
    worksheet1.write(0,1,'marker_name',format1)
    worksheet1.write(0,2,'classgroup',format1)
    worksheet1.write(0,3,'num_samples',format1)
    worksheet1.write(0,4,'comment',format1)
  
    for i in range(5):
        worksheet1.set_column(i,i,12)
    worksheet1.set_column(0,0,60)
    worksheet1.set_column(4,4,100)
    worksheet1.freeze_panes(1,0)

    i = 1
    for mf in recs:
        for marker in mf.markers:
            worksheet1.write(i,0,mf.dataset,format2)
            worksheet1.write(i,1,marker.name, format2)
            worksheet1.write(i,2,marker.classgroup, format2)
            worksheet1.write_number(i,3,marker.num_samples, format2)
            worksheet1.write(i,4,marker.comment, format2)
            i += 1
        i += 1
    workbook.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(usage="parsemarks_report.py [-h] [-v] studydir",
        description="Reports on the marker set from every dataset directory under a study directory.  MEG \
        studies consists of a collection of datasets, each with its own MarkerFile.mrk, organized under a top level \
        (studydir) directory. Output is an excel file stored in your ~/excel folder.")
    parser.add_argument('studydir', help="path to a toplevel directory holding a set of dataset directories (required)")
    if len(sys.argv) == 1:
        sys.exit(parser.print_help())
    args = parser.parse_args()

    studydir = args.studydir
    if studydir.endswith('/'):
        studydir = studydir[:-1]
    study_name = os.path.basename(studydir)

    if not os.path.isdir(studydir):
        sys.exit('Error opening studydir: {}'.format(studydir))

    datasets = get_datasets(studydir)
    recs = get_allmarkers(datasets)

    # xlsdir = studydir    # modify for your site
    xlsdir = os.path.join(os.environ['HOME'], 'excel')
    xlsname = '{}_{}.xlsx'.format(study_name, 'markers_report')
    mkExcel(xlsname, xlsdir, recs)