# -*- coding: utf-8 -*-
#This code is adaped from
# https://github.com/dfm/corner.py
# git hash 5c2cd63  on May 25
# Modifications by John Baker NASA-GSFC (2016)
# Modifications by Sylvain Marsat AEI (2017)
#Copyright (c) 2013-2016 Daniel Foreman-Mackey
#All rights reserved.
#
#Redistribution and use in source and binary forms, with or without
#modification, are permitted provided that the following conditions are met:
#
#1. Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#2. Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
#THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#The views and conclusions contained in the software and documentation are those
#of the authors and should not be interpreted as representing official policies,
#either expressed or implied, of the FreeBSD Project.

from __future__ import absolute_import, division, print_function
import sys
if sys.version_info[0] == 2:
    from future_builtins import map, filter


import logging
import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import LinearSegmentedColormap, colorConverter
from matplotlib.ticker import ScalarFormatter
from matplotlib.patches import Ellipse

try:
    from scipy.ndimage import gaussian_filter
except ImportError:
    gaussian_filter = None

__all__ = ["corner", "hist2d", "quantile"]


def corner(xs, bins=20, params_range=None, weights=None, cov=None, color="k",
           smooth=None, smooth1d=None,
           labels=None, label_kwargs=None,
           show_histograms=True, show_titles=False, title_fmt=".2f", title_kwargs=None,
           truths=None, add_truths=None, truth_color="#4682b4", add_truth_colors=None, cov_color=None,
           scale_hist=False, quantiles=None, show_quantiles=True, verbose=False, figaxes=None,
           max_n_ticks=5, top_ticks=False, use_math_text=False, resize_fig=1.,
           hist_kwargs=None, **hist2d_kwargs):
    """
    Make a *sick* corner plot showing the projections of a data set in a
    multi-dimensional space. kwargs are passed to hist2d() or used for
    `matplotlib` styling.

    Parameters
    ----------
    xs : array_like[nsamples, ndim]
        The samples. This should be a 1- or 2-dimensional array. For a 1-D
        array this results in a simple histogram. For a 2-D array, the zeroth
        axis is the list of samples and the next axis are the dimensions of
        the space.

    bins : int or array_like[ndim,]
        The number of bins to use in histograms, either as a fixed value for
        all dimensions, or as a list of integers for each dimension.

    weights : array_like[nsamples,]
        The weight of each sample. If `None` (default), samples are given
        equal weight.

    color : str
        A ``matplotlib`` style color for all histograms.

    smooth, smooth1d : float
       The standard deviation for Gaussian kernel passed to
       `scipy.ndimage.gaussian_filter` to smooth the 2-D and 1-D histograms
       respectively. If `None` (default), no smoothing is applied.

    labels : iterable (ndim,)
        A list of names for the dimensions. If a ``xs`` is a
        ``pandas.DataFrame``, labels will default to column names.

    label_kwargs : dict
        Any extra keyword arguments to send to the `set_xlabel` and
        `set_ylabel` methods.

    show_histograms : bool
        Displays 1-D histogram showing marginalized distributions and quantiles.

    show_titles : bool
        Displays a title above each 1-D histogram showing the 0.5 quantile
        with the upper and lower errors supplied by the quantiles argument.

    title_fmt : string
        The format string for the quantiles given in titles. If you explicitly
        set ``show_titles=True`` and ``title_fmt=None``, the labels will be
        shown as the titles. (default: ``.2f``)

    title_kwargs : dict
        Any extra keyword arguments to send to the `set_title` command.

    params_range : iterable (ndim,)
        A list where each element is either a length 2 tuple containing
        lower and upper bounds or a float in range (0., 1.)
        giving the fraction of samples to include in bounds, e.g.,
        [(0.,10.), (1.,5), 0.999, etc.].
        If a fraction, the bounds are chosen to be equal-tailed.

    truths : iterable (ndim,)
        A list of reference values to indicate on the plots.  Individual
        values can be omitted by using ``None``.

    truth_color : str
        A ``matplotlib`` style color for the ``truths`` makers.

    scale_hist : bool
        Should the 1-D histograms be scaled in such a way that the zero line
        is visible?

    quantiles : iterable
        A list of fractional quantiles to show on the 1-D histograms as
        vertical dashed lines.

    verbose : bool
        If true, print the values of the computed quantiles.

    plot_contours : bool
        Draw contours for dense regions of the plot.

    use_math_text : bool
        If true, then axis tick labels for very large or small exponents will
        be displayed as powers of 10 rather than using `e`.

    max_n_ticks: int
        Maximum number of ticks to try to use

    top_ticks : bool
        If true, label the top ticks of each axis

    fig : matplotlib.Figure
        Overplot onto the provided figure object.

    hist_kwargs : dict
        Any extra keyword arguments to send to the 1-D histogram plots.

    **hist2d_kwargs
        Any remaining keyword arguments are sent to `corner.hist2d` to generate
        the 2-D histogram plots.

    """
    if quantiles is None:
        quantiles = []
    if title_kwargs is None:
        title_kwargs = dict()
    if label_kwargs is None:
        label_kwargs = dict()

    # Try filling in labels from pandas.DataFrame columns.
    if labels is None:
        try:
            labels = xs.columns
        except AttributeError:
            pass

    # Deal with 1D sample lists.
    xs = np.atleast_1d(xs)
    if len(xs.shape) == 1:
        xs = np.atleast_2d(xs)
    else:
        assert len(xs.shape) == 2, "The input sample array must be 1- or 2-D."
        xs = xs.T
    assert xs.shape[0] <= xs.shape[1], "I don't believe that you want more " \
                                       "dimensions than samples!"

    # Parse the weight array.
    if weights is not None:
        weights = np.asarray(weights)
        if weights.ndim != 1:
            raise ValueError("Weights must be 1-D")
        if xs.shape[1] != weights.shape[0]:
            raise ValueError("Lengths of weights must match number of samples")

    # Parse the parameter ranges.
    if params_range is None:
        if "extents" in hist2d_kwargs:
            logging.warn("Deprecated keyword argument 'extents'. "
                         "Use 'params_range' instead.")
            params_range = hist2d_kwargs.pop("extents")
        else:
            params_range = [[x.min(), x.max()] for x in xs]
            # Check for parameters that never change.
            m = np.array([e[0] == e[1] for e in params_range], dtype=bool)
            if np.any(m):
                raise ValueError(("It looks like the parameter(s) in "
                                  "column(s) {0} have no dynamic range. "
                                  "Please provide a `params_range` argument.")
                                 .format(", ".join(map(
                                     "{0}".format, np.arange(len(m))[m]))))

    else:
        # If any of the extents are percentiles, convert them to ranges.
        # Also make sure it's a normal list.
        params_range = list(params_range)
        for i, pr in enumerate(params_range):
            if pr is None:
                params_range[i] = [xs[i].min(), xs[i].max()]
            else:
                try:
                    emin, emax = params_range[i]
                except TypeError:
                    q = [0.5 - 0.5*params_range[i], 0.5 + 0.5*params_range[i]]
                    params_range[i] = quantile(xs[i], q, weights=weights)

    if len(params_range) != xs.shape[0]:
        raise ValueError("Dimension mismatch between samples and params_range")

    # Parse the bin specifications.
    try:
        bins = [int(bins) for _ in params_range]
    except TypeError:
        if len(bins) != len(params_range):
            raise ValueError("Dimension mismatch between bins and params_range")

    # Some magic numbers for pretty axis layout.
    K = len(xs)
    if show_histograms:    # total number of plots: grid nplot*nplot
        nplot = K
    else:
        nplot = K-1
    factor = 2.0           # size of one side of one panel
    lbdim = 0.5 * factor   # size of left/bottom margin
    trdim = 0.2 * factor   # size of top/right margin
    whspace = 0.05         # w/hspace size
    plotdim = factor * nplot + factor * (nplot - 1.) * whspace
    dim = lbdim + plotdim + trdim

    # Create a new figure if one wasn't provided.
    # if fig is None:
    #     fig, axes = plt.subplots(nplot, nplot, figsize=(dim, dim))
    # else:
    #     print(type(fig.axes))
    #     try:
    #         axes = np.array(fig.axes).reshape((nplot, nplot))
    #     except:
    #         raise ValueError("Provided figure has {0} axes, but data has "
    #                          "dimensions K={1}".format(len(fig.axes), nplot))
    if figaxes is None:
        fig, axes = plt.subplots(nplot, nplot, figsize=(resize_fig*dim, resize_fig*dim))
    else:
        fig, axes = figaxes


    #idea is to pass in covariance, otherwise concoct something from the 1-sigma range.
    # if(cov==None):
    #     print("concocting covar elements from 1-sigma ranges")
    #     cov=np.zeros((K,K))
    #     for k in np.arange(K):
    #         q_16, q_50, q_84 = quantile(xs[k], [0.16, 0.5, 0.84],weights=weights)
    #         deltax=(q_84-q_16)/2.0
    #         cov[k,k]=deltax**2
    # print("cov=",cov)

    # Format the figure.
    lb = lbdim / dim
    tr = (lbdim + plotdim) / dim
    fig.subplots_adjust(left=lb, bottom=lb, right=tr, top=tr,
                        wspace=whspace, hspace=whspace)

    # If asked to show the 1D histograms
    if show_histograms:
        # Set up the default histogram keywords.
        if hist_kwargs is None:
            hist_kwargs = dict()
        hist_kwargs["color"] = hist_kwargs.get("color", color)
        if smooth1d is None:
            hist_kwargs["histtype"] = hist_kwargs.get("histtype", "step")
        for i, x in enumerate(xs):
            # Deal with masked arrays.
            if hasattr(x, "compressed"):
                x = x.compressed()

            if np.shape(xs)[0] == 1:
                ax = axes
            else:
                ax = axes[i, i]
            # Plot the histograms.
            if smooth1d is None:
                n, _, _ = ax.hist(x, bins=bins[i], weights=weights, density=True,
                                  range=np.sort(params_range[i]), **hist_kwargs)
            else:
                if gaussian_filter is None:
                    raise ImportError("Please install scipy for smoothing")
                n, b = np.histogram(x, bins=bins[i], weights=weights, density=True,
                                    range=np.sort(params_range[i]))
                n = gaussian_filter(n, smooth1d)
                x0 = np.array(list(zip(b[:-1], b[1:]))).flatten()
                y0 = np.array(list(zip(n, n))).flatten()
                ax.plot(x0, y0, **hist_kwargs)

            if truths is not None and truths[i] is not None:
                ax.axvline(truths[i], color=truth_color)

            if add_truths is not None:
                for add_truth, add_truth_color in zip(add_truths, add_truth_colors):
                     if add_truth[i] is not None:
                         ax.axvline(add_truth[i], color=add_truth_color)

            # Plot quantiles if wanted.
            if show_quantiles:
                if len(quantiles) > 0:
                    qvalues = quantile(x, quantiles, weights=weights)
                    for q in qvalues:
                        ax.axvline(q, ls="dashed", color=color)

                    if verbose:
                        print("Quantiles:")
                        for item in zip(quantiles, qvalues):
                            print(item)

            if show_titles:
                title = None
                if title_fmt is not None:
                    # Compute the quantiles for the title. This might redo
                    # unneeded computation but who cares.
                    q_16, q_50, q_84 = quantile(x, [0.16, 0.5, 0.84],
                                                weights=weights)
                    q_m, q_p = q_50-q_16, q_84-q_50

                    # Format the quantile display.
                    fmt = "{{0:{0}}}".format(title_fmt).format
                    title = r"${{{0}}}_{{-{1}}}^{{+{2}}}$"
                    title = title.format(fmt(q_50), fmt(q_m), fmt(q_p))

                    # Add in the column name if it's given.
                    if labels is not None:
                        title = "{0} = {1}".format(labels[i], title)

                elif labels is not None:
                    title = "{0}".format(labels[i])

                if title is not None:
                    ax.set_title(title, **title_kwargs)

            # Set up the axes.
            ax.set_xlim(params_range[i])
            if scale_hist:
                maxn = np.max(n)
                ax.set_ylim(-0.1 * maxn, 1.1 * maxn)
            else:
                ax.set_ylim(0, 1.1 * np.max(n))
            ax.set_yticklabels([])
            ax.xaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower"))

            if i < K - 1:
                if top_ticks:
                    ax.xaxis.set_ticks_position("top")
                    [l.set_rotation(45) for l in ax.get_xticklabels()]
                else:
                    ax.set_xticklabels([])
            else:
                [l.set_rotation(45) for l in ax.get_xticklabels()]
                if labels is not None:
                    ax.set_xlabel(labels[i], **label_kwargs)
                    ax.xaxis.set_label_coords(0.5, -0.3)

                # use MathText for axes ticks
                ax.xaxis.set_major_formatter(
                    ScalarFormatter(useMathText=use_math_text))

            for j, y in enumerate(xs):
                if np.shape(xs)[0] == 1:
                    ax = axes
                else:
                    ax = axes[i, j]
                if j > i:
                    ax.set_frame_on(False)
                    ax.set_xticks([])
                    ax.set_yticks([])
                    continue
                elif j == i:
                    continue

                # Deal with masked arrays.
                if hasattr(y, "compressed"):
                    y = y.compressed()

                hist2d(y, x, ax=ax, range=[params_range[j], params_range[i]], weights=weights,
                       color=color, smooth=smooth, bins=[bins[j], bins[i]],
                       **hist2d_kwargs)

                # Add an error ellipse based on the provided covariance matrix
                if cov is not None:
                    # #center
                    # cx=truths[j]#need to add checking for availability of truths?
                    # cy=truths[i]
                    # #ang=math.acos(cov[0,1]/math.sqrt(cov[0,0]*cov[1,1]))*180/math.pi
                    # print (j,i,labels[j],labels[i],"center=",cx,cy)
                    # N_thetas=60
                    # dtheta=2.0*math.pi/(N_thetas-1)
                    # thetas=np.arange(0,(2.0*math.pi+dtheta),dtheta)
                    # Cplus=(cov[i,i]+cov[j,j])/2.0
                    # Cminus=(-cov[i,i]+cov[j,j])/2.0
                    # print("cov[ii],cov[ij],cov[jj],Cplus,Cminus:",cov[i,i],cov[i,j],cov[j,j],Cplus,Cminus)
                    # ang=-math.pi/4.
                    # root=cov[i,j]/math.sqrt(cov[i,i]*cov[j,j])
                    # acoeff=math.sqrt(1-root)
                    # bcoeff=math.sqrt(1+root)
                    # xcoeff=math.sqrt(cov[j,j])
                    # ycoeff=math.sqrt(cov[i,i])
                    # print("a2,b2",acoeff*acoeff,bcoeff*bcoeff)
                    # print("a,b,ang, xcoeff,ycoeff, root=",acoeff,bcoeff,ang,xcoeff,ycoeff,root)
                    # elxs=[cx+xcoeff*(acoeff*math.cos(th)*math.cos(ang)-bcoeff*math.sin(th)*math.sin(ang)) for th in thetas]
                    # elys=[cy+ycoeff*(acoeff*math.cos(th)*math.sin(ang)+bcoeff*math.sin(th)*math.cos(ang)) for th in thetas]
                    # #print (thetas)
                    # #print (elxs)
                    # #print (elys)
                    # ax.plot(elxs,elys,color='r')
                    #
                    #center
                    if truths is None:
                        raise ValueError('Error: truths is None but is required to plot Fisher ellipses.')
                    cx = truths[j] # need to add checking for availability of truths?
                    cy = truths[i]
                    N_thetas = 60
                    dtheta = 2.0*math.pi/(N_thetas-1)
                    thetas = np.arange(0,(2.0*math.pi+dtheta),dtheta)
                    ang = -math.pi/4.
                    root = cov[i,j]/math.sqrt(cov[i,i]*cov[j,j])
                    if(root>1): root = 1
                    if(root<-1): root = -1
                    acoeff = math.sqrt(1-root)
                    bcoeff = math.sqrt(1+root)
                    xcoeff = math.sqrt(cov[j,j])
                    ycoeff = math.sqrt(cov[i,i])
                    if "levels" in hist2d_kwargs:
                        levels= hist2d_kwargs["levels"]
                    else:
                        levels== 1.0 - np.exp(-0.5 * np.arange(0.5, 2.1, 0.5) ** 2)
                    for xlev in levels:
                        # in the next line we convert the credibility limit
                        # to a "sigma" limit for a 2-d normal
                        # this becomes a scale-factor for the error ellipse
                        # 1-exp(x^2/(-2)=y
                        # -2*log(1-y)=x^2
                        lev_fac = math.sqrt( -2 * math.log( 1 - xlev ) )
                        elxs = [cx+lev_fac*xcoeff*(acoeff*math.cos(th)*math.cos(ang)-bcoeff*math.sin(th)*math.sin(ang)) for th in thetas]
                        elys = [cy+lev_fac*ycoeff*(acoeff*math.cos(th)*math.sin(ang)+bcoeff*math.sin(th)*math.cos(ang)) for th in thetas]
                        if cov_color is None:
                            cov_color='k'
                        ax.plot(elxs, elys, color=cov_color)

                if truths is not None:
                    if truths[i] is not None and truths[j] is not None:
                        ax.plot(truths[j], truths[i], "s", color=truth_color)
                    if truths[j] is not None:
                        ax.axvline(truths[j], color=truth_color)
                    if truths[i] is not None:
                        ax.axhline(truths[i], color=truth_color)

                if add_truths is not None:
                    for add_truth, add_truth_color in zip(add_truths, add_truth_colors):
                        if add_truth[i] is not None and add_truth[j] is not None:
                            ax.plot(add_truth[j], add_truth[i], "s", color=add_truth_color)
                        if add_truth[j] is not None:
                            ax.axvline(add_truth[j], color=add_truth_color)
                        if add_truth[i] is not None:
                            ax.axhline(add_truth[i], color=add_truth_color)

                ax.xaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower"))
                ax.yaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower"))

                if i < K - 1:
                    ax.set_xticklabels([])
                else:
                    [l.set_rotation(45) for l in ax.get_xticklabels()]
                    if labels is not None:
                        ax.set_xlabel(labels[j], **label_kwargs)
                        ax.xaxis.set_label_coords(0.5, -0.3)

                    # use MathText for axes ticks
                    ax.xaxis.set_major_formatter(
                        ScalarFormatter(useMathText=use_math_text))

                if j > 0:
                    ax.set_yticklabels([])
                else:
                    [l.set_rotation(45) for l in ax.get_yticklabels()]
                    if labels is not None:
                        ax.set_ylabel(labels[i], **label_kwargs)
                        ax.yaxis.set_label_coords(-0.3, 0.5)

                    # use MathText for axes ticks
                    ax.yaxis.set_major_formatter(
                        ScalarFormatter(useMathText=use_math_text))

    # If not asking for the histograms to be plotted
    else:
        for i, x in enumerate(xs):
            if i==0:
                continue
            # Deal with masked arrays.
            if hasattr(x, "compressed"):
                x = x.compressed()

            for j, y in enumerate(xs):
                if j==K-1:
                    continue
                if np.shape(xs)[0] == 2:
                    ax = axes
                else:
                    ax = axes[i-1, j]
                if j > i-1 :
                    ax.set_frame_on(False)
                    ax.set_xticks([])
                    ax.set_yticks([])
                    continue
                elif j == i:
                    continue

                # Deal with masked arrays.
                if hasattr(y, "compressed"):
                    y = y.compressed()

                hist2d(y, x, ax=ax, range=[params_range[j], params_range[i]], weights=weights,
                       color=color, smooth=smooth, bins=[bins[j], bins[i]],
                       **hist2d_kwargs)

                # Add an error ellipse based on the provided covariance matrix
                if cov is not None:
                    # #center
                    # cx=truths[j]#need to add checking for availability of truths?
                    # cy=truths[i]
                    # #ang=math.acos(cov[0,1]/math.sqrt(cov[0,0]*cov[1,1]))*180/math.pi
                    # print (j,i,labels[j],labels[i],"center=",cx,cy)
                    # N_thetas=60
                    # dtheta=2.0*math.pi/(N_thetas-1)
                    # thetas=np.arange(0,(2.0*math.pi+dtheta),dtheta)
                    # Cplus=(cov[i,i]+cov[j,j])/2.0
                    # Cminus=(-cov[i,i]+cov[j,j])/2.0
                    # print("cov[ii],cov[ij],cov[jj],Cplus,Cminus:",cov[i,i],cov[i,j],cov[j,j],Cplus,Cminus)
                    # ang=-math.pi/4.
                    # root=cov[i,j]/math.sqrt(cov[i,i]*cov[j,j])
                    # acoeff=math.sqrt(1-root)
                    # bcoeff=math.sqrt(1+root)
                    # xcoeff=math.sqrt(cov[j,j])
                    # ycoeff=math.sqrt(cov[i,i])
                    # print("a2,b2",acoeff*acoeff,bcoeff*bcoeff)
                    # print("a,b,ang, xcoeff,ycoeff, root=",acoeff,bcoeff,ang,xcoeff,ycoeff,root)
                    # elxs=[cx+xcoeff*(acoeff*math.cos(th)*math.cos(ang)-bcoeff*math.sin(th)*math.sin(ang)) for th in thetas]
                    # elys=[cy+ycoeff*(acoeff*math.cos(th)*math.sin(ang)+bcoeff*math.sin(th)*math.cos(ang)) for th in thetas]
                    # #print (thetas)
                    # #print (elxs)
                    # #print (elys)
                    # ax.plot(elxs,elys,color='r')
                    #
                    #center
                    cx = truths[j] # need to add checking for availability of truths?
                    cy = truths[i]
                    N_thetas = 60
                    dtheta = 2.0*math.pi/(N_thetas-1)
                    thetas = np.arange(0,(2.0*math.pi+dtheta),dtheta)
                    ang = -math.pi/4.
                    root = cov[i,j]/math.sqrt(cov[i,i]*cov[j,j])
                    if(root>1): root = 1
                    if(root<-1): root = -1
                    acoeff = math.sqrt(1-root)
                    bcoeff = math.sqrt(1+root)
                    xcoeff = math.sqrt(cov[j,j])
                    ycoeff = math.sqrt(cov[i,i])
                    if "levels" in hist2d_kwargs:
                        levels= hist2d_kwargs["levels"]
                    else:
                        levels== 1.0 - np.exp(-0.5 * np.arange(0.5, 2.1, 0.5) ** 2)
                    for xlev in levels:
                        # in the next line we convert the credibility limit
                        # to a "sigma" limit for a 2-d normal
                        # this becomes a scale-factor for the error ellipse
                        # 1-exp(x^2/(-2)=y
                        # -2*log(1-y)=x^2
                        lev_fac = math.sqrt( -2 * math.log( 1 - xlev ) )
                        elxs = [cx+lev_fac*xcoeff*(acoeff*math.cos(th)*math.cos(ang)-bcoeff*math.sin(th)*math.sin(ang)) for th in thetas]
                        elys = [cy+lev_fac*ycoeff*(acoeff*math.cos(th)*math.sin(ang)+bcoeff*math.sin(th)*math.cos(ang)) for th in thetas]
                        if cov_color is None:
                            cov_color='k'
                        ax.plot(elxs, elys, color=cov_color)

                if truths is not None:
                    if truths[i] is not None and truths[j] is not None:
                        ax.plot(truths[j], truths[i], "s", color=truth_color)
                    if truths[j] is not None:
                        ax.axvline(truths[j], color=truth_color)
                    if truths[i] is not None:
                        ax.axhline(truths[i], color=truth_color)

                if add_truths is not None:
                    for add_truth, add_truth_color in zip(add_truths, add_truth_colors):
                        if add_truth[i] is not None and add_truth[j] is not None:
                            ax.plot(add_truth[j], add_truth[i], "s", color=add_truth_color)
                        if add_truth[j] is not None:
                            ax.axvline(add_truth[j], color=add_truth_color)
                        if add_truth[i] is not None:
                            ax.axhline(add_truth[i], color=add_truth_color)

                ax.xaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower"))
                ax.yaxis.set_major_locator(MaxNLocator(max_n_ticks, prune="lower"))

                if i < K - 1:
                    ax.set_xticklabels([])
                else:
                    [l.set_rotation(45) for l in ax.get_xticklabels()]
                    if labels is not None:
                        ax.set_xlabel(labels[j], **label_kwargs)
                        ax.xaxis.set_label_coords(0.5, -0.3)

                    # use MathText for axes ticks
                    ax.xaxis.set_major_formatter(
                        ScalarFormatter(useMathText=use_math_text))

                if j > 0:
                    ax.set_yticklabels([])
                else:
                    [l.set_rotation(45) for l in ax.get_yticklabels()]
                    if labels is not None:
                        ax.set_ylabel(labels[i], **label_kwargs)
                        ax.yaxis.set_label_coords(-0.3, 0.5)

                    # use MathText for axes ticks
                    ax.yaxis.set_major_formatter(
                        ScalarFormatter(useMathText=use_math_text))

    return fig, axes


