# 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
# March 24 2020 See notes at line 74.
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()
        catalog_path = self.saveAll(gamma, quality)
        #it.app.Info('Images saved to "%s"' % catalog_path)
        it.app.RaiseLogWindow()
        self.writeFFjpegScripts(catalog_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
        # No longer using cwdpath = os.getcwd() because it returns different
        # paths for different OS's. For example, on OSX it returns an empty path!
        cwdpath = os.environ['HOME']
  
        if len(cwdpath) == 0:
            it.app.Error('Unable to determine the HOME 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, catalog_path):
        isWindows = False
        if os.name == "posix":
            script_path = os.path.join(catalog_path,'ffmpeg')
            
            if sys.platform == 'darwin':
                path = catalog_path + '/'
            else:
                # linux does not require the full path to the source jpegs or the mp4
                path = ''
        else:
            script_path = os.path.join(catalog_path,'a_ffmpeg.bat')
            path = catalog_path + '\\'
            isWindows = True
            
        ffmpegscript = """
# 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.
#    GENERALLY BEST NOT TO CHANGE THIS VALUE.
#
# -framerate 1/3
#     Use this for slide shows. For example, 1/3 will display each image for 3 secoonds.
#     Add this flag and its value immediately before the -i flag. For example,
#         ffmpeg -pattern_type glob -framerate 1/3 -i
# 
#  aout.mp4
#    Name of the output MP4 video file. THE NAME MAY BE CHANGED
#
# Prepared by M.Kesson April 7 2018
# Modified to include -framerate Jan 17 2019
#
ffmpeg -pattern_type glob -i "%s*.jpg" \\
-s 960x540 -codec:v libx264 -r 24 -pix_fmt yuv420p \\
"%saout.mp4"
""" % (path, path)
        if isWindows:
            ffmpegscript = ffmpegscript.replace('#', '::')
        f = open(script_path, 'w')
        f.write(ffmpegscript)
        f.close()
        if isWindows == False:
            os.chmod(script_path, 0777)
        it.app.Info('ffmpeg script saved to "%s"' % catalog_path)
it.commands.append(SaveAll)