#!/usr/bin/python26

import sys
import os
import string
import re
import optparse
import csv
import StringIO
import Image
import nibabel
import pyxnat

try:
	set
except NameError:
	import sets
	set = sets.Set

############################################################################
# defaults and definitions
#

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

scan_chars = '_' + string.ascii_letters + string.digits

usage = """usage: %prog [options] <project> <subjs file> <sess file> <scan file> <files file> <data directory>
       %prog [options] -c <subjs file> <sess file> <scan file> <files file> <data directory>"""
description = """Uploads data to an XNAT.

Data must be prepared according to the instructions on NITRC.  Positional arguments are the (existing but empty) XNAT project, subjects CSV file, the sessions CSV file, the scans CSV file, the files CSV file, and the base data directory (relative to which paths in the scans file are generated).

The arguments for the base XNAT URL and the XNAT user and password override the envrionment variables XNAT_URL, XNAT_USER, XNAT_PASSWORD, respectively.
"""

parser = optparse.OptionParser(usage=usage, description=description)
parser.add_option('-c', '--check', dest='check_only', default=False, action='store_true', help='check sources, don\'t upload')
parser.add_option('-t', '--test', dest='test', default=False, action='store_true', help='upload only the test subjects')
parser.add_option('-x', '--xnat', dest='xnat_url', help='the base XNAT URL')
parser.add_option('-u', '--user', dest='user', help='the XNAT user')
parser.add_option('-p', '--password', dest='password', help='XNAT password')

class CSVError(Exception):

    """error in a CSV file"""

    def __init__(self, msg, fname=None, line=None):
        self.msg = msg
        self.fname = fname
        self.line = line
        return

    def __str__(self):
        if self.line:
            return '%s on line %d of in %s' % (self.msg, self.line, self.fname)
        if self.fname:
            return '%s in %s' % (self.msg, self.fname)
        return self.msg

############################################################################
# functions
#

#---------------------------------------------------------------------------
# validators
#

def string_not_null_val(value):
    value = value.strip()
    if not value:
        raise ValueError
    return value

def string_val(value):
    value = value.strip()
    if not value:
        return None
    return value

def gender_val(value):
    value = value.strip().lower()
    if value not in ('female', 'male', 'unknown'):
        raise ValueError
    return value

def handedness_val(value):
    value = value.strip().lower()
    if value not in ('left', 'right', 'ambidextrous', 'unknown'):
        raise ValueError
    return value

def int_val(value):
    value = value.strip()
    if not value:
        return None
    return int(value)

def float_val(value):
    value = value.strip()
    if not value:
        return None
    return float(value)

def diagnosis_val(value):
    value = value.strip()
    if not value:
        raise ValueError
    return value

def test_val(value):
    value = value.strip().lower()
    if value == 'true':
        return True
    if value == 'false':
        return False
    raise ValueError

def quality_val(value):
    value = value.strip().lower()
    if value not in ('usable', 'questionable', 'unusable'):
        raise ValueError
    return value

date_re = re.compile('^[0-9]{4}-[0-9]{2}-[0-9]{2}$')
def date_val(value):
    value = value.strip()
    if not date_re.search(value):
        raise ValueError
    return value

def format_val(value):
    value = value.strip()
    if ' ' in value:
        raise ValueError
    return value

#---------------------------------------------------------------------------
# CSV reader
#

def read_csv(fname, required_headers, optional_headers):
    rv = []
    cased_headers = {}
    for h in required_headers:
        cased_headers[h.lower()] = h
    for h in optional_headers:
        cased_headers[h.lower()] = h
    fo = open(fname)
    try:
        reader = csv.reader(fo)
        headers = []
        for header in reader.next():
            header = header.strip().lower()
            if header not in cased_headers:
                raise CSVError('unknown column "%s"'% header, fname)
            headers.append(cased_headers[header])
        for header in required_headers:
            if header not in headers:
                raise CSVError('no column "%s"' % header, fname)
        line = 1
        for row in reader:
            line += 1
            if len(row) != len(headers):
                raise CSVError('wrong number of values', fname, line)
            row_dict = dict(zip(headers, row))
            for header in headers:
                if header in required_headers:
                    validator = required_headers[header]
                else:
                    validator = optional_headers[header]
                try:
                    value = validator(row_dict[header])
                except ValueError:
                    raise CSVError('bad value for %s' % header, fname, line)
                row_dict[header] = value
            rv.append(row_dict)
    finally:
        fo.close()
    return rv