def quantile(x, q, weights=None):
    """
    Compute sample quantiles with support for weighted samples.

    Note
    ----
    When ``weights`` is ``None``, this method simply calls numpy's percentile
    function with the values of ``q`` multiplied by 100.

    Parameters
    ----------
    x : array_like[nsamples,]
       The samples.

    q : array_like[nquantiles,]
       The list of quantiles to compute. These should all be in the range
       ``[0, 1]``.

    weights : Optional[array_like[nsamples,]]
        An optional weight corresponding to each sample. These

    Returns
    -------
    quantiles : array_like[nquantiles,]
        The sample quantiles computed at ``q``.

    Raises
    ------
    ValueError
        For invalid quantiles; ``q`` not in ``[0, 1]`` or dimension mismatch
        between ``x`` and ``weights``.

    """
    x = np.atleast_1d(x)
    q = np.atleast_1d(q)

    if np.any(q < 0.0) or np.any(q > 1.0):
        raise ValueError("Quantiles must be between 0 and 1")

    if weights is None:
        return np.percentile(x, 100.0 * q)
    else:
        weights = np.atleast_1d(weights)
        if len(x) != len(weights):
            raise ValueError("Dimension mismatch: len(weights) != len(x)")
        idx = np.argsort(x)
        sw = weights[idx]
        cdf = np.cumsum(sw)[:-1]
        cdf /= cdf[-1]
        cdf = np.append(0, cdf)
        return np.interp(q, cdf, x[idx]).tolist()


