#!python

__author__ = 'SungHo Lee (shlee@unc.edu)'
__version_info__ = ('2017', '11', '02')

import os, argparse, glob, re
from pynit.handler.images import BrukerRawData

parser = argparse.ArgumentParser(prog='checkscans', description="Check the Scans in the Subjects of Bruker's Raw files")
parser.add_argument("input", help="Input Image", type=str)

args = parser.parse_args()

def check_subject(subject_path):
    # load subject information
    header = BrukerRawData(subject_path).subject

    cur_path = os.getcwd()
    path = os.path.join(cur_path, subject_path)
    regex = re.compile('^[0-9]+$')
    scan_list = sorted([int(n) for n in os.listdir(path) if (os.path.isdir(os.path.join(path, n)) and regex.match(n))])
    context = []
    context.append('*' * 90)
    context.append('Folder: {}'.format(subject_path))
    context.append('Scan date: {:<22}\tResearcher: {}'.format(header['SUBJECT_date'], header['SUBJECT_name_string']))
    context.append('Protocol: {:<22}\tSession: {:<20}\tStudyNo: {}'.format(header['SUBJECT_location'],
                                                                   header['SUBJECT_study_name'],
                                                                   header['SUBJECT_study_nr']))
    context.append('AnimalID: {:<22}\tSex: {:<19}\tWeight: {}kg'.format(header['SUBJECT_id'],
                                                              header['SUBJECT_sex'],
                                                              header['SUBJECT_weight']))
    context.append('*' * 90)
    context.append('Acquired scans:')
    for num in sorted(scan_list):
        try:
            _h = BrukerRawData(subject_path, num=num)
            _shape = 'x'.join([str(i) for i in _h.get_shape() if i != 1][::-1])
            scan_path = os.path.join(path, str(num))
            has_data = '- No data'
            reco = None
            if 'fid' in os.listdir(scan_path):
                has_data = ''
                pdata_path = os.path.join(scan_path, 'pdata')
                reco = [str(n) for n in os.listdir(pdata_path) if os.path.isdir(os.path.join(pdata_path, n))]
            try:
                te = "TE:{}".format(_h.acqp['ACQ_echo_time'])
            except:
                te = ''
            context.append(' [{}] {} ({}) <{}> TR:{} {} {}'.format(str(num), _h.acqp['ACQ_protocol_name'],
                                                                   _h.method['Method'], _shape,
                                                                   _h.acqp['ACQ_repetition_time'],
                                                                   te, has_data))
            if reco:
                _context = []
                for r in reco:
                    _r = BrukerRawData(subject_path, num=num, pid=r).reco
                    _context.append('    pid={}, mode={}, range={}'.format(r, _r['RECO_map_mode'],
                                                                           _r['RECO_map_range']))
                context.extend(_context)
        except:
            context.append(' [{}] -No data'.format(str(num)))
    context.append('\n')
    return "\n".join(context)

list = glob.glob(args.input+"*")
if list == []:
    print('There is no available folder to meet your input')
else:
    for folder in list:
        if os.path.isdir(folder):
            print(check_subject(folder))
        else:
            pass

