#!env python
#-----------------------------------------------------------------------------
# Copyright (c) 2013, NeXpy Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING, distributed with this software.
#-----------------------------------------------------------------------------

import os, getopt, glob, sys, timeit
import numpy as np
from nexpy.api.nexus import *

def get_files(directory, prefix, extension):
    os.chdir(directory)
    if not extension.startswith('.'):
        extension = '.'+extension
    filenames = glob.glob(prefix+'*'+extension)
    return sorted(filenames,key=natural_sort)

def open_nexus_file(directory, prefix, filenames):
    try:
        from libtiff import TIFF
        im = TIFF.open(filenames[0])
        v0 = im.read_image()
        x = NXfield(range(v0.shape[1]), dtype=np.uint16, name='x')
        y = NXfield(range(v0.shape[0]), dtype=np.uint16, name='y')
        z = NXfield(range(1,len(filenames)+1), dtype=np.uint16, name='z')
        v = NXfield(name='v',shape=(len(filenames),v0.shape[0],v0.shape[1]),
                    dtype=v0.dtype)
    except ImportError:
        import Image
        im = Image.open(filenames[0])
        dtype = np.dtype(np.uint16)
        if im.mode == "I;32" or im.mode == "I":
            dtype=np.dtype(np.uint32)
        z = NXfield(range(1,len(filenames)+1), dtype=np.uint16, name='z')
        y = NXfield(range(im.size[1]), name='y')
        x = NXfield(range(im.size[0]), name='x')
        v = NXfield(name='v', shape=(len(filenames),im.size[1],im.size[0]),
                    dtype=dtype)
    data = NXdata(v, (z,y,x))
    nexus_file = os.path.join(directory, prefix+'.nxs')
    NXroot(NXentry(data)).save(nexus_file)
    return nexus_file

def write_data(nexus_file, filenames):
    root = nxload(nexus_file, 'rw')
    for i in range(len(filenames)):
        try:
            im = TIFF.open(filenames[i])
            v = im.read_image()
        except NameError:
            im = Image.open(filenames[i])
            dtype = np.dtype(np.uint16)
            if im.mode == "I;32" or im.mode == "I":
                dtype=np.dtype(np.uint32)
            v = np.array(im.getdata(),dtype=dtype).reshape(im.size[1],im.size[0])
        root.entry.data.v.put(v, (i, 0, 0))
    root.save()

def natural_sort(key):
    import re
    return [int(t) if t.isdigit() else t for t in re.split(r'(\d+)', key)]    

if __name__=="__main__":
    try:
        opts, args = getopt.getopt(sys.argv[1:],"hd:e:",["directory=","ext="])
    except getopt.GetoptError:
        print 'merge_tiffs.py -d <directory> -e <extension> prefix'
        sys.exit(2)
    directory = './'
    extension = 'tif'
    first = None
    last = None
    for opt, arg in opts:
        if opt == '-h':
            print 'merge_tiffs.py prefix -d <directory> -e <extension>'
            sys.exit()
        elif opt in ('-d', '--directory'):
            directory = arg
        elif opt in ('-e', '--extension'):
            extension = arg
        elif opt == '-f':
            first = arg
        elif opt == '-l':
            last = arg
    prefix = args[0]
    filenames = get_files(directory, prefix, extension)
    nexus_file = open_nexus_file(directory, prefix, filenames)
    tic=timeit.default_timer()
    write_data(nexus_file, filenames)
    toc=timeit.default_timer()
    print toc-tic, 'seconds'