#---------------------------------------------------------------------------
# sorted(), in case we're running Python < 2.4
#

try:
	sorted
except NameError:
	def sorted(iterable):
		list = [ item for item in iterable ]
		list.sort()
		return list

############################################################################
# command line handling
#

(opts, args) = parser.parse_args()

if opts.check_only:
    if len(args) < 5:
        parser.error('not enough arguments')
    if len(args) > 5:
        parser.error('too many arguments')
    (subjects_csv, sessions_csv, scans_csv, files_csv, data_dir) = args
else:
    if len(args) < 6:
        parser.error('not enough arguments')
    if len(args) > 6:
        parser.error('too many arguments')
    (project, subjects_csv, sessions_csv, scans_csv, files_csv, data_dir) = args
    if not opts.xnat_url:
        try:
            opts.xnat_url = os.environ['XNAT_URL']
        except KeyError:
            parser.error('XNAT base url undefined')
    if not opts.user:
        try:
            opts.user = os.environ['XNAT_USER']
        except KeyError:
            parser.error('XNAT user undefined')
    if not opts.password:
        try:
            opts.password = os.environ['XNAT_PASSWORD']
        except KeyError:
            parser.error('XNAT password undefined')

data_dir = data_dir.rstrip('/')
opts.xnat_url = opts.xnat_url.rstrip('/')

############################################################################
# read and validate the CSV files
#

#---------------------------------------------------------------------------
# subjects
#

required_headers = {'ID': string_not_null_val}
optional_headers = {'Gender': gender_val, 
                    'Handedness': handedness_val, 
                    'Age': int_val, 
                    'Diagnosis': diagnosis_val, 
                    'Test': test_val}
try:
    subjects = read_csv(subjects_csv, required_headers, optional_headers)
except Exception, data:
    sys.stderr.write('%s: %s\n' % (progname, str(data)))
    sys.exit(1)
if not subjects:
    sys.stderr.write('%s: no subjects in %s\n' % (progname, subjects_csv))
    sys.exit(1)

if opts.test:
    if not subjects[0].has_key('Test'):
        sys.stderr.write('%s: test upload requested, but no Test column in subjects file\n' % progname)
        sys.exit(1)

#---------------------------------------------------------------------------
# sessions
#

required_headers = {'ID': string_not_null_val, 
                    'Subject ID': string_not_null_val}
optional_headers = {'Date': date_val, 
                    'Field Strength': string_val, 
                    'Resting TR': float_val}
try:
    sessions = read_csv(sessions_csv, required_headers, optional_headers)
except Exception, data:
    sys.stderr.write('%s: %s\n' % (progname, str(data)))
    sys.exit(1)
if not sessions:
    sys.stderr.write('%s: no sessions in %s\n' % (progname, sessions_csv))
    sys.exit(1)

#---------------------------------------------------------------------------
# scans
#

required_headers = {'Name': string_not_null_val, 
                    'Session ID': string_not_null_val}
optional_headers = {'Type': string_val, 
                    'Quality': quality_val, 
                    'Flip Angle': int_val, 
                    'TR': float_val, 
                    'TE': float_val, 
                    'TI': float_val}
try:
    scans = read_csv(scans_csv, required_headers, optional_headers)
except Exception, data:
    sys.stderr.write('%s: %s\n' % (progname, str(data)))
    sys.exit(1)
if not scans:
    sys.stderr.write('%s: no scans in %s\n' % (progname, scans_csv))
    sys.exit(1)

#---------------------------------------------------------------------------
# files
#

required_headers = {'Session ID': string_not_null_val, 
                    'Scan Name': string_not_null_val, 
                    'Path': string_not_null_val, 
                    'Format': format_val}
optional_headers = {}

try:
    files = read_csv(files_csv, required_headers, optional_headers)
except Exception, data:
    sys.stderr.write('%s: %s\n' % (progname, str(data)))
    sys.exit(1)
if not files:
    sys.stderr.write('%s: no files in %s\n' % (progname, scans_csv))
    sys.exit(1)

############################################################################
# check data
#

#---------------------------------------------------------------------------
# summarize data
#

subject_ids = set()
duplicate_subjects = set()
for subject in subjects:
    if subject['ID'] in subject_ids:
        duplicate_subjects.add(subject['ID'])
    subject_ids.add(subject['ID'])

session_ids = set()
duplicate_sessions = set()
referenced_subjects = set()
bad_subjects = set()
for session in sessions:
    if session['ID'] in session_ids:
        duplicate_sessions.add(session['ID'])
    session_ids.add(session['ID'])
    if session['Subject ID'] not in subject_ids:
        bad_subjects.add(session['Subject ID'])
    else:
        referenced_subjects.add(session['Subject ID'])
