#
# The Python Imaging Library.
#
# SPIDER image file handling
#
# History:
# 2004-08-02    Created BB
# 2006-03-02    added save method
# 2006-03-13    added support for stack images
#
# Copyright (c) 2004 by Health Research Inc. (HRI) RENSSELAER, NY 12144.
# Copyright (c) 2004 by William Baxter.
# Copyright (c) 2004 by Secret Labs AB.
# Copyright (c) 2004 by Fredrik Lundh.
#

##
# Image plugin for the Spider image format.  This format is is used
# by the SPIDER software, in processing image data from electron
# microscopy and tomography.
##

#
# SpiderImagePlugin.py
#
# The Spider image format is used by SPIDER software, in processing
# image data from electron microscopy and tomography.
#
# Spider home page:
# http://www.wadsworth.org/spider_doc/spider/docs/spider.html
#
# Details about the Spider image format:
# http://www.wadsworth.org/spider_doc/spider/docs/image_doc.html
#

import Image, ImageFile
import os, struct, sys

def isInt(f):
    try:
        i = int(f)
        if f-i == 0: return 1
        else:        return 0
    except:
        return 0

iforms = [1,3,-11,-12,-21,-22]

# There is no magic number to identify Spider files, so just check a
# series of header locations to see if they have reasonable values.
# Returns no.of bytes in the header, if it is a valid Spider header,
# otherwise returns 0

def isSpiderHeader(t):
    h = (99,) + t   # add 1 value so can use spider header index start=1
    # header values 1,2,5,12,13,22,23 should be integers
    for i in [1,2,5,12,13,22,23]:
        if not isInt(h[i]): return 0
    # check iform
    iform = int(h[5])
    if not iform in iforms: return 0
    # check other header values
    labrec = int(h[13])   # no. records in file header
    labbyt = int(h[22])   # total no. of bytes in header
    lenbyt = int(h[23])   # record length in bytes
    #print "labrec = %d, labbyt = %d, lenbyt = %d" % (labrec,labbyt,lenbyt)
    if labbyt != (labrec * lenbyt): return 0
    # looks like a valid header
    return labbyt

def isSpiderImage(filename):
    fp = open(filename,'rb')
    f = fp.read(92)   # read 23 * 4 bytes
    fp.close()
    bigendian = 1
    t = struct.unpack('>23f',f)    # try big-endian first
    hdrlen = isSpiderHeader(t)
    if hdrlen == 0:
        bigendian = 0
        t = struct.unpack('<23f',f)  # little-endian
        hdrlen = isSpiderHeader(t)
    return hdrlen


class SpiderImageFile(ImageFile.ImageFile):

    format = "SPIDER"
    format_description = "Spider 2D image"

    def _open(self):
        # check header
        n = 27 * 4  # read 27 float values
        f = self.fp.read(n)

        try:
            self.bigendian = 1
            t = struct.unpack('>27f',f)    # try big-endian first
            hdrlen = isSpiderHeader(t)
            if hdrlen == 0:
                self.bigendian = 0
                t = struct.unpack('<27f',f)  # little-endian
                hdrlen = isSpiderHeader(t)
            if hdrlen == 0:
                raise SyntaxError("not a valid Spider file")
        except struct.error:
            raise SyntaxError("not a valid Spider file")

        h = (99,) + t   # add 1 value : spider header index starts at 1
        iform = int(h[5])
        if iform != 1:
            raise SyntaxError("not a Spider 2D image")

        self.size = int(h[12]), int(h[2]) # size in pixels (width, height)
        self.istack = int(h[24])
        self.imgnumber = int(h[27])

        if self.istack == 0 and self.imgnumber == 0:
            # stk=0, img=0: a regular 2D image
            offset = hdrlen
            self.nimages = 1
        elif self.istack > 0 and self.imgnumber == 0:
            # stk>0, img=0: Opening the stack for the first time
            self.imgbytes = int(h[12]) * int(h[2]) * 4
            self.hdrlen = hdrlen
            self.nimages = int(h[26])
            # Point to the first image in the stack
            offset = hdrlen * 2
            self.imgnumber = 1
        elif self.istack == 0 and self.imgnumber > 0:
            # stk=0, img>0: an image within the stack
            offset = hdrlen + self.stkoffset
            self.istack = 2  # So Image knows it's still a stack
        else:
            raise SyntaxError("inconsistent stack header values")

        if self.bigendian:
            self.rawmode = "F;32BF"
        else:
            self.rawmode = "F;32F"
        self.mode = "F"

        self.tile = [("raw", (0, 0) + self.size, offset,
                    (self.rawmode, 0, 1))]
        self.__fp = self.fp # FIXME: hack

    # 1st image index is zero (although SPIDER imgnumber starts at 1)
    def tell(self):
        if self.imgnumber < 1:
            return 0
        else:
            return self.imgnumber - 1

    def seek(self, frame):
        if self.istack == 0:
            return
        if frame >= self.nimages:
            raise EOFError("attempt to seek past end of file")
        self.stkoffset = self.hdrlen + frame * (self.hdrlen + self.imgbytes)
        self.fp = self.__fp
        self.fp.seek(self.stkoffset)
        self._open()

    # returns a byte image after rescaling to 0..255
    def convert2byte(self, depth=255):
        (min, max) = self.getextrema()
        m = 1
        if max != min:
            m = depth / (max-min)
        b = -m * min
        return self.point(lambda i, m=m, b=b: i * m + b).convert("L")

    # returns a ImageTk.PhotoImage object, after rescaling to 0..255
    def tkPhotoImage(self):
        import ImageTk
        return ImageTk.PhotoImage(self.convert2byte(), palette=256)

