"""
Generates plots / figures when run as a script.
Plot files are placed in the :file:`plots` directory.

By default, simply running ``python -m src.plots`` generates **ALL** plots,
which may not be desired.  Instead, one can pass a list of plots to generate:
``python -m src.plots plot1 plot2 ...``.  The full list of plots is shown in
the usage information ``python -m src.plots --help``.

Typing can be reduced by using shell brace expansion, e.g. ``python -m
src.plots observables_{design,posterior}`` for both ``observables_design`` and
``observables_posterior``.  In addition, plots may be given as paths to plot
filenames, which enables shell globbing, e.g. ``python -m src.plots
plots/observables_*``.

In the code, each plot is generated by a function tagged with the ``@plot``
decorator.
"""

import itertools
import logging
from pathlib import Path
import subprocess
import tempfile
import warnings

from sklearn.externals import joblib
import h5py
import hsluv
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import lines
from matplotlib import patches
from matplotlib import ticker
from scipy import special
from scipy.interpolate import PchipInterpolator
from sklearn.decomposition import PCA
from sklearn.gaussian_process import GaussianProcessRegressor as GPR
from sklearn.gaussian_process import kernels
from sklearn.mixture import GaussianMixture
from functools import wraps

from . import workdir, mcmc, find_obs_index, data_list, exp_data_param_list, exp_data_list#, data_list_val#, model, expt
from .design import Design
from .emulator import emulators

fontsmall, fontnormal, fontlarge = 5, 6, 7
offblack = '#262626'
aspect = 1/1.618
resolution = 72.27
textwidth = 307.28987/resolution
textheight = 261.39864/resolution
fullwidth = 350/resolution
fullheight = 270/resolution

plt.rcdefaults()
plt.rcParams.update({
#    'font.family': 'sans-serif',
#    'font.sans-serif': ['Lato'],
#    'mathtext.fontset': 'custom',
#    'mathtext.default': 'it',
#    'mathtext.rm': 'sans',
#    'mathtext.it': 'sans:italic:medium',
#    'mathtext.cal': 'sans',
    'font.size': fontnormal,
    'legend.fontsize': fontnormal,
    'axes.labelsize': fontnormal,
    'axes.titlesize': fontlarge,
    'xtick.labelsize': fontsmall,
    'ytick.labelsize': fontsmall,
    'font.weight': 400,
    'axes.labelweight': 400,
    'axes.titleweight': 400,
    'lines.linewidth': .5,
    'lines.markersize': 3,
    'lines.markeredgewidth': 0,
    'patch.linewidth': .5,
    'axes.linewidth': .4,
    'xtick.major.width': .4,
    'ytick.major.width': .4,
    'xtick.minor.width': .4,
    'ytick.minor.width': .4,
    'xtick.major.size': 1.2,
    'ytick.major.size': 1.2,
    'xtick.minor.size': .8,
    'ytick.minor.size': .8,
    'xtick.major.pad': 1.5,
    'ytick.major.pad': 1.5,
    'axes.formatter.limits': (-5, 5),
    'axes.spines.top': False,
    'axes.spines.right': False,
    'axes.labelpad': 3,
    'text.color': offblack,
    'axes.edgecolor': offblack,
    'axes.labelcolor': offblack,
    'xtick.color': offblack,
    'ytick.color': offblack,
    'legend.numpoints': 1,
    'legend.scatterpoints': 1,
    'legend.frameon': False,
    'image.cmap': 'Blues',
    'image.interpolation': 'none',
    'pdf.fonttype': 42
})


plotdir = workdir / 'plots'
plotdir.mkdir(exist_ok=True)

plot_functions = {}


def plot(f):
    """
    Plot function decorator.  Calls the function, does several generic tasks,
    and saves the figure as the function name.

    """
    @wraps(f)
    def wrapper(*args, **kwargs):
        logging.info('generating plot: %s', f.__name__)
        f(*args, **kwargs)

        fig = plt.gcf()

        if not fig.get_tight_layout():
            set_tight(fig)

        plotfile = plotdir / '{}.pdf'.format(f.__name__)
        fig.savefig(str(plotfile))
        logging.info('wrote %s', plotfile)
        plt.close(fig)

    plot_functions[f.__name__] = wrapper

    return wrapper


def set_tight(fig=None, **kwargs):
    """
    Set tight_layout with a better default pad.

    """
    if fig is None:
        fig = plt.gcf()

    kwargs.setdefault('pad', .1)
    fig.set_tight_layout(kwargs)


def auto_ticks(ax, axis='both', minor=False, **kwargs):
    """
    Convenient interface to matplotlib.ticker locators.

    """
    axis_list = []

    if axis in {'x', 'both'}:
        axis_list.append(ax.xaxis)
    if axis in {'y', 'both'}:
        axis_list.append(ax.yaxis)

    for axis in axis_list:
        axis.get_major_locator().set_params(**kwargs)
        if minor:
            axis.set_minor_locator(ticker.AutoMinorLocator(minor))

def format_ci(samples, ci=.9):
    """
    Compute the median and a credible interval for an array of samples and
    return a TeX-formatted string.

    """
    cil, cih = mcmc.credible_interval(samples, ci=ci)
    median = np.median(samples)
    ul = median - cil
    uh = cih - median

    # decide precision for formatting numbers
    # this is NOT general but it works for the present data
    if abs(median) < .2 and ul < .02:
        precision = 3
    elif abs(median) < 1:
        precision = 2
    else:
        precision = 1

    fmt = str(precision).join(['{:#.', 'f}'])

    return ''.join([
        '$', fmt.format(median),
        '_{-', fmt.format(ul), '}',
        '^{+', fmt.format(uh), '}$'
    ])

