# March 8 2018. Added removal of existing numeric extension (line 87).
# April 4 2018. Corrected a bug whereby numeric extension was duplicated (line 91).
# April 5 2018. Added ffmpeg script so that a catelog of images can be converted
#               to an MP4 movie on linux, MacOSX and Windows - see https://www.ffmpeg.org
#               To install ffmpeg on MacOSX see,
#               https://github.com/fluent-ffmpeg/node-fluent-ffmpeg/wiki/Installing-ffmpeg-on-Mac-OS-X
  
import it, time, os
from it.It3Command import It3Command
from PythonQt.QtGui import QDialogButtonBox
from PythonQt.QtGui import QHBoxLayout
from PythonQt.QtGui import QVBoxLayout
from PythonQt.QtGui import QComboBox
  
class SaveAll(It3Command):
    def __init__(self):
        self.m_menuPath = 'Commands/Save All...'
        self.m_dialog = None
        self.m_stdButtons = QDialogButtonBox.Apply|QDialogButtonBox.Cancel
        
    def Invoke(self):
        if self.m_dialog == None:
            self.m_dialog = self.makeUI()
        self.m_dialog.show()
        self.m_dialog.raise_()
        self.m_dialog.activateWindow()
                
    def doit(self):
        gamma = float(self.m_gamma.currentText)
        quality = float(self.m_quality.currentText) * 100.0
        self.m_dialog.hide()
        catelog_path = self.saveAll(gamma, quality)
        #it.app.Info('Images saved to "%s"' % catelog_path)
        it.app.RaiseLogWindow()
        self.writeFFjpegScripts(catelog_path)
        
    def makeUI(self):
        dlg = self.CreateDialog('Save All')
        contents = dlg.findChild(QVBoxLayout, 'contents')
  
        layout = QHBoxLayout()
        contents.addLayout(layout)
        label_gamma = QLabel("Gamma ")
        layout.addWidget(label_gamma)
        
        self.m_gamma = QComboBox()
        self.m_gamma.addItems(['1.0','1.8','2.0','2.2'])
        self.m_gamma.setCurrentIndex(3)
        layout.addWidget(self.m_gamma)
  
        layout = QHBoxLayout()
        contents.addLayout(layout)
        label_quality = QLabel("Quality ")
        layout.addWidget(label_quality)
        
        self.m_quality = QComboBox()
        self.m_quality.addItems(['0.5','0.75','1.0'])
        self.m_quality.setCurrentIndex(2)
        layout.addWidget(self.m_quality)
  
        bbox = dlg.findChild(QDialogButtonBox, 'bbox')
        doItButton = bbox.button(QDialogButtonBox.Apply)
        doItButton.setText('Save All')
        doItButton.connect('clicked()', self.doit)
        return dlg
            
    def saveAll(self, gamma=1.0, quality=100):
        localtime = time.asctime( time.localtime(time.time()) )
        localtime = localtime.replace(' ', '_')
        localtime = localtime.replace(':', '_')    
        outname = 'catalog_%s' % localtime
        cwdpath = os.getcwd()
        if len(cwdpath) < 3:
            cwdpath = os.environ['HOME']
            if len(cwdpath) < 3:    
                it.app.Error('Unable to determine the current working directory.')
                it.app.Error('Unable to save the images.')
                it.app.RaiseLogWindow()
                return ''
        out_dirpath = os.path.join(cwdpath,outname)
        if not os.path.exists(out_dirpath):
            os.mkdir(out_dirpath)
        aovDict = {}
        cat = it.app.GetCurrentCatalog()
        img_counter = 1
        all_counter = 0
        for i in range (0, cat.GetChildCount()):
            element = cat.GetChild(i)
            imgname = it.os.path.basename( element.GetFilename() )
            if imgname == '_preview':
                continue
            # Remove and existing numeric extension
            parts = imgname.rsplit('.')
            if len(parts) == 3 and parts[1].isdigit():
                imgname = '%s.%04d.jpg' % (parts[0],img_counter)
            elif len(parts) == 2 and parts[1].isdigit:
                imgname = '%s.%04d.jpg' % (parts[0],img_counter)
            else:
                # Add a padded numeric extension
                imgname = '%s.%04d.jpg' % (imgname,img_counter)
            it.app.Info('    saving primary image: "%s"' % imgname)
            out_imgpath = os.path.join(out_dirpath, imgname)
            self.saveImage(element, out_imgpath, gamma, quality)
            img_counter += 1
            all_counter += 1
            # Save any AOV's
            for j in range (0, element.GetChildCount()):
                aov = element.GetChild(j)
                aovname = it.os.path.basename( aov.GetFilename() )
                # Remove the extension
                aovname = os.path.splitext(aovname)[0]
                # September 16th 2017 don't save the mask for an IPR render
                if aovname.startswith('IPR_id'):
                    continue
                # Maintain unique counters for each AOV
                aov_counter = 1
                if aovDict.has_key(aovname):
                    aov_counter = aovDict[aovname]
                    aov_counter += 1
                    all_counter += 1
                aovDict[aovname] = aov_counter
                
                # Add a padded numeric extension
                aovname = aovname + ('.%04d' % aov_counter) + '.jpg'
                it.app.Info('    saving aov image: "%s"' % aovname)
                out_aovpath = os.path.join(out_dirpath, aovname)
                self.saveImage(aov, out_aovpath, gamma, quality)
        it.app.Info('Saved %d images to "%s".' % (all_counter,cwdpath))
        return out_dirpath
    def saveImage(self, element,path,gamma,quality):
        image = element.GetImage()
        image = image.Gamma(gamma);
        image.SetMetaDataItem('JPEG_QUALITY', quality)
        image.Save(path, ice.constants.FMT_JPEG)
        
    # Added April 5 2018
    def writeFFjpegScripts(self, catelog_path):
        unix_script = """
# The flags in use are:
# -pattern_type glob -i '*.jpg'
#    Ensures the input (-i) jpg image names can be (almost) any format. For example, 
#    foo_0001.jpg or 
#    goo.01.jpg
#
# -s 960x540
#    Forces all input images to be resized to 960 pixels wide by 540 pixels height.
#    WIDTH and HEIGHT MAY BE CHANGED BUT BE CAREFUL NOT TO INCLUDE SPACES. For example,
#    this is wrong 960 x 540.
#
# -codec:v libx264
#    Use H264 video compression. DO NOT CHANGE. 
#
# -pix_fmt yuv420p
#    Color encoding. DO NOT CHANGE. 
#    See https://en.wikipedia.org/wiki/YUV for details.
#
# -r 24
#    Frame rate per second playback speed.
#    MAY BE CHANGED.
#
#  aout.mp4
#    Name of the output MP4 video file. THE NAME MAY BE CHANGED
#
# Prepared by M.Kesson April 7 2018
#
ffmpeg -pattern_type glob -i '%s/*.jpg' \\
-s 960x540 -codec:v libx264 -r 24 -pix_fmt yuv420p \\
'%s/aout.mp4'
""" % (catelog_path,catelog_path)
        win_script = unix_script.replace('#', '::')
        unixpath = os.path.join(catelog_path,'a_ffjpeg')
        winpath = os.path.join(catelog_path,'a_ffjpeg.bat')
        uf = open(unixpath, 'w')
        uf.write(unix_script)
        uf.close()
        os.chmod(unixpath, 0777)
        wf = open(winpath, 'w')
        wf.write(win_script)
        wf.close()
        it.app.Info('ffjpeg scripts saved to "%s"' % catelog_path)
it.commands.append(SaveAll)