"""
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.
"""

from collections import OrderedDict, Counter
import itertools
import logging
import multiprocessing
from pathlib import Path
import pickle
import subprocess
import tempfile
import warnings

import h5py
import hsluv
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib import lines
from matplotlib import patches
from matplotlib import ticker
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
from scipy import special
from scipy.interpolate import PchipInterpolator, UnivariateSpline, interp2d
from scipy.optimize import curve_fit
from sklearn.decomposition import PCA
from sklearn.gaussian_process import GaussianProcessRegressor as GPR
from sklearn.gaussian_process import kernels
from sklearn.model_selection import KFold

from . import cachedir, workdir, systems, parse_system, expt, model, mcmc
from .design import Design
from .emulator import Emulator, emulators


# golden ratio
aspect = 1/1.618

# font sizes
fontsize = dict(
    large=11,
    normal=10,
    small=9,
    tiny=8,
)

# new tableau colors
# https://www.tableau.com/about/blog/2016/7/colors-upgrade-tableau-10-56782
colors = OrderedDict([
    ('blue', '#4e79a7'),
    ('orange', '#f28e2b'),
    ('green', '#59a14f'),
    ('red', '#e15759'),
    ('cyan', '#76b7b2'),
    ('purple', '#b07aa1'),
    ('brown', '#9c755f'),
    ('yellow', '#edc948'),
    ('pink', '#ff9da7'),
    ('gray', '#bab0ac')
])

offblack = '.15'

plt.rcdefaults()
plt.rcParams.update({
    'font.family': 'sans-serif',
    'font.sans-serif': ['Lato'],
    'mathtext.fontset': 'custom',
    'mathtext.default': 'it',
    'mathtext.rm': 'sans',
    'mathtext.cal': 'sans',
    'font.size': fontsize['normal'],
    'legend.fontsize': fontsize['normal'],
    'axes.labelsize': fontsize['normal'],
    'axes.titlesize': fontsize['large'],
    'xtick.labelsize': fontsize['small'],
    'ytick.labelsize': fontsize['small'],
    'font.weight': 400,
    'axes.labelweight': 400,
    'axes.titleweight': 400,
    'axes.prop_cycle': plt.cycler('color', list(colors.values())),
    'lines.linewidth': .8,
    'lines.markersize': 3,
    'lines.markeredgewidth': 0,
    'patch.linewidth': .8,
    'hatch.linewidth': .8,
    'axes.linewidth': .6,
    'xtick.major.width': .6,
    'ytick.major.width': .6,
    'xtick.minor.width': .4,
    'ytick.minor.width': .4,
    'xtick.major.size': 3.,
    'ytick.major.size': 3.,
    'xtick.minor.size': 2.,
    'ytick.minor.size': 2.,
    'xtick.major.pad': 3.5,
    'ytick.major.pad': 3.5,
    'axes.labelpad': 4.,
    'axes.formatter.limits': (-5, 5),
    'axes.spines.top': False,
    'axes.spines.right': False,
    'text.color': offblack,
    'axes.edgecolor': offblack,
    'axes.labelcolor': offblack,
    'xtick.color': offblack,
    'ytick.color': offblack,
    'legend.frameon': False,
    'image.cmap': 'Blues',
    'image.interpolation': 'none',
})


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.

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

        fig = plt.gcf()

        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 figsize(relwidth=1, aspect=.618, refwidth=6.5):
    """
    Return figure dimensions from a relative width (to a reference width) and
    aspect ratio (default: 1/golden ratio).

    """
    width = relwidth * refwidth

    return width, width*aspect


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 cmap_to_alpha(cmap=plt.cm.inferno, fraction=.2):
    """
    Fade the bottom of a colormap to white.

    Currently only works with ListedColormap objects, of which the new cmaps
    (inferno, magma, plasma, viridis) are examples.

    """
    try:
        colors = cmap.colors.copy()
        n = int(fraction*len(colors))

        for i, rgb in enumerate(colors[:n]):
            colors[i] = rgb + [i/n]

        return type(cmap)(colors, cmap.name + '_mod')
    except AttributeError:
        cmin = 25 if cmap == plt.cm.Oranges_r else 0
        colors = [list(cmap(n/256)) for n in range(cmin, 200)]
        n = int(fraction*len(colors))

        for i, rgb in enumerate(colors[:n]):
            rgb[-1] = i/n
            colors[i] = rgb

        return ListedColormap(colors, cmap.name + '_mod')


def format_system(system):
    """
    Format a system string into a display name, e.g.:

    >>> format_system('PbPb2760')
    'Pb-Pb 2.76 TeV'

    >>> format_system('AuAu200')
    'Au-Au 200 GeV'

    """
    proj, energy = parse_system(system)

    if energy > 1000:
        energy /= 1000
        prefix = 'T'
    else:
        prefix = 'G'

    return '{} {} {}eV'.format('-'.join(proj), energy, prefix)


def darken(color_hex, amount=.5):
    """
    Darken a color by the given amount in HSLuv space.

    """
    return color_hex
    #H, S, L = hsluv.hex_to_hsluv(color_hex)
    #return hsluv.hsluv_to_rgb((H, S, (1 - amount)*L))


def obs_color_hsluv(obs, subobs):
    """
    Return a nice color for the given observable in HSLuv space.
    Use obs_color() to obtain an RGB color.

    """
    if obs in {'dNch_deta', 'pT_fluct'}:
        return 250, 90, 55

    if obs == 'mean_pT':
        return 230, 90, 65

    if obs == 'dET_deta':
        return 10, 65, 55

    if obs in {'iden_dN_dy', 'iden_mean_pT'}:
        return dict(
            charged=(250, 90, 55),
            pion=(210, 85, 70),
            kaon=(130, 88, 68),
            proton=(30, 90, 62),
        )[subobs]

    if obs == 'vnk':
        return {
            (2, 2): (250, 90, 65),
            (3, 2): (150, 90, 67),
            (4, 2): (20, 90, 62),
            (2, 4): (310, 70, 50),
        }[subobs]

    raise ValueError('unknown observable: {} {}'.format(obs, subobs))


def obs_color(obs, subobs):
    """
    Return a nice color for the given observable.

    """
    return hsluv.hsluv_to_rgb(obs_color_hsluv(obs, subobs))


def obs_label(obs, subobs, differentials=False, full_cumulants=False):
    """
    Return a formatted label for the given observable.

    """
    if obs.startswith('d') and obs.endswith('_deta'):
        return (r'$d{}/d\eta$' if differentials else '${}$').format(
            {'Nch': r'N_\mathrm{ch}', 'ET': r'E_T'}[obs[1:-5]])

    id_parts_labels = {
        'dN_dy': '$dN_{}/dy$' if differentials else '$N_{}$',
        'mean_pT': r'$\langle p_T^{} \rangle$'
    }
    if obs in id_parts_labels:
        return id_parts_labels[obs].format(
                {'pion': '\pi', 'kaon': 'K', 'proton': 'p', None: '{}'}[subobs]
        )

    if obs == 'pT_fluct':
        return r'$\delta p_T/\langle p_T \rangle$'

    if obs == 'vnk':
        n, k = subobs
        return '$v_{}{}$'.format(
            n,
            (r'\{' + str(k) + r'\}') if full_cumulants else ''
        )


def _observables_plots():
    """
    Metadata for observables plots.

    """
    def id_parts_plots(obs):
        return [(obs, species, dict(label=label)) for species, label in [
            ('pion', '$\pi$'), ('kaon', '$K$'), ('proton', '$p$')
        ]]

    return [
        dict(
            title='Yields',
            xlabel=dict(
                pPb5020='Centrality %',
                PbPb5020='Centrality %',
            ),
            ylabel=(
                r'$dN_\mathrm{ch}/d\eta,\ dN/dy,\ dE_T/d\eta\ [\mathrm{GeV}]$'
            ),
            xlim=dict(
                pPb5020=(0, 60),
                PbPb5020=(0, 80),
            ),
            ylim=(1e2, 1e5),
            yscale='log',
            height_ratio=1.5,
            subplots=[
                ('dNch_deta', None, dict(label=r'$N_\mathrm{ch}^{\times 25}$', scale=25)),
                ('dET_deta', None, dict(label=r'$E_T^{\times 5}$', scale=5)),
                *id_parts_plots('iden_dN_dy')
            ]
        ),
        dict(
            title='Mean $p_T$',
            xlabel=dict(
                pPb5020=r'$n_\mathrm{ch} / \langle n_\mathrm{ch} \rangle$',
                PbPb5020='Centrality %'
            ),
            ylabel=r'$\langle p_T \rangle$ [GeV]',
            xlim=dict(
                pPb5020=(1, 6),
                PbPb5020=(0, 80),
            ),
            ylim=(0, 1.5),
            subplots=[
                ('mean_pT', None, dict(label='ch')),
                *id_parts_plots('iden_mean_pT')
            ]
        ),
        dict(
            title='Mean $p_T$ fluctuations',
            xlabel=dict(
                pPb5020=r'$n_\mathrm{ch} / \langle n_\mathrm{ch} \rangle$',
                PbPb5020='Centrality %'
            ),
            ylabel=r'$\delta p_T/\langle p_T \rangle$',
            xlim=dict(
                pPb5020=(1, 6),
                PbPb5020=(0, 80),
            ),
            ylim=(0, 0.05),
            subplots=[('pT_fluct', None, dict())]
        ),
        dict(
            title='Flow cumulants',
            xlabel=dict(
                pPb5020=(r'$n_\mathrm{trk}^\mathrm{offline}/'
                         r'\langle n_\mathrm{trk}^\mathrm{offline} \rangle$'),
                PbPb5020='Centrality %',
            ),
            ylabel=r'$v_n\{2\}$',
            xlim=dict(
                pPb5020=(1, 6),
                PbPb5020=(0, 80),
            ),
            ylim=(0, .15),
            subplots=[
                ('vnk', (n, 2), dict(label='$v_{}$'.format(n)))
                for n in [2, 3, 4]
            ]
        )
    ]