def _posterior(
        params=None, ignore=None,
        scale=1, padr=.99, padt=.98,
        cmap=None
):
    """
    Triangle plot of posterior marginal and joint distributions.

    """
    chain = mcmc.Chain()

    if params is None and ignore is None:
        params = set(chain.keys)
    elif params is not None:
        params = set(params)
    elif ignore is not None:
        params = set(chain.keys) - set(ignore)

    keys, labels, ranges = map(list, zip(*(
        i for i in zip(chain.keys, chain.labels, chain.range)
        if i[0] in params
    )))
    ndim = len(params)

    data = chain.load(*keys).T

    cmap = plt.get_cmap(cmap)
    cmap.set_bad('white')

    line_color = cmap(.8)
    fill_color = cmap(.5, alpha=.1)

    fig, axes = plt.subplots(
        nrows=ndim, ncols=ndim,
        sharex='col', sharey='row',
        figsize=2*(scale*fullheight,)
    )

    for samples, key, lim, ax in zip(data, keys, ranges, axes.diagonal()):
        counts, edges = np.histogram(samples, bins=50, range=lim)
        x = (edges[1:] + edges[:-1]) / 2
        y = .85 * (lim[1] - lim[0]) * counts / counts.max() + lim[0]
        # smooth histogram with monotonic cubic interpolation
        interp = PchipInterpolator(x, y)
        x = np.linspace(x[0], x[-1], 10*x.size)
        y = interp(x)
        ax.plot(x, y, lw=.5, color=line_color)
        ax.fill_between(x, lim[0], y, color=fill_color, zorder=-10)

        ax.set_xlim(lim)
        ax.set_ylim(lim)

        if key == 'dmin3':
            samples = samples**(1/3)

        ax.annotate(
            format_ci(samples), (.62, .92), xycoords='axes fraction',
            ha='center', va='bottom', fontsize=4.5
        )

    for ny, nx in zip(*np.tril_indices_from(axes, k=-1)):
        axes[ny][nx].hist2d(
            data[nx], data[ny], bins=100,
            range=(ranges[nx], ranges[ny]),
            cmap=cmap, cmin=1
        )
        axes[nx][ny].set_axis_off()

    for key, label, axb, axl in zip(keys, labels, axes[-1], axes[:, 0]):
        for axis in [axb.xaxis, axl.yaxis]:
            axis.set_label_text(label.replace(r'\ [', '$\n$['), fontsize=4)
            axis.set_tick_params(labelsize=3)
            if key == 'dmin3':
                ticks = [0., 1.2, 1.5, 1.7]
                axis.set_ticklabels(list(map(str, ticks)))
                axis.set_ticks([t**3 for t in ticks])
            else:
                axis.set_major_locator(ticker.LinearLocator(3))
                if (
                        axis.axis_name == 'x'
                        and scale / ndim < .13
                        and any(len(str(x)) > 4 for x in axis.get_ticklocs())
                ):
                    for t in axis.get_ticklabels():
                        t.set_rotation(30)

        axb.get_xticklabels()[0].set_horizontalalignment('left')
        axb.get_xticklabels()[-1].set_horizontalalignment('right')
        axl.get_yticklabels()[0].set_verticalalignment('bottom')
        axl.get_yticklabels()[-1].set_verticalalignment('top')

    set_tight(fig, pad=.05, h_pad=.1, w_pad=.1, rect=[0., 0., padr, padt])


@plot
def posterior():
    """
    Pairplot of posteriors for all calibration inputs.
    Diagonal displays marginal density.
    Lower off-diagonal displays pairwise scatter plot.

    """
    _posterior( scale=1.6, padr=1., padt=.99)

@plot
def design_exp_compare(obs,**kwargs):
    index=find_obs_index(obs, **kwargs)
    x=exp_data_param_list[index]
    print(x)
    y=exp_data_list[index,0]
    print(y)
    plt.plot(x, y, color='black', marker='o', linewidth=2, markersize=12)
    design_y = data_list[:,index]
    for y in design_y:
        plt.plot(x,y)
    plt.ylim((0,1))

@plot
def model_exp_compare(obs,**kwargs):
    index=find_obs_index(obs, **kwargs)
    x=exp_data_param_list[index]
    print(x)
    y=exp_data_list[index,0]
    print(y)
    plt.plot(x, y, color='black', marker='o', linewidth=2, markersize=12)
    samples = mcmc.Chain().samples(100)
    for y in samples:
        plt.plot(x,y)
    plt.ylim((0,1))

if __name__ == '__main__':
    import argparse
    from matplotlib.mathtext import MathTextWarning

    warnings.filterwarnings(
        'ignore',
        category=MathTextWarning,
        message='Substituting with a symbol from Computer Modern.'
    )

    choices = list(plot_functions)

    def arg_to_plot(arg):
        arg = Path(arg).stem
        if arg not in choices:
            raise argparse.ArgumentTypeError(arg)
        return arg

    parser = argparse.ArgumentParser(description='generate plots')
    parser.add_argument(
        'plots', nargs='*', type=arg_to_plot, metavar='PLOT',
        help='{} (default: all)'.format(', '.join(choices).join('{}'))
    )
    args = parser.parse_args()

    if args.plots:
        for p in args.plots:
            plot_functions[p]()
    else:
        for f in plot_functions.values():
            f()