empty_subjects = subject_ids - referenced_subjects

# scan names need only be unique within a session, so I call the (unique) 
# (session ID, scan name) then scan ID
scan_ids = set()
bad_scan_names = set()
duplicate_scan_ids = set()
referenced_sessions = set()
bad_sessions = set()
for scan in scans:
    for c in scan['Name']:
        if c not in scan_chars:
            bad_scan_names.add(scan['Name'])
            break
    scan_id = (scan['Session ID'], scan['Name'])
    if scan_id in scan_ids:
        duplicate_scan_ids.add(scan_id)
    scan_ids.add(scan_id)
    if scan['Session ID'] not in session_ids:
        bad_sessions.add(scan['Session ID'])
    else:
        referenced_sessions.add(scan['Session ID'])
empty_sessions = session_ids - referenced_sessions

file_paths = set()
duplicate_file_paths = set()
referenced_scan_ids = set()
bad_scan_ids = set()
for file in files:
    if file['Path'] in file_paths:
        duplicate_file_paths.add(file['Path'])
    file_paths.add(file['Path'])
    scan_id = (file['Session ID'], file['Scan Name'])
    if scan_id not in scan_ids:
        bad_scan_ids.add(scan_id)
    else:
        referenced_scan_ids.add(scan_id)
empty_scans = scan_ids - referenced_scan_ids

empty_directories = set()
files_on_disk = set()
for (path, dir_names, file_names) in os.walk(data_dir):
    relpath = path[len(data_dir)+1:]
    if not file_names and not dir_names:
        empty_directories.add(path)
    for fname in file_names:
        files_on_disk.add('%s/%s' % (relpath, fname))

missing_files = file_paths - files_on_disk
extra_files = files_on_disk - file_paths

#---------------------------------------------------------------------------
# report warnings
#

for subject_id in sorted(empty_subjects):
    print 'WARNING: empty subject %s' % subject_id

for session_id in sorted(empty_sessions):
    print 'WARNING: empty session %s' % session_id

for (session_id, scan_name) in sorted(empty_scans):
    print 'WARNING: no files for scan %s in session %s' % \
          (scan_name, session_id)

for directory in sorted(empty_directories):
    print 'WARNING: empty directory %s' % directory

for fname in sorted(extra_files):
    print 'WARNING: unreferenced file %s' % fname

#---------------------------------------------------------------------------
# report errors
#

have_error = False

for subject_id in sorted(duplicate_subjects):
    sys.stderr.write('ERROR: duplicate subject %s\n' % subject_id)
    have_error = True

for subject_id in sorted(bad_subjects):
    sys.stderr.write('ERROR: unknown subject %s in sessions file\n' % 
                     subject_id)
    have_error = True

for session_id in sorted(duplicate_sessions):
    sys.stderr.write('ERROR: duplicate session %s\n' % session_id)
    have_error = True

for scan_name in sorted(bad_scan_names):
    sys.stderr.write('ERROR: bad scan name %s in scans file\n' % scan_name)
    have_error = True

for session_id in sorted(bad_sessions):
    sys.stderr.write('ERROR: unknown session %s in scans file\n' % 
                     session_id)
    have_error = True

for (session_id, scan_name) in sorted(duplicate_scan_ids):
    sys.stderr.write('ERROR: duplicate scan %s in session %s\n' % 
                     (scan_name, session_id))
    have_error = True

for (session_id, scan_name) in sorted(bad_scan_ids):
    sys.stderr.write('ERROR: unknown scan %s for session %s in files file\n' % 
                     (scan_name, session_id))
    have_error = True

for fname in missing_files:
    sys.stderr.write('ERROR: missing file %s in files file\n' % fname)
    have_error = True

if have_error:
    sys.exit(1)

if opts.check_only:
    sys.exit(0)

#---------------------------------------------------------------------------
# clean up variables used for checking
#

del subject_ids
del duplicate_subjects
del session_ids
del duplicate_sessions
del referenced_subjects
del bad_subjects
del empty_subjects
del scan_ids
del bad_scan_names
del duplicate_scan_ids
del referenced_sessions
del bad_sessions
del empty_sessions
del file_paths
del duplicate_file_paths
del referenced_scan_ids
del bad_scan_ids
del empty_scans
del empty_directories
del files_on_disk
del missing_files
del extra_files