def _observables(posterior=False):
    """
    Model observables at all design points or drawn from the posterior with
    experimental data points.
    """
    plots = _observables_plots()
    plots.pop(2)

    plot_fixes = [
        dict(
            ylabel=r'$dN_\mathrm{ch}/d\eta$',
            ylim=(1, 1e4),
            height_ratio=1,
            subplots=[
                ('dNch_deta', None, dict(label=r'$N_\mathrm{ch}$', scale=1))
            ]
        ),
        dict(
            subplots=[('mean_pT', None, dict(label='ch'))]
        ),
        dict()
    ]

    for a, b in zip(plots, plot_fixes):
        for k, v in b.items():
            a[k] = v

    fig, axes = plt.subplots(
        nrows=len(plots), ncols=len(systems),
        figsize=figsize(.8, aspect=1.1),
    )

    if posterior:
        samples = mcmc.Chain().samples(100)

    for (plot, system), ax in zip(itertools.product(plots, systems), axes.flat):
        for obs, subobs, opts in plot['subplots']:
            color = obs_color(obs, subobs)
            scale = opts.get('scale')

            try:
                model_data = model.data[system][obs][subobs]
                x = model_data['x']
                Y = (samples[system][obs][subobs]
                     if posterior else model_data['Y'])
            except KeyError:
                continue

            if scale is not None:
                Y = Y*scale

            for y in Y:
                ax.plot(x, y, color=color, alpha=.2, lw=.3)

            try:
                dset = expt.data[system][obs][subobs]
            except KeyError:
                print(system, obs, subobs, 'not found')
                pass
            else:
                x = dset['x']
                y = dset['y']
                yerr = np.sqrt(sum(
                    e**2 for e in dset['yerr'].values()
                ))

                if scale is not None:
                    y = y*scale
                    yerr = yerr*scale

                ax.errorbar(
                    x, y, yerr=yerr, fmt='o',
                    capsize=0, mfc='.25', mec='.25', mew=.2, zorder=1000
                )

                xmin, xmax = plot['xlim'][system]

                ax.text(
                    x[-1] + .03*(xmax - xmin), y[-1], opts['label'],
                    color=darken(color), ha='left', va='center'
                )

        auto_ticks(ax, 'x', nbins=4, minor=2)

        if plot.get('yscale') == 'log':
            ax.set_yscale('log')
            ax.minorticks_off()
        else:
            auto_ticks(ax, 'y', nbins=4, minor=2)

        if ax.is_first_row():
            ax.set_title(format_system(system), va='top')

        if ax.is_first_col():
            ax.set_ylabel(plot['ylabel'])
        else:
            ax.set_yticklabels([])

        if ax.is_last_col():
            ax.text(
                1.02, .5, plot['title'],
                transform=ax.transAxes, ha='left', va='center',
                size=plt.rcParams['axes.labelsize'], rotation=-90
            )

        ax.set_xlabel(plot['xlabel'][system])
        ax.set_xlim(*plot['xlim'][system])
        ax.set_ylim(plot['ylim'])

    set_tight(fig, rect=[0, 0, .97, 1])


@plot
def observables_design():
    _observables(posterior=False)


@plot
def observables_posterior():
    _observables(posterior=True)


def observables(system):
    """
    Model observables at all design points or drawn from the posterior with
    experimental data points.

    """
    plots = _observables_plots()
    plots.pop(2)

    plot_fixes = [
        dict(
            title=r'Yields $dN_\mathrm{ch}/d\eta$',
            ylabel=r'$dN_\mathrm{ch}/d\eta$',
            ylim=(1e1, 1e4) if system == 'PbPb5020' else (1, 1e3),
            height_ratio=1,
            subplots=[
                ('dNch_deta', None, dict(label=''))
            ]
        ),
        dict(
            title=r'Mean $p_T$ [GeV]',
            subplots=[('mean_pT', None, dict(label=''))]),
        dict(
            title=r'Flow cumulants $v_n\{2\}$',
        )
    ]

    for a, b in zip(plots, plot_fixes):
        for k, v in b.items():
            a[k] = v

    fig, axes = plt.subplots(
        nrows=2, ncols=len(plots),
        figsize=figsize(1, aspect=.6),
    )

    title = dict(
        pPb5020=r'$p$-Pb 5.02 TeV',
        PbPb5020=r'Pb-Pb 5.02 TeV',
    )

    for (posterior, plot), ax in zip(
            itertools.product([False, True], plots), axes.flat):

        if posterior:
            samples = mcmc.Chain().samples(100)

        for obs, subobs, opts in plot['subplots']:
            color = obs_color(obs, subobs)
            scale = opts.get('scale')

            try:
                model_data = model.data[system][obs][subobs]
                x = model_data['x']
                Y = (samples[system][obs][subobs]
                     if posterior else model_data['Y'])
            except KeyError:
                continue

            if scale is not None:
                Y = Y*scale

            for y in Y:
                alpha = .2 if system == 'pPb5020' else 0.1
                ax.plot(x, y, color=color, alpha=alpha, lw=.3)

            try:
                dset = expt.data[system][obs][subobs]
            except KeyError:
                print(system, obs, subobs, 'not found')
                pass
            else:
                x = dset['x']
                y = dset['y']
                yerr = np.sqrt(sum(
                    e**2 for e in dset['yerr'].values()
                ))

                if scale is not None:
                    y = y*scale
                    yerr = yerr*scale

                ax.errorbar(
                    x, y, yerr=yerr, fmt='o',
                    capsize=0, mfc='.25', mec='.25', mew=.2, zorder=1000
                )

                xmin, xmax = plot['xlim'][system]

                ax.text(
                    x[-1] + .03*(xmax - xmin), y[-1], opts['label'],
                    color=darken(color), ha='left', va='center'
                )

        ax.set_xlim(*plot['xlim'][system])
        ax.set_ylim(plot['ylim'])

        auto_ticks(ax, 'x', nbins=4, minor=2)
        auto_ticks(ax, 'y', nbins=4, minor=2)

        if plot.get('yscale') == 'log':
            ax.set_yscale('log')
            ax.minorticks_off()

        if ax.is_first_row():
            ax.set_title(plot['title'], fontsize=fontsize['normal'])
            ax.set_xticklabels([])
            if ax.is_first_col():
                ax.set_ylabel('Training data')
        else:
            ax.set_xlabel(plot['xlabel'][system])
            if ax.is_first_col():
                ax.set_ylabel('Posterior samples')
                ax.annotate(title[system], xy=(.1, .1),
                            xycoords='axes fraction',ha='left', va='bottom')
    set_tight()


@plot
def observables_ppb():
    observables('pPb5020')


@plot
def observables_pbpb():
    observables('PbPb5020')


@plot
def observables_map():
    """
    Model observables and ratio to experiment at the maximum a posteriori
    (MAP) estimate.

    """
    systems = ['pPb5020', 'PbPb5020']
    plots = _observables_plots()

    ylim = {
        'Yields': (1e-1, 1e5),
        'Mean $p_T$': (0, 1.7),
        'Mean $p_T$ fluctuations': (0, 0.05),
        'Flow cumulants': (0, .12),
    }

    for n, p in enumerate(plots):
        p['ylim'] = ylim[p['title']]
        if p['title'] == 'Flow cumulants':
            move_index = n
            p.update(
                ylabel=r'$v_n\{k\}$',
                subplots=[
                    ('vnk', nk, dict(label='$v_{}\{{{}\}}$'.format(*nk)))
                    for nk in [(2, 2), (2, 4), (3, 2), (4, 2)]
                ],
                legend=True
            )
    fig = plt.figure(figsize=figsize(0.88, 1.5))

    yields, mean_pT, mean_pT_fluct, flows = [
        gridspec.GridSpecFromSubplotSpec(
            2, 2, gs, height_ratios=[5, 1] if n == 0 else [3, 1],
            hspace=0.1, wspace=.14
        ) for n, gs in enumerate(
            gridspec.GridSpec(4, 1, height_ratios=[6, 4, 4, 4])
        )
    ]

    gridspecs = [yields, mean_pT, mean_pT_fluct, flows]
    rows = zip(plots, gridspecs)

    for nrow, (plot, gs) in enumerate(rows):
        axes = [fig.add_subplot(ax) for ax in gs]
        cols = zip(systems, axes[:2], axes[2:])

        for ncol, (system, ax, ratio_ax) in enumerate(cols):
            for obs, subobs, opts in plot['subplots']:
                color = obs_color(obs, subobs)
                scale = opts.get('scale')

                try:
                    x = model.map_data[system][obs][subobs]['x']
                    y = model.map_data[system][obs][subobs]['Y']
                except KeyError:
                    continue

                if scale is not None:
                    y = y*scale

                ax.plot(x, y, color=color)

                if 'label' in opts:
                    xmin, xmax = plot['xlim'][system]
                    ax.text(
                        x[-1] + .03*(xmax - xmin), y[-1], opts['label'],
                        color=darken(color), ha='left', va='center'
                    )

                try:
                    dset = expt.data[system][obs][subobs]
                except KeyError:
                    continue

                x = dset['x']
                yexp = dset['y']
                yerr = dset['yerr']
                yerrstat = yerr.get('stat')
                yerrsys = yerr.get('sys', yerr.get('sum'))

                if scale is not None:
                    yexp = yexp*scale
                    if yerrstat is not None:
                        yerrstat = yerrstat*scale
                    if yerrsys is not None:
                        yerrsys = yerrsys*scale

                ax.errorbar(
                    x, yexp, yerr=yerrstat, fmt='o',
                    capsize=0, mfc='.25', mec='.25', mew=.2, zorder=1000
                )

                ax.fill_between(
                    x, yexp - yerrsys, yexp + yerrsys,
                    color='.9', zorder=-10
                )

                ratio_ax.plot(x, y/yexp, color=color)

            # main axes

            if plot.get('yscale') == 'log':
                ax.set_yscale('log')
                ax.minorticks_off()
            else:
                auto_ticks(ax, 'y', nbins=4, minor=2)

            ax.set_xticklabels([])
            ax.set_xlim(*plot['xlim'][system])
            ax.set_ylim(plot['ylim'])

            if nrow == 0:
                ax.set_title(format_system(system))

            if ncol == 0:
                ax.set_ylabel(plot['ylabel'])
                ratio_ax.set_ylabel('Ratio')
            else:
                ax.text(
                    1.02, .5, plot['title'],
                    transform=ax.transAxes, ha='left', va='center',
                    size=plt.rcParams['axes.labelsize'], rotation=-90
                )
                ax.set_yticklabels([])
                ratio_ax.set_yticklabels([])

            # ratio axes

            ratio_ax.axhline(1, lw=.5, color='0.5', zorder=-100)
            ratio_ax.axhspan(0.9, 1.1, color='0.95', zorder=-200)

            ratio_ax.set_xlim(plot['xlim'][system])
            ratio_ax.set_ylim(0.8, 1.2)

            ratio_ax.set_xlabel(plot['xlabel'][system])
            ratio_ax.set_yticks(np.arange(80, 121, 20)/100)

            ratio_ax.get_yticklabels()[0].set_verticalalignment('bottom')
            ratio_ax.get_yticklabels()[-1].set_verticalalignment('top')

    set_tight(fig, h_pad=.5, rect=[0, 0, .97, 1])


