#!/usr/bin/python

import sys
import os
import csv

# csv name -> (site, scan_params_key)
csv_tr = {'phenotypic_CALTECH.csv': 'Caltech', 
          'phenotypic_CMU.csv': 'CMU', 
          'phenotypic_KKI.csv': 'KKI', 
          'phenotypic_LEUVEN_1.csv': 'Leuven', 
          'phenotypic_LEUVEN_2.csv': 'Leuven', 
          'phenotypic_MAX_MUN.csv': 'MaxMun', 
          'phenotypic_NYU.csv': 'NYU', 
          'phenotypic_OHSU.csv': 'OHSU', 
          'phenotypic_OLIN.csv': 'Olin', 
          'phenotypic_PITT.csv': 'Pitt', 
          'phenotypic_SBL.csv': 'SBL', 
          'phenotypic_SDSU.csv': 'SDSU', 
          'phenotypic_STANFORD.csv': 'Stanford', 
          'phenotypic_TRINITY.csv': 'Trinity', 
          'phenotypic_UCLA_1.csv': 'UCLA', 
          'phenotypic_UCLA_2.csv': 'UCLA', 
          'phenotypic_UM_1.csv': 'UM', 
          'phenotypic_UM_2.csv': 'UM', 
          'phenotypic_USM.csv': 'USM', 
          'phenotypic_YALE.csv': 'Yale'}

field_strengths = {'Caltech': '3T', 
                   'CMU_a': '3T', 
                   'CMU_b': '3T', 
                   'KKI': None, 
                   'Leuven_1': '3T', 
                   'Leuven_2': '3T', 
                   'MaxMun_a': '3T', 
                   'MaxMun_b': '3T', 
                   'MaxMun_c': '3T', 
                   'MaxMun_d': '3T', 
                   'NYU': '3T', 
                   'OHSU': '3T', 
                   'Olin': '3T', 
                   'Pitt': '3T', 
                   'SBL': '3T', 
                   'SDSU': '3T', 
                   'Stanford': '3T', 
                   'Trinity': '3T', 
                   'UCLA_1': '3T', 
                   'UCLA_2': '3T', 
                   'UM_1': '3T', 
                   'UM_2': '3T', 
                   'USM': '3T', 
                   'Yale': '3T'}