############################################################################
# create XNAT entities and upload
#

interface = pyxnat.Interface(opts.xnat_url, opts.user, opts.password)
xnat_project = interface.select.project(project)
if not xnat_project.exists():
    sys.stderr.write('%s: XNAT project %s does not exist\n' % \
                     (progname, project))
    sys.exit(1)

if opts.test:
    n_subjects = 0
    for subject in subjects:
        if subject['Test']:
            n_subjects += 1
else:
    n_subjects = len(subjects)

i = 0
for subject in subjects:
    if opts.test and not subject['Test']:
        continue
    i += 1
    print 'Subject %s (%d of %d)' % (subject['ID'], i, n_subjects)
    xnat_subject = xnat_project.subject(subject['ID'])
    if xnat_subject.exists():
        print '    exists (skipping)'
        continue
    xnat_subject.create()
    if subject.has_key('Gender'):
        if subject['Gender'] in ('male', 'female'):
            xnat_subject.attrs.set('gender', subject['Gender'])
    if subject.has_key('Age'):
        if subject['Age'] is not None:
            xnat_subject.attrs.set('age', str(subject['Age']))
    if subject.has_key('Handedness'):
        if subject['Handedness'] in ('left', 'right', 'ambidextrous'):
            xnat_subject.attrs.set('handedness', subject['Handedness'])
    if subject.has_key('Diagnosis'):
        if subject['Diagnosis']:
            xnat_subject.attrs.set('group', subject['Diagnosis'])
    for session in sessions:
        if session['Subject ID'] != subject['ID']:
            continue
        print '    Session %s' % session['ID']
        xnat_session = xnat_subject.experiment(session['ID'])
        xnat_session.create()
        if session.has_key('Field Strength'):
            if session['Field Strength']:
                xnat_session.attrs.set('xnat:mrSessionData/fieldStrength', 
                                       session['Field Strength'])
        if session.has_key('Resting TR'):
            if session['Resting TR']:
                xnat_session.attrs.set('xnat:mrSession/RestingTR', 
                                       str(session['Resting TR']))
        for scan in scans:
            if scan['Session ID'] != session['ID']:
                continue
            print '        Scan %s' % scan['Name']
            xnat_scan = xnat_session.scan(scan['Name'])
            xnat_scan.create()
            for (key, xpath) in (('Quality', 'quality'), 
                                 ('TR', 'xnat:mrScanData/parameters/tr'), 
                                 ('TE', 'xnat:mrScanData/parameters/te'), 
                                 ('TI', 'xnat:mrScanData/parameters/ti'), 
                                 ('Flip Angle', 'xnat:mrScanData/parameters/flip')):
                if scan.has_key(key):
                    if scan[key] is not None:
                        xnat_scan.attrs.set(xpath, str(scan[key]))



            for file in files:
                if file['Session ID'] != session['ID'] \
                   or file['Scan Name'] != scan['Name']:
                    continue
                print '            File %s' % file['Path']
                full_path = '%s/%s' % (data_dir, file['Path'])
                file_contents = open(full_path).read()
                xnat_resource = xnat_scan.resource(file['Format'])
                file_name = os.path.basename(file['Path'])
                xnat_resource.file(file_name).create(file_contents, 
                                                     file['Format'], 
                                                     'RAW')
                if file['Format'].lower() != 'nifti':
                    continue
                # --- if we have a NIfTI file, we can create a snapshot
                print '            Thumbnail'
                vol = nibabel.load(full_path)
                vol_data = vol.get_data()
                if len(vol_data.shape) == 3:
                    # 3D (anatomical)
                    slice = vol_data.shape[0]/2
                    image_data = vol_data[slice,:,:]
                else:
                    # 4D (resting state)
                    slice = vol_data.shape[2]/2
                    image_data = vol_data[:,:,slice,0]
                image_data = image_data.transpose()
                image_data = image_data.astype('float32')
                scale = 255.0 / image_data.max()
                image_data = scale * image_data
                image_data = image_data.astype('uint8')
                im = Image.fromarray(image_data, 'L')
                im = im.transpose(Image.FLIP_TOP_BOTTOM)
                gif_buffer = StringIO.StringIO()
                im.save(gif_buffer, format='GIF')
                gif_data = gif_buffer.getvalue()
                gif_buffer.close()
                xnat_resource = xnat_scan.resource('SNAPSHOTS')
                xnat_resource.create()
                xnat_resource.file('qc.gif').put(gif_data, 'GIF', 'THUMBNAIL')

sys.exit(0)

# eof