#@plot
def find_map():
    """
    Find the maximum a posteriori (MAP) point and compare emulator predictions
    to experimental data.
    """
    from scipy.optimize import minimize

    chain = mcmc.Chain()

    fixed_params = {
        'parton_number': 6,
        'Tswitch': 0.151,
    }

    opt_params = [k for k in chain.keys if k not in fixed_params]

    def full_x(x):
        x = dict(zip(opt_params, x), **fixed_params)
        return [x[k] for k in chain.keys]

    res = minimize(
        lambda x: -chain.log_posterior(full_x(x))[0],
        x0=np.median(chain.load(*opt_params, thin=1), axis=0),
        tol=1e-8,
        bounds=[
            (a + 1e-6*(b - a), b - 1e-6*(b - a))
            for (a, b), k in zip(chain.range, chain.keys)
            if k in opt_params
        ]
    )

    logging.debug('optimization result:\n%s', res)
    width = max(map(len, chain.keys)) + 2
    logging.info(
        'MAP params:\n%s',
        '\n'.join(
            k.ljust(width) + str(x) for k, x in zip(chain.keys, full_x(res.x))
        )
    )

    pred = chain._predict(np.atleast_2d(full_x(res.x)))

    plots = _observables_plots()

    ylim = {
        'Yields': (1e2, 1e5),
        'Mean $p_T$': (0, 1.7),
        'Flow cumulants': (0, .12),
    }

    for p in plots:
        p['ylim'] = ylim[p['title']]

    fig = plt.figure(figsize=figsize(1, 1.5))
    gs = gridspec.GridSpec(6, 2, height_ratios=[3, 1, 3, 1, 3, 1])

    axes = []
    ratio_axes = []

    nrow = len(plots)
    ncol = len(systems)

    for row, col in itertools.product(range(0, 2*nrow, 2), range(ncol)):
        ax = fig.add_subplot(gs[row, col])
        ratio_ax = fig.add_subplot(gs[row + 1, col], sharex=ax)
        axes.append(ax)
        ratio_axes.append(ratio_ax)

    for (plot, system), ax, ratio_ax in zip(
            itertools.product(plots, systems), axes, ratio_axes
    ):
        for obs, subobs, opts in plot['subplots']:
            color = obs_color(obs, subobs)
            scale = opts.get('scale')

            try:
                x = model.data[system][obs][subobs]['x']
            except KeyError:
                continue
            y = pred[system][obs][subobs][0]

            if scale is not None:
                y = y*scale

            ax.plot(x, y, color=color)

            if 'label' in opts and obs == 'vnk':
                ax.text(
                    x[-1] + (.1 if system == 'pPb5020' else 3), y[-1],
                    opts['label'],
                    color=darken(color), ha='left', va='center'
                )

            try:
                dset = expt.data[system][obs][subobs]
            except KeyError:
                continue

            x = dset['x']
            yexp = dset['y']
            yerr = dset['yerr']
            yerrstat = yerr.get('stat')
            yerrsys = yerr.get('sys', yerr.get('sum'))

            if scale is not None:
                yexp = yexp*scale
                if yerrstat is not None:
                    yerrstat = yerrstat*scale
                if yerrsys is not None:
                    yerrsys = yerrsys*scale

            ax.errorbar(
                x, yexp, yerr=yerrstat, fmt='o', ms=1.7,
                capsize=0, color='.25', zorder=1000
            )

            ax.fill_between(
                x, yexp - yerrsys, yexp + yerrsys,
                color='.9', zorder=-10
            )

            ratio_ax.plot(x, y/yexp, color=color)

        if plot.get('yscale') == 'log':
            ax.set_yscale('log')
            ax.minorticks_off()
        else:
            auto_ticks(ax, 'y', nbins=4, minor=2)

        ax.set_xticklabels([])
        ax.set_ylim(plot['ylim'])

        if ax.is_first_row():
            ax.set_title(format_system(system))

        if ax.is_first_col():
            ax.set_ylabel(plot['ylabel'])
            ratio_ax.set_ylabel('Ratio')

        if ax.is_last_col():
            ax.text(
                1.08, .5, plot['title'],
                transform=ax.transAxes, ha='left', va='center',
                size=plt.rcParams['axes.labelsize'], rotation=-90
            )

        ratio_ax.axhline(1, lw=.5, color='0.5', zorder=-100)
        ratio_ax.axhspan(0.9, 1.1, color='0.95', zorder=-200)
        ratio_ax.set_xlabel(xlabel(system, obs))
        ratio_ax.set_ylim(0.8, 1.2)
        ratio_ax.set_yticks(np.arange(80, 121, 20)/100)

    set_tight(fig, rect=(0, 0, .95, 1))


@plot
def flow_corr():
    """
    Symmetric cumulants SC(m, n) at the MAP point compared to experiment.

    """
    fig, axes = plt.subplots(
        figsize=figsize(0.5, 1.2), sharex=True,
        nrows=2, gridspec_kw=dict(height_ratios=[4, 5])
    )

    observables = ['sc', 'sc_normed']
    ylims = [(-2.5e-6, 2.5e-6), (-.9, .8)]
    labels = ['(4,2)', '(3,2)']
    system = 'PbPb5020'

    def label(*mn, normed=False):
        fmt = r'\mathrm{{SC}}({0}, {1})'
        if normed:
            fmt += r'/\langle v_{0}^2 \rangle\langle v_{1}^2 \rangle'
        return fmt.format(*mn).join('$$')

    for obs, ylim, ax in zip(observables, ylims, axes.flat):
        for (mn, cmap), lbl in zip([((4, 2), 'Blues'), ((3, 2), 'Oranges')], labels):
            x = model.map_data[system][obs][mn]['x']
            y = model.map_data[system][obs][mn]['Y']

            ax.plot(x, y, color=getattr(plt.cm, cmap)(.7))
            ax.text(1.02*x[-1], y[-1], lbl, va='center', ha='left')

        ax.axhline(
            0, color='.5', lw=plt.rcParams['xtick.major.width'],
            zorder=-100
        )

        ax.set_xlim(0, 80)
        ax.set_ylim(*ylim)

        auto_ticks(ax, nbins=7, minor=2)

        if ax.is_first_col():
            ax.set_ylabel(label('m', 'n', normed='normed' in obs))

        if ax.is_first_row():
            ax.set_title('Pb-Pb 5.02 TeV')
        else:
            ax.set_xlabel('Centrality %')


    # MAP estimate for Pb-Pb collisions at 5.02 TeV, calibrated to Pb-Pb
    # data at 2.76 and 5.02 TeV using a model without nucleon substructure.

    # symmetric cumulants
    SC = np.array([
        [2.5e+00, 5.8591e-09,  5.9204e-09],
        [7.5e+00, 2.1582e-08, -2.1367e-08],
        [1.5e+01, 1.2228e-07, -1.3942e-07],
        [2.5e+01, 4.3989e-07, -5.4267e-07],
        [3.5e+01, 9.4414e-07, -1.0677e-06],
        [4.5e+01, 1.4138e-06, -1.4616e-06],
        [5.5e+01, 1.4456e-06, -1.2317e-06],
        [6.5e+01, 7.3726e-07, -3.3222e-07],
    ])

    # normalized symmetric cumulants
    NSC = np.array([
        [2.5e+00, 7.3202e-02,  2.1091e-02],
        [7.5e+00, 7.6282e-02, -2.0918e-02],
        [1.5e+01, 1.5216e-01, -4.7261e-02],
        [2.5e+01, 2.4814e-01, -8.6423e-02],
        [3.5e+01, 3.4423e-01, -1.1640e-01],
        [4.5e+01, 4.5614e-01, -1.4251e-01],
        [5.5e+01, 6.1072e-01, -1.5021e-01],
    ])

    for ax, obs in zip(axes, [SC, NSC]):
        x, y42, y32 = obs.T
        ax.plot(x, y42, color=plt.cm.Blues(.7), linestyle='dashed')
        ax.plot(x, y32, color=plt.cm.Oranges(.7), linestyle='dashed')

    solid_line = lines.Line2D([], [], color=offblack)
    dashed_line = lines.Line2D([], [], linestyle='dashed', color=offblack)

    handles = [solid_line, dashed_line]
    labels = ["p-Pb, Pb-Pb 5.02 TeV", "Pb-Pb 2.76, 5.02 TeV"]

    plt.legend(handles, labels, loc=8, title='Bayesian calibration on:')

    set_tight(fig)


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) < .05 or (uh + ul) < abs(median) < .5:
        precision = 3
    elif abs(median) < 5:
        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, pad_subplots=-.1, rect_r=1, rect_t=.99,
        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=figsize(.15*scale*ndim, aspect=1)
    )

    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, linewidth=1, 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=fontsize['large']
        )

    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 ax in axes.flat:
        ax.tick_params(length=2/3*plt.rcParams['xtick.major.size'])

    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$['),
            )
            axis.set_tick_params(labelsize=fontsize['tiny'])
            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 any(
                        len(str(round(x, 5))) > 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=0, w_pad=pad_subplots, h_pad=pad_subplots,
        rect=(0, 0, rect_r, rect_t)
    )


