# Create figures in Python that handle LaTeX, and save images to files in my
# preferred formatting. I typically place this code in the root of each of my
# projects, and import using:
# from latexify import *
# which will also run the latexify() function on the import.

# Based on code from https://nipunbatra.github.io/blog/2014/latexify.html

import matplotlib
import matplotlib.pyplot as plt
from math import sqrt

# Back-end to use depends on the system
from matplotlib.backends.backend_pgf import FigureCanvasPgf

matplotlib.backend_bases.register_backend("pdf", FigureCanvasPgf)
# matplotlib.use('pgf')
# from matplotlib.backends.backend_pgf import FigureCanvasPgf
# matplotlib.backend_bases.register_backend('ps', FigureCanvasPgf)

import seaborn as sns

sns.set_style("white", {"legend.frameon": False})

# my preferred palette. From
# https://seaborn.pydata.org/tutorial/color_palettes.html: "The cubehelix color
# palette system makes sequential palettes with a linear increase or decrease in
# brightness and some variation in hue. This means that the information in your
# colormap will be preserved when converted to black and white (for printing) or
# when viewed by a colorblind individual."

# I typically set the number of colors (below, 8) to the distinct colors I need
# in a given plot, so as to use the full range.
sns.set_palette(sns.color_palette("cubehelix", 5))


# The following is the latexify function. It allows you to create 2 column or 1
# column figures. You may also wish to alter the height or width of the figure.
# The default settings are good for most cases. You may also change the
# parameters such as labelsize and fontsize based on your classfile.
def latexify(fig_width=None, fig_height=None, columns=1):
    """Set up matplotlib's RC params for LaTeX plotting.
    Call this before plotting a figure.
    Parameters
    ----------
    fig_width : float, optional, inches
    fig_height : float,  optional, inches
    columns : {1, 2}
    """

    # code adapted from http://www.scipy.org/Cookbook/Matplotlib/LaTeX_Examples
    # Width and max height in inches for IEEE journals taken from
    # computer.org/cms/Computer.org/Journal%20templates/transactions_art_guide.pdf

    assert columns in [1, 2]

    if fig_width is None:
        fig_width = 6.9 if columns == 1 else 13.8  # width in inches #3.39

    if fig_height is None:
        golden_mean = (sqrt(5) - 1.0) / 2.0  # Aesthetic ratio
        fig_height = fig_width * golden_mean  # height in inches

    MAX_HEIGHT_INCHES = 16.0
    if fig_height > MAX_HEIGHT_INCHES:
        print(("WARNING: fig_height too large:" + fig_height + "so will reduce to" + MAX_HEIGHT_INCHES + "inches."))
        fig_height = MAX_HEIGHT_INCHES

    params = {
        # 'backend': 'ps',
    # 'backend': 'ps',
        "pgf.rcfonts": False,
        # "pgf.preamble": '\n'.join([
        #     "\\PassOptionsToPackage{dvipsnames}{xcolor}",
        #     "\\usepackage{gensymb,amsmath}", 
        #     "\\usepackage{lmodern}"
        # ]),
        # "pgf.texsystem": "pdflatex",
        # "text.latex.preamble": '\n'.join([
        #     "\\PassOptionsToPackage{dvipsnames}{xcolor}",
        #     "\\usepackage{gensymb,amsmath}", 
        #     "\\usepackage{lmodern}"
        # ]),
        "pgf.preamble": '\n'.join(["\\usepackage{gensymb,amsmath}", "\\usepackage{lmodern}"]),  # Removed xcolor
        "pgf.texsystem": "pdflatex",
        "text.latex.preamble": '\n'.join(["\\usepackage{gensymb,amsmath}", "\\usepackage{lmodern}"]),  # Removed xcolor

        # values below are useful defaults. individual plot fontsizes are
        # modified as necessary.
        "axes.labelsize": 13,  # fontsize for x and y labels
        "axes.titlesize": 13,
        "font.size": 10,
        "legend.fontsize": 15,
        "xtick.labelsize": 15,
        "ytick.labelsize": 15,
        # 'text.usetex': True,
        "figure.figsize": [fig_width, fig_height],
        # "font.sans-serif": ["Arial", "Liberation Sans",
        # "Bitstream Vera Sans", "sans-serif"],
        # 'font.family': 'serif',
        "lines.linewidth": 1,
        "lines.markersize": 1,
        "xtick.major.pad": 2,
        "ytick.major.pad": 2,
        "axes.unicode_minus": False,
    }

    matplotlib.rcParams.update(params)


def saveimage(name, fig=plt, extension="pdf", folder="plots/", close=True):
    sns.despine()

    # Minor ticks off by default in matplotlib
    # plt.minorticks_off()

    # grid being off is the default for seaborn white style, so not needed.
    plt.grid(False, axis="x")
    plt.grid(False, axis="y")

    fig.savefig("{}{}.{}".format(folder, name, extension), bbox_inches="tight", dpi=500)
    if close:
        plt.close()


latexify()
