#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""
Python script to read DMP file.
Written by David Haberthür <david.haberthuer@psi.ch> based on code by Martin
Nyvlt <martin.nyvlt@gmail.com>
"""

import numpy
import os
import matplotlib.pyplot as plt
import urllib2


def readDMP(filename):
    """
    Opens and reads DMP file.
    Returns a numpy ndarray which can be displayed/written/handled further.
    """
    # Open the file
    dmpfile = open(filename, 'rb')
    # The DMP file format is a very simple file format; the header consists of
    # 3 unsigned 16bit integers specifying the image dimensions, followed by a
    # stream of 32-bit floats in little-endian byte order. It's first two
    # values are the image size.
    # Load Header.
    header = numpy.fromfile(dmpfile, numpy.int16, 3)
    imageSize = (header[1], header[0])
    # Load the whole file and reshape to the size given in the header
    image = numpy.fromfile(dmpfile, numpy.float32)
    image = image.reshape(imageSize)
    dmpfile.close()
    return image


def writeDMP(array, outputfilename):
    """
    Writes numpy ndarray to DMP file.
    The array.dtype should be numpy.float32.
    If it's not, it will be converted before saving.
    """
    # Open the file for writing (binary)
    dmpfile = open(outputfilename, 'wb')
    # Create header. *Must* be dtype numpy.uint16!
    width = numpy.array(array.shape[1])
    height = numpy.array(array.shape[0])
    header = numpy.array([width, height, 0], numpy.uint16)
    header.tofile(dmpfile)
    # Write Image data. dtype has to be np.float32
    if array.dtype != 'float32':
        array = array.astype(numpy.float32)
    array.tofile(dmpfile)
    dmpfile.close()

if __name__ == "__main__":
    # Read example DMP file with the function defined above
    DMPFileName = os.path.join(os.getcwd(), '..', 'example.DMP')
    DMP = readDMP(DMPFileName)
    # Show example image
    plt.subplot(121)
    plt.imshow(DMP, cmap='gray')
    plt.title('example.DMP')
    # Load some image from the internet and save its green channel as DMP.
    URL = urllib2.urlopen('http://www.psi.ch/sls/tomcat/HomeEN/tomcat_banner.png')
    image = plt.imread(URL)[:, :, 1]
    writeDMP(image, 'randomimage.DMP')
    # Read the image again as DMP and display it
    plt.subplot(122)
    plt.imshow(readDMP('randomimage.DMP'), cmap='gray')
    plt.title('Some random image')

    plt.show()

    print 'Now open randomimage.DMP in Fiji!'