@plot
def posterior():
    _posterior(ignore={'etas_hrg'})


@plot
def posterior_shear():
    _posterior(
        scale=1.35, pad_subplots=.1, rect_t=.97,
        params={'etas_min', 'etas_slope', 'etas_crv'}
    )


@plot
def posterior_bulk():
    _posterior(
        scale=1.35, pad_subplots=.1, rect_t=.97,
        params={'zetas_max', 'zetas_width', 'zetas_t0'}
    )


@plot
def posterior_p():
    """
    Distribution of trento p parameter with annotations for other models.

    """
    plt.figure(figsize=figsize(.5, .5))
    ax = plt.axes()

    data = mcmc.Chain().load('trento_p').ravel()

    counts, edges = np.histogram(data, bins=50)
    x = (edges[1:] + edges[:-1]) / 2
    y = counts / counts.max()
    interp = PchipInterpolator(x, y)
    x = np.linspace(x[0], x[-1], 10*x.size)
    y = interp(x)
    ax.plot(x, y, color=plt.cm.Blues(0.8))
    ax.fill_between(x, y, color=plt.cm.Blues(0.15), zorder=-10)

    ax.set_xlabel('$p$')

    for spine in ax.spines.values():
        spine.set_visible(False)

    for label, x, err in [
            ('KLN', -.67, .01),
            ('EKRT /\nIP-Glasma', 0, .1),
            ('Wounded\nnucleon', 1, None),
    ]:
        args = ([x], [0], 'o') if err is None else ([x - err, x + err], [0, 0])
        ax.plot(*args, lw=4, ms=4, color=offblack, alpha=.58, clip_on=False)

        if label.startswith('EKRT'):
            x -= .29

        ax.text(x, .05, label, va='bottom', ha='center')

    ax.text(.1, .8, format_ci(data))
    ax.set_xticks(np.arange(-10, 11, 5)/10)
    ax.set_xticks(np.arange(-75, 76, 50)/100, minor=True)

    xm = 1.2
    ax.set_xlim(-xm, xm)
    ax.add_artist(
        patches.FancyArrowPatch(
            (-xm, 0), (xm, 0),
            linewidth=plt.rcParams['axes.linewidth'],
            arrowstyle=patches.ArrowStyle.CurveFilledAB(
                head_length=3, head_width=1.5
            ),
            facecolor=offblack, edgecolor=offblack,
            clip_on=False, zorder=100
        )
    )

    ax.set_yticks([])
    ax.set_ylim(0, 1.01*y.max())

    set_tight(pad=0)


def posterior_parameter(parameter, label, xticks, bins=50):
    """
    Marginal distribution of a single parameter.

    """
    plt.figure(figsize=figsize(.5, .75))
    ax = plt.axes()

    data = mcmc.Chain().load(parameter).ravel()

    counts, edges = np.histogram(data, bins=bins)
    x = (edges[1:] + edges[:-1]) / 2
    y = counts / counts.max()
    interp = PchipInterpolator(x, y)
    x = np.linspace(x[0], x[-1], 100*x.size)
    y = interp(x)
    ax.plot(x, y, color=plt.cm.Blues(0.8))
    ax.fill_between(x, y, color=plt.cm.Blues(0.15), zorder=-10)
    ax.spines['left'].set_visible(False)

    ax.set_xlabel(label)
    ax.set_xticks(xticks)
    ax.set_yticks([])
    ax.set_ylim(0, 1.01*y.max())

    set_tight(pad=0)


@plot
def posterior_parton_number():
    """
    Posterior distribution on the number of constituents.

    """
    plt.figure(figsize=figsize(.5, .75))
    ax = plt.axes()

    data = mcmc.Chain().load('parton_number').ravel()

    counts, edges = np.histogram(data, bins=50)
    x = (edges[1:] + edges[:-1]) / 2
    y = counts / counts.max()

    spline = UnivariateSpline(x, y, s=1e-3)
    x = np.linspace(data.min(), data.max(), 1000)
    y = spline(x)

    plt.plot([x[0], x[0]], [0, y[0]], linestyle='dashed',
             color=offblack, clip_on=False)

    plt.plot([x[-1], x[-1]], [0, y[-1]], linestyle='dashed',
             color=offblack, clip_on=False)

    ax.plot(x, y, color=plt.cm.Blues(0.8))
    ax.fill_between(x, y, color=plt.cm.Blues(0.15), zorder=-10)
    ax.spines['left'].set_visible(False)

    ax.set_xlabel('Constituent number $n_c$')
    ax.set_xticks([1, 3, 5, 7, 9])
    ax.set_yticks([])
    ax.set_ylim(0, 1.01*y.max())

    set_tight(pad=0)


@plot
def posterior_freestreaming():
    posterior_parameter(
        'tau_fs',
        'Free streaming time [fm/$c$]',
        [.1, .8, 1.5]
    )


def _region(ax, name, chain, cmap=plt.cm.Blues, legend=False, title=False):
    """
    Visual estimate (posterior median and credible region) of
    temperature-dependent shear or bulk viscosity.
    """
    var, keys, function, ymax = dict(
        shear=(
            'eta',
            ['min', 'slope', 'crv'],
            lambda T, m, s, c: m + s*(T - Tc)*(T/Tc)**c,
            .4
        ),
        bulk=(
            'zeta',
            ['max', 'width', 't0'],
            lambda T, m, w, T0: m / (1 + ((T - T0)/w)**2),
            .08
        ),
    )[name]

    Tmin, Tmax = .150, .300
    Tc = .154

    samples = chain.load(
        *['{}s_{}'.format(var, k) for k in keys], thin=1
    )

    T = np.linspace(Tc if name == 'shear' else Tmin, Tmax, 1000)
    ax.plot(
        T, function(T, *np.median(samples, axis=0)),
        color=cmap(.75), label='Posterior median'
    )

    Tsparse = np.linspace(T[0], T[-1], 25)
    intervals = [
        PchipInterpolator(Tsparse, y)(T)
        for y in np.array([
            mcmc.credible_interval(function(t, *samples.T))
            for t in Tsparse
        ]).T
    ]
    ax.fill_between(
        T, *intervals,
        color=cmap(.3), label='90% credible region'
    )

    ax.set_xlim(Tmin, Tmax)
    ax.set_ylim(0, ymax)
    auto_ticks(ax, nbins=5)
    ax.xaxis.set_major_formatter(
        ticker.FuncFormatter(lambda x, pos: int(1000*x))
    )

    ax.set_xlabel('Temperature [MeV]')
    ax.set_ylabel(r'$\{}/s$'.format(var))

    if title:
        ax.set_title(name.capitalize() + ' viscosity')

    if legend:
        ax.legend(loc=legend if isinstance(legend, str) else 'best')

    if name == 'shear':
        ax.axhline(
            1/(4*np.pi),
            color='.5', linewidth=plt.rcParams['ytick.major.width']
        )
        ax.text(Tmax, .07, r'$1/4\pi$', va='top', ha='right', color='.3')


@plot
def region_shear():
    """
    Region plot for eta/s.
    """
    chain = mcmc.Chain()
    fig, ax = plt.subplots(figsize=figsize(.5, .65))
    _region(ax, 'shear', chain, legend='upper left')
    set_tight(fig)


@plot
def region_bulk():
    """
    Region plot for zeta/s.
    """
    chain = mcmc.Chain()
    fig, ax = plt.subplots(figsize=figsize(.5, .65))
    _region(ax, 'bulk', chain, legend='upper right')
    set_tight(fig)


@plot
def region_shear_bulk(cmap=plt.cm.Blues):
    """
    Visual estimates (posterior median and credible region) of the
    temperature-dependent shear and bulk viscosity.

    """
    fig, axes = plt.subplots(ncols=2, figsize=figsize(1, .4))
    ax_shear, ax_bulk = axes

    Tmin, Tmax = .150, .300
    Tc = .154

    prj_path = Path('/home/morelandjs/research/chains/jonah_nature.hdf')
    energies = (mcmc.Chain(prj_path), plt.cm.Blues, .6, 'Pb-Pb 2.76, 5.02 TeV')
    nuclei = (mcmc.Chain(), plt.cm.Oranges, .25, 'p-Pb, Pb-Pb 5.02 TeV')
    handles = []

    for zorder, (chain, cmap, darkness, label) in enumerate([nuclei, energies]):
        for (name,  var, keys, function, ymax), ax in zip([
                ('shear', 'eta', ['min', 'slope', 'crv'],
                lambda T, m, s, c: m + s*(T - Tc)*(T/Tc)**c,
                .4),
                ('bulk', 'zeta', ['max', 'width', 't0'],
                lambda T, m, w, T0: m / (1 + ((T - T0)/w)**2),
                .08)
        ], axes):
            samples = chain.load(*['{}s_{}'.format(var, k) for k in keys], thin=1)

            T = np.linspace(Tc if name == 'shear' else Tmin, Tmax, 1000)

            ax.plot(
                T, function(T, *np.median(samples, axis=0)),
                color=cmap(.75), label='Posterior median', zorder=zorder
            )

            Tsparse = np.linspace(T[0], T[-1], 25)

            intervals = [
                PchipInterpolator(Tsparse, y)(T)
                for y in np.array([
                    mcmc.credible_interval(function(t, *samples.T))
                    for t in Tsparse
                ]).T
            ]

            ax.fill_between(
                T, *intervals,
                color=cmap(darkness), label='90% credible region',
                alpha=.6 if zorder == 1 else 1, lw=0, zorder=zorder
            )

            ax.set_xlim(Tmin, Tmax)
            ax.set_ylim(0, ymax)
            auto_ticks(ax, nbins=5)
            ax.xaxis.set_major_formatter(
                ticker.FuncFormatter(lambda x, pos: int(1000*x))
            )

            ax.set_xlabel('Temperature [MeV]')
            ax.set_ylabel(r'$\{}/s$'.format(var))
            ax.set_title(name.capitalize() + ' viscosity')

            if name == 'shear':
                ax.axhline(
                    1/(4*np.pi),
                    color='.5', linewidth=plt.rcParams['ytick.major.width']
                )
                ax.text(Tmax, .07, r'$1/4\pi$', va='top', ha='right', color='.3')

        line = lines.Line2D([], [], color=cmap(.8))
        band = patches.Patch(color=cmap(.3))
        handles.append((band, line))

    line = lines.Line2D([], [], color=offblack, label='Posterior median')
    band = patches.Patch(color='.85', label='90% credible region')
    ax_shear.legend(handles=[band, line], loc='upper left')
    labels = ["p-Pb, Pb-Pb 5.02 TeV", "Pb-Pb 2.76, 5.02 TeV"]
    ax_bulk.legend(
        handles, labels, loc='upper right', markerfirst=False,
        bbox_to_anchor=(1, 1.05), title='Bayesian calibration on:'
    )

    set_tight(w_pad=.2)


