#!/usr/bin/env python
''' Quick 3D SPLEEM visualization script. Requires mayavi '''

import sys
import pylab as pl
from enthought.mayavi import mlab
import Image

w = 100
X, Y = pl.meshgrid(range(-w/2, w/2), range(-w/2, w/2))
R = pl.sqrt(X**2+Y**2)

def im2arr(path, resize=False):
    ''' Loads an image as a 2d array, resized to a square, w by w '''
    image = Image.open(path)
    image = image.convert("I")
    if resize:
        image = image.resize(resize)
    newArr = pl.fromstring(image.tostring(), pl.uint8).astype(pl.int16)
    newArr = pl.reshape(newArr, (100, 100))
    return newArr

def load_spleem(up_path, dn_path):
    ''' Loads spleem data from two leem images '''
    up = im2arr(up_path, resize=(w, w))
    down = im2arr(dn_path, resize=(w, w))
    spleem = (up-down) / (up+down)
    return spleem

def touchup(im):
    ''' removes integer offset (naiive interpretation,
    subtracts mean) and crops to circle '''
    image = im - im[pl.where(im != im.max())].mean()
    image[pl.where(R>(w/2-5))] = 0
    return image

def loadimages(p):
    ''' Loads an arbitrary number of images,
    touches them up and returns them as a tuple '''
    ret = []
    for path in p: ret.append(touchup(im2arr(path, resize=(w, w))))
    return tuple(ret)

def loadspleems(p):
    ''' Load several spleems at once '''
    ret = []
    for i in range(0, len(p), 2): ret.append(load_spleem(p[i], p[i+1]))
    return tuple(ret)

if __name__ == '__main__':
    x, y, z = (False, False, False)

    if len(sys.argv) > 1: f = sys.argv[1:]
    else: f = raw_input('Which files would you like to evaluate?: ').split()

    if len(f) == 3: x, y, z = loadimages(f)
    elif len(f) == 6: x, y, z = loadspleems(f)

    good = x.any() and y.any() and z.any()
    assert good, 'Images have not been initialized.'
    if not good: sys.exit(1)

    #z *= 0.2 # Scale z coordinate differently to enhance in-plane contrast.
    mlab.figure(bgcolor=(1, 1, 1), size=(600, 400))
    y = -1 * y
    pts = mlab.quiver3d(x, y, z, scalars=y, colormap='RdYlBu', mode='cone', resolution=60, mask_points=3)
    pts.glyph.color_mode = 'color_by_scalar'
    mlab.view(*(95, 53, 105, pl.array([50, 70, 0])))
    mlab.savefig('test_output.png', magnification=3)
    #mlab.show()
    #while True:
    #    print mlab.veiw()
    #mlab.savefig('snapshot.png', size=(2400, 1600))