'''
Minimum working example
version 1.3.1
Plots a descriptive flowchart showing error
@Author: wronk
'''

from os import path as op
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.image import BboxImage
from matplotlib.offsetbox import TextArea, AnnotationBbox
from matplotlib.transforms import Bbox
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.backends.backend_pdf import PdfPages

###############################################################################
#EDIT SAVE FOLDER TO AUTOSAVE
saveFig = True
save_fname = '/home/wronk/'

mpl.rcParams['mathtext.default'] = 'rm'

#Box properties
bboxProp = dict(boxstyle='round,pad=0.3', fill=False, ec='w', linewidth=8)

#Colors for annotation arrows
senCol = (0.35, 0.35, 0.35)
srcCol = (.8, .216, 0.0)

#Define color map
cm = 'autumn'
colRange = np.atleast_2d(np.arange(256)/256.)

###############################################################################
#Initialize Plot
plt.ion()
plt.close('all')
mpl.rcParams['pdf.fonttype'] = 42

figSize = (12, 6)
ftSize = 32
dpi = 80
rowYVals = [0.085, 0.55, 0.85]
rowXVals = [0, .075, .2, .34, .5, .65, .8, .875]

fig = plt.figure(figsize=figSize, facecolor='white', dpi=dpi)
#fig = plt.figure(figsize=figSize, facecolor='white')
#ax = fig.gca()
ax = plt.subplot(111)

###############################################################################
#Mid level of annotation boxes
#Text for boxes in mathematical font
flowMid = r'$j_{N-1}$'

midBox3 = ax.annotate(flowMid, xy=(0, .5), xycoords='axes fraction',
                      xytext=(rowXVals[3], rowYVals[1]),
                      textcoords='axes fraction',
                      size=ftSize, ha='left', va='center', bbox=bboxProp,
                      arrowprops=None, color='black')

#######################################
#HEREIN LIES THE PROBLEM
#Add color (gradients) behind the boxes

#Something about the bbox window extent property must get shifted when saving
#as a vector format (pdf, svg, eps) but not common image types. In my experience,
#every bboximage was translated and scaled uniformly.

#I tried changing DPI, image size, subplot parameters as well as sychronizing
#plotting and saving figures in matplotlibrc to no avail. DPI does seem to
#have some sort of effect in terms of how the bboximage gets shifted though.

gradient = BboxImage(midBox3.get_bbox_patch().get_window_extent,
                     data=np.zeros_like(colRange), cmap=cm, norm=None,
                     origin=None)
ax.add_artist(gradient)

#######################################
plt.draw()

#Save figure as pdf and png to highlight difference
#Same results when saving from GUI window
if saveFig:
    pdfFile = PdfPages(save_fname + 'flowchart_PDF.pdf')
    plt.savefig(pdfFile, format='pdf', dpi=fig.dpi)
    pdfFile.close()
    plt.savefig(save_fname + 'flowchart_PNG.png', dpi=fig.dpi)
plt.show()