"""Module containing a transformer that extracts all of the figures from the
notebook file.  The extracted figures are returned in the 'resources' dictionary.
"""
#-----------------------------------------------------------------------------
# Copyright (c) 2013, the IPython Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

from IPython.utils.traitlets import (Dict, List, Unicode)
from .activatable import ActivatableTransformer

#-----------------------------------------------------------------------------
# Constants
#-----------------------------------------------------------------------------

FIGURES_KEY = "figures"
BINARY_KEY = "binary"
TEXT_KEY = "text"

#-----------------------------------------------------------------------------
# Classes
#-----------------------------------------------------------------------------

class ExtractFigureTransformer(ActivatableTransformer):
    """
    Extracts all of the figures from the notebook file.  The extracted 
    figures are returned in the 'resources' dictionary.
    """

    extra_extension_map =  Dict({},
        config=True,
        help="""Extra map to override extension based on type.
        Useful for latex where SVG will be converted to PDF before inclusion
        """)
    
    key_format_map =  Dict({}, config=True,)
    figure_name_format_map =  Dict({}, config=True)

    display_data_priority = List(['svg', 'png', 'latex', 'jpg', 'jpeg','text'])

    #TODO: Change this to .format {} syntax
    default_key_template = Unicode('_fig_{index:02d}.{ext}', config=True)
        
    def cell_transform(self, cell, resources, index):
        """
        Apply a transformation on each cell,
        
        Parameters
        ----------
        cell : NotebookNode cell
            Notebook cell being processed
        resources : dictionary
            Additional resources used in the conversion process.  Allows
            transformers to pass variables into the Jinja engine.
        index : int
            Index of the cell being processed (see base.py)
        """
        
        if resources.get(FIGURES_KEY, None) is None :
            resources[FIGURES_KEY] = {TEXT_KEY:{},BINARY_KEY:{}}
            
        for out in cell.get('outputs', []):
            for out_type in self.display_data_priority:
                
                if out.hasattr(out_type):
                    figname, key, data, binary = self._new_figure(out[out_type], out_type)
                    out['key_'+out_type] = figname
                    
                    if binary :
                        resources[FIGURES_KEY][BINARY_KEY][key] = data
                    else :
                        resources[FIGURES_KEY][TEXT_KEY][key] = data
                        
                    index += 1
        return cell, resources


    def _get_override_extension(self, extension):
        """Gets the overriden extension if it exists, else returns extension. 

        Parameters
        ----------
        extension : str
            File extension.
        """
        
        if extension in self.extra_extension_map :
            return self.extra_extension_map[extension]
    
        return extension


    def _new_figure(self, data, format):
        """Create a new figure file in the given format.

        Parameters
        ----------
        data : str
            Cell data (from Notebook node cell)
        format : str
            Figure format
        index : int
            Index of the figure being extracted
        """
        
        figure_name_template = self.figure_name_format_map.get(format, self.default_key_template)
        key_template = self.key_format_map.get(format, self.default_key_template)
        
        #Make sure the figure counter has been initiated 
        if self.index_generator is None:
            self.index_generator = self._index_generator()
            
        #TODO: option to pass the hash as data?
        index = self.index_generator.next()
        figure_name = figure_name_template.format(index=index, ext=self._get_override_extension(format))
        key = key_template.format(index=index, ext=self._get_override_extension(format))

        #Binary files are base64-encoded, SVG is already XML
        binary = False
        if format in ('png', 'jpg', 'pdf'):
            data = data.decode('base64')
            binary = True

        return figure_name, key, data, binary

    
    def _index_generator(self):
        """
        Generates a unique index for association with extracted figures.
        """
        index = 0
        while True:
            yield index
            index+=1
    
@@ -676,16 +676,33 @@
 --------
+%0Aimport itertools
 %0A%0Afrom I
@@ -1920,24 +1920,468 @@
 g=True)%0A
-        
+%0A    def __init__(self, config=None, **kw):%0A        %22%22%22%0A        Public constructor%0A        %0A        Parameters%0A        ----------%0A        config : Config%0A            Configuration file structure%0A        **kw : misc%0A            Additional arguments%0A        %22%22%22%0A%0A        super(ExtractFigureTransformer, self).__init__(config=config, **kw)%0A%0A        # A unique index for association with extracted figures%0A        self.index_generator = itertools.count(1)%0A
 %0A    def
@@ -4531,179 +4531,8 @@
    %0A
-        #Make sure the figure counter has been initiated %0A        if self.index_generator is None:%0A            self.index_generator = self._index_generator()%0A            %0A
     
@@ -4591,16 +4591,21 @@
 index = 
+next(
 self.ind
@@ -4620,14 +4620,8 @@
 ator
-.next(
 )%0A  
@@ -5059,226 +5059,4 @@
 ary%0A
-%0A    %0A    def _index_generator(self):%0A        %22%22%22%0A        Generates a unique index for association with extracted figures.%0A        %22%22%22%0A        index = 0%0A        while True:%0A            yield index%0A            index+=1%0A    

Fix problem with missing .index_generator