#!/usr/bin/python

# Version 1.0
# Updated Sat Oct 15 13:42:51 PDT 2016 -- spl

import argparse
import sys
import os
import glob
import csv
import re
from PIL import Image
import math
import threading
from xml.etree.ElementTree import Element, SubElement
from xml.dom import minidom
from xml.etree import ElementTree

def processCSV( file_name ):
    
    # Read the CSV file and find the #img elements
    
    f = open( file_name, 'rb' )
    reader = csv.reader( f )
    
    # Scan the header row for the "#img" element
    
    header = reader.next()
    
    img_index = -1
    for index, item in enumerate( header ):
        
        if ( item == "#img" ):
            
            img_index = index       # We win!  There's an "#img" element
            break
        
    if ( img_index == -1 ):         # Oh, noes!  No "#img" element.  Die.
            
        print "No '#img' item in '" + file_name + "'"
        sys.exit()
            
    # Now grab the image object names from the CSV
    
    images = {}                     # Make associative to cull duplicates

    for row in reader:

        image = row[img_index]
        images[image] = image

    f.close()

    return images

# End processCSV

def processXML( file_name ):
    
    # Read the XML file

    root = ElementTree.parse( file_name ).getroot()
    root_tag = root.tag
    xmlns = re.sub( "[A-Za-z]*$", "", root_tag )
    items = root.find( xmlns + "Items" )

    children = list( items.iter() )
    t = children[1].tag
    images = {}
    for i in items.iter( t ):

        if "Source" in i.keys():
            source = os.path.basename( i.attrib["Source"].replace( ".xml", "" ) )
        elif "Href" in i.keys():
            source = os.path.basename( i.attrib["Href"].replace( ".html", "" ) )

        id = i.attrib["Id"]
        
        images[id] = source

    return images

# End processXML

# Worker thread.  Spawned to parallel process the various levels.

def worker( level,
            lock,
            dirpath,
            outpath,
            images_wide,
            images_high,
            sprite_index ):

    ps = os.path.sep            # Get the path separator

    # Calculate the size of the tile at this level and derive the
    # sheet size.

    # Note that we create a basically square sheet mostly for quality
    # assurance purposes, since long sheet tend to confuse some image
    # viewers for some unknown reason. It adds a small amount of
    # complexity to the code for writing and reading the sprite sheets
    # but basically it's no big deal.

    tile_size = 1 << level
    sprite_sheet_width = images_wide * tile_size
    sprite_sheet_height = images_high * tile_size

    # Create the sprite sheet
    
    sprite_sheet = Image.new( "RGBA",
                              ( sprite_sheet_width, sprite_sheet_height ),
                              ( 0, 0, 0, 0 ) )

    # Fish up the DZC file for this collection if it exists and get
    # the default file format.

    dzcbasename = os.path.basename( dirpath )
    dzcpath = dirpath + ps + dzcbasename + ".dzc"
    
    if os.path.exists( dzcpath ):

        root = ElementTree.parse( dzcpath ).getroot()
        root_tag = root.tag
        xmlns = re.sub( "[A-Za-z]*$", "", root_tag )
        default_format = root.attrib["Format"]

    else:                       # If there's no DZC, punt.
        default_format = ""
        
    # Iterate over the input XML for each image object, read the
    # image, and stuff it into the sprite sheet.

    for sprite in sprite_index.iter( "Tile" ):

        # Get the object name.
        
        obj = sprite.attrib["file"];

        dzipath = dirpath + ps + obj + ".dzi"

        if os.path.exists( dzipath ):

            root = ElementTree.parse( dzipath ).getroot()
            root_tag = root.tag
            xmlns = re.sub( "[A-Za-z]*$", "", root_tag )
            image_format = "." + root.attrib["Format"]
            
        else:                       # If there's no DZI, punt.
            image_format = default_format

        # Find the file in the level's object director

        objpath = dirpath + ps + \
                  obj + "_files" + ps + \
                  str( level ) + ps + \
                  "*" + image_format
        
        files = glob.glob( objpath )

        # Make sure we're not processing a multitiled image, which
        # would probably create too large of a sprite sheet to be
        # useful, anyhow.

        if ( len( files ) > 1 ):

            print "Can't process multitiled images."
            sys.exit()

        elif ( len( files ) == 0 ):

            print "No such file or directory:", objpath
            sys.exit()
            
        image_file = files[0]

        # Open and read the file

        with Image.open( image_file ) as image:
            
            # Get the image
            
            width, height = image.size

            with lock:          # Lock to avoid trashing the XML tree.
                
                tile_params = SubElement( sprite, "TileParams" )

            # Set essential XML attributes.
            
            tile_params.set( "width", str( width ) );
            tile_params.set( "height", str( height ) );
            tile_params.set( "level", str( level ) );

            # Fetch the position of image within sprite sheet from the
            # XML.

            tilex = int( sprite.attrib["x"] ) * tile_size
            tiley = int ( sprite.attrib["y"] ) * tile_size
            
            # Copy the image into the sprite sheet.

            sprite_sheet.paste( image.convert( "RGBA" ),
                                ( tilex, tiley, tilex + width, tiley + height ) )

    # We're done processing so write the sprite sheet to disk and bail.

    sprite_sheet_file = outpath + ps + "sprite_" + str( level ) + ".png"

    sprite_sheet.save( sprite_sheet_file )

    return