region_style = dict(color='.93', zorder=-100)
Tc = .154


@plot
def design():
    """
    Projection of a LH design into two dimensions.

    """
    fig = plt.figure(figsize=figsize(.5, 1))
    ratio = 5
    gs = plt.GridSpec(ratio + 1, ratio + 1)

    ax_j = fig.add_subplot(gs[1:, :-1])
    ax_x = fig.add_subplot(gs[0, :-1], sharex=ax_j)
    ax_y = fig.add_subplot(gs[1:, -1], sharey=ax_j)

    d = Design(systems[0])

    keys = ('etas_min', 'etas_slope')
    indices = tuple(d.keys.index(k) for k in keys)

    x, y = (d.array[:, i] for i in indices)
    ax_j.plot(x, y, 'o', color=plt.cm.Blues(0.75), mec='white', mew=.3)

    hist_kw = dict(bins=30, color=plt.cm.Blues(0.4), edgecolor='white', lw=.5)
    ax_x.hist(x, **hist_kw)
    ax_y.hist(y, orientation='horizontal', **hist_kw)

    for ax in fig.axes:
        ax.tick_params(top=False, right=False)
        spines = ['top', 'right']
        if ax is ax_x:
            spines += ['left']
        elif ax is ax_y:
            spines += ['bottom']
        for spine in spines:
            ax.spines[spine].set_visible(False)
        for ax_name in 'xaxis', 'yaxis':
            getattr(ax, ax_name).set_ticks_position('none')

    auto_ticks(ax_j)

    for ax in ax_x, ax_y:
        ax.tick_params(labelbottom=False, labelleft=False)

    for i, xy in zip(indices, 'xy'):
        for f, l in [('lim', d.range), ('label', d.labels)]:
            getattr(ax_j, 'set_{}{}'.format(xy, f))(l[i])

    set_tight(fig)


@plot
def gp():
    """
    Conditioning a Gaussian process.

    """
    fig, axes = plt.subplots(
        figsize=figsize(.5, 2*aspect),
        nrows=2, sharex='col'
    )

    def dummy_optimizer(obj_func, initial_theta, bounds):
        return initial_theta, 0.

    gp = GPR(1.*kernels.RBF(.8), optimizer=dummy_optimizer)

    def sample_y(*args, **kwargs):
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', RuntimeWarning)
            return gp.sample_y(*args, **kwargs)

    x = np.linspace(0, 5, 1000)
    X = x[:, np.newaxis]

    x_train = np.linspace(.5, 4.5, 4)
    X_train = x_train[:, np.newaxis]

    for title, ax in zip(['Random functions', 'Conditioned on data'], axes):
        if title.startswith('Conditioned'):
            y = sample_y(X_train, random_state=23158).squeeze()
            y -= .5*(y.max() + y.min())
            gp.fit(X_train, y)
            training_data, = plt.plot(x_train, y, 'o', color='.3', zorder=50)

        for s, c in zip(
                sample_y(X, n_samples=4, random_state=34576).T,
                ['Blues', 'Greens', 'Oranges', 'Purples']
        ):
            ax.plot(x, s, color=getattr(plt.cm, c)(.6))

        mean, std = gp.predict(X, return_std=True)
        std = ax.fill_between(x, mean - std, mean + std, color='.92')
        mean, = ax.plot(x, mean, color='.42', dashes=(3.5, 1.5))

        ax.set_ylim(-2, 2)
        ax.set_ylabel('Output')
        auto_ticks(ax)

        ax.set_title(title, y=.9)

    ax.set_xlabel('Input')
    ax.legend(*zip(*[
        (mean, 'Mean prediction'),
        (std, 'Uncertainty'),
        (training_data, 'Training data'),
    ]), loc='lower left')

    set_tight(fig, h_pad=1)


@plot
def pca():
    fig = plt.figure(figsize=figsize(.5, aspect=1))
    ratio = 5
    gs = plt.GridSpec(ratio + 1, ratio + 1)

    ax_j = fig.add_subplot(gs[1:, :-1])
    ax_x = fig.add_subplot(gs[0, :-1], sharex=ax_j)
    ax_y = fig.add_subplot(gs[1:, -1], sharey=ax_j)

    x, y = (
        model.data['PbPb5020'][obs][subobs]['Y'][:, 3]
        for obs, subobs in [('dNch_deta', None), ('vnk', (2, 2))]
    )
    xlabel = r'$dN_\mathrm{ch}/d\eta$'
    ylabel = r'$v_2\{2\}$'
    xlim = .5, 6
    ylim = -6.5, 0

    cmap = plt.cm.Blues

    ax_j.plot(x, y, 'o', color=cmap(.75), mec='white', mew=.25, zorder=10)

    for d, ax, orientation in [(x, ax_x, 'vertical'), (y, ax_y, 'horizontal')]:
        ax.hist(
            d, bins=20,
            orientation=orientation, color=cmap(.4), edgecolor='white'
        )

    xy = np.column_stack([x, y])
    xymean = xy.mean(axis=0)
    xystd = xy.std(axis=0)
    xy -= xymean
    xy /= xystd
    pca = PCA().fit(xy)
    pc = (
        6 * xystd *
        pca.explained_variance_ratio_[:, np.newaxis] *
        pca.components_
    )

    for w, p in zip(pca.explained_variance_ratio_, pc):
        if np.all(p < 0):
            p *= -1
        ax_j.annotate(
            '', xymean + p, xymean, zorder=20,
            arrowprops=dict(
                arrowstyle='->', shrinkA=0, shrinkB=0,
                color=offblack, lw=.7
            )
        )
        ax_j.text(
            *(xymean + p + (.8, .002)*np.sign(p)), s='{:.0f}%'.format(100*w),
            color=offblack, ha='center', va='top' if p[1] < 0 else 'bottom',
            zorder=20
        )

    for ax in fig.axes:
        ax.tick_params(top=False, right=False)
        spines = ['top', 'right']
        if ax is ax_x:
            spines += ['left']
        elif ax is ax_y:
            spines += ['bottom']
        for spine in spines:
            ax.spines[spine].set_visible(False)
        for ax_name in 'xaxis', 'yaxis':
            getattr(ax, ax_name).set_ticks_position('none')

    for ax in ax_x, ax_y:
        ax.tick_params(labelbottom=False, labelleft=False)

    auto_ticks(ax_j, nbins=5, prune='upper')

    ax_j.set_xlim(xlim)
    ax_j.set_ylim(ylim)

    ax_j.set_xlabel(xlabel)
    ax_j.set_ylabel(ylabel)

    set_tight(pad=.1, h_pad=.3, w_pad=.3)


default_system = 'PbPb5020'


@plot
def pca_vectors_variance(system=default_system):
    """
    PCA vectors and explained variance.

    """
    fig, axes = plt.subplots(
        figsize=figsize(1.2, aspect=.4),
        ncols=2, gridspec_kw=dict(width_ratios=[5, 1])
    )

    emu = emulators[system]
    pca = emu.pca

    ax = axes[0]

    for n, (pc, var) in enumerate(zip(
            pca.components_[:3], pca.explained_variance_ratio_
    ), start=1):
        ax.plot(pc, 'o', label='PC {} ({:.0f}%)'.format(n, 100*var))

    ax.axhline(
        0,
        color='.5', linewidth=plt.rcParams['ytick.major.width'],
        zorder=-100
    )

    x = -.5
    ticks = []
    ticklabels = []

    for obs, subobslist in emu.pPb5020:
        for subobs in subobslist:
            i = model.data[system][obs][subobs]['Y'].shape[1]
            ticks.append(x + .5*i)
            ticklabels.append(obs_label(obs, subobs))
            x += i

    ax.set_xticks(ticks)
    ax.set_xticklabels(ticklabels)
    ax.tick_params(
        'x',
        bottom=False, labelbottom=False,
        labeltop=True, pad=1
    )
    for t in ax.get_xticklabels():
        t.set_verticalalignment('baseline')

    ax.set_ylabel('PCA coefficient', labelpad=1)
    auto_ticks(ax, 'y', nbins=4, minor=2)
    ax.legend(loc='best', handletextpad=0)

    ax = axes[1]

    npc = 10
    ax.plot(
        np.arange(1, 1 + npc),
        pca.explained_variance_ratio_.cumsum()[:npc],
        '-o',
    )

    ax.set_xlim(.5, npc + .5)
    ax.set_ylim(0, 1)

    majorticks = [1, 4, 7, 10]
    ax.set_xticks(majorticks)
    ax.set_xticks(sorted(set(range(1, npc)) - set(majorticks)), minor=True)
    auto_ticks(ax, 'y', nbins=5, minor=2)
    ax.xaxis.set_ticks_position('top')

    ax.set_xlabel('Number of PC')
    ax.set_ylabel('Cumulative explained variance fraction')
    ax.xaxis.set_ticks_position('top')
    ax.xaxis.set_label_position('top')

    for ax in axes:
        for s in ax.spines.values():
            s.set_visible(True)

    set_tight(w_pad=.5)