# --------------------------------------------------------------------
# Image series

# given a list of filenames, return a list of images
def loadImageSeries(filelist=None):
    " create a list of Image.images for use in montage "
    if filelist == None or len(filelist) < 1:
        return

    imglist = []
    for img in filelist:
        if not os.path.exists(img):
            print "unable to find %s" % img
            continue
        try:
            im = Image.open(img).convert2byte()
        except:
            if not isSpiderImage(img):
                print img + " is not a Spider image file"
            continue
        im.info['filename'] = img
        imglist.append(im)
    return imglist

# --------------------------------------------------------------------
# For saving images in Spider format

def makeSpiderHeader(im):
    nsam,nrow = im.size
    lenbyt = nsam * 4  # There are labrec records in the header
    labrec = 1024 / lenbyt
    if 1024%lenbyt != 0: labrec += 1
    labbyt = labrec * lenbyt
    hdr = []
    nvalues = labbyt / 4
    for i in range(nvalues):
        hdr.append(0.0)

    if len(hdr) < 23:
        return []

    # NB these are Fortran indices
    hdr[1]  = 1.0           # nslice (=1 for an image)
    hdr[2]  = float(nrow)   # number of rows per slice
    hdr[5]  = 1.0           # iform for 2D image
    hdr[12] = float(nsam)   # number of pixels per line
    hdr[13] = float(labrec) # number of records in file header
    hdr[22] = float(labbyt) # total number of bytes in header
    hdr[23] = float(lenbyt) # record length in bytes

    # adjust for Fortran indexing
    hdr = hdr[1:]
    hdr.append(0.0)
    # pack binary data into a string
    hdrstr = []
    for v in hdr:
        hdrstr.append(struct.pack('f',v))
    return hdrstr

def _save(im, fp, filename):
    if im.mode[0] != "F":
        im = im.convert('F')

    hdr = makeSpiderHeader(im)
    if len(hdr) < 256:
        raise IOError("Error creating Spider header")

    # write the SPIDER header
    try:
        fp = open(filename, 'wb')
    except:
        raise IOError("Unable to open %s for writing" % filename)
    fp.writelines(hdr)

    rawmode = "F;32NF"  #32-bit native floating point
    ImageFile._save(im, fp, [("raw", (0,0)+im.size, 0, (rawmode,0,1))])

    fp.close()

def _save_spider(im, fp, filename):
    # get the filename extension and register it with Image
    fn, ext = os.path.splitext(filename)
    Image.register_extension("SPIDER", ext)
    _save(im, fp, filename)

# --------------------------------------------------------------------

Image.register_open("SPIDER", SpiderImageFile)
Image.register_save("SPIDER", _save_spider)

if __name__ == "__main__":

    if not sys.argv[1:]:
        print "Syntax: python SpiderImagePlugin.py Spiderimage [outfile]"
        sys.exit()

    filename = sys.argv[1]
    if not isSpiderImage(filename):
        print "input image must be in Spider format"
        sys.exit()

    outfile = ""
    if len(sys.argv[1:]) > 1:
        outfile = sys.argv[2]

    im = Image.open(filename)
    print "image: " + str(im)
    print "format: " + str(im.format)
    print "size: " + str(im.size)
    print "mode: " + str(im.mode)
    print "max, min: ",
    print im.getextrema()

    if outfile != "":
        # perform some image operation
        im = im.transpose(Image.FLIP_LEFT_RIGHT)
        print "saving a flipped version of %s as %s " % (os.path.basename(filename), outfile)
        im.save(outfile, "SPIDER")