scan_params = {
    'Caltech': {
        'anat':  {'tr': '1590',  'te': '2.73', 'ti': '800',  'flip': '10'}, 
        'rest':  {'tr': '2000',  'te': '30',   'ti': None,   'flip': '75'}},
    'CMU_a': {
        'anat':  {'tr': '1870',  'te': '2.48', 'ti': '1100', 'flip': '8'}, 
        'rest':  {'tr': '2000',    'te': None,   'ti': None,   'flip': None}}, 
    'CMU_b': {
        'anat':  {'tr': '1870',  'te': '2.48', 'ti': '1100', 'flip': '8'}, 
        'rest':  {'tr': '1500',    'te': None,   'ti': None,   'flip': None}}, 
    'KKI': {
        'anat':  {'tr': None,    'te': None,   'ti': None,   'flip': '8'},  
        'rest':  {'tr': '2500',  'te': '30',   'ti': None,   'flip': '75'}},
    'Leuven_1': {
        'anat':  {'tr': None,    'te': '4.6',  'ti': None,   'flip': '8'}, 
        'rest':  {'tr': None,    'te': '33',   'ti': None,   'flip': '90'}},
    'Leuven_2': {
        'anat':  {'tr': None,    'te': '4.6',  'ti': None,   'flip': '8'},
        'rest':  {'tr': None,    'te': '33',   'ti': None,   'flip': '90'}},
    'MaxMun_a': {
        'anat':  {'tr': '1800',  'te': '3.06', 'ti': '900',  'flip': '9'}, 
        'rest':  {'tr': '3000',  'te': '30',   'ti': None,   'flip': '80'}}, 
    'MaxMun_b': {
        'anat':  {'tr': '1800',  'te': '3.06', 'ti': '900',  'flip': '9'}, 
        'rest':  {'tr': '3000',  'te': '30',   'ti': None,   'flip': '80'}}, 
    'MaxMun_c': {
        'anat':  {'tr': '1800',  'te': '3.06', 'ti': '900',  'flip': '9'}, 
        'rest':  {'tr': '3000',  'te': '30',   'ti': None,   'flip': '80'}}, 
    'MaxMun_d': {
        'anat':  {'tr': '1800',  'te': '3.06', 'ti': '900',  'flip': '9'}, 
        'rest':  {'tr': '3000',  'te': '30',   'ti': None,   'flip': '80'}}, 
    'NYU': {
        'anat':  {'tr': '2530',  'te': '3.25', 'ti': '1100', 'flip': '7'},  
        'rest':  {'tr': '2000',  'te': '15',   'ti': None,   'flip': '90'}},
    'OHSU': {
        'anat':  {'tr': '2300',  'te': '3.58', 'ti': '900',  'flip': '10'}, 
        'rest':  {'tr': '2500',  'te': '30',   'ti': None,   'flip': '90'}}, 
    'Olin': {
        'anat':  {'tr': '2500',  'te': '2.74', 'ti': '900',  'flip': '8'}, 
        'rest':  {'tr': '1500',  'te': '27',   'ti': None,   'flip': '60'}}, 
    'Pitt': {
        'anat':  {'tr': '2100',  'te': '3.93', 'ti': '1000', 'flip': '7'}, 
        'rest':  {'tr': '1500',  'te': '25',   'ti': None,   'flip': '70'}}, 
    'SBL': {
        'anat':  {'tr': None,    'te': None,   'ti': None,   'flip': '8'}, 
        'rest':  {'tr': '2200',  'te': '30',   'ti': None,   'flip': '80'}}, 
    'SDSU': {
        'anat':  {'tr': None,    'te': None,   'ti': '600',  'flip': '8'}, 
        'rest':  {'tr': '2000',  'te': '30',   'ti': None,   'flip': '90'}}, 
    'Stanford': {
        'anat':  {'tr': None,    'te': None,   'ti': None,   'flip': '15'}, 
        'rest':  {'tr': '2000',  'te': '30',   'ti': None,   'flip': '80'}}, 
    'Trinity': {
        'anat':  {'tr': None,    'te': None,   'ti': None,   'flip': '8'}, 
        'rest':  {'tr': '2000',  'te': '28',   'ti': None,   'flip': '90'}}, 
    'UCLA_1': {
        'anat':  {'tr': '2300',  'te': '2.84', 'ti': '853',  'flip': '9'}, 
        'rest':  {'tr': '3000',  'te': '28',   'ti': None,   'flip': '90'}, 
        'hires': {'tr': '5000',  'te': '34',   'ti': '',     'flip': '90'}}, 
    'UCLA_2': {
        'anat':  {'tr': '2300',  'te': '2.84', 'ti': '853',  'flip': '9'}, 
        'rest':  {'tr': '3000',  'te': '28',   'ti': None,   'flip': '90'}, 
        'hires': {'tr': '5000',  'te': '34',   'ti': '',     'flip': '90'}}, 
    'UM_1': {
        'anat':  {'tr': None,    'te': None,   'ti': None,   'flip': None}, 
        'rest':  {'tr': '2000',  'te': '30',   'ti': None,   'flip': '90'}}, 
    'UM_2': {
        'anat':  {'tr': None,    'te': None,   'ti': None,   'flip': None}, 
        'rest':  {'tr': '2000',  'te': '30',   'ti': None,   'flip': '90'}}, 
    'USM': {
        'anat':  {'tr': '2300',  'te': '2.91', 'ti': '900',  'flip': '9'}, 
        'rest':  {'tr': '2000',  'te': '28',   'ti': None,   'flip': '90'}}, 
    'Yale': {
        'anat':  {'tr': '1230',  'te': '1.73', 'ti': '624',  'flip': '9'}, 
        'rest':  {'tr': '2000',  'te': '25',   'ti': None,   'flip': '60'}}}

# one from each base data directory (to catch all the scan parameters)
test_ids = ['51471', # Caltech
            '50664', # CMU_a
            '50650', # CMU_b
            '50778', # KKI
            '50711', # Leuven_1
            '50732', # Leuven_2
            '51366', # MaxMun_a
            '51324', # MaxMun_b
            '51334', # MaxMun_c
            '51351', # MaxMun_d
            '50989', # NYU
            '50147', # OHSU
            '50109', # Olin
            '50010', # Pitt
            '51556', # SBL
            '50216', # SDSU
            '51163', # Stanford
            '50237', # Trinity
            '51252', # UCLA_1
            '51304', # UCLA_2
            '50275', # UM_1
            '50388', # UM_2
            '50433', # USM
            '50552' # Yale
           ]

