#!/usr/bin/python26

import sys
import os
import tempfile
import numpy
import PIL.Image
import nibabel
import pyxnat

border = 5

colors = {}

for line in open('FreeSurferColorLUT.txt'):
    line = line.strip()
    if not line or line.startswith('#'):
        continue
    (index, label, r, g, b, a) = line.split()
    index = int(index)
    if index == 0:
        continue
    r = int(r)
    g = int(g)
    b = int(b)
    colors[index] = (r, g, b, 255)
max_index = max(colors)
color_table = numpy.zeros((max_index+1, 4), dtype='uint8')
for i in xrange(max_index+1):
    if i in colors:
        color_table[i] = colors[i]

i = pyxnat.Interface(os.environ['XNAT_URI'].rstrip('/'), 
                     os.environ['XNAT_USER'], 
                     os.environ['XNAT_PASSWORD'])

for subject in i.select.project('cs_schizbull08').subjects():
    print subject.label()
    experiment = subject.experiment('%s_MR' % subject.label())
    scan = experiment.scan('anat')
    seg = experiment.assessor('%s_MR_seg' % subject.label())
    anat_fname = scan.resource('NIfTI').file('anat.nii.gz').get()
    seg_fname = seg.out_resource('NIfTI').file('seg.nii.gz').get()
    anat = nibabel.load(anat_fname)
    seg = nibabel.load(seg_fname)
    anat_data = anat.get_data()
    seg_data = seg.get_data()
    assert anat_data.shape == seg_data.shape
    slice = anat_data.shape[2] / 2
    # determine the border of zeroed voxels so we can cut them off
    nonzeros_0 = anat_data.max(0).nonzero()[0]
    nonzeros_1 = anat_data.max(1).nonzero()[0]
    # height
    range_0 = [nonzeros_0.min()-border, nonzeros_0.max()+border]
    range_0[0] = max(0, range_0[0])
    range_0[1] = min(255, range_0[1])
    side_0 = range_0[1] - range_0[0] + 1
    # width
    range_1 = [nonzeros_1.min()-border, nonzeros_1.max()+border]
    range_1[0] = max(0, range_1[0])
    range_1[1] = min(255, range_1[1])
    side_1 = range_1[1] - range_1[0] + 1
    slice_data = anat_data[range_1[0]:range_1[1]+1,range_0[0]:range_0[1]+1,slice].transpose()
    scale = 255.0 / slice_data.max()
    slice_data = slice_data * scale
    slice_data = slice_data.astype('uint8')
    anat_im = PIL.Image.fromarray(slice_data, 'L')
    anat_im = anat_im.convert('RGB')
    seg_slice_data = seg_data[range_1[0]:range_1[1]+1,range_0[0]:range_0[1]+1,slice].transpose()
    seg_slice_rgb = color_table[seg_slice_data]
    seg_im = PIL.Image.fromarray(seg_slice_rgb, 'RGBA')
    comp_im = PIL.Image.new('RGB', (2*anat_im.size[0], anat_im.size[1]), 0x00)
    comp_im.paste(anat_im, (0, 0))
    comp_im.paste(anat_im, (anat_im.size[0], 0))
    comp_im.paste(seg_im, (anat_im.size[0], 0), seg_im)
    comp_im = comp_im.convert('P', palette=PIL.Image.ADAPTIVE)
    fname = tempfile.mktemp('.gif')
    comp_im.save(fname)
    qc_fname = '%s_MR_anat_qc_t.gif' % subject.label()
    scan.resource('SNAPSHOTS').file(qc_fname).put(fname, 'GIF', 'THUMBNAIL')
    os.unlink(fname)

sys.exit(0)

# eof