def hist2d(x, y, bins=20, range=None, weights=None, levels=None, smooth=None,
           ax=None, color=None, plot_datapoints=True, plot_density=True,
           plot_contours=True, no_fill_contours=False, fill_contours=False,
           contour_kwargs=None, contourf_kwargs=None, data_kwargs=None,
           **kwargs):
    """
    Plot a 2-D histogram of samples.

    Parameters
    ----------
    x : array_like[nsamples,]
       The samples.

    y : array_like[nsamples,]
       The samples.

    levels : array_like
        The contour levels to draw.

    ax : matplotlib.Axes
        A axes instance on which to add the 2-D histogram.

    plot_datapoints : bool
        Draw the individual data points.

    plot_density : bool
        Draw the density colormap.

    plot_contours : bool
        Draw the contours.

    no_fill_contours : bool
        Add no filling at all to the contours (unlike setting
        ``fill_contours=False``, which still adds a white fill at the densest
        points).

    fill_contours : bool
        Fill the contours.

    contour_kwargs : dict
        Any additional keyword arguments to pass to the `contour` method.

    contourf_kwargs : dict
        Any additional keyword arguments to pass to the `contourf` method.

    data_kwargs : dict
        Any additional keyword arguments to pass to the `plot` method when
        adding the individual data points.

    """
    if ax is None:
        ax = plt.gca()

    # Set the default range based on the data range if not provided.
    if range is None:
        if "extent" in kwargs:
            logging.warn("Deprecated keyword argument 'extent'. "
                         "Use 'range' instead.")
            range = kwargs["extent"]
        else:
            range = [[x.min(), x.max()], [y.min(), y.max()]]

    # Set up the default plotting arguments.
    if color is None:
        color = "k"

    # Choose the default "sigma" contour levels.
    if levels is None:
        levels = 1.0 - np.exp(-0.5 * np.arange(0.5, 2.1, 0.5) ** 2)

    # This is the color map for the density plot, over-plotted to indicate the
    # density of the points near the center.
    density_cmap = LinearSegmentedColormap.from_list(
        "density_cmap", [color, (1, 1, 1, 0)])

    # This color map is used to hide the points at the high density areas.
    white_cmap = LinearSegmentedColormap.from_list(
        "white_cmap", [(1, 1, 1), (1, 1, 1)], N=2)

    # This "color map" is the list of colors for the contour levels if the
    # contours are filled.
    rgba_color = colorConverter.to_rgba(color)
    contour_cmap = [list(rgba_color) for l in levels] + [rgba_color]
    for i, l in enumerate(levels):
        contour_cmap[i][-1] *= float(i) / (len(levels)+1)

    # We'll make the 2D histogram to directly estimate the density.
    try:
        H, X, Y = np.histogram2d(x.flatten(), y.flatten(), bins=bins,
                                 range=list(map(np.sort, range)),
                                 weights=weights)
    except ValueError:
        raise ValueError("It looks like at least one of your sample columns "
                         "have no dynamic range. You could try using the "
                         "'range' argument.")

    if smooth is not None:
        if gaussian_filter is None:
            raise ImportError("Please install scipy for smoothing")
        H = gaussian_filter(H, smooth)

    # Compute the density levels.
    Hflat = H.flatten()
    inds = np.argsort(Hflat)[::-1]
    Hflat = Hflat[inds]
    sm = np.cumsum(Hflat)
    sm /= sm[-1]
    V = np.empty(len(levels))
    for i, v0 in enumerate(levels):
        try:
            V[i] = Hflat[sm <= v0][-1]
        except:
            V[i] = Hflat[0]
    V.sort()
    m = np.diff(V) == 0
    if np.any(m):
        logging.warning("Too few points to create valid contours")
    while np.any(m):
        V[np.where(m)[0][0]] *= 1.0 - 1e-4
        m = np.diff(V) == 0
    V.sort()

    # Compute the bin centers.
    X1, Y1 = 0.5 * (X[1:] + X[:-1]), 0.5 * (Y[1:] + Y[:-1])

    # Extend the array for the sake of the contours at the plot edges.
    H2 = H.min() + np.zeros((H.shape[0] + 4, H.shape[1] + 4))
    H2[2:-2, 2:-2] = H
    H2[2:-2, 1] = H[:, 0]
    H2[2:-2, -2] = H[:, -1]
    H2[1, 2:-2] = H[0]
    H2[-2, 2:-2] = H[-1]
    H2[1, 1] = H[0, 0]
    H2[1, -2] = H[0, -1]
    H2[-2, 1] = H[-1, 0]
    H2[-2, -2] = H[-1, -1]
    X2 = np.concatenate([
        X1[0] + np.array([-2, -1]) * np.diff(X1[:2]),
        X1,
        X1[-1] + np.array([1, 2]) * np.diff(X1[-2:]),
    ])
    Y2 = np.concatenate([
        Y1[0] + np.array([-2, -1]) * np.diff(Y1[:2]),
        Y1,
        Y1[-1] + np.array([1, 2]) * np.diff(Y1[-2:]),
    ])

    if plot_datapoints:
        if data_kwargs is None:
            data_kwargs = dict()
        data_kwargs["color"] = data_kwargs.get("color", color)
        data_kwargs["ms"] = data_kwargs.get("ms", 2.0)
        data_kwargs["mec"] = data_kwargs.get("mec", "none")
        data_kwargs["alpha"] = data_kwargs.get("alpha", 0.1)
        ax.plot(x, y, "o", zorder=-1, rasterized=True, **data_kwargs)

    # Plot the base fill to hide the densest data points.
    if (plot_contours or plot_density) and not no_fill_contours:
        ax.contourf(X2, Y2, H2.T, [V.min(), H.max()],
                    cmap=white_cmap, antialiased=False)

    if plot_contours and fill_contours:
        if contourf_kwargs is None:
            contourf_kwargs = dict()
        contourf_kwargs["colors"] = contourf_kwargs.get("colors", contour_cmap)
        contourf_kwargs["antialiased"] = contourf_kwargs.get("antialiased",
                                                             False)
        ax.contourf(X2, Y2, H2.T, np.concatenate([[0], V, [H.max()*(1+1e-4)]]),
                    **contourf_kwargs)

    # Plot the density map. This can't be plotted at the same time as the
    # contour fills.
    elif plot_density:
        ax.pcolor(X, Y, H.max() - H.T, cmap=density_cmap)

    # Plot the contour edge colors.
    if plot_contours:
        if contour_kwargs is None:
            contour_kwargs = dict()
        contour_kwargs["colors"] = contour_kwargs.get("colors", color)
        ax.contour(X2, Y2, H2.T, V, **contour_kwargs)

    ax.set_xlim(range[0])
    ax.set_ylim(range[1])
