"""Handles running example scripts and saving the generated plots."""
import os
import glob
import sys
import io
import warnings

from PIL import Image
import matplotlib as mpl
mpl.use('Agg')
import matplotlib._pylab_helpers as pylab_helpers

SAVED_DIRECTORY = os.path.join(os.path.dirname(__file__), 'test_files')

def example_files():
    """Returns an iterator over the *.py files in the example directory."""
    current_dir = os.path.dirname(__file__)
    example_dir = os.path.join(current_dir, '..', 'examples')
    for filename in glob.iglob(example_dir + '/*.py'):
        yield filename

def current_figures():
    """Returns an iterator over the current matplotlib figures. After iterating
    through, all figures are destroyed!"""
    for manager in pylab_helpers.Gcf.get_all_fig_managers():
        yield manager.canvas.figure
    pylab_helpers.Gcf.destroy_all()

def image_filename(example_filename, fig):
    """Given the filename of the source code that generated the figure and a
    matplotlib figure object, calculate the filename of the saved image in
    the saved image directory."""
    ext = '_{:02d}.png'.format(fig.number)
    return _process_path(example_filename, ext)

def text_filename(example_filename):
    """Given the filename of the example source code, calculate the filename of
    the saved output (i.e. printed to stdout) in the saved image directory."""
    return _process_path(example_filename, '.txt')

def _process_path(orig_filename, ext):
    """Generate the approriate output filename for a given python script's
    filename and an extension."""
    orig_filename = os.path.basename(orig_filename)
    new_name, _ = os.path.splitext(orig_filename)
    return os.path.join(SAVED_DIRECTORY, new_name + ext)

def make_pil_image(fig):
    """Make a PIL image object from a matplotlib figure instance."""
    fig.canvas.draw()
    size = fig.canvas.get_width_height()
    im = Image.frombytes('RGB', size, fig.canvas.tostring_rgb())
    return im

def run(example_filename):
    """Run an example and return the output text (or an empty string if nothing
    was printed) and a list of figure objects generated by the script."""
    # Ignore the fact that the examples call show() and this warns with Agg..
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        output_text = run_file(example_filename)
        figures = list(current_figures())
    return output_text, figures

def run_file(example_filename):
    """Run the python file specified by "example_filename" in the current global
    namespace and return anything it prints to stdout."""
    # All of this appears to be overkill, but is necessary to emulate the file
    # being run in its directory and have the __file__, etc attributes work
    # correctly.
    import __main__
    # Get current environment
    prev_stdout = sys.stdout
    prev_cwd = os.getcwd()
    prev_file = __main__.__dict__.get('__file__', '')

    # Direction that we'll execute the file in
    example_dir = os.path.dirname(example_filename)
    if not example_dir:
        example_dir = '.'

    # File-like object that stdout will be captured in
    # I can't find a better way to do this than version conditionals...
    if sys.version_info.major >= 3:
        printed_output = io.StringIO()
    else:
        printed_output = io.BytesIO()

    # Setup the environment for execution
    sys.path.append(example_dir)
    os.chdir(example_dir)
    __main__.__file__ = example_filename
    sys.stdout = printed_output

    try:
        # Run the file and capture anything it prints to stdout
        with open(example_filename) as infile:
            code = compile(infile.read(), example_filename, 'exec')
            exec(code, __main__.__dict__)
    finally:
        # Return the environment to "normal"
        os.chdir(prev_cwd)
        __main__.__file__ = prev_file
        sys.path.pop()
        sys.stdout = prev_stdout

    return printed_output.getvalue()

def save_output(filename):
    """Save the figures generated by the specified python script to the
    saved images directory."""
    output_text, figures = run(filename)
    if output_text:
        with open(text_filename(filename), 'w') as outfile:
            outfile.write(output_text)
    for fig in figures:
        im = make_pil_image(fig)
        im.save(image_filename(filename, fig))
