#!/usr/bin/python
 
import mmap
from struct import unpack_from
from numpy import frombuffer
 
types = { # for both numpy and struct
    2: "h",    3: "i",    4: "H",
    5: "I",    6: "f",    7: "d",
    8: "?",    #9: "c",   
   10: "b",   11: "q",   12: "Q",
}
 
class Value:
    def __init__(self, name, mm, offset):
        self.name = name
        self.mm = mm
        self.offset = offset
 
    def read(self):
        offset = self.offset
 
        head, ninfo = unpack_from(">4sQ", self.mm, offset); offset += 12
        info = unpack_from(">" + "Q"*ninfo, self.mm, offset); offset += ninfo*8
 
        if ninfo == 1: # single value
            fmt = types[info[0]]
            return unpack_from(fmt, self.mm, offset)[0]
 
        elif info[0] == 15: # tuple
            fmt = "".join(types[i] for i in info[4::2])
            return unpack_from(fmt, self.mm, offset)
 
        elif info[:2] == (20,4): # unicode string
            return self.mm[offset:offset+2*info[2]].decode("utf16")
 
        elif info[:2] == (20,10): # byte string
            return str(self.mm[offset:offset+info[2]])
 
        elif info[:2] == (20,15): # array of groups
            return (info, "not implemented")
 
        elif info[0] == 20: # array of simple values
            dtype = types[info[1]]
            return frombuffer(self.mm, dtype, info[2], offset)
 
        else:
            raise ValueError("unknown tag {} {!r}".format( info, self.mm[offset:offset+128]))
 
    def value(self):
        if not hasattr(self, "_value"):
            self._value = self.read()
        return self._value
 
    def dump(self, indent=""):
        print indent, self.name, "=", repr(self.value())
 
class Directory:
    def __init__(self, name, mm, offset):
        self.name = name
        self.mm = mm
        self.offset = offset
 
    def read(self):
        offset = self.offset
        sortf, closef, ntags = unpack_from(">BBQ", self.mm, offset); offset += 10
        for i in xrange(ntags):
            tag, ltname = unpack_from(">BH", self.mm, offset); offset += 3
            tname = str(self.mm[offset:offset+ltname]); offset += ltname
            if not tname:
                tname = i
            tlen, = unpack_from(">Q", self.mm, offset); offset += 8
            if tag == 20:
                yield Directory(tname, self.mm, offset)
            elif tag == 21:
                yield Value(tname, self.mm, offset)
            else:
                raise ValueError("tag != 20 or 21")
            offset += tlen
 
    def __getitem__(self, key):
        if not hasattr(self, "_dict"):
            self._dict = {i.name:i for i in self}
        return self._dict[key]
 
    def __iter__(self):
        if not hasattr(self, "_list"):
            self._list = list(self.read())
        return iter(self._list)
 
    def dump(self, indent=""):
        print indent, self.name
        for i in self:
            i.dump(indent+"   ")
 
class DM4:
    def __init__(self, fname):
        with open(fname) as fh:
            mm = mmap.mmap(fh.fileno(), 0, prot=mmap.PROT_READ)
        version, rootlen, byteord = unpack_from(">LQL", mm)
        assert rootlen == mm.size()-24
        assert version == 4
        assert byteord == 1
        #assert mm[rootlen+16:] == '\0'*8
        self.root = Directory("root", mm, 16)
 
    def data(self):
        if not hasattr(self, "_data"):
            for img in self.root['ImageList']:
                try:
                    img['ImageTags']['Acquisition']
                except KeyError:
                    continue
                else:
                    break
            else:
                raise ValueError("Acquisition not found")
            self._data = img['ImageData']['Data'].value()
            self._data.shape = [d.value() for d in img['ImageData']['Dimensions']]
        return self._data
 
if __name__ == "__main__":
    import sys
    d = DM4(sys.argv[1])
    d.root.dump()
    from pylab import imshow, show
    imshow(d.data(), cmap='gray')
    show()