def boxplot(
        ax, percentiles, x=0, y=0, box_width=1,
        line_width=plt.rcParams['lines.linewidth'],
        color=(0, 0, 0), alpha=.6, zorder=10, vert=True
):
    """
    Draw a minimal boxplot.

    `percentiles` must be a np.array of five numbers:

        whisker_low, quartile_1, median, quartile_3, whisker_high

    """
    pl, q1, q2, q3, ph = percentiles + y

    # IQR box
    ax.add_patch(patches.Rectangle(
        xy=(x - .5*box_width, q1),
        width=box_width, height=(q3 - q1),
        color=color, alpha=alpha, lw=0, zorder=zorder
    ))

    # median line
    ax.plot(
        [x - .5*box_width, x + .5*box_width], 2*[q2],
        lw=line_width, solid_capstyle='butt', color=color,
        zorder=zorder + 1
    )

    # whisker lines
    for y in [[q1, pl], [q3, ph]]:
        ax.plot(
            2*[x], y, lw=line_width, solid_capstyle='butt',
            color=color, alpha=alpha, zorder=zorder
        )


def validation_data(system, n_splits=20):
    """
    Partition the design into training and test data using K-fold
    cross validation. Train the emulator on each fold (subset of the design)
    and return the emulator prediction (mean and cov) for each fold as a list.

    """
    design = Design(system)
    kf = KFold(n_splits=n_splits)
    npc = {'pPb5020': 4, 'PbPb5020': 8}[system]

    mean_folds = []
    cov_folds = []

    cachefile = Path(cachedir, 'validation', '{}.pkl'.format(system))
    cachefile.parent.mkdir(parents=True, exist_ok=True)

    if cachefile.exists():
        return pickle.load(cachefile.open(mode='rb'))

    for train_index, test_index in kf.split(design.array):
        test_points, train_points = [
            [design.points[index] for index in indices]
            for indices in (test_index, train_index)
        ]

        emu = Emulator(system, exclude_points=test_points, npc=npc)
        test_mean, test_cov = emu.predict(
            design.array[test_index], return_cov=True
        )

        mean_folds.append(test_mean)
        cov_folds.append(test_cov)

    pickle.dump((mean_folds, cov_folds), cachefile.open(mode='wb'))
    return mean_folds, cov_folds


def validation_all(system):
    """
    Emulator validation: normalized residuals and RMS error for each
    observable.

    """
    fig, (ax_box, ax_rms) = plt.subplots(
        nrows=2, figsize=figsize(1, aspect=.4),
        gridspec_kw=dict(height_ratios=[1.5, 1])
    )

    index = 1
    ticks = []
    ticklabels = []

    plots = _observables_plots()
    model_data  = model.data[system]
    mean_folds, cov_folds = validation_data(system)

    for plot in plots:
        for (obs, subobs, opts) in plot['subplots']:
            color = obs_color(obs, subobs)

            try:
                # model data
                Y = model_data[obs][subobs]['Y']

                # emulator predictions
                Y_ = np.concatenate(
                    [mean[obs][subobs] for mean in mean_folds], axis=0
                )
                S_ = np.concatenate(
                    [np.sqrt(cov[(obs, subobs), (obs, subobs)].T.diagonal())
                    for cov in cov_folds], axis=0
                )
            except KeyError:
                continue

            Z = (Y_ - Y)/S_

            for i, percentiles in enumerate(
                    np.percentile(Z, [10, 25, 50, 75, 90], axis=0).T,
                    start=index
            ):
                boxplot(ax_box, percentiles, x=i, box_width=.8, color=color)

            Ymin, Ymax = np.percentile(Y, (.5, 99.5))
            Yerr = (Y_ - Y)/(Ymax - Ymin)
            rms = 100*np.sqrt(np.square(Yerr).mean(axis=0))
            ax_rms.plot(
                np.arange(index, index + rms.size), rms, 'o', color=color
            )

            ticks.append(.5*(index + i))
            ticklabels.append(obs_label(obs, subobs))

            index = i + 2

    ax_box.set_xticks(ticks)
    ax_box.tick_params(axis='x', pad=-6)
    ax_box.set_xticklabels(ticklabels)
    ax_box.tick_params('x', bottom=False, labelsize=plt.rcParams['font.size'])

    ax_box.set_ylim(-2.25, 2.25)
    ax_box.set_ylabel(r'Normalized residuals')
    nuclei, roots = parse_system(system)
    ax_box.set_title('{}-{} {} TeV'.format(*nuclei, roots/1000))

    q, p = np.sqrt(2) * special.erfinv(2*np.array([.75, .90]) - 1)
    ax_box.axhspan(-q, q, color='.85', zorder=-20)
    for s in [-1, 0, 1]:
        ax_box.axhline(s*p, color='.5', zorder=-10)

    ax_q = ax_box.twinx()
    ax_q.set_ylim(ax_box.get_ylim())
    ax_q.set_yticks([-p, -q, 0, q, p])
    ax_q.set_yticklabels([10, 25, 50, 75, 90])
    ax_q.tick_params('y', right=False)
    ax_q.set_ylabel(
        'Normal quantiles',
        fontdict=dict(rotation=-90),
        labelpad=4*plt.rcParams['axes.labelpad']
    )

    ax_rms.set_xticks([])
    ax_rms.set_yticks(np.arange(0, 21, 5))
    ax_rms.set_ylim(0, 20)
    ax_rms.set_ylabel('Frac. error')

    for y in ax_rms.get_yticks():
        ax_rms.axhline(y, color='.5', zorder=-10)

    for ax in fig.axes:
        ax.set_xlim(0, index - 1)
        ax.spines['bottom'].set_visible(False)


@plot
def validation_pPb5020():
    validation_all('pPb5020')


@plot
def validation_PbPb5020():
    validation_all('PbPb5020')


@plot
def validation_example(
        system='PbPb5020',
        obs='dNch_deta', subobs=None,
        label=r'$dN_\mathrm{ch}/d\eta$',
        cent=(20, 30)
):
    """
    Example of emulator validation for a single observable.  Scatterplot of
    model calculations vs emulator predictions with histogram and boxplot of
    normalized residuals.

    """
    fig, axes = plt.subplots(
        ncols=2, figsize=figsize(.8, aspect=.6),
        gridspec_kw=dict(width_ratios=[3, 1])
    )

    ax_scatter, ax_hist = axes

    # model data
    model_data  = model.data[system]
    vdata = model_data[obs][subobs]
    cent_slc = (slice(None), vdata['cent'].index(cent))
    y = vdata['Y'][cent_slc]

    # emulator predictions
    mean_folds, cov_folds = validation_data(system)
    y_ = np.concatenate(
        [mean[obs][subobs][cent_slc] for mean in mean_folds], axis=0
    )
    std_ = np.concatenate(
        [np.sqrt(cov[(obs, subobs), (obs, subobs)].T.diagonal()[cent_slc])
            for cov in cov_folds], axis=0
    )

    color = obs_color(obs, subobs)
    alpha = .6

    ax_scatter.set_aspect('equal')
    ax_scatter.errorbar(
        y_, y, xerr=std_,
        fmt='o', mew=.2, mec='white',
        color=color, alpha=alpha
    )
    dy = .03*y.ptp()
    x = [y.min() - dy, y.max() + dy]
    ax_scatter.plot(x, x, color='.4')
    ax_scatter.set_xlabel('Emulator prediction')
    ax_scatter.set_ylabel('Model calculation')
    ax_scatter.text(
        .04, .96, '{} {}–{}'.format(label, *cent),
        horizontalalignment='left', verticalalignment='top',
        transform=ax_scatter.transAxes
    )

    zmax = 3.5
    zrange = (-zmax, zmax)

    z = (y_ - y)/std_

    ax_hist.hist(
        z, bins=30, range=zrange, density=True, histtype='stepfilled',
        orientation='horizontal', color=color, alpha=alpha
    )
    x = np.linspace(-zmax, zmax, 1000)
    ax_hist.plot(np.exp(-.5*x*x)/np.sqrt(2*np.pi), x, color='.25')

    box_x = .75
    box_width = .1

    boxplot(
        ax_hist, np.percentile(z, [10, 25, 50, 75, 90]),
        x=box_x, box_width=box_width,
        line_width=2*plt.rcParams['lines.linewidth'],
        color=color, alpha=alpha
    )

    guide_width = 2.5*box_width

    q, p = np.sqrt(2) * special.erfinv(2*np.array([.75, .90]) - 1)
    ax_hist.add_patch(patches.Rectangle(
        xy=(box_x - .5*guide_width, -q),
        width=guide_width, height=2*q,
        color='.85', zorder=-20
    ))
    for s in [-1, 0, 1]:
        ax_hist.plot(
            [box_x - .5*guide_width, box_x + .5*guide_width], 2*[s*p],
            color='.5', zorder=-10
        )

    ax_hist.set_ylim(zrange)
    ax_hist.spines['bottom'].set_visible(False)
    ax_hist.tick_params('x', bottom=False, labelbottom=False)
    ax_hist.set_ylabel('Normalized residuals')

    ax_q = ax_hist.twinx()
    ax_q.spines['bottom'].set_visible(False)
    ax_q.set_ylim(ax_hist.get_ylim())
    ax_q.set_yticks([-p, -q, 0, q, p])
    ax_q.set_yticklabels([10, 25, 50, 75, 90])
    ax_q.tick_params('y', right=False)
    ax_q.set_ylabel(
        'Normal quantiles',
        fontdict=dict(rotation=-90),
        labelpad=4*plt.rcParams['axes.labelpad']
    )

    set_tight(fig, rect=[.02, 0, 1, 1])


