#!/usr/bin/env python

import sys
import os
import re
import numpy
from PIL import Image
import StringIO

# you need to install this library yourself
# recent versions handle bigtiff too...
import tifffile

"""
Extract pyramidal TIFF files with JPEG tiled storage into a tree of
separate JPEG files into DZI compliant directory that is usable 
by openseadragon. Multiple channels info are combined into single
jpeg.

usage: extract2dzi_rgb.py pyramid-file-dir dest-dir
         
The pyramid-file must be a multi-page TIFF with each page having an
image scaled by 1/2 from the previous page.  All pages must be tiled
with the same tile size, and tiles must be stored using the new-style
JPEG compression format, i.e. TIFF compression == 7.

The lowest resolution page must have 4 or fewer tiles.  If it has
more than 1, this script will leave space for the user to decide whether
final lowest zoom tile 0/0_0.jpg that is 1/2 scaled version of the image 
represented by that last page should be generated or not.

File directory generated
   
    dest-dir
      ImageProperties.xml
      0 
        0_0.jpg
      1
        0_0.jpg
        1_0.jpg
      ...

Since the tiled tiff kept padded tiles and openseadragon expected its
jpeg files to be cropped but not padded, the border tiles are cropped
and the width and height of image uses the actual image dimension
"""

try:
    srcloc = sys.argv[1]
    outloc = sys.argv[2]

    if not os.path.exists(srcloc) or not os.path.isdir(srcloc):
        sys.stderr.write('Pyramid directory must be given and exist')
        sys.stderr.write('\nusage: extract2dzi_rgb.py pyramid-file-directory dest-dir\n\n')
        sys.exit(1)

    if not os.path.exists(outloc):
      os.makedirs(outloc)
except:
    sys.stderr.write('\nusage: extract2dzi_rgb.py pyramid-file-directory dest-dir\n\n')
    raise


## 20140403-R26-Tdt-JJG-0-38-000-DAPI-Z3.tif
## 20140403-R26-Tdt-JJG-0-38-000-FITC-Z3.tif
## 20140403-R26-Tdt-JJG-0-38-000-Rhodamine-Z3.tif
## iterate through the files,
## if valid tiff file, then change the outdir to 
##     outdir/DAPI/.xml,0,1..
## essentialy like calling extract2dzi.py filename outdir/color

infile=None
txsize=0
tysize=0
pxsize=0
pysize=0
zoomno=0
outinfo=[]
total_tiles=0
topdir_template = '%(outdir)s'
dir_template = topdir_template +'/%(zoomno)d'
tile_template = dir_template + '/%(tcolno)d_%(trowno)d.jpg'
image_template = '%(outdir)s/ImageProperties.xml'

################# helper functions ###################
# http://www.w3.org/Graphics/JPEG/jfif3.pdf
def jpeg_assemble(jpeg_tables_bytes, jpeg_bytes):
    return jpeg_bytes[0:2] + jpeg_tables_bytes + jpeg_bytes[2:]

def load_tile(infile, tile_offset, tile_length):
    infile.seek(tile_offset)
    return infile.read(tile_length)

def getTile(page, infile, jpeg_tables_bytes, tileno):
    jpeg = jpeg_assemble(jpeg_tables_bytes, load_tile(infile, page.tags.tile_offsets.value[tileno], page.tags.tile_byte_counts.value[tileno]))
    outfile = StringIO.StringIO()
    outfile.write( jpeg )
    outfile.seek(0)
    image = Image.open(outfile)
    ret = numpy.asarray(image)
    outfile.close()
    return ret

def maxTile(page, infile):
    pxsize, pysize, txsize, tysize, tcols, trows, jpeg_tables_bytes = get_page_info(page)
    maxval = 0
    for tileno in range(0, len(page.tags.tile_offsets.value)):
        tile = getTile(page, infile, jpeg_tables_bytes, tileno)
        maxval = max(maxval, tile.max())
    return maxval


def write_tile(tileno, trow, trows, tcol, tcols, rgb_image):
    """Output one tile.  Note this manages global state for tile grouping in subdirs."""
    global zoomno
    global total_tiles

    cropIt = False
    if (trow+1 == trows) or (tcol+1 == tcols) : 
        #this is a border tile, crop it if need to
        if tcol+1 == tcols :
           cpxsize= (pxsize-(txsize * tcol))
        else:
           cpxsize=txsize
        if trow+1 == trows :
           cpysize= (pysize-(tysize * trow))
        else:
           cpysize=tysize
        cropIt = True
    
    total_tiles += 1

    topdir = topdir_template % dict(
        outdir = outloc
    )
    if not os.path.exists(topdir):
        os.makedirs(topdir, mode=0755)

    dirname = dir_template % dict(