#!/usr/bin/env python
# coding=utf-8
#
# Copyright 2012 Vicomtech-IK4
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

print """
This code was created by Luis Kabongo, Vicomtech-IK4 Copyright 2012-2013.
This application converts the slices found in a folder into a tiled 2D texture 
image in PNG format (it assumes all files in the folder are of the same type 
and dimensions). It uses Python with PIL, numpy and pydicom packages are 
recommended for other formats.
Information links:
http://www.volumerc.org
http://demos.vicomtech.org
Contact mailto:volumerendering@vicomtech.org
"""

import os, errno
import sys
import math
from PIL import Image  # this is required to manage the images

# This is the default size when loading a Raw image
sizeOfRaw = (512, 512)
# This determines if the endianness should be reversed
rawByteSwap = True
# Standard deviation for Gaussian kernel
sigmaValue = 2


def normalize(inputData):
    old_min = inputData.min()
    old_range = inputData.max() - old_min
    return (inputData - old_min) / old_range


# This function calculates the gradient from a 3 dimensional numpy array
def calculateGradient(arr):
    r = np.zeros(arr.shape)
    g = np.zeros(arr.shape)
    b = np.zeros(arr.shape)
    ndimage.gaussian_filter1d(arr, sigma=sigmaValue, axis=1, order=1, output=r)
    ndimage.gaussian_filter1d(arr, sigma=sigmaValue, axis=0, order=1, output=g)
    ndimage.gaussian_filter1d(arr, sigma=sigmaValue, axis=2, order=1, output=b)
    return normalize(np.concatenate((r[..., np.newaxis], g[..., np.newaxis], b[..., np.newaxis]), axis=3))


# Write here your own load data version here ...
def loadMyData(filename):
    myImageSize = (32, 32)
    im = Image.new("L", myImageSize)
    putpix = im.im.putpixel

    for y in range(myImageSize[1]):
        for x in range(myImageSize[0]):
            val = 0  # ... here get the value from your file (make sure that 0 <= val <= 255 )
            putpix((x, y), val)

    return im


# This function uses the images retrieved with loadImgFunction (whould return a PIL.Image) and
#	writes them as tiles within a new square Image. 
#	Returns a set of Image, size of a slice, number of slices and number of slices per axis
def ImageSlices2TiledImage(filenames, loadImgFunction=loadMyData):
    filenames = sorted(filenames)
    print "Desired load function=", loadImgFunction.__name__
    size = loadImgFunction(filenames[0]).size
    numberOfSlices = len(filenames)
    slicesPerAxis = int(math.ceil(math.sqrt(numberOfSlices)))
    imout = Image.new("L", (size[0] * slicesPerAxis, size[1] * slicesPerAxis))

    i = 0
    for filename in filenames:
        im = loadImgFunction(filename)

        row = int((math.floor(i / slicesPerAxis)) * size[0])
        col = int((i % slicesPerAxis) * size[1])

        box = (int(col), int(row), int(col + size[0]), int(row + size[1]))
        imout.paste(im, box)
        i += 1
        print "processed slice  : " + str(i) + "/" + str(numberOfSlices)  # filename

    return imout, size, numberOfSlices, slicesPerAxis


# This functions takes a (tiled) image and writes it to a png file with base filename outputFilename.
#	It also writes several versions in different sizes determined by dimensions
def WriteVersions(tileImage, outputFilename, dimensions=[8192, 4096, 2048, 1024]):
    try:
        print 'Creating folder', os.path.dirname(outputFilename), '...',
        os.makedirs(os.path.dirname(outputFilename))
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(os.path.dirname(outputFilename)):
            print 'was already there.'
        else:
            print ', folders might not be created, trying to write anyways...'
    except:
        print "Could not create folders, trying to write anyways..."

    print "Writing complete image: " + outputFilename + "_full.png"
    try:
        tileImage.save(outputFilename + "_full.png", "PNG")
    except:
        print "Failed writing ", outputFilename + "_full.png"
    for dim in dimensions:
        if tileImage.size[0] > dim:
            print "Writing " + str(dim) + "x" + str(dim) + " version: " + outputFilename + "_" + str(dim) + ".png"
            tmpImage = tileImage.resize((dim, dim))
            try:
                tmpImage.save(outputFilename + "_" + str(dim) + ".png", "PNG")
            except:
                print "Failed writing ", outputFilename, "_", str(dim), ".png"


# This function lists the files within a given directory dir
def listdir_fullpath(d):
    return [os.path.join(d, f) for f in os.listdir(d)]


# This is the main program, it takes at least 2 arguments <InputFolder> and <OutputFilename>
def main(argv=None):
    print "Parsing arguments..."
    if argv is None:
        argv = sys.argv

    if len(argv) < 3:
        print "Usage: command <InputFolder> <OutputFilename>"
        print "	<InputFolder> must contain only one series of DATA files to be processed"
        print "	<OutputFilename> must contain the path and base name of the desired output, extensions will be added automatically"
        print "Note1: this version requires you to define the MyLoadData function."
        print "Note2: this version does not process several DATA folders recursively."
        print "You typed:", argv
        return 2

    filenames = listdir_fullpath(argv[1])

    # Convert into a tiled image
    if len(filenames) > 0:
        try:
            global ndimage, misc
            global np
            import numpy as np
            from scipy import ndimage, misc

            gradient = True
        except ImportError:
            print "You need SciPy and Numpy (http://numpy.scipy.org/) to also calculate the gradient!"
            gradient = False

        imgTile, sliceResolution, numberOfSlices, slicesPerAxis = ImageSlices2TiledImage(filenames, loadMyData)
    else:
        print "No files found in that folder, check your parameters and your load function."
        return 2

    # Write a text file containing the number of slices for reference
    try:
        try:
            print 'Creating folder', os.path.dirname(argv[2]), '...',
            os.makedirs(os.path.dirname(argv[2]))
        except OSError as exc:
            if exc.errno == errno.EEXIST and os.path.isdir(os.path.dirname(argv[2])):
                print 'was already there.'
            else:
                print ', folders might not be created, trying to write anyways...'
        except:
            print ", could not create folders, trying to write anyways..."
        with open(argv[2] + "_AtlasDim.txt", 'w') as f:
            f.write(str((numberOfSlices, (slicesPerAxis, slicesPerAxis))))
    except:
        print "Could not write a text file", argv[2] + "_AtlasDim.txt", "containing dimensions (total slices, slices per axis):",\
            (numberOfSlices, (slicesPerAxis, slicesPerAxis))
    else:
        print "Created", argv[2] + "_AtlasDim.txt", "containing dimensions (total slices, slices per axis):", (
        numberOfSlices, (slicesPerAxis, slicesPerAxis))

    # Output is written in different sizes
    WriteVersions(imgTile, argv[2])


if __name__ == "__main__":
    sys.exit(main())