@plot
def validation_example_vert(
        system='PbPb5020',
        obs='dNch_deta', subobs=None,
        label=r'$dN_\mathrm{ch}/d\eta$',
        cent=(20, 30)
):
    """
    Example of emulator validation for a single observable.  Scatterplot of
    model calculations vs emulator predictions with histogram and boxplot of
    normalized residuals.

    """
    fig, axes = plt.subplots(
        nrows=2, figsize=figsize(.5, aspect=1.3),
        gridspec_kw=dict(height_ratios=[1, 3]),
    )

    ax_hist, ax_scatter = axes

    # model data
    model_data  = model.data[system]
    vdata = model_data[obs][subobs]
    cent_slc = (slice(None), vdata['cent'].index(cent))
    y = vdata['Y'][cent_slc]

    # emulator predictions
    mean_folds, cov_folds = validation_data(system)
    y_ = np.concatenate(
        [mean[obs][subobs][cent_slc] for mean in mean_folds], axis=0
    )
    std_ = np.concatenate(
        [np.sqrt(cov[(obs, subobs), (obs, subobs)].T.diagonal()[cent_slc])
            for cov in cov_folds], axis=0
    )

    color = obs_color(obs, subobs)
    alpha = .6

    ax_scatter.set_aspect('equal')
    ax_scatter.errorbar(
        y_, y, xerr=std_,
        fmt='o', mew=.2, mec='white',
        color=color, alpha=alpha
    )
    dy = .03*y.ptp()
    x = [y.min() - dy, y.max() + dy]
    ax_scatter.plot(x, x, color='.4')
    ax_scatter.set_xlabel('Emulator prediction')
    ax_scatter.set_ylabel('Model calculation')
    ax_scatter.text(
        .96, .06, '{} {}–{}%'.format(label, *cent),
        horizontalalignment='right', verticalalignment='bottom',
        transform=ax_scatter.transAxes
    )

    zmax = 3.5
    zrange = (-zmax, zmax)

    z = (y_ - y)/std_

    ax_hist.hist(
        z, bins=30, range=zrange, density=True, histtype='stepfilled',
        color=color, alpha=alpha
    )
    x = np.linspace(-zmax, zmax, 1000)
    ax_hist.plot(x, np.exp(-.5*x*x)/np.sqrt(2*np.pi), color='.25')

    box_y = .75
    box_width = .1

    # percentiles
    pl, q1, q2, q3, ph = np.percentile(z, [10, 25, 50, 75, 90])

    # IQR box
    ax_hist.add_patch(patches.Rectangle(
        xy=(q1, box_y - .5*box_width),
        width=(q3 - q1), height=box_width,
        color=color, alpha=alpha, lw=0, zorder=1
    ))

    # median line
    ax_hist.plot(
        2*[q2], [box_y - .5*box_width, box_y + .5*box_width],
        lw=2*plt.rcParams['lines.linewidth'], solid_capstyle='butt',
        color=color, zorder=2
    )

    # whisker lines
    for x in [[pl, q1], [q3, ph]]:
        ax_hist.plot(
            x, 2*[box_y], lw=2*plt.rcParams['lines.linewidth'],
            solid_capstyle='butt', color=color, alpha=alpha, zorder=1
        )

    guide_width = 2.5*box_width

    q, p = np.sqrt(2) * special.erfinv(2*np.array([.75, .90]) - 1)
    ax_hist.add_patch(patches.Rectangle(
        xy=(-q, box_y - .5*guide_width),
        width=2*q, height=guide_width,
        color='.85', zorder=-20
    ))
    for s in [-1, 0, 1]:
        ax_hist.plot(
           2*[s*p], [box_y - .5*guide_width, box_y + .5*guide_width],
            color='.5', zorder=-10
        )

    ax_hist.set_xlim(zrange)
    ax_hist.spines['left'].set_visible(False)
    ax_hist.tick_params('y', left=False, labelleft=False)
    ax_hist.set_xlabel('Normalized residuals')

    ax_q = ax_hist.twiny()
    ax_q.spines['left'].set_visible(False)
    ax_q.set_xlim(ax_hist.get_xlim())
    ax_q.set_xticks([-p, -q, 0, q, p])
    ax_q.set_xticklabels([10, 25, 50, 75, 90])
    ax_q.tick_params('x', top=False)
    ax_q.set_xlabel(
        'Normal quantiles',
        labelpad=2*plt.rcParams['axes.labelpad']
    )

    set_tight(fig)


@plot
def correlation_matrices(system=default_system):
    """
    Correlation (normalized covariance) matrices for model and experiment.

    """
    chain = mcmc.Chain()

    emu = emulators[system]
    emu_slices = [
        (obs, subobs, slc)
        for obs, subobs_slc in emu._slices.items()
        for subobs, slc in subobs_slc.items()
    ]

    design = Design(system)
    X = np.random.uniform(design.min, design.max).reshape(1, -1)
    emu_cov = emu.predict(X, return_cov=True)[1].array[0]

    fig, axes = plt.subplots(
        ncols=3, figsize=figsize(1, .47),
        gridspec_kw=dict(width_ratios=[1, 1, .02])
    )

    for (cov, slices, title), ax in zip([
            (emu_cov, emu_slices, 'Model (emulator)'),
            (chain._expt_cov[system], chain._slices[system], 'Experiment'),
    ], axes):
        s = np.sqrt(cov.diagonal())
        img = ax.imshow(
            cov / np.outer(s, s), vmin=-1, vmax=1,
            interpolation='nearest', cmap='RdBu'
        )

        ticks = []
        ticklabels = []

        for obs, subobs, slc in slices:
            ticks.append(.5*(slc.start + slc.stop - 1))
            ticklabels.append(obs_label(obs, subobs))

        ax.set_xticks(ticks)
        ax.set_yticks(ticks)
        ax.set_xticklabels(ticklabels)
        ax.set_yticklabels(ticklabels)
        ax.set_title(title, y=1.05)

        ax.tick_params(
            bottom=False, top=False, left=False, right=False,
            labelbottom=False, labeltop=True,
            pad=0
        )

        for s in ax.spines.values():
            s.set_visible(False)

    axes[0].tick_params(labelleft=True)
    for t in axes[1].get_yticklabels():
        t.set_horizontalalignment('center')
        t.set_x(-.05)

    cax = axes[-1]
    fig.colorbar(img, cax=cax, ticks=[-1, -.5, 0, .5, 1])
    cax.set_aspect(40)
    cax.yaxis.set_ticks_position('left')
    cax.set_title('Corr', y=1.02, fontsize=fontsize['normal'])

    set_tight(fig, rect=(0, 0, 1, .96))


@plot
def diag_pca(system=default_system):
    """
    Diagnostic: histograms of principal components and scatterplots of pairs.

    """
    Y = [g.y_train_ for g in emulators[default_system].gps]
    n = len(Y)
    ymax = np.ceil(max(np.fabs(y).max() for y in Y))
    lim = (-ymax, ymax)

    fig, axes = plt.subplots(nrows=n, ncols=n, figsize=2*(n,))

    for y, ax in zip(Y, axes.diagonal()):
        ax.hist(y, bins=30)
        ax.set_xlim(lim)

    for ny, nx in zip(*np.tril_indices_from(axes, k=-1)):
        ax = axes[ny][nx]
        ax.scatter(Y[nx], Y[ny])
        ax.set_xlim(lim)
        ax.set_ylim(lim)
        axes[nx][ny].set_axis_off()

    for i in range(n):
        label = 'PC {}'.format(i)
        axes[-1][i].set_xlabel(label)
        axes[i][0].set_ylabel(label)


def _diag_emu(system=default_system, pcs=None, params=None, label_all=True):
    """
    Diagnostic: plots of each principal component vs each input parameter,
    overlaid by emulator predictions at several points in design space.

    """
    gps = emulators[system].gps
    pcs = (
        range(len(gps)) if pcs is None else
        [p if p >= 0 else (len(gps) + p) for p in pcs]
    )
    nrows = len(pcs)

    design = Design(system)
    if params is None:
        params = design.keys
    ncols = len(params)

    fig, axes = plt.subplots(
        nrows=nrows, ncols=ncols,
        figsize=figsize((.5 if label_all else .375)*ncols, .62*nrows/ncols),
        sharex=(False if label_all else 'col'),
        sharey=(False if label_all else 'row')
    )

    ymax = np.ceil(2*max(np.fabs(gps[pc].y_train_).max() for pc in pcs))/2
    ylim = (-ymax, ymax)

    tmax = int(ymax)
    yticksmajor = [-tmax, 0, tmax]
    yticksminor = list(range(-tmax + 1, 0)) + list(range(1, tmax))

    for pc, row in zip(pcs, axes):
        gp = gps[pc]
        y = gp.y_train_

        for param, ax in zip(params, row):
            i = design.keys.index(param)
            x = gp.X_train_[:, i]
            ax.plot(
                x, y, 'o',
                markersize=.4*plt.rcParams['lines.markersize'],
                color='.7',
                zorder=-30
            )

            xlim = design.range[i]
            x = np.linspace(xlim[0], xlim[1], 100)
            X = np.empty((x.size, gp.X_train_.shape[1]))

            for r, c in [(.2, 'purple'), (.5, 'blue'), (.8, 'green')]:
                X[:] = r*design.min + (1 - r)*design.max
                X[:, i] = x
                mean, std = gp.predict(X, return_std=True)

                color = colors[c]
                ax.plot(
                    x, mean,
                    linewidth=.8*plt.rcParams['lines.linewidth'],
                    color=color,
                    zorder=-10
                )
                ax.fill_between(
                    x, mean - std, mean + std,
                    lw=0, color=color, alpha=.3, zorder=-20
                )

            if param == 'parton_number':
                xlim = (0, 10)
            ax.set_xlim(xlim)
            ax.set_ylim(ylim)

            auto_ticks(ax, 'x', nbins=3, minor=2)
            ax.set_yticks(yticksmajor)
            ax.set_yticks(yticksminor, minor=True)

            if label_all or ax.is_last_row():
                ax.set_xlabel(design.labels[i])
            if label_all or ax.is_first_col():
                ax.set_ylabel('PC {}'.format(pc + 1))

            ax.set_ylim(-4, 4)

    set_tight(fig, w_pad=.5, h_pad=.25)