def extract_site(full_site):
    if full_site.endswith('_1') \
       or full_site.endswith('_2') \
       or full_site.endswith('_a') \
       or full_site.endswith('_b') \
       or full_site.endswith('_c') \
       or full_site.endswith('_d'):
        return full_site[:-2]
    return full_site

# subject information, indexed by subject ID
subjects = {}

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

if len(sys.argv) != 3:
    print 'usage: %s <csv dir> <data dir>' % progname
    sys.exit(1)

csv_dir = sys.argv[1]
data_dir = sys.argv[2]

for fname in ('subjects.csv', 'sessions.csv', 'scans.csv', 'files.csv'):
    if os.path.exists(fname):
        sys.stderr.write('%s: %s exists; bailing out\n' % (progname, fname))
        sys.exit(1)

print 'reading subject information...'

subjects_fo = open('subjects.csv', 'w')
subjects_writer = csv.writer(subjects_fo)
subjects_writer.writerow(['ID', 
                          'Gender', 
                          'Age', 
                          'Handedness', 
                          'Diagnosis', 
                          'Test'])

for fname in sorted(os.listdir(csv_dir)):
    try:
        site = csv_tr[fname]
    except KeyError:
        print ' ** WARNING: skipping %s' % fname
        continue
    print '    %s' % fname
    reader = csv.reader(open('%s/%s' % (csv_dir, fname), 'rU'))
    headers = reader.next()
    for row in reader:
        row_dict = dict(zip(headers, row))
        if row_dict['SUB_ID'] == '':
            print '  ** WARNING: skipping row in %s: %s' % (fname, 
                                                            ','.join(row))
            continue
        if row_dict['SUB_ID'] in test_ids:
            test_flag = 'true'
            test_ids.remove(row_dict['SUB_ID'])
        else:
            test_flag = 'false'
        subject_id = '%s_%s' % (site, row_dict['SUB_ID'])
        if row_dict['DX_GROUP'] == '1':
            diagnosis = 'Autism'
        elif row_dict['DX_GROUP'] == '2':
            diagnosis = 'Control'
        else:
            sys.stderr.write('bad diagnosis group "%s" in %s\n' \
                             % (row_dict['DX_GROUP'], fname))
            sys.exit(1)
        age = int(float(row_dict['AGE_AT_SCAN']))
        if row_dict['SEX'] == '1':
            gender = 'Male'
        elif row_dict['SEX'] == '2':
            gender = 'Female'
        else:
            raise ValueError('bad gender %s for %s' % (row_dict['SEX'], 
                                                       subject_id))
        if row_dict['HANDEDNESS_CATEGORY'] == '-9999' \
            or row_dict['HANDEDNESS_SCORES'] == '-9999':
                handedness = 'Unknown'
        elif row_dict['HANDEDNESS_CATEGORY'] == 'L':
            handedness = 'Left'
        elif row_dict['HANDEDNESS_CATEGORY'] == 'R':
            handedness = 'Right'
        elif row_dict['HANDEDNESS_CATEGORY'] == 'Ambi':
            handedness = 'Ambidextrous'
        elif row_dict['HANDEDNESS_CATEGORY'] == 'NA':
            handedness = 'Unknown'
        elif row_dict['HANDEDNESS_CATEGORY'] == 'L->R':
            handedness = 'Left'
        elif row_dict['HANDEDNESS_CATEGORY'] not in ('', 'Mixed'):
            fmt = 'handedness category for %s is %s'
            raise ValueError(fmt % (subject_id, 
                                    row_dict['HANDEDNESS_CATEGORY']))
        else:
            score = float(row_dict['HANDEDNESS_SCORES'])
            if score == 0.0:
                handedness = 'Ambidextrous'
            elif score < 0.0:
                handedness = 'Left'
            else:
                handedness = 'Right'
        subjects_writer.writerow([subject_id, 
                                  gender, 
                                  age, 
                                  handedness, 
                                  diagnosis, 
                                  test_flag])
        first_subject = False

subjects_fo.close()

