#!/usr/bin/python

import sys
import os
import csv

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

if len(sys.argv) != 2:
    print 'usage: %s <base directory>' % progname
    sys.exit(1)

base_dir = sys.argv[1]

csv_files = ('BeijingEnhanced.csv', 
             'BeijingShortTR.001.001.csv', 
             'BeijingShortTR.002.001.csv', 
             'BeijingShortTR.003.001.csv', 
             'BeijingShortTR.004.001.csv', 
             'BeijingShortTR.005.001.csv', 
             'BeijingShortTR.006.001.csv', 
             'BeijingShortTR.007.001.csv', 
             'BeijingShortTR.008.001.csv', 
             'BeijingShortTR.009.001.csv')

subjects = {}

for fname in csv_files:
    fo = open('%s/%s' % (base_dir, fname), 'rU')
    reader = csv.reader(fo)
    header = reader.next()
    if header[-1] == '':
        header.pop()
    for row in reader:
        row_dict = dict(zip(header, row))
        subject_id = int(row_dict['SubjectID'])
        assert subject_id not in subjects
        age = int(row_dict['Age'])
        gender = {'F': 'female', 'M': 'male'}[row_dict['Sex']]
        if row_dict.has_key('Handedness'):
            handedness = {'L': 'left', 'R': 'right'}[row_dict['Handedness']]
        else:
            handedness = 'unknown'
        test = 'true'
        subjects[subject_id] = {'age': age, 
                                'gender': gender, 
                                'handedness': handedness, 
                                'test': test, 
                                'sessions': {}}
    fo.close()

for dir in os.listdir(base_dir):
    if not os.path.isdir('%s/%s' % (base_dir, dir)):
        continue
    subject_id = int(dir)
    assert subject_id in subjects
    subject = subjects[subject_id]
    for (dirpath, dirnames, filenames) in os.walk('%s/%s' % (base_dir, dir)):
        if dirpath.endswith('/phenotypic'):
            continue
        if not filenames:
            continue
        (session, scan) = dirpath.split('/')[-2:]
        assert session.startswith('session_')
        session_num = int(session[8:])
#        assert scan[-2] == '_'
#        scan_type = scan[:-2]
#        scan_num = int(scan[-1])
        session_id = '%d_%d' % (subject_id, session_num)
        d = subject['sessions'].setdefault(session_id, {'scans': {}})
        d['scans'][scan] = {'files': {}}
        files = d['scans'][scan]['files']
        for fname in filenames:
            path = '%s/%s' % (dirpath, fname)
            # strip the path to the base directory
            path = path[len(base_dir)+1:]
            if fname.endswith('.nii.gz'):
                files.setdefault('NIfTI', []).append(path)
            elif path.endswith('bval'):
                files.setdefault('bval', []).append(path)
            elif path.endswith('bvec'):
                files.setdefault('bvec', []).append(path)
            else:
                raise TypeError('uncaught type for file %s' % fname)

assert not os.path.exists('subjects.csv')
fo = open('subjects.csv', 'w')
writer = csv.writer(fo)
writer.writerow(['ID', 'Gender', 'Age', 'Handedness', 'Test'])
for subject_id in sorted(subjects):
    row = [str(subject_id), 
           subjects[subject_id]['age'], 
           subjects[subject_id]['gender'], 
           subjects[subject_id]['handedness'], 
           subjects[subject_id]['test']]
    writer.writerow(row)
fo.close()

assert not os.path.exists('sessions.csv')
fo = open('sessions.csv', 'w')
writer = csv.writer(fo)
writer.writerow(['Subject ID', 'ID'])
for subject_id in sorted(subjects):
    for session_id in sorted(subjects[subject_id]['sessions']):
        writer.writerow([str(subject_id), session_id])
fo.close()

assert not os.path.exists('scans.csv')
fo = open('scans.csv', 'w')
writer = csv.writer(fo)
writer.writerow(['Session ID', 'Name'])
for subject_id in sorted(subjects):
    subject = subjects[subject_id]
    for session_id in sorted(subject['sessions']):
        session = subject['sessions'][session_id]
        for scan_name in sorted(session['scans']):
            writer.writerow([session_id, scan_name])
fo.close()

assert not os.path.exists('files.csv')
fo = open('files.csv', 'w')
writer = csv.writer(fo)
writer.writerow(['Session ID', 'Scan Name', 'Format', 'Path'])
for subject_id in sorted(subjects):
    subject = subjects[subject_id]
    for session_id in sorted(subject['sessions']):
        session = subject['sessions'][session_id]
        for scan_name in sorted(session['scans']):
            scan = session['scans'][scan_name]
            for format in sorted(scan['files']):
                for path in scan['files'][format]:
                    writer.writerow([session_id, scan_name, format, path])
fo.close()

sys.exit(0)

# eof
