

# from models_withLik import interpolate_experiment
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

import impala.superCal as impala

plt.rcParams.update({"figure.max_open_warning": 0})
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.lines import Line2D
from numpy.random import uniform
from scipy.interpolate import interp1d
from scipy.stats import gaussian_kde


# diagnostic plots to include:
#    - traces of theta, s2, but for hier which ones? All?
#   - tempering and MCMC counts
class Diagnostics(object):
    def __init__(self, setup, out):
        """ """
        self.setup = setup
        self.out = out
        return

    def pooled_trace_plots():
        pass


class PTW_Plotter(object):
    """PTW Prediction and Pairwise Plots"""

    pooled = None

    def __init__(self, setup, out):
        """ """
        self.setup = setup
        self.out = out
        return

    def ptw_prediction_plot_single(
        self,
        actual_x,
        actual_y,
        pred_y_lb,
        pred_y_ub,
        hier_y_lb=None,
        hier_y_ub=None,
        pnew_y_lb=None,
        pnew_y_ub=None,
        edot=None,
        temp=None,
        pdf=None,
        model_type=None,
    ):
        """PTW Prediction Single Plot"""
        ymaxs = [actual_y.max(), pred_y_ub.max()]
        if hier_y_ub is not None:
            ymaxs.append(hier_y_ub.max())
        if pnew_y_ub is not None:
            ymaxs.append(pnew_y_ub.max())
        ymax = max(ymaxs)
        xmax = actual_x.max()
        text_coords = (0.1 * xmax, 0.85 * ymax)

        fig = plt.figure(figsize=(4, 3))
        if pnew_y_lb is not None:
            plt.fill_between(
                actual_x,
                pnew_y_lb,
                pnew_y_ub,
                color="lightgrey",
                label=r"$\theta^*$",
            )
        if hier_y_lb is not None:
            plt.fill_between(
                actual_x,
                hier_y_lb,
                hier_y_ub,
                color="lightblue",
                label=r"$\theta_0$",
            )
        plt.fill_between(
            actual_x,
            pred_y_lb,
            pred_y_ub,
            color="lightgreen",
            label=r"$\theta_i$",
        )
        plt.scatter(actual_x, actual_y, color="blue", s=0.5, label="y")
        plt.ylim(0.0, 0.025)
        if edot is not None and temp is not None:
            plt.text(*text_coords, "edot: {:.3E}\n temp: {}".format(edot, temp))
        if pdf:
            pdf.savefig(fig)
        return

    def ptw_prediction_plots_hier(self, path, uu):
        """PTW Prediction Hierarchical Plots (no input)"""
        sel = uu  # np.arange(20000, self.setup.nmcmc, 10) # need to script this in
        pred_theta_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]
        pred_theta0_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]
        pred_thetap_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]

        theta_parent = impala.chol_sample_1per_constraints(
            self.out.theta0[sel, 0],
            self.out.Sigma0[sel, 0],
            self.setup.checkConstraints,
            self.setup.bounds_mat,
            self.setup.bounds.keys(),
            self.setup.bounds,
            self.setup.constants,
        )

        for i in range(self.setup.nexp):
            self.setup.models[i].pool = False
            pred_theta_raw[i] = self.setup.models[i].eval(
                impala.tran_unif(
                    self.out.theta[i][sel, 0, :, :].reshape([-1, self.setup.p]),
                    self.setup.bounds_mat,
                    self.setup.bounds.keys(),
                ),
                pool=False,
                nugget=True,
            )
            pred_theta0_raw[i] = self.setup.models[i].eval(
                impala.tran_unif(
                    self.out.theta0[sel, 0],
                    self.setup.bounds_mat,
                    self.setup.bounds.keys(),
                ),
                pool=True,
                nugget=True,
            )
            pred_thetap_raw[i] = self.setup.models[i].eval(
                impala.tran_unif(
                    theta_parent,
                    self.setup.bounds_mat,
                    self.setup.bounds.keys(),
                ),
                pool=True,
                nugget=True,
            )
            self.setup.models[i].pool = True
        real_strain = []
        real_stress = []
        pred_theta = []
        pred_theta0 = []
        pred_thetap = []
        pred_theta_quant_lb = []
        pred_theta_quant_ub = []
        pred_theta0_quant_lb = []
        pred_theta0_quant_ub = []
        pred_thetap_quant_lb = []
        pred_thetap_quant_ub = []
        edots = []
        temps = []
        model_type = []

        for i in range(self.setup.nexp):
            for j in range(self.setup.ntheta[i]):
                pred_theta.append(
                    pred_theta_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_theta_quant_lb.append(
                    np.quantile(pred_theta[-1], 0.005, 0)
                )
                pred_theta_quant_ub.append(
                    np.quantile(pred_theta[-1], 0.995, 0)
                )
                pred_theta0.append(
                    pred_theta0_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_theta0_quant_lb.append(
                    np.quantile(pred_theta0[-1], 0.005, 0)
                )
                pred_theta0_quant_ub.append(
                    np.quantile(pred_theta0[-1], 0.995, 0)
                )
                pred_thetap.append(
                    pred_thetap_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_thetap_quant_lb.append(
                    np.quantile(pred_thetap[-1], 0.005, 0)
                )
                pred_thetap_quant_ub.append(
                    np.quantile(pred_thetap[-1], 0.995, 0)
                )
                model_type.append(type(self.setup.models[i]).__name__)
                if model_type[-1] == "ModelMaterialStrength":
                    edots.append(self.setup.models[i].edots[j])
                    temps.append(self.setup.models[i].temps[j])
                    real_strain.append(
                        self.setup.models[i].meas_strain_histories[j]
                    )
                    real_stress.append(
                        self.setup.ys[i][self.setup.s2_ind[i] == j]
                    )
                else:
                    edots.append(None)
                    temps.append(None)
                    real_strain.append(
                        np.linspace(
                            0,
                            1,
                            len(self.setup.ys[i][self.setup.s2_ind[i] == j]),
                        )
                    )
                    real_stress.append(
                        self.setup.ys[i][self.setup.s2_ind[i] == j]
                    )

        rows = zip(
            real_strain,
            real_stress,
            pred_theta_quant_lb,
            pred_theta_quant_ub,
            pred_theta0_quant_lb,
            pred_theta0_quant_ub,
            pred_thetap_quant_lb,
            pred_thetap_quant_ub,
            edots,
            temps,
            model_type,
        )
        keys = [
            "actual_x",
            "actual_y",
            "pred_y_lb",
            "pred_y_ub",
            "hier_y_lb",
            "hier_y_ub",
            "pnew_y_lb",
            "pnew_y_ub",
            "edot",
            "temp",
            "model_type",
        ]
        plot_param_list = [dict(zip(keys, row)) for row in rows]

        pdf = PdfPages(path)
        for plot_params in plot_param_list:
            if plot_params["model_type"] == "ModelMaterialStrength":
                self.ptw_prediction_plot_single(**plot_params, pdf=pdf)
            else:
                if len(plot_params["actual_x"]) == 1:
                    pass
                else:
                    self.ptw_prediction_plot_single(**plot_params, pdf=pdf)
        pdf.close()
        return

    def ptw_prediction_plots_pool(self, path, uu):
        """PTW Prediction Hierarchical Plots (no input)"""
        sel = uu  # np.arange(20000, self.setup.nmcmc, 10) # need to script this in
        pred_theta_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]
        for i in range(self.setup.nexp):
            pred_theta_raw[i] = self.setup.models[i].eval(
                impala.tran_probit(
                    self.out.theta[
                        sel, 0
                    ],  # .repeat(self.setup.ns2[i], axis = 0),
                    self.setup.bounds_mat,
                    self.setup.bounds.keys(),
                ),
            )

        real_strain = []
        real_stress = []
        pred_theta = []
        pred_theta_quant_lb = []
        pred_theta_quant_ub = []
        edots = []
        temps = []
        model_type = []

        for i in range(self.setup.nexp):
            for j in range(self.setup.ns2[i]):
                pred_theta.append(
                    pred_theta_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_theta_quant_lb.append(
                    np.quantile(pred_theta[-1], 0.025, 0)
                )
                pred_theta_quant_ub.append(
                    np.quantile(pred_theta[-1], 0.975, 0)
                )
                model_type.append(type(self.setup.models[i]).__name__)
                if model_type[-1] == "ModelMaterialStrength":
                    edots.append(self.setup.models[i].edots[j])
                    temps.append(self.setup.models[i].temps[j])
                    real_strain.append(
                        self.setup.models[i].meas_strain_histories[j]
                    )
                    real_stress.append(
                        self.setup.ys[i][self.setup.s2_ind[i] == j]
                    )
                else:
                    edots.append(None)
                    temps.append(None)
                    real_strain.append(
                        np.linspace(
                            0,
                            1,
                            len(self.setup.ys[i][self.setup.s2_ind[i] == j]),
                        )
                    )
                    real_stress.append(
                        self.setup.ys[i][self.setup.s2_ind[i] == j]
                    )

        rows = zip(
            real_strain,
            real_stress,
            pred_theta_quant_lb,
            pred_theta_quant_ub,
            edots,
            temps,
            model_type,
        )
        keys = [
            "actual_x",
            "actual_y",
            "pred_y_lb",
            "pred_y_ub",
            "edot",
            "temp",
            "model_type",
        ]
        plot_param_list = [dict(zip(keys, row)) for row in rows]

        pdf = PdfPages(path)
        for plot_params in plot_param_list:
            if plot_params["model_type"] == "ModelMaterialStrength":
                self.ptw_prediction_plot_single(**plot_params, pdf=pdf)
            else:
                if len(plot_params["actual_x"]) == 1:
                    pass
                else:
                    self.ptw_prediction_plot_single(**plot_params, pdf=pdf)
        pdf.close()
        return

    def ptw_prediction_plots_cluster(self, path):
        sel = np.arange(20000, self.setup.nmcmc, 10)  # need to script this in
        pred_theta_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]
        pred_theta0_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]
        pred_thetap_raw = [
            np.empty([sel.shape[0], self.setup.ys[i].shape[0]])
            for i in range(self.setup.nexp)
        ]

        thetas = self.out.theta[sel, 0]
        deltas = [self.out.delta[i][sel] for i in range(self.setup.nexp)]
        nclustmax = (
            max(self.out.delta[i].max() for i in range(self.setup.nexp)) + 1
        )
        dcounts = np.zeros((sel.shape[0], nclustmax))
        for it, s in enumerate(sel):
            for i in range(self.setup.nexp):
                dcounts[it] += np.bincount(
                    self.out.delta[i][s, 0], minlength=nclustmax
                )
        etas = self.out.eta[sel, 0]
        nclust = (dcounts > 0).sum(axis=1)
        prob = dcounts + (dcounts == 0) * (
            etas / (nclustmax - nclust + 1e-9)
        ).reshape(-1, 1)
        prob[:] /= prob.sum(axis=1).reshape(-1, 1)
        cumprob = np.cumsum(prob, axis=1)
        unif = uniform(size=(sel.shape[0], 1))
        dnew = (unif > cumprob).sum(axis=1)
        theta_parent = thetas[np.arange(sel.shape[0]), dnew]
        for i in range(self.setup.nexp):
            for j in range(sel.shape[0]):
                pred_theta_raw[i][j] = self.setup.models[i].eval(
                    impala.tran_unif(
                        self.out.theta_hist[i][sel[j], 0],
                        self.setup.bounds_mat,
                        self.setup.bounds.keys(),
                    ),
                )
                pred_theta0_raw[i][j] = self.setup.models[i].eval(
                    impala.tran_unif(
                        np.repeat(
                            self.out.theta0[sel[j], 0].reshape(1, -1),
                            self.setup.ns2[i],
                            axis=0,
                        ),
                        self.setup.bounds_mat,
                        self.setup.bounds.keys(),
                    ),
                )
                pred_thetap_raw[i][j] = self.setup.models[i].eval(
                    impala.tran_unif(
                        np.repeat(
                            theta_parent[j].reshape(1, -1),
                            self.setup.ns2[i],
                            axis=0,
                        ),
                        self.setup.bounds_mat,
                        self.setup.bounds.keys(),
                    ),
                )

        real_strain = []
        real_stress = []
        pred_theta = []
        pred_theta0 = []
        pred_thetap = []
        pred_theta_quant_lb = []
        pred_theta_quant_ub = []
        pred_theta0_quant_lb = []
        pred_theta0_quant_ub = []
        pred_thetap_quant_lb = []
        pred_thetap_quant_ub = []
        edots = []
        temps = []

        for i in range(self.setup.nexp):
            for j in range(self.setup.ntheta[i]):
                pred_theta.append(
                    pred_theta_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_theta_quant_lb.append(
                    np.quantile(pred_theta[-1], 0.025, 0)
                )
                pred_theta_quant_ub.append(
                    np.quantile(pred_theta[-1], 0.975, 0)
                )
                pred_theta0.append(
                    pred_theta0_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_theta0_quant_lb.append(
                    np.quantile(pred_theta0[-1], 0.025, 0)
                )
                pred_theta0_quant_ub.append(
                    np.quantile(pred_theta0[-1], 0.975, 0)
                )
                pred_thetap.append(
                    pred_thetap_raw[i].T[self.setup.s2_ind[i] == j].T
                )
                pred_thetap_quant_lb.append(
                    np.quantile(pred_thetap[-1], 0.025, 0)
                )
                pred_thetap_quant_ub.append(
                    np.quantile(pred_thetap[-1], 0.975, 0)
                )
                edots.append(self.setup.models[i].edots[j])
                temps.append(self.setup.models[i].temps[j])
                real_strain.append(
                    self.setup.models[i].meas_strain_histories[j]
                )
                real_stress.append(self.setup.ys[i][self.setup.s2_ind[i] == j])

        rows = zip(
            real_strain,
            real_stress,
            pred_theta_quant_lb,
            pred_theta_quant_ub,
            pred_theta0_quant_lb,
            pred_theta0_quant_ub,
            pred_thetap_quant_lb,
            pred_thetap_quant_ub,
            edots,
            temps,
        )
        keys = [
            "actual_x",
            "actual_y",
            "pred_y_lb",
            "pred_y_ub",
            "hier_y_lb",
            "hier_y_ub",
            "pnew_y_lb",
            "pnew_y_ub",
            "edot",
            "temp",
        ]
        plot_param_list = [dict(zip(keys, row)) for row in rows]

        pdf = PdfPages(path)
        for plot_params in plot_param_list:
            self.ptw_prediction_plot_single(**plot_params, pdf=pdf)
        pdf.close()
        return

    def ptw_prediction_plots(self, path):
        """PTW Prediction Plots against model"""
        if type(self.out) is impala.OutCalibPool:
            return self.ptw_prediction_plots_pool(path)
        elif type(self.out) is impala.OutCalibHier:
            return self.ptw_prediction_plots_hier(path)
        elif type(self.out) is impala.OutCalibClust:
            return self.ptw_prediction_plots_cluster(path)
        else:
            raise ValueError("Improper out type")
        return

    @staticmethod
    def kde_contour(x1, x2, percentile):
        density = gaussian_kde([x1, x2], bw_method="silverman")
        X, Y = np.mgrid[min(x1) : max(x1) : 100j, min(x2) : max(x2) : 100j]
        positions = np.vstack([X.ravel(), Y.ravel()])
        Z = density(positions)
        Z /= Z.sum()
        t = np.linspace(0, Z.max(), 1000)
        integral = ((Z >= t[:, None, None]) * Z).sum(axis=(1, 2))
        f = interp1d(integral, t)
        t_contours = f(np.array([percentile]))
        return {"X": X, "Y": Y, "Z": Z.reshape([100, 100]), "conts": t_contours}

    def pairwise_theta_plot_hier(self, path, uu, highlight=None):
        """Pairwise Theta scatterplot"""
        if highlight is None:
            highlight = [
                range(self.setup.ntheta[k]) for k in range(self.setup.nexp)
            ]
        sel = uu  # np.arange(20000, self.setup.nmcmc, 10)
        theta_parent = impala.chol_sample_1per_constraints(
            self.out.theta0[sel, 0],
            self.out.Sigma0[sel, 0],
            self.setup.checkConstraints,
            self.setup.bounds_mat,
            self.setup.bounds.keys(),
            self.setup.bounds,
            self.setup.constants,
        )
        theta_names = list(self.setup.bounds.keys())
        theta0_unst = impala.unnormalize(
            self.out.theta0[sel, 0, :], self.setup.bounds_mat
        )
        theta_parent_unst = impala.unnormalize(
            theta_parent, self.setup.bounds_mat
        )
        theta_unst = [
            self.out.theta[k][sel, 0, :, :] for k in range(self.setup.nexp)
        ]
        for k in range(self.setup.nexp):
            for s in range(self.setup.ntheta[k]):
                theta_unst[k][:, s, :] = impala.unnormalize(
                    self.out.theta[k][sel, 0, s, :], self.setup.bounds_mat
                )
        plt.figure(figsize=(15, 15))
        for i in range(self.setup.p):
            for j in range(self.setup.p):
                if i == j:
                    plt.subplot2grid((self.setup.p, self.setup.p), (i, j))
                    for k in range(self.setup.nexp):
                        # for s in range(self.setup.ntheta[k]):
                        if highlight[k] is not None:
                            for s in highlight[k]:
                                sns.kdeplot(
                                    theta_unst[k][:, s, i], color="lightgreen"
                                )
                    sns.kdeplot(theta0_unst[:, i], color="blue")
                    sns.kdeplot(theta_parent_unst[:, i], color="grey")
                    # plt.xlim(0,1)

                    plt.xlim(
                        self.setup.bounds_mat[i, 0], self.setup.bounds_mat[i, 1]
                    )

                    ax = plt.gca()
                    ax.axes.yaxis.set_visible(False)
                    plt.xlabel(theta_names[i])
                    ax.tick_params(axis="x", which="major", labelsize=8)
                    plt.setp(
                        ax.get_xticklabels(),
                        rotation=30,
                        horizontalalignment="right",
                    )
                elif i < j:
                    plt.subplot2grid((self.setup.p, self.setup.p), (i, j))
                    for k in range(self.setup.nexp):
                        # for s in range(self.setup.ntheta[k]):
                        if highlight[k] is not None:
                            for s in highlight[k]:
                                contour = self.kde_contour(
                                    theta_unst[k][:, s, j],
                                    theta_unst[k][:, s, i],
                                    0.9,
                                )
                                plt.contour(
                                    contour["X"],
                                    contour["Y"],
                                    contour["Z"],
                                    contour["conts"],
                                    colors="lightgreen",
                                )
                    contour = self.kde_contour(
                        theta0_unst[:, j], theta0_unst[:, i], 0.9
                    )
                    plt.contour(
                        contour["X"],
                        contour["Y"],
                        contour["Z"],
                        contour["conts"],
                        colors="blue",
                    )

                    contour = self.kde_contour(
                        theta_parent_unst[:, j], theta_parent_unst[:, i], 0.9
                    )
                    plt.contour(
                        contour["X"],
                        contour["Y"],
                        contour["Z"],
                        contour["conts"],
                        colors="grey",
                    )
                    # plt.xlim(0,1)
                    # plt.ylim(0,1)

                    plt.xlim(
                        self.setup.bounds_mat[j, 0], self.setup.bounds_mat[j, 1]
                    )
                    plt.ylim(
                        self.setup.bounds_mat[i, 0], self.setup.bounds_mat[i, 1]
                    )

                    ax = plt.gca()
                    ax.axes.xaxis.set_visible(False)
                    ax.axes.yaxis.set_visible(False)
                else:
                    pass
        plt.subplots_adjust(wspace=0.05, hspace=0.05)
        plt.subplot2grid((self.setup.p, self.setup.p), (2, 0))
        colors = ["lightgreen", "blue", "grey"]
        lines = [Line2D([0], [0], color=c, linewidth=2) for c in colors]
        labels = [r"$\theta_i$", r"$\theta_0$", r"$\theta^*$"]
        plt.legend(lines, labels)
        plt.axis("off")
        if path:
            plt.savefig(path, bbox_inches="tight")
        else:
            plt.show()
        return

    def pairwise_theta_plot_pool(self, path, uu):
        sel = uu  # np.arange(20000, self.setup.nmcmc, 10)
        plt.figure(figsize=(15, 15))
        for i in range(self.setup.p):
            for j in range(self.setup.p):
                if i == j:
                    plt.subplot2grid((self.setup.p, self.setup.p), (i, j))
                    sns.distplot(
                        impala.invprobit(self.out.theta[sel, 0, i]),
                        hist=False,
                        kde=True,
                        color="blue",
                    )
                    plt.xlim(0, 1)
                    ax = plt.gca()
                    ax.axes.yaxis.set_visible(False)
                    ax.tick_params(axis="x", which="major", labelsize=8)
                    plt.setp(
                        ax.get_xticklabels(),
                        rotation=30,
                        horizontalalignment="right",
                    )
                elif i < j:
                    plt.subplot2grid((self.setup.p, self.setup.p), (i, j))
                    contour = self.kde_contour(
                        impala.invprobit(self.out.theta[sel, 0, j]),
                        impala.invprobit(self.out.theta[sel, 0, i]),
                        0.9,
                    )
                    plt.contour(
                        contour["X"],
                        contour["Y"],
                        contour["Z"],
                        contour["conts"],
                        colors="blue",
                    )
                    plt.xlim(0, 1)
                    plt.ylim(0, 1)
                    ax = plt.gca()
                    ax.axes.xaxis.set_visible(False)
                    ax.axes.yaxis.set_visible(False)
                else:
                    pass
        plt.subplots_adjust(wspace=0.05, hspace=0.05)
        plt.axis("off")
        if path:
            plt.savefig(path, bbox_inches="tight")
        else:
            plt.show()
        return

    def pairwise_theta_plot_cluster(self, path=None):
        sel = np.arange(20000, self.setup.nmcmc, 10)
        thetas = self.out.theta[sel, 0]
        deltas = [self.out.delta[i][sel] for i in range(self.setup.nexp)]
        nclustmax = (
            max(self.out.delta[i].max() for i in range(self.setup.nexp)) + 1
        )
        dcounts = np.zeros((sel.shape[0], nclustmax))
        for it, s in enumerate(sel):
            for i in range(self.setup.nexp):
                dcounts[it] += np.bincount(
                    self.out.delta[i][s, 0], minlength=nclustmax
                )
        etas = self.out.eta[sel, 0]
        nclust = (dcounts > 0).sum(axis=1)
        prob = dcounts + (dcounts == 0) * (
            etas / (nclustmax - nclust + 1e-9)
        ).reshape(-1, 1)
        prob /= prob.sum(axis=1).reshape(-1, 1)
        cumprob = np.cumsum(prob, axis=1)
        unif = uniform(size=(sel.shape[0], 1))
        dnew = (unif > cumprob).sum(axis=1)
        theta_parent = thetas[np.arange(sel.shape[0]), dnew]
        plt.figure(figsize=(15, 15))
        for i in range(self.setup.p):
            for j in range(self.setup.p):
                if i == j:
                    plt.subplot2grid((self.setup.p, self.setup.p), (i, j))
                    for k in range(self.setup.nexp):
                        for s in range(self.setup.ns2[k]):
                            sns.distplot(
                                impala.invprobit(
                                    self.out.theta_hist[k][sel, 0, s, i]
                                ),
                                hist=False,
                                kde=True,
                                color="lightgreen",
                            )
                    sns.distplot(
                        impala.invprobit(self.out.theta0[sel, 0, i]),
                        hist=False,
                        kde=True,
                        color="blue",
                    )
                    sns.distplot(
                        impala.invprobit(theta_parent[:, i]),
                        hist=False,
                        kde=True,
                        color="grey",
                    )
                    plt.xlim(0, 1)
                    ax = plt.gca()
                    ax.axes.yaxis.set_visible(False)
                    ax.tick_params(axis="x", which="major", labelsize=8)
                    plt.setp(
                        ax.get_xticklabels(),
                        rotation=30,
                        horizontalalignment="right",
                    )
                elif i < j:
                    plt.subplot2grid((self.setup.p, self.setup.p), (i, j))
                    for k in range(self.setup.nexp):
                        for s in range(self.setup.ns2[k]):
                            contour = self.kde_contour(
                                impala.invprobit(
                                    self.out.theta_hist[k][sel, 0, s, j]
                                ),
                                impala.invprobit(
                                    self.out.theta_hist[k][sel, 0, s, i]
                                ),
                                0.9,
                            )
                            plt.contour(
                                contour["X"],
                                contour["Y"],
                                contour["Z"],
                                contour["conts"],
                                colors="lightgreen",
                            )
                    contour = self.kde_contour(
                        impala.invprobit(self.out.theta0[sel, 0, j]),
                        impala.invprobit(self.out.theta0[sel, 0, i]),
                        0.9,
                    )
                    plt.contour(
                        contour["X"],
                        contour["Y"],
                        contour["Z"],
                        contour["conts"],
                        colors="blue",
                    )

                    contour = self.kde_contour(
                        impala.invprobit(theta_parent[:, j]),
                        impala.invprobit(theta_parent[:, i]),
                        0.9,
                    )
                    plt.contour(
                        contour["X"],
                        contour["Y"],
                        contour["Z"],
                        contour["conts"],
                        colors="grey",
                    )
                    plt.xlim(0, 1)
                    plt.ylim(0, 1)
                    ax = plt.gca()
                    ax.axes.xaxis.set_visible(False)
                    ax.axes.yaxis.set_visible(False)
                else:
                    pass
        plt.subplots_adjust(wspace=0.05, hspace=0.05)
        plt.subplot2grid((self.setup.p, self.setup.p), (2, 0))
        colors = ["lightgreen", "blue", "grey"]
        lines = [Line2D([0], [0], color=c, linewidth=2) for c in colors]
        labels = [r"$\theta_i$", r"$\theta_0$", r"$\theta^*$"]
        plt.legend(lines, labels)
        plt.axis("off")
        try:
            plt.subplot2grid((self.setup.p, self.setup.p), (4, 0))
            sns.distplot(nclust, kde=True, color="blue")
            plt.xlim(0, nclustmax)
        except IndexError:
            pass
        if path:
            plt.savefig(path, bbox_inches="tight")
        else:
            plt.show()
        pass

    def pairwise_theta_plot(self, path=None):
        if type(self.out) is impala.OutCalibPool:
            return self.pairwise_theta_plot_pool(path)
        elif type(self.out) is impala.OutCalibHier:
            return self.pairwise_theta_plot_hier(path)
        elif type(self.out) is impala.OutCalibClust:
            return self.pairwise_theta_plot_cluster(path)
        else:
            raise ValueError("Improper out type")
        return

    @staticmethod
    def cluster_matrix(delta_list, ns2, nclustmax, nburn=20000, nthin=10):
        # subset delta to post burn-in
        delta_relist = [d[nburn::nthin] for d in delta_list]
        # Declare constants
        nsamp = delta_relist[0].shape[0]
        nexp = len(delta_relist)
        # create a combined delta array (for all experiments/vectorized experiments)
        # Boolean array, so (True iff member of cluster)
        breaks = np.hstack((0, np.cumsum(ns2)))
        bounds = [
            (breaks[i], breaks[i + 1]) for i in range(breaks.shape[0] - 1)
        ]
        delta_mat = np.empty(
            (delta_relist[0].shape[0], breaks[-1], nclustmax), dtype=bool
        )
        # Fill the combined delta array
        for i in range(breaks.shape[0] - 1):
            delta_mat[:, bounds[i][0] : bounds[i][1]] = delta_relist[i][
                :, 0, :, None
            ] == np.arange(nclustmax)
        # X = matrix of incidence (nsamp x nclust) ->  XX^t matrix of coincidence (nsamp x nsamp)
        # Average over all iterations -> matrix of average coincidence / shared cluster membership
        out = np.einsum("icp,iqp->icq", delta_mat, delta_mat).mean(axis=0)
        return out, breaks

    def cluster_matrix_plot(self, path=None, **kwargs):
        cmat, breaks = self.cluster_matrix(
            self.out.delta, self.setup.ns2, self.out.nclustmax, **kwargs
        )
        plt.matshow(cmat)
        if breaks.shape[0] > 1:
            for breakpoint in breaks[1:-1] - 0.5:
                plt.axhline(breakpoint, color="red", linestyle="--")
                plt.axvline(breakpoint, color="green", linestyle="--")
        plt.legend()
        if path:
            plt.savefig(path, bbox_inches="tight")
        else:
            plt.show()
        return

    pass


# EOF