# End worker

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

# Main program starts here.

# Parse the arguments

parser = argparse.ArgumentParser( description =
                                  "Creates sprite sheets for SuAVE." );

parser.add_argument( "-m", "--maxdepth",
                     metavar = "maxdepth",
                     nargs = 1,
                     default = [-1],
                     dest = "maxdepth",
                     help = "Maximum depth sprite sheet to be created. " + \
                     "Default: the maximum depth will be automatically " + \
                     "computed.",
                     type = int )
parser.add_argument( "-M", "--Microsoft",
                     dest = "microsoft_mode",
                     help = "Process Microsoft XML file rather than CSV.",
                     action = "store_true" )
parser.set_defaults( microsoft_mode = False )
parser.add_argument( "file_name",
                     nargs = 1,
                     help = "The file to be parsed (assumed CSV)." )

ps = os.path.sep                # Get the system's path separator

options = parser.parse_args()

# Mostly for laziness sake, so we can avoid typing "options." . . .

maxdepth = options.maxdepth[0]
file_name = options.file_name[0]

# Chastise the user if there's something bogus

if ( not os.path.exists( file_name ) ):

    print "No such file or directory:", file_name
    sys.exit()

if ( maxdepth != -1 ) and ( math.pow( 2, maxdepth ) > 128 ):

    print "maxdepth is too large."
    print "You don't want to create sprite sheets with cells greater than"
    print "128 x 128"
    sys.exit()

if ( options.microsoft_mode ):
    images = processXML( file_name )
else:
    images = processCSV( file_name )

# If the user has not set a maxdepth, we will calculate one. Arbitrary
# maximum spritesheet size of 10 megapixels.  We'll assume that this
# will compress as a PNG to something fairly manageable.

if ( maxdepth == -1 ):          
                                
    ten_mp = 10 * 1024 * 1024   
    maxdepth = int( math.log( math.sqrt( ten_mp / len( images ) ), 2 ) )
    if ( maxdepth > 7 ):
        maxdepth = 7

# Munge the file name until it looks like a basename.  Enormous kludge
# here.

basename = os.path.basename( file_name )
basename = re.sub( "\\.[xc][ms][lv]$", "", basename )
basename = re.sub( "\\.cxml$", "", basename )
dirname = os.path.dirname( file_name )

if ( dirname == "" ):
    dirname = "."

# Build the path were we're going to store the output

# Some versions of the Microsoft format apparently store files in an
# images subdirectory

if os.path.exists( dirname + ps + "images" ):

    dirpath = dirname + ps + "images"
    outpath = dirname
    
else:

    dirpath = dirname + ps + basename
    outpath = dirpath

# process the images into sprites

nimages = float( len( images ) )
images_wide = int( math.floor( math.sqrt( nimages ) ) )
images_high = int( math.ceil( nimages / images_wide ) )

print "Sprite sheet will be", images_wide, "images x", images_high, "images"

# Insert the image position indices into the XML.

sprite_index = Element( "SpriteIndex" );
sprite_index.set( "maxdepth", str( maxdepth ) );

imageno = 0
for image in images:

    image_element = SubElement( sprite_index, "Tile" )
    image_element.set( "x", str( imageno % images_wide ) )
    image_element.set( "y", str( int( math.floor( imageno / images_wide ) ) ) )
    image_element.set( "object", image )
    image_element.set( "file", images[image] )
    imageno = imageno + 1

# Create sprite sheet generating threads
    
lock = threading.Lock()         # A lock to maintain the XML sanctity.

threads = []                    # Thread list.

for level in range( maxdepth + 1 ): # +1 so that the maxdepth is
                                    # included
    # While we're at it, create the correspondence between the index
    # (depth) and the actual sprite sheet.  This is probably overkill,
    # since we can easily generate the file name but it avoids baking
    # a file name structure into the Javascript in SuAVE deepzoom.

    sprite_sheet = SubElement( sprite_index, "SpriteSheet" )
    sprite_sheet.set( "sprite", "sprite_" + str( level ) + ".png" )
    sprite_sheet.set( "index", str( level ) )

    # Okey dokey, let's create a worker thread.
    
    threads.append( threading.Thread( target = worker,
                                      args = ( level,
                                               lock,
                                               dirpath,
                                               outpath,
                                               images_wide,
                                               images_high,
                                               sprite_index ) ) )

# Now launch them

for thread in threads:

    thread.start()

# Now sit back, relax, enjoy, and wait for the threads to complete.

for thread in threads:

    thread.join()

# Write the sprite index file.  Pretty print it so that humans can
# enjoy it, too.

sprite_idx = outpath + ps + "sprite.idx"
print "Writing", sprite_idx

# This is a rather hideous hack, since it emits and reparses the XML,
# which can take a fair amount of memory in a large XML tree.  We may
# want to revert to writing plain old compact XML if this becomes a
# problem.

text = minidom.parseString( ElementTree.tostring( sprite_index,
                                                  "UTF-8" ) ).toprettyxml( indent = "  " )

# Write it out and take a hike.

try:
    
    f = open( sprite_idx, "wb" )
    f.write( text )
    f.close()

except IOError:

    print "Couldn't open sprite index file", sprite_idx, "or write failed."
    sys.exit()

sys.exit()