@plot
def diag_emu_all():
    _diag_emu()


@plot
def diag_emu_partial():
    _diag_emu(
        pcs=[0, 2, -1],
        params=['trento_p', 'tau_fs', 'etas_min'],
        label_all=False
    )


#@plot
def grid_error():
    """
    Scatter plot observables calculated on a grid with grid scale = 0.2 against
    observables calculated on a grid with grid scale = 0.1.

    """
    obs_list = [
        ('dNch_deta', r'$dN_\mathrm{ch}/d\eta$'),
        ('mean_pT', r'mean $p_T$ [GeV]'),
        ('v2', r'$|Q_2|/M$'),
        ('v3', r'$|Q_3|/M$'),
    ]

    fig, axes = plt.subplots(ncols=4, nrows=1, figsize=figsize(1, aspect=.3))

    fine_design_points, coarse_design_points = (
        [Path(
            '/var/phy/project/nukeserv/jsm55',
            'hic-events/qm18-grid-scale',
            'grid-scale-{}'.format(gs),
            'events/{}.dat'.format(p)
        ) for p in Design(default_system).points if int(p) != 228]
        for gs in (.1, .2)
    )

    fine_events, coarse_events = [
        [ev for p, ev in model.ModelData(default_system, *design_points).design_events]
        for design_points in (fine_design_points, coarse_design_points)
    ]

    def obs(event, name):
        flow = event['flow']['cms']
        TINY = 1e-12
        return dict(
            dNch_deta=event['dNch_deta'],
            mean_pT=event['mean_pT']['pT'],
            v2=np.absolute(flow['Qn'][1])/(flow['N'] + TINY),
            v3=np.absolute(flow['Qn'][2])/(flow['N'] + TINY),
        )[name]

    for ax, (name, label) in zip(axes.flat, obs_list):
        x = [obs(ev, name) for ev in itertools.chain(*fine_events)]
        y = [obs(ev, name) for ev in itertools.chain(*coarse_events)]
        ax.scatter(x, y, s=15, edgecolors='white', linewidths=.2)

        xy_max = np.nanmax(np.append(x, y))
        ax.plot((0, xy_max), (0, xy_max), color=offblack)

        if ax.is_last_row():
            ax.set_xlabel('grid scale 0.1', fontsize=fontsize['tiny'])
        if ax.is_first_col():
            ax.set_ylabel('grid scale 0.2', fontsize=fontsize['tiny'])

        ax.set_title(label, y=.9, fontsize=fontsize['tiny'])
        ax.set_aspect('equal')
        auto_ticks(ax, 'x', nbins=3)
        auto_ticks(ax, 'y', nbins=3)

    set_tight(fig)


#@plot
def entropy_scaling():
    """
    Plot initial entropy vs final dNch/deta

    """
    ncols= 10
    nrows = 5

    fig, axes = plt.subplots(
        nrows=nrows, ncols=ncols,
        figsize=figsize(relwidth=3, aspect=nrows/ncols)
    )

    data = Path(workdir, 'model_output', 'main', default_system)

    design = Design(default_system, validation=False)
    files = [data / '{}.dat'.format(p) for p in design.points]
    model_data = model.ModelData(default_system, *files)

    def powerlaw(x, a, b, c):
        return a*x**b + c

    for ax, (pt, ev), in zip(axes.flat, model_data.design_events):
        x, y = [ev[k] for k in ('init_entropy', 'dNch_deta')]

        y = y[x.argsort()]
        x = np.sort(x)

        try:
            nonzero = (y > 0)
            popt, pcov = curve_fit(powerlaw, x[nonzero], y[nonzero])

            xval = np.linspace(0, max(x[nonzero]), 100)
            ax.plot(xval, powerlaw(xval, *popt), color=offblack, zorder=1)
        except RuntimeError:
            pass

        ax.plot(x, y, 'o', zorder=0)
        ax.set_title(str(pt))

    set_tight(fig)


@plot
def posterior_proton_shape():
    """
    Draw samples from the proton sampling radius and parton width,
    marginalized over the parton number.

    Plot the resulting joint posterior distribution.

    """
    chain = mcmc.Chain()
    sampling_radius, parton_width = chain.load(
        'sampling_radius', 'parton_width'
    ).T

    fig = plt.figure(figsize=figsize(.5, aspect=.833))
    ax = plt.gca()

    plt.hist2d(sampling_radius, parton_width, bins=50, cmap=plt.cm.Blues)

    plt.xlabel('Constituent position radius $r_{cp}$ [fm]')
    plt.ylabel('Constituent width $w_c$ [fm]')

    ax.set_aspect('equal')
    set_tight(fig, pad=.2)


#@plot
def statistics(system='pPb5020'):
    """
    Event statistics in each p-Pb trigger bin

    """
    plt.figure(figsize=figsize())
    ax = plt.gca()

    files = [
        Path(workdir, 'model_output', 'main', system, '{}.dat'.format(p))
        for p in Design(system, validation=False).points
    ]

    events = [
        tuple(event)
        for point, events in model.ModelData(system, *files).design_events
        for event in events['trigger']
    ]

    counter = Counter(events).most_common()
    shift = 0

    labels = ['flow bins', 'minimum bias', 'mean $p_T$ bins']

    for label, counts in zip(labels, np.split(counter, [8, 9])):
        lbl, height = zip(*counts)
        loc = np.arange(len(lbl))
        ax.bar(loc + shift, height, label=label)
        shift += len(lbl)

    plt.ylabel('Events')
    plt.xticks([])
    plt.legend()

    set_tight()


@plot
def trim_design_points():
    """
    Identify design points with observables which fall outside a reasonable
    prior range.

    """
    emu_obs = [Emulator.pPb5020, Emulator.PbPb5020]

    def bad(obs, y):
        """
        Flag 'bad" design points with extreme observable values

        """
        if any(np.isnan(y)):
            return True
        elif obs == 'dNch_deta':
            return any(y < 1)
        else:
            return False

    bad_points = set()

    for system, obs_data in zip(systems, emu_obs):
        model_data = model._data(system)
        for obs, subobs_data in obs_data:
            for subobs in subobs_data:
                Y = model_data[obs][subobs]['Y']
                for point, y in enumerate(Y):
                    if bad(obs, y):
                        bad_points.add(point)

    print('{} bad points\n{}'.format(len(bad_points), bad_points))


def run_cmd(*args):
    """
    Run and log a subprocess.

    """
    cmd = ' '.join(args)
    logging.info('running command: %s', cmd)

    try:
        proc = subprocess.run(
            cmd.split(), check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True
        )
    except subprocess.CalledProcessError as e:
        logging.error(
            'command failed with status %d:\n%s',
            e.returncode, e.output.strip('\n')
        )
        raise
    else:
        logging.debug(
            'command completed successfully:\n%s',
            proc.stdout.strip('\n')
        )
        return proc


def trento_pcoll(impact_parameter):
    """
    Returns collision probability at the specified
    impact parameter.

    """
    proc = run_cmd(
            'trento p p',
            '--number-events 100000',
            '--reduced-thickness 0',
            '--fluctuation 0.19',
            '--nucleon-width 0.855',
            '--constit-width 0.43',
            '--constit-number 6',
            '--cross-section 7.0',
            '--b-min {}'.format(impact_parameter),
            '--b-max {}'.format(impact_parameter),
            '--ncoll'
    )

    ncoll = np.loadtxt(proc.stdout.splitlines(), usecols=2)
    return np.count_nonzero(ncoll) / ncoll.size


@plot
def inelasticity_profile():
    """
    proton-proton inelasticity density as a function of collision impact
    parameter.

    """
    fig = plt.figure(figsize=figsize(.5, aspect=.7))
    ncpu = multiprocessing.cpu_count()

    # inelasticity density at 7 TeV by Alba Ontoso et. al.
    fname = workdir / 'expt/inelasticity_density_lhc.txt'
    b_values, Pcoll = np.loadtxt(fname).T
    plt.plot(b_values, Pcoll, label=r'$7$ TeV, Fit to $d(\sigma/dt)(t)$')

    # trento inelasticity density at 5.02 TeV
    b_values = np.linspace(0, 5, 20)
    p = multiprocessing.Pool(ncpu)
    Pcoll = np.array(p.map(trento_pcoll, b_values))
    plt.plot(b_values, Pcoll/Pcoll[0],
             label='$5.02$ TeV, MAP estimate')

    plt.xlabel(r'Impact parameter $b$ [fm]')
    plt.ylabel(r'$G_\mathrm{inel}(b) / G_\mathrm{inel}(0)$')
    plt.xlim(0.0, 4.0)
    plt.ylim(0.0, 1.2)
    plt.legend(title='LHC $p$-$p$ collisions',
               loc=1, bbox_to_anchor=(1.05,1.05))
    set_tight(fig)


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

    warnings.filterwarnings(
        'ignore',
        category=MathTextWarning,
        message='Substituting with a symbol from Computer Modern.'
    )
    warnings.filterwarnings(
        'ignore',
        category=UserWarning,
        message=r"'[\w\.]+' can not be subsetted into a Type 3 font."
    )

    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(
        '--paper', action='store_true',
        help='use paper style: cm serif font, true black text + axes'
    )
    parser.add_argument(
        'plots', nargs='*', type=arg_to_plot, metavar='PLOT',
        help='{} (default: all)'.format(', '.join(choices).join('{}'))
    )
    args = parser.parse_args()

    if args.paper:
        plt.rcParams.update({
            'font.family': 'serif',
            'font.serif': ['CMU Serif'],
            'mathtext.fontset': 'cm',
            'text.color': 'black',
            'axes.edgecolor': 'black',
            'axes.labelcolor': 'black',
            'xtick.color': 'black',
            'ytick.color': 'black',
        })

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