if test_ids:
    print 'ERROR: missed test subject(s):'
    for tid in test_ids:
        print '    %s' % tid
    sys.exit(1)

print 'descending into data directories...'

sessions_fo = open('sessions.csv', 'w')
sessions_writer = csv.writer(sessions_fo)
sessions_writer.writerow(['ID', 'Subject ID', 'Field strength', 'Resting TR'])

scans_fo = open('scans.csv', 'w')
scans_writer = csv.writer(scans_fo)
scans_writer.writerow(['Name', 
                       'Session ID', 
                       'Type', 
                       'Quality', 
                       'Flip angle', 
                       'TR', 
                       'TE', 
                       'TI'])

files_fo = open('files.csv', 'w')
files_writer = csv.writer(files_fo)
files_writer.writerow(['Session ID', 'Scan name', 'Format', 'Path'])

for site_dir in os.listdir(data_dir):
    site = extract_site(site_dir)
    print '    %s: %s' % (site_dir, site)
    for subject_dir in os.listdir('%s/%s' % (data_dir, site_dir)):
        if not os.path.isdir('%s/%s/%s' % (data_dir, site_dir, subject_dir)):
            continue
        subject_id = '%s_%s' % (site, subject_dir.lstrip('0'))
        session_id = subject_id
        params = scan_params[site_dir]
        resting_tr = None
        session_dir = '%s/%s/%s/session_1' % (data_dir, site_dir, subject_dir)
        scan_dirs = []
        for f in os.listdir(session_dir):
            if os.path.isdir('%s/%s' % (session_dir, f)):
                scan_dirs.append(f)
        if 'anat_1' in scan_dirs:
            scans_writer.writerow(['anat', 
                                   session_id, 
                                   'anat', 
                                   'Usable', 
                                   params['anat']['flip'], 
                                   params['anat']['tr'], 
                                   params['anat']['te'], 
                                   params['anat']['ti']])
            nifti_fname = '%s/anat_1/mprage.nii.gz' % session_dir
            files_writer.writerow([session_id, 
                                   'anat', 
                                   'NIfTI', 
                                   nifti_fname[len(data_dir)+1:]])
            scan_dirs.remove('anat_1')
        if 'hires_1' in scan_dirs:
            scans_writer.writerow(['hires', 
                                   session_id, 
                                   'hires', 
                                   'Usable', 
                                   params['hires']['flip'], 
                                   params['hires']['tr'], 
                                   params['hires']['te'], 
                                   params['hires']['ti']])
            nifti_fname = '%s/hires_1/hires.nii.gz' % session_dir
            files_writer.writerow([session_id, 
                                   'hires', 
                                   'NIfTI', 
                                   nifti_fname[len(data_dir)+1:]])
            scan_dirs.remove('hires_1')
        if len(scan_dirs) == 1:
            scans_writer.writerow(['rest', 
                                   session_id, 
                                   'rest', 
                                   'Usable', 
                                   params['rest']['flip'], 
                                   params['rest']['tr'], 
                                   params['rest']['te'], 
                                   params['rest']['ti']])
            nifti_fname = '%s/rest_1/rest.nii.gz' % session_dir
            files_writer.writerow([session_id, 
                                   'rest', 
                                   'NIfTI', 
                                   nifti_fname[len(data_dir)+1:]])
            resting_tr = params['rest']['tr']
        else:
            for scan in scan_dirs:
                scans_writer.writerow([scan, 
                                       session_id, 
                                       'rest', 
                                       'Usable', 
                                        params['rest']['flip'], 
                                        params['rest']['tr'], 
                                        params['rest']['te'], 
                                        params['rest']['ti']])
                if resting_tr:
                    # make sure resting TR is the same across scans
                    assert resting_tr == params['rest']['tr']
                else:
                    resting_tr = params['rest']['tr']
                nifti_fname = '%s/%s/rest.nii.gz' % (session_dir, scan)
                files_writer.writerow([session_id, 
                                       scan, 
                                       'NIfTI', 
                                       nifti_fname[len(data_dir)+1:]])
        if not resting_tr:
            resting_tr = ''
        sessions_writer.writerow([session_id, 
                                  subject_id, 
                                  field_strengths[site_dir], 
                                  resting_tr])

files_fo.close()
scans_fo.close()
sessions_fo.close()

sys.exit(0)

# eof
