# plots.py

# Imports
import os
import re
import math
import h5py
import glob
import logging
import hashlib
import operator
from textwrap import wrap
from functools import reduce
from datetime import datetime
from multiprocessing import Pool
from itertools import islice, product
from collections import Counter, OrderedDict, defaultdict
from typing import Iterable, DefaultDict, Dict, List, Tuple, Optional, Union, Callable
import statsmodels.api as sm

import numpy as np
import pandas as pd
from tensorflow.keras.optimizers.schedules import LearningRateSchedule

import matplotlib

from ml4h.tensor_generators import _sample_csv_to_set

matplotlib.use("Agg")  # Need this to write images from the GSA servers.  Order matters:
import matplotlib.pyplot as plt  # First import matplotlib, then use Agg, then import plt
from matplotlib.ticker import NullFormatter
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib.ticker import AutoMinorLocator, MultipleLocator

from sklearn import manifold
from sklearn.decomposition import PCA
from sklearn.metrics import (
    roc_curve,
    auc,
    precision_recall_curve,
    average_precision_score,
)
from sklearn.metrics import (
    brier_score_loss,
    precision_score,
    recall_score,
    f1_score,
    roc_auc_score,
)
from sklearn.calibration import calibration_curve

import seaborn as sns
from biosppy.signals import ecg
from scipy.ndimage.filters import gaussian_filter
from scipy import stats

import ml4h.tensormap.ukb.ecg
import ml4h.tensormap.mgb.ecg
from ml4h.tensormap.mgb.dynamic import make_waveform_maps
from ml4h.TensorMap import TensorMap
from ml4h.metrics import concordance_index, concordance_index_censored, coefficient_of_determination
from ml4h.defines import (
    IMAGE_EXT,
    JOIN_CHAR,
    PDF_EXT,
    TENSOR_EXT,
    ECG_REST_LEADS,
    ECG_REST_MEDIAN_LEADS,
    PARTNERS_DATETIME_FORMAT,
    PARTNERS_DATE_FORMAT,
)

RECALL_LABEL = "Recall | Sensitivity | True Positive Rate | TP/(TP+FN)"
FALLOUT_LABEL = "Fallout | 1 - Specificity | False Positive Rate | FP/(FP+TN)"
PRECISION_LABEL = "Precision | Positive Predictive Value | TP/(TP+FP)"

SUBPLOT_SIZE = 7

COLOR_ARRAY = [
    "tan",
    "indigo",
    "cyan",
    "pink",
    "purple",
    "blue",
    "chartreuse",
    "deepskyblue",
    "green",
    "salmon",
    "aqua",
    "magenta",
    "aquamarine",
    "red",
    "coral",
    "tomato",
    "grey",
    "black",
    "maroon",
    "hotpink",
    "steelblue",
    "orange",
    "papayawhip",
    "wheat",
    "chocolate",
    "darkkhaki",
    "gold",
    "orange",
    "crimson",
    "slategray",
    "violet",
    "cadetblue",
    "midnightblue",
    "darkorchid",
    "paleturquoise",
    "plum",
    "lime",
    "teal",
    "peru",
    "silver",
    "darkgreen",
    "rosybrown",
    "firebrick",
    "saddlebrown",
    "dodgerblue",
    "orangered",
]

ECG_REST_PLOT_DEFAULT_YRANGE = 3.0
ECG_REST_PLOT_MAX_YRANGE = 10.0
ECG_REST_PLOT_LEADS = [
    ["strip_I", "strip_aVR", "strip_V1", "strip_V4"],
    ["strip_II", "strip_aVL", "strip_V2", "strip_V5"],
    ["strip_III", "strip_aVF", "strip_V3", "strip_V6"],
]
ECG_REST_PLOT_MEDIAN_LEADS = [
    ["median_I", "median_aVR", "median_V1", "median_V4"],
    ["median_II", "median_aVL", "median_V2", "median_V5"],
    ["median_III", "median_aVF", "median_V3", "median_V6"],
]
ECG_REST_PLOT_AMP_LEADS = [
    [0, 3, 6, 9],
    [1, 4, 7, 10],
    [2, 5, 8, 11],
]


def evaluate_predictions(
    tm: TensorMap,
    y_predictions: np.ndarray,
    y_truth: np.ndarray,
    protected: Dict[TensorMap, np.ndarray],
    title: str,
    folder: str,
    test_paths: List[str] = None,
    max_melt: int = 10000000,
    rocs: List[Tuple[np.ndarray, np.ndarray, Dict[str, int]]] = [],
    scatters: List[Tuple[np.ndarray, np.ndarray, str, List[str]]] = [],
    dpi: int = 300,
    width: int = 7,
    height: int = 7,
) -> Dict[str, float]:
    """Evaluate predictions for a given TensorMap with truth data and plot the appropriate metrics.
    Accumulates data in the rocs and scatters lists to facilitate subplotting.

    :param tm: The TensorMap predictions to evaluate
    :param y_predictions: The predictions
    :param y_truth: The truth
    :param title: A title for the plots
    :param folder: The folder to save the plots at
    :param test_paths: The tensor paths that were predicted
    :param max_melt: For multi-dimensional prediction the maximum number of prediction to allow in the flattened array
    :param protected: TensorMaps and tensors sensitive to bias
    :param rocs: (output) List of Tuples which are inputs for ROC curve plotting to allow subplotting downstream
    :param scatters: (output) List of Tuples which are inputs for scatter plots to allow subplotting downstream
    :param dpi: Dots per inch
    :param width: Figure width in inches
    :param height: Figure height in inches
    :return: Dictionary of performance metrics with string keys for labels and float values
    """
    print("y_predictions",y_predictions.shape)
    print("y_truth",y_truth.shape)
    performance_metrics = {}
    if tm.is_categorical() and tm.axes() == 1:
        logging.info(
            f"For tm:{tm.name} with channel map:{tm.channel_map} examples:{y_predictions.shape[0]}",
        )
        logging.info(
            f"\nSum Truth:{np.sum(y_truth, axis=0)} \nSum pred :{np.sum(y_predictions, axis=0)}",
        )
        plot_precision_recall_per_class(
            y_predictions, y_truth, tm.channel_map, title, folder, dpi, width, height,
        )
        plot_prediction_calibration(
            y_predictions, y_truth, tm.channel_map, title, folder, 10,  dpi, width, height,
        )
        performance_metrics.update(
            subplot_roc_per_class(
                y_predictions, y_truth, tm.channel_map, protected, title, folder, dpi, width, height,
            ),
        )
        rocs.append((y_predictions, y_truth, tm.channel_map))
    elif tm.is_categorical() and tm.axes() == 2:
        melt_shape = (
            y_predictions.shape[0] * y_predictions.shape[1],
            y_predictions.shape[2],
        )
        idx = np.random.choice(
            np.arange(melt_shape[0]), min(melt_shape[0], max_melt), replace=False,
        )
        y_predictions = y_predictions.reshape(melt_shape)[idx]
        y_truth = y_truth.reshape(melt_shape)[idx]
        performance_metrics.update(
            subplot_roc_per_class(
                y_predictions, y_truth, tm.channel_map, protected, title, folder, dpi, width, height,
            ),
        )
        performance_metrics.update(
            plot_precision_recall_per_class(
                y_predictions, y_truth, tm.channel_map, title, folder, dpi, width, height,
            ),
        )
        plot_prediction_calibration(
            y_predictions, y_truth, tm.channel_map, title, folder, 10, dpi, width, height,
        )
        rocs.append((y_predictions, y_truth, tm.channel_map))
    elif tm.is_categorical() and tm.axes() == 3:
        melt_shape = (
            y_predictions.shape[0] * y_predictions.shape[1] * y_predictions.shape[2],
            y_predictions.shape[3],
        )
        idx = np.random.choice(
            np.arange(melt_shape[0]), min(melt_shape[0], max_melt), replace=False,
        )
        y_predictions = y_predictions.reshape(melt_shape)[idx]
        y_truth = y_truth.reshape(melt_shape)[idx]
        performance_metrics.update(
            subplot_roc_per_class(
                y_predictions, y_truth, tm.channel_map, protected, title, folder, dpi, width, height,
            ),
        )
        performance_metrics.update(
            plot_precision_recall_per_class(
                y_predictions, y_truth, tm.channel_map, title, folder, dpi, width, height,
            ),
        )
        plot_prediction_calibration(
            y_predictions, y_truth, tm.channel_map, title, folder, 10, dpi, width, height,
        )
        rocs.append((y_predictions, y_truth, tm.channel_map))
    elif tm.is_categorical() and tm.axes() == 4:
        melt_shape = (
            y_predictions.shape[0]
            * y_predictions.shape[1]
            * y_predictions.shape[2]
            * y_predictions.shape[3],
            y_predictions.shape[4],
        )
        idx = np.random.choice(
            np.arange(melt_shape[0]), min(melt_shape[0], max_melt), replace=False,
        )
        y_predictions = y_predictions.reshape(melt_shape)[idx]
        y_truth = y_truth.reshape(melt_shape)[idx]
        performance_metrics.update(
            subplot_roc_per_class(
                y_predictions, y_truth, tm.channel_map, protected, title, folder, dpi, width, height,
            ),
        )
        performance_metrics.update(
            plot_precision_recall_per_class(
                y_predictions, y_truth, tm.channel_map, title, folder, dpi, width, height,
            ),
        )
        plot_prediction_calibration(
            y_predictions, y_truth, tm.channel_map, title, folder, 10, dpi, width, height,
        )
        rocs.append((y_predictions, y_truth, tm.channel_map))
    elif tm.is_survival_curve():
        performance_metrics.update(
            plot_survival(
                y_predictions, y_truth, title, days_window=tm.days_window,
                prefix=folder, dpi=dpi, width=width, height=height,
            ),
        )
        plot_survival_curves(
            y_predictions, y_truth, title, days_window=tm.days_window,
            prefix=folder, paths=test_paths, dpi=dpi, width=width, height=height,
        )
        time_steps = tm.shape[-1] // 2
        days_per_step = 1 + tm.days_window // time_steps
        predictions_at_end = (
            1 - np.cumprod(y_predictions[:, :time_steps], axis=-1)[:, -1]
        )
        events_at_end = np.cumsum(y_truth[:, time_steps:], axis=-1)[:, -1]
        follow_up = np.cumsum(y_truth[:, :time_steps], axis=-1)[:, -1] * days_per_step
        logging.info(
            f"Event {events_at_end.shape}, preds:{predictions_at_end.shape} end {events_at_end[:, np.newaxis].shape}",
        )
        calibration_title = f"{title}_at_{tm.days_window}_days"
        plot_prediction_calibration(
            predictions_at_end[:, np.newaxis], events_at_end[:, np.newaxis], {tm.name: 0},
            calibration_title, folder, dpi, width, height,
        )
        plot_survivorship(
            events_at_end, follow_up, predictions_at_end, tm.name,
            folder, tm.days_window, dpi, width, height,
        )
    elif tm.is_time_to_event():
        c_index = concordance_index_censored(y_truth[:, 0] == 1.0, y_truth[:, 1], y_predictions[:, 0])
        concordance_return_values = [
            "C-Index",
            "Concordant Pairs",
            "Discordant Pairs",
            "Tied Predicted Risk",
            "Tied Event Time",
        ]
        logging.info(
            f"{[f'{label}: {value:.3f}' for label, value in zip(concordance_return_values, c_index)]}",
        )
        new_title = f"{title}_C_Index_{c_index[0]:0.3f}"
        performance_metrics.update(
            subplot_roc_per_class(
                y_predictions, y_truth[:, 0, np.newaxis], {f"{new_title}_vs_ROC": 0}, protected,
                new_title, folder, dpi, width, height,
            ),
        )
        calibration_title = f"{title}_at_{tm.days_window}_days"
        plot_prediction_calibration(
            y_predictions,
            y_truth[:, 0, np.newaxis],
            {tm.name: 0},
            calibration_title,
            folder,
            dpi,
            width,
            height,
        )
        plot_survivorship(
            y_truth[:, 0],
            y_truth[:, 1],
            y_predictions[:, 0],
            tm.name,
            folder,
            tm.days_window,
            dpi,
            width,
            height,
        )
    elif tm.is_language():
        prediction_1hot = y_predictions.reshape((y_predictions.shape[0]*y_predictions.shape[1], y_predictions.shape[2]))
        truth_1hot = make_one_hot(y_truth.flatten()[:max_melt], len(tm.channel_map))
        logging.info(f"shapes are: {prediction_1hot.shape} {truth_1hot.shape} {y_predictions.shape}, {y_truth.shape}")
        performance_metrics.update(
            subplot_roc_per_class(
                prediction_1hot, truth_1hot, tm.channel_map, protected, title, folder, dpi, width, height,
            ),
        )
        performance_metrics.update(
            plot_precision_recall_per_class(
                prediction_1hot, truth_1hot, tm.channel_map, title, folder, dpi, width, height,
            ),
        )
    elif tm.axes() > 1 or tm.is_mesh():
        prediction_flat = tm.rescale(y_predictions).flatten()[:max_melt]
        truth_flat = tm.rescale(y_truth).flatten()[:max_melt]
        protected_repeated = {}
        for ptm in protected:
            repeat_axis = 0 if ptm.shape[0] > 1 else None
            protected_repeated[ptm] = np.repeat(protected[ptm], np.prod(tm.shape), axis=repeat_axis)
            logging.info(f'Protected TM {ptm} had shape {protected[ptm].shape} now has {protected_repeated[ptm].shape} ')
            protected_repeated[ptm] = protected_repeated[ptm][:max_melt]
        if tm.sentinel is not None:
            y_predictions = y_predictions[y_truth != tm.sentinel]
            y_truth = y_truth[y_truth != tm.sentinel]
        plot_reconstruction(tm, y_truth, y_predictions, folder, test_paths)
        if prediction_flat.shape[0] == truth_flat.shape[0]:
            performance_metrics.update(
                subplot_pearson_per_class(
                    prediction_flat,
                    truth_flat,
                    tm.channel_map,
                    protected_repeated,
                    title,
                    prefix=folder,
                    dpi=dpi,
                    width=width,
                    height=height,
                ),
            )
    elif tm.is_continuous():
        if tm.sentinel is not None:
            y_predictions = y_predictions[y_truth != tm.sentinel, np.newaxis]
            y_truth = y_truth[y_truth != tm.sentinel, np.newaxis]
        performance_metrics.update(
            subplot_pearson_per_class(
                tm.rescale(y_predictions),
                tm.rescale(y_truth),
                tm.channel_map,
                protected,
                title,
                folder,
                dpi,
                width,
                height,
            ),
        )
        scatters.append(
            (tm.rescale(y_predictions), tm.rescale(y_truth), title, test_paths),
        )
    else:
        logging.warning(f"No evaluation clause for tensor map {tm.name}")

    if tm.name == "median":
        plot_waves(y_predictions, y_truth, "median_waves_" + title, folder)

    return performance_metrics


def make_one_hot(y, num_labels):
    ohy = np.zeros((y.shape[-1], num_labels))
    for i in range(0, y.shape[-1]):
        ohy[i, int(y[i])] = 1.0
    return ohy


def plot_metric_history(history, training_steps: int, title: str, prefix="./figures/", dpi=300, width=3, height=3):
    row = 0
    col = 0
    total_plots = int(
        math.ceil(len(history.history) / 2),
    )  # divide by 2 because we plot validation and train histories together
    cols = max(2, int(math.ceil(math.sqrt(total_plots))))
    rows = max(2, int(math.ceil(total_plots / cols)))
    f, axes = plt.subplots(
        rows, cols, figsize=(int(cols * width), int(rows * height)), dpi=dpi,
    )

    for k in sorted(history.history.keys()):
        if not k.startswith("val_"):
            if isinstance(history.history[k][0], LearningRateSchedule):
                history.history[k] = [
                    history.history[k][0](i * training_steps)
                    for i in range(len(history.history[k]))
                ]
            axes[row, col].plot(history.history[k])
            k_split = str(k).replace("output_", "").split("_")
            k_title = " ".join(OrderedDict.fromkeys(k_split))
            axes[row, col].set_title(k_title.capitalize())
            axes[row, col].set_xlabel("Epoch")
            if "val_" + k in history.history:
                axes[row, col].plot(history.history["val_" + k])
                labels = ["Train", "Validation"]
            else:
                labels = [k]
            axes[row, col].legend(labels)

            row += 1
            if row == rows:
                row = 0
                col += 1
                if col >= cols:
                    break

    plt.tight_layout()
    now_string = datetime.now().strftime('%Y-%m-%d_%H-%M')
    figure_path = os.path.join(prefix, f'metrics_{now_string}_{title}{IMAGE_EXT}')
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    if 'loss' in history.history:
        logging.info(f'Starting training loss:   {history.history["loss"][0]:0.3f}, Final training loss:   {history.history["loss"][-1]:0.4f}')
    if 'val_loss' in history.history:
        logging.info(
            f'Starting validation loss: {history.history["val_loss"][0]:0.3f}, Final validation loss: {history.history["val_loss"][-1]:0.4f}, '
            f'Minimum validation loss: {min(history.history["val_loss"]):0.4f}',
        )
    logging.info(f"Saved learning curves at:{figure_path}")


def plot_rocs(
    predictions: Dict[str, np.ndarray],
    truth: np.ndarray,
    labels: Dict[str, int],
    title: str,
    prefix: str = "./figures/",
    dpi: int = 300,
    width: int = 7,
    height: int = 7,
    line_width = 2,
):
    """Plot Receiver Operating Characteristic (ROC) curves from a dictionary of predictions

    Typically, this function is used to compare several models predictions across multiple labels.
    As a hack to avoid repetitive ROC curves for binary classification label string containing 'no_' are skipped.

    :param predictions: The keys are strings identifying the model the values are numpy arrays
                        The arrays have shape (num_samples, num_classes)
    :param truth: The true classifications of each class, one hot encoded of shape (num_samples, num_classes)
    :param labels: Dictionary mapping strings describing each class to their corresponding index in the arrays
    :param title: The name of this plot
    :param prefix: Optional path prefix where the plot will be saved
    :param dpi: Dots per inch of the figure
    :param width: Width in inches of the figure
    :param height: Height in inches of the figure
    :param line_width: Line width in pixels
    """

    true_sums = np.sum(truth, axis=0)
    plt.figure(figsize=(width, height), dpi=dpi)
    for p in predictions:
        fpr, tpr, roc_auc = get_fpr_tpr_roc_pred(predictions[p], truth, labels)
        for key in labels:
            if "no_" in str(key) and len(labels) == 2:
                continue
            color = _hash_string_to_color(p + str(key))
            label_text = f"{p}_{key} area:{roc_auc[labels[key]]:.3f} n={true_sums[labels[key]]:.0f}"
            plt.plot(
                fpr[labels[key]], tpr[labels[key]], color=color, lw=line_width, label=label_text,
            )
            logging.info(f"ROC Label {label_text}")

    plt.xlim([0.0, 1.0])
    plt.ylim([-0.02, 1.03])
    plt.ylabel(RECALL_LABEL)
    plt.xlabel(FALLOUT_LABEL)
    plt.legend(loc="lower right")
    plt.plot([0, 1], [0, 1], "k:", lw=0.5)
    plt.title(f"ROC {title} n={np.sum(true_sums):.0f}")

    figure_path = os.path.join(prefix, "per_class_roc_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved ROC curve at: {figure_path}")


def plot_prediction_calibrations(
    predictions: Dict[str, np.ndarray],
    truth: np.ndarray,
    labels: Dict[str, int],
    title: str,
    prefix: str = "./figures/",
    n_bins: int = 10,
    dpi: int = 300,
    width: int = 7,
    height: int = 7,
):
    """Plot calibration performance and compute Brier Score.

    Typically this function is used to compare several models predictions across multiple labels.

    :param predictions: The keys are strings identifying the model the values are numpy arrays
                        The arrays have shape (num_samples, num_classes)
    :param truth: The true classifications of each class, one hot encoded of shape (num_samples, num_classes)
    :param labels: Dictionary mapping strings describing each class to their corresponding index in the arrays
    :param title: The name of this plot
    :param prefix: Optional path prefix where the plot will be saved
    :param n_bins: Number of bins to quantize predictions into
    :param dpi: Dots per inch of the figure
    :param width: Width in inches of the figure
    :param height: Height in inches of the figure
    """
    _, (ax1, ax2) = plt.subplots(3, figsize=(width, height * 2), dpi=dpi)
    #ax1 = plt.subplot2grid((3, 1), (0, 0), rowspan=2)
    #ax2 = plt.subplot2grid((3, 1), (2, 0))

    true_sums = np.sum(truth, axis=0)
    ax1.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated Brier score: 0.0")

    for p in predictions:
        for k in labels:
            color = _hash_string_to_color(p + k)
            brier_score = brier_score_loss(
                truth[..., labels[k]], predictions[p][..., labels[k]], pos_label=1,
            )
            fraction_of_positives, mean_predicted_value = calibration_curve(
                truth[..., labels[k]], predictions[p][..., labels[k]], n_bins=n_bins,
            )
            ax1.plot(
                mean_predicted_value,
                fraction_of_positives,
                "s-",
                label=f"{p} {k} Brier: {brier_score:0.3f}",
                color=color,
            )
            ax2.hist(
                predictions[p][..., labels[k]],
                range=(0, 1),
                bins=10,
                label=f"{p} {k} n={true_sums[labels[k]]:.0f}",
                histtype="step",
                lw=2,
                color=color,
            )
            logging.info(
                f"{p} {k} n={true_sums[labels[k]]:.0f}\nBrier score: {brier_score:0.3f}",
            )
    ax1.set_ylabel("Fraction of positives")
    ax1.set_ylim([-0.05, 1.05])
    ax1.legend(loc="lower right")
    ax1.set_title("Calibrations plots  (reliability curve)")

    ax2.set_xlabel("Mean predicted value")
    ax2.set_ylabel("Count")
    ax2.legend(loc="upper center")
    plt.tight_layout()

    figure_path = os.path.join(prefix, "calibrations_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    logging.info(f"Try to save calibration comparison plot at: {figure_path}")
    plt.savefig(figure_path)


def plot_prediction_calibration(
    prediction: np.ndarray,
    truth: np.ndarray,
    labels: Dict[str, int],
    title: str,
    prefix: str = "./figures/",
    n_bins: int = 10,
    dpi: int = 300,
    width: int = 6,
    height: int = 6,
):
    """Plot calibration performance and compute Brier Score.

    :param prediction: Array of probabilistic predictions with shape (num_samples, num_classes)
    :param truth: The true classifications of each class, one hot encoded of shape (num_samples, num_classes)
    :param labels: Dictionary mapping strings describing each class to their corresponding index in the arrays
    :param title: The name of this plot
    :param prefix: Optional path prefix where the plot will be saved
    :param n_bins: Number of bins to quantize predictions into
    :param dpi: Dots per inch of the figure
    :param width: Width in inches of the figure
    :param height: Height in inches of the figure
    """
    fig, (ax1, ax3, ax2) = plt.subplots(3, 1, figsize=(width*2, height*4), dpi=dpi)

    true_sums = np.sum(truth, axis=0)
    ax1.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated Brier score: 0.0")
    ax3.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated Brier score: 0.0")

    for k in labels:
        y_true = truth[..., labels[k]]
        y_prob = prediction[..., labels[k]]
        color = _hash_string_to_color(k)
        brier_score = brier_score_loss(y_true, prediction[..., labels[k]], pos_label=1)
        fraction_of_positives, mean_predicted_value = calibration_curve(
            y_true, y_prob, n_bins=n_bins,
        )
        ax3.plot(
            mean_predicted_value,
            fraction_of_positives,
            "s-",
            label=f"{k} Brier score: {brier_score:0.3f}",
            color=color,
        )
        ax2.hist(
            y_prob,
            range=(0, 1),
            bins=n_bins,
            label=f"{k} n={true_sums[labels[k]]:.0f}",
            histtype="step",
            lw=2,
            color=color,
        )

        bins = stats.mstats.mquantiles(y_prob, np.arange(0.0, 1.0, 1.0 / n_bins))
        binids = np.digitize(y_prob, bins) - 1

        bin_sums = np.bincount(binids, weights=y_prob, minlength=len(bins))
        bin_true = np.bincount(binids, weights=y_true, minlength=len(bins))
        bin_total = np.bincount(binids, minlength=len(bins))

        nonzero = bin_total != 0
        prob_true = bin_true[nonzero] / bin_total[nonzero]
        prob_pred = bin_sums[nonzero] / bin_total[nonzero]
        ax1.plot(
            prob_pred,
            prob_true,
            "s-",
            label=f"{k} Brier score: {brier_score:0.3f}",
            color=color,
        )
    ax1.set_ylabel("Fraction of positives")
    ax1.set_ylim([-0.05, 1.05])
    ax1.legend(loc="lower right")
    ax1.set_title(f'{title.replace("_", " ")}\nCalibration plot (equally sized bins)')
    ax2.set_xlabel("Mean predicted value")
    ax2.set_ylabel("Count")
    ax2.legend(loc="upper center", ncol=2)
    ax3.set_title("Calibration plot (equally spaced bins)")
    plt.tight_layout()

    figure_path = os.path.join(prefix, "calibrations_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    logging.info(f"Try to save calibrations plot at: {figure_path}")
    plt.savefig(figure_path)


def _pearson_wrapper(prediction, truth):
    return np.corrcoef(prediction.flatten(), truth.flatten())[1, 0]


def bootstrap_confidence_interval(
        prediction, truth, n_boot: int = 1000, bottom: float = 2.5, top: float = 97.5,
        max_n: int = 1000000, metric_fxn: Callable = _pearson_wrapper,
) -> Tuple[float, Tuple[float, float]]:
    n = min(max_n, len(truth))
    sample_idxs = np.random.randint(n, size=(n_boot, n_boot))
    r2s = [metric_fxn(truth[idx], prediction[idx]) for idx in sample_idxs]
    return np.mean(r2s), np.percentile(r2s, [bottom, top])


def plot_scatter(
    prediction, truth, title, prefix="./figures/", paths=None, top_k=3, alpha=0.5,
    bootstrap=True, dpi=300, width=4, height=4,
):
    margin = float((np.max(truth) - np.min(truth)) / 100)
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(width, 2 * height), dpi=dpi)
    ax1.plot(
        [np.min(truth), np.max(truth)], [np.min(truth), np.max(truth)], linewidth=2,
    )
    ax1.plot(
        [np.min(prediction), np.max(prediction)],
        [np.min(prediction), np.max(prediction)],
        linewidth=4,
    )
    pearson = np.corrcoef(prediction.flatten(), truth.flatten())[
        1, 0,
    ]  # corrcoef returns full covariance matrix
    big_r_squared = coefficient_of_determination(truth, prediction)

    if bootstrap:
        pearson, ci = bootstrap_confidence_interval(prediction, truth)
        label = f'Pearson:{pearson:0.4f} $R^2$:{big_r_squared:0.4f}\n 95% Confidence:({ci[0]:0.4f}, {ci[1]:0.4f})'
    else:
        label = f"Pearson:{pearson:0.3f} $r^2$:{pearson * pearson:0.3f} $R^2$:{big_r_squared:0.3f}"
    logging.info(f"{label}")
    ax1.scatter(
        prediction,
        truth,
        label=label,
        marker=".",
        alpha=alpha,
    )
    if paths is not None:
        diff = np.abs(prediction - truth)
        arg_sorted = diff[:, 0].argsort()
        # The path of the best prediction, ie the inlier
        _text_on_plot(
            ax1,
            prediction[arg_sorted[0]] + margin,
            truth[arg_sorted[0]] + margin,
            os.path.basename(paths[arg_sorted[0]]),
        )
        # Plot the paths of the worst predictions ie the outliers
        for idx in arg_sorted[-top_k:]:
            _text_on_plot(
                ax1,
                prediction[idx] + margin,
                truth[idx] + margin,
                os.path.basename(paths[idx]),
            )

    ax1.set_xlabel("Predictions")
    ax1.set_ylabel("Actual")
    ax1.set_title(title)
    ax1.legend(loc="lower right")

    sns.distplot(prediction, label="Predicted", color="r", ax=ax2)
    sns.distplot(truth, label="Truth", color="b", ax=ax2)
    ax2.legend(loc="upper left")

    figure_path = os.path.join(prefix, "scatter_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    logging.info(f"Try to save scatter plot at: {figure_path}")
    plt.savefig(figure_path)
    return {title + "_pearson": pearson}


def plot_scatters(
    predictions, truth, title, prefix="./figures/", paths=None, top_k=3, alpha=0.5, dpi=300, width=4, height=4,
):
    margin = float((np.max(truth) - np.min(truth)) / 100)
    plt.figure(figsize=(width, height), dpi=dpi)
    plt.plot([np.min(truth), np.max(truth)], [np.min(truth), np.max(truth)])
    for k in predictions:
        color = _hash_string_to_color(k)
        pearson = np.corrcoef(predictions[k].flatten(), truth.flatten())[
            1, 0,
        ]  # corrcoef returns full covariance matrix
        r2 = pearson * pearson
        big_r2 = coefficient_of_determination(truth.flatten(), predictions[k].flatten())
        plt.plot(
            [np.min(predictions[k]), np.max(predictions[k])],
            [np.min(predictions[k]), np.max(predictions[k])],
            color=color,
        )
        plt.scatter(
            predictions[k],
            truth,
            color=color,
            label=str(k) + f" Pearson:{pearson:0.3f} r^2:{r2:0.3f} R^2:{big_r2:0.3f}",
            marker=".",
            alpha=alpha,
        )
        if paths is not None:
            diff = np.abs(predictions[k] - truth)
            arg_sorted = diff[:, 0].argsort()
            _text_on_plot(
                plt,
                predictions[k][arg_sorted[0]] + margin,
                truth[arg_sorted[0]] + margin,
                os.path.basename(paths[arg_sorted[0]]),
            )
            for idx in arg_sorted[-top_k:]:
                _text_on_plot(
                    plt,
                    predictions[k][idx] + margin,
                    truth[idx] + margin,
                    os.path.basename(paths[idx]),
                )
    plt.xlabel("Predictions")
    plt.ylabel("Actual")
    plt.title(title)
    plt.legend()

    figure_path = os.path.join(prefix, "scatters_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved scatter plot at: {figure_path}")


def subplot_pearson_per_class(
    prediction: np.ndarray,
    truth: np.ndarray,
    labels: Dict[str, int],
    protected: Dict[TensorMap, np.ndarray],
    title: str,
    prefix: str = "./figures/",
    dpi: int = 300,
    width: int = 6,
    height: int = 6,
) -> Dict[str, float]:
    lw = 2
    alpha = 0.5
    labels_to_areas = {}
    total_plots = len(protected) + 1
    if total_plots == 1:
        return plot_scatter(prediction, truth, title, prefix, dpi=dpi, width=width, height=height)
    cols = max(2, int(math.ceil(math.sqrt(total_plots))))
    rows = max(2, int(math.ceil(total_plots / cols)))
    fig, axes = plt.subplots(
        rows, cols, figsize=(cols * width, rows * height), dpi=dpi,
    )
    _protected_subplots(
        prediction,
        truth,
        protected,
        axes,
        metric=coefficient_of_determination,
        metric_name="$R^2$",
    )

    color = _hash_string_to_color(title)
    axes[-1, -1].plot(
        [np.min(truth), np.max(truth)], [np.min(truth), np.max(truth)], linewidth=2,
    )
    axes[-1, -1].plot(
        [np.min(prediction), np.max(prediction)],
        [np.min(prediction), np.max(prediction)],
        linewidth=4,
    )
    pearson = np.corrcoef(prediction.flatten(), truth.flatten())[1, 0]
    big_r_squared = coefficient_of_determination(truth, prediction)
    label_text = f"Pearson:{pearson:0.3f} r^2:{pearson * pearson:0.3f} R^2:{big_r_squared:0.3f} n={truth.shape[0]:.0f}"
    axes[-1, -1].scatter(
        prediction, truth, color=color, lw=lw, label=label_text, marker=".", alpha=alpha,
    )
    axes[-1, -1].legend(loc="lower right")
    axes[-1, -1].set_title(f"Pearson {title}")

    figure_path = os.path.join(prefix, "per_class_pearson_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path, bbox_inches="tight")
    logging.info(
        f"{label_text} saved at: {figure_path}{f' with {len(protected)} protected TensorMaps.' if len(protected) else '.'}",
    )
    return labels_to_areas


def subplot_scatters(
    scatters: List[Tuple[np.ndarray, np.ndarray, str, Optional[List[str]]]],
    prefix: str = "./figures/",
    top_k: int = 3,
    alpha: float = 0.5,
    dpi: int = 300,
    width: int = 6,
    height: int = 6,
):
    fig, axes = _figure_and_subplot_axes_from_total(len(scatters), dpi=dpi, width=width, height=height)
    for scatter_data, ax in zip(scatters, axes.ravel()):
        prediction, truth, title, paths = scatter_data
        ax.plot([np.min(truth), np.max(truth)], [np.min(truth), np.max(truth)])
        ax.plot(
            [np.min(prediction), np.max(prediction)],
            [np.min(prediction), np.max(prediction)],
        )
        ax.scatter(prediction, truth, marker=".", alpha=alpha)
        margin = float((np.max(truth) - np.min(truth)) / 100)
        if (
            paths is not None
        ):  # If tensor paths are provided we plot the file names of top_k outliers and the #1 inlier
            diff = np.abs(prediction - truth)
            arg_sorted = diff[:, 0].argsort()
            # The path of the best prediction, ie the inlier
            _text_on_plot(
                ax,
                prediction[arg_sorted[0]] + margin,
                truth[arg_sorted[0]] + margin,
                os.path.basename(paths[arg_sorted[0]]),
            )
            # Plot the paths of the worst predictions ie the outliers
            for idx in arg_sorted[-top_k:]:
                _text_on_plot(
                    ax,
                    prediction[idx] + margin,
                    truth[idx] + margin,
                    os.path.basename(paths[idx]),
                )
        ax.set_xlabel("Predictions")
        ax.set_ylabel("Actual")
        ax.set_title(title + "\n")
        pearson = np.corrcoef(prediction.flatten(), truth.flatten())[
            1, 0,
        ]  # corrcoef returns full covariance matrix
        r2 = pearson * pearson
        big_r2 = coefficient_of_determination(truth.flatten(), prediction.flatten())
        ax.text(
            0,
            1,
            f"Pearson:{pearson:0.3f} r^2:{r2:0.3f} R^2:{big_r2:0.3f}",
            verticalalignment="bottom",
            transform=ax.transAxes,
        )

    figure_path = prefix + "scatters_together" + IMAGE_EXT
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved scatters together at: {figure_path}")


def subplot_comparison_scatters(
    scatters: List[Tuple[Dict[str, np.ndarray], np.ndarray, str, Optional[List[str]]]],
    prefix: str = "./figures/",
    top_k: int = 3,
    alpha: float = 0.5,
    dpi: int = 300,
    width: int = 6,
    height: int = 6,
):
    fig, axes = _figure_and_subplot_axes_from_total(len(scatters), dpi=dpi, width=width, height=height)
    for scatter_data, ax in zip(scatters, axes.ravel()):
        predictions, truth, title, paths = scatter_data
        for k in predictions:
            c = _hash_string_to_color(title + k)
            pearson = np.corrcoef(predictions[k].flatten(), truth.flatten())[
                1, 0,
            ]  # corrcoef returns full covariance matrix
            r2 = pearson * pearson
            big_r2 = coefficient_of_determination(
                truth.flatten(), predictions[k].flatten(),
            )
            ax.plot(
                [np.min(predictions[k]), np.max(predictions[k])],
                [np.min(predictions[k]), np.max(predictions[k])],
                color=c,
            )
            ax.scatter(
                predictions[k],
                truth,
                color=c,
                label=f"{k} r:{pearson:0.3f} r^2:{r2:0.3f} R^2:{big_r2:0.3f}",
                marker=".",
                alpha=alpha,
            )
            ax.legend(loc="upper left")
            if (
                paths is not None
            ):  # If tensor paths are provided we plot the file names of top_k outliers and the #1 inlier
                margin = float((np.max(truth) - np.min(truth)) / 100)
                diff = np.abs(predictions[k] - truth)
                arg_sorted = diff[:, 0].argsort()
                _text_on_plot(
                    ax,
                    predictions[k][arg_sorted[0]] + margin,
                    truth[arg_sorted[0]] + margin,
                    os.path.basename(paths[arg_sorted[0]]),
                )
                for idx in arg_sorted[-top_k:]:
                    _text_on_plot(
                        ax,
                        predictions[k][idx] + margin,
                        truth[idx] + margin,
                        os.path.basename(paths[idx]),
                    )
        ax.set_xlabel("Predictions")
        ax.set_ylabel("Actual")
        ax.set_title(title + "\n")

    figure_path = os.path.join(prefix, "scatters_compared_together" + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved scatter comparisons together at: {figure_path}")


def plot_survivorship(
    events: np.ndarray, days_follow_up: np.ndarray, predictions: np.ndarray, title: str,
    prefix: str = './figures/', days_window: int = 1825, dpi: int = 300, width: int = 6, height: int = 6,
):
    """Plot Kaplan-Meier survivorship curves and stratify by median model prediction.
    All input arrays have the same shape: (num_samples,)

    :param events: Array indicating if each sample had an event (1) or not (0) by the end of follow up
    :param days_follow_up: Array with the total days of follow up for each sample
    :param predictions: Array with model predictions of an event before the end of follow up.
    :param title: Title for the plot
    :param prefix: Path prefix where plot will be saved
    :param days_window: Maximum days of follow up
    :param dpi: Dots per inch of the figure
    :param width: Width in inches of the figure
    :param height: Height in inches of the figure
    """
    plt.figure(figsize=(width, height), dpi=dpi)
    days_sorted_index = np.argsort(days_follow_up)
    days_sorted = days_follow_up[days_sorted_index]
    alive_per_step = len(events)

    sick_per_step = 0
    censored = 0
    survivorship = [1.0]
    real_survivorship = [1.0]
    for cur_day, day_index in enumerate(days_sorted_index):
        if days_follow_up[day_index] > days_window:
            break
        sick_per_step += events[day_index]
        censored += 1 - events[day_index]
        alive_per_step -= events[day_index]
        survivorship.append(1 - (sick_per_step / (alive_per_step+sick_per_step)))
        real_survivorship.append(real_survivorship[cur_day] * (1 - (events[day_index] / alive_per_step)))
    plt.plot([0]+days_sorted[:cur_day+1], real_survivorship[:cur_day+1], marker='.', label='Survivorship')
    groups = ['High risk', 'Low risk']
    predicted_alive = {g: len(events) // 2 for g in groups}

    predicted_sick = {g: 0 for g in groups}
    predicted_days = defaultdict(list)
    predicted_survival = defaultdict(list)
    threshold = np.median(predictions)
    for cur_day, day_index in enumerate(days_sorted_index):
        if days_follow_up[day_index] > days_window:
            break
        group = 'High risk' if predictions[day_index] > threshold else 'Low risk'
        predicted_sick[group] += events[day_index]
        predicted_survival[group].append(1 - (predicted_sick[group] / (predicted_alive[group]+predicted_sick[group])))
        predicted_alive[group] -= events[day_index]
        predicted_days[group].append(days_follow_up[day_index])

    for group in groups:
        plt.plot(
            [0]+predicted_days[group], [1]+predicted_survival[group], color='r' if 'High' in group else 'g', marker='o',
            label=f'{group} group had {predicted_sick[group]} events',
        )

    plt.title(
        f'{title}\nEnrolled: {len(events)}, Censored: {censored:.0f}, {100 * (censored / len(events)):2.1f}%, Events: {sick_per_step:.0f}, '
        f'{100 * (sick_per_step / len(events)):2.1f}%\nMax follow up: {days_window} days, {days_window // 365} years.',
    )
    plt.xlabel('Follow up time (days)')
    plt.ylabel('Proportion Surviving')
    plt.legend(loc="lower left")

    figure_path = os.path.join(prefix, f'survivorship_fu_{days_window}_{title}{IMAGE_EXT}')
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    logging.info(f'Try to save survival plot at: {figure_path}')
    plt.savefig(figure_path)
    return {}


def plot_survival(
    prediction: np.ndarray, truth: np.ndarray, title: str, days_window: int,
    prefix: str = './figures/', dpi: int = 300, width: int = 7, height: int = 7,
) -> Dict[str, float]:
    """Plot Kaplan-Meier survivorship and predicted proportion surviving, calculate and return C-Index

    :param prediction: Array with model predictions of an event at each time step, with shape (num_samples, intervals*2).
    :param truth: Array with survival at each time step followed by events, shape is (num_samples, intervals*2)
    :param title: Title for the plot
    :param days_window: Maximum days of follow up
    :param prefix: Path prefix where plot will be saved
    :param dpi: Dots per inch of the figure
    :param width: Width in inches of the figure
    :param height: Height in inches of the figure

    :return: Dictionary mapping metric names to their floating point values
    """
    c_index, concordant, discordant, tied_risk, tied_time = concordance_index(prediction, truth)

    intervals = truth.shape[-1] // 2
    plt.figure(figsize=(width, height), dpi=dpi)

    cumulative_sick = np.cumsum(np.sum(truth[:, intervals:], axis=0))
    cumulative_censored = (truth.shape[0]-np.sum(truth[:, :intervals], axis=0)) - cumulative_sick
    alive_per_step = np.sum(truth[:, :intervals], axis=0)
    sick_per_step = np.sum(truth[:, intervals:], axis=0)
    survivorship = np.cumprod(1 - (sick_per_step / alive_per_step))
    predicted_proportion = np.sum(np.cumprod(prediction[:, :intervals], axis=1), axis=0) / truth.shape[0]

    plt.plot(range(0, days_window, 1 + days_window // intervals), predicted_proportion, marker='o', label=f'Predicted Proportion C-Index:{c_index:0.3f}')
    plt.plot(range(0, days_window, 1 + days_window // intervals), survivorship, marker='o', label='Survivorship')
    plt.xlabel('Follow up time (days)')
    plt.ylabel('Proportion Surviving')
    full_title = f"""{title} C-Index: {c_index:.4f}
        Enrolled: {truth.shape[0]}, Censored: {cumulative_censored[-1]:.0f}, {100 * (cumulative_censored[-1] / truth.shape[0]):2.1f}%,
        Events: {cumulative_sick[-1]:.0f}, {100 * (cumulative_sick[-1] / truth.shape[0]):2.1f}%\nMax follow up: {days_window} days, {days_window // 365} years."""
    plt.title(full_title)
    plt.legend(loc="upper right")
    logging.info(f"Concordant:{concordant} discordant:{discordant} tied_risk:{tied_risk} tied_time:{tied_time}")
    logging.info(full_title)

    figure_path = os.path.join(prefix, f'survivorship_{title}_c_{c_index:.3f}{IMAGE_EXT}')
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    logging.info(f'Survival plot: {figure_path}')
    plt.savefig(figure_path)
    return {'c_index': c_index, 'concordant': concordant, 'discordant': discordant, 'tied_risk': tied_risk, 'tied_time': tied_time}


def plot_survival_curves(
    prediction,
    truth,
    title,
    days_window,
    prefix="./figures/",
    num_curves=30,
    paths=None,
    dpi=300,
    width: int = 6,
    height: int = 6,
):
    intervals = truth.shape[-1] // 2
    plt.figure(figsize=(width * 2, height * 2), dpi=dpi)
    predicted_survivals = np.cumprod(prediction[:, :intervals], axis=1)
    sick = np.sum(truth[:, intervals:], axis=-1)
    censor_periods = np.argmin(truth[:, :intervals], axis=-1)
    x_days = range(0, days_window, 1 + days_window // intervals)
    cur_sick = 0
    cur_healthy = 0
    min_sick = num_curves * 0.1
    for i in range(truth.shape[0]):
        p = os.path.basename(paths[i]).replace(TENSOR_EXT, "")
        if sick[i] == 1:
            sick_period = np.argmax(truth[i, intervals:])
            sick_day = sick_period * (days_window // intervals)
            plt.plot(
                x_days[: sick_period + 2],
                predicted_survivals[i, : sick_period + 2],
                label=f"Failed:{p} p:{predicted_survivals[i, sick_period]:0.2f}",
                color="red",
            )
            plt.text(
                sick_day,
                predicted_survivals[i, sick_period],
                f"Diagnosed day:{sick_day} id:{p}",
            )
            cur_sick += 1
            if cur_sick >= min_sick and i >= num_curves:
                break
        elif censor_periods[i] != 0:  # individual was censored before failure
            plt.plot(
                x_days[: censor_periods[i]],
                predicted_survivals[i, : censor_periods[i]],
                label=f"Censored:{p} p:{predicted_survivals[i, censor_periods[i]]:0.2f}",
                color="blue",
            )
        elif cur_healthy < num_curves:
            plt.plot(
                x_days,
                predicted_survivals[i],
                label=f"Survived:{p} p:{predicted_survivals[i, -1]:0.2f}",
                color="green",
            )
            cur_healthy += 1
    plt.title(title + "\n")
    plt.legend(loc="lower left")
    plt.xlabel("Follow up time (days)")
    plt.ylabel("Survival Curve Prediction")
    figure_path = os.path.join(prefix, "survival_curves_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    logging.info("Try to save survival plot at: {}".format(figure_path))
    plt.savefig(figure_path)
    return {}


def plot_noise(noise):
    samples = 240
    real_weight = 2.0
    real_bias = 0.5
    x = np.linspace(10, 100, samples)
    y1_real = real_weight * x + real_bias
    y2_real = 4.0 * x + 0.8
    y1 = y1_real + (np.random.randn(*x.shape) * noise)
    y2 = y2_real + (np.random.randn(*x.shape) * noise)
    y_ratio = (y2 - y1) / y2
    y_ratio_real = (y2_real - y1_real) / y2_real
    pearson = np.corrcoef(y1.flatten(), y1_real.flatten())[1, 0]
    pearson2 = np.corrcoef(y2.flatten(), y2_real.flatten())[1, 0]
    ratio_pearson = np.corrcoef(y_ratio.flatten(), y_ratio_real.flatten())[1, 0]
    return pearson, pearson2, ratio_pearson


def plot_noisy():
    samples = 140
    p1s = []
    p2s = []
    prats = []
    noises = np.linspace(0.0, 0.01, samples)
    for n in noises:
        p1, p2, prat = plot_noise(n)
        p1s.append(1.0 - p1)
        p2s.append(1.0 - p2)
        prats.append(1.0 - prat)

    plt.figure(figsize=(28, 42))
    matplotlib.rcParams.update({"font.size": 36})
    plt.xlabel("Noise")
    plt.ylabel("Error")
    plt.scatter(noises, p1s, color="cyan", label="p1")
    plt.scatter(noises, p2s, color="green", label="p2")
    plt.scatter(noises, prats, color="red", label="p_ratio")
    plt.legend(loc="lower right")
    plt.savefig("./figures/noise_fxn.png")


def plot_value_counter(categories, counts, title, prefix="./figures/"):
    matplotlib.rcParams.update({"font.size": 14})
    counters = defaultdict(Counter)
    for k in categories:
        parts = k.split(JOIN_CHAR)
        group = parts[0]
        label = parts[1]
        counters[group][label] = counts[k]

    rows = int(math.ceil(math.sqrt(len(counters))))
    fig, axes = plt.subplots(rows, rows, figsize=(28, 24))
    for i, group in enumerate(counters):
        ax = plt.subplot(rows, rows, i + 1)
        ax.set_title(group)
        idxs = np.arange(len(counters[group]))
        ax.barh(idxs, list(counters[group].values()))
        ax.set_yticks(idxs)
        ax.set_yticklabels(list(counters[group].keys()))

    plt.tight_layout()

    figure_path = os.path.join(prefix, "counter_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved counter plot at: {figure_path}")


def plot_histograms(continuous_stats, title, prefix="./figures/", num_bins=50):
    matplotlib.rcParams.update({"font.size": 14})

    rows = int(math.ceil(math.sqrt(len(continuous_stats))))
    fig, axes = plt.subplots(rows, rows, figsize=(28, 24))
    for i, group in enumerate(continuous_stats):
        a = np.array(continuous_stats[group])
        ax = plt.subplot(rows, rows, i + 1)
        ax.set_title(group + "\n Mean:%0.3f STD:%0.3f" % (np.mean(a), np.std(a)))
        ax.hist(continuous_stats[group], bins=num_bins)
    plt.tight_layout()

    figure_path = os.path.join(prefix, "histograms_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved histograms plot at: {figure_path}")


def plot_histograms_in_pdf(
    stats: Dict[str, Dict[str, List[float]]],
    all_samples_count: int,
    output_file_name: str,
    output_folder_path: str = "./figures",
    num_rows: int = 4,
    num_cols: int = 6,
    num_bins: int = 50,
    title_line_width: int = 50,
) -> None:
    """
    Plots histograms of field values given in 'stats' in pdf
    :param stats: field names extracted from hd5 dataset names to list of values, one per sample_instance_arrayidx
    :param all_samples_count: total number of samples fields were drawn from; samples don't necessarily have values for each field
    :param output_file_name: name of output file in pdf
    :param output_folder_path: directory that output file will be written to
    :param num_rows: number of histograms that will be plotted vertically per pdf page
    :param num_cols: number of histograms that will be plotted horizontally per pdf page
    :param num_bins: number of histogram bins
    :param title_line_width: max number of characters that a plot title line will span; longer lines will be wrapped into multiple lines
    :return: None
    """

    def _sorted_chunks(
        d: Dict[str, Dict[str, List[float]]], size: int,
    ) -> Iterable[DefaultDict[str, List[float]]]:
        """
        :param d: dictionary to be chunked
        :param size: size of chunks
        :return: iterator of dictionary chunks with keys alphabetically sorted
        """
        it = iter(dict(sorted(d.items())))
        for i in range(0, len(d), size):
            yield {k: d[k] for k in islice(it, size)}

    subplot_width = 7.4 * num_cols
    subplot_height = 6 * num_rows
    matplotlib.rcParams.update(
        {"font.size": 14, "figure.figsize": (subplot_width, subplot_height)},
    )

    figure_path = os.path.join(output_folder_path, output_file_name + PDF_EXT)
    with PdfPages(figure_path) as pdf:
        for stats_chunk in _sorted_chunks(stats, num_rows * num_cols):
            plt.subplots(num_rows, num_cols)
            for i, field in enumerate(stats_chunk):
                field_values = reduce(operator.concat, stats_chunk[field].values())
                field_sample_count = len(stats_chunk[field].keys())
                missingness = int(
                    100 * (all_samples_count - field_sample_count) / all_samples_count,
                )
                ax = plt.subplot(num_rows, num_cols, i + 1)
                title_text = "\n".join(wrap(field, title_line_width))
                title_stats = "\n".join(
                    wrap(
                        f"Mean:{np.mean(field_values):.2f} STD:{np.std(field_values):.2f} Missing:{missingness}% #Samples:{field_sample_count}",
                        title_line_width,
                    ),
                )
                ax.set_title(title_text + "\n" + title_stats)
                ax.hist(field_values, bins=min(num_bins, len(set(field_values))))
            plt.tight_layout()
            pdf.savefig()

    logging.info(f"Saved histograms plot at: {figure_path}")


def plot_heatmap(
    stats: Dict[str, Dict[str, List[float]]],
    output_file_name: str,
    min_samples: int,
    output_folder_path: str,
) -> None:

    """
    Plot heatmap of correlations between field pairs derived from 'stats'
    :param stats: field names extracted from hd5 dataset names to list of values, one per sample_instance_arrayidx
    :param output_file_name: name of output file in pdf
    :param output_folder_path: directory that output file will be written to
    :param min_samples: calculate correlation coefficient only if both fields have values from that many common samples
    :return: None
    """
    fields = stats.keys()
    num_fields = len(fields)
    field_pairs = product(fields, fields)
    correlations_by_field_pairs: DefaultDict[Tuple[str, str], float] = defaultdict(
        float,
    )
    logging.info(f"There are {int(num_fields * (num_fields - 1) / 2)} field pairs.")
    processed_field_pair_count = 0
    nan_counter = Counter()  # keep track of if we've seen a field have NaNs
    for field1, field2 in field_pairs:
        if field1 not in nan_counter.keys() and field2 not in nan_counter.keys():
            if field1 == field2:
                correlations_by_field_pairs[(field1, field2)] = 1
            elif (field2, field1) in correlations_by_field_pairs:
                correlations_by_field_pairs[
                    (field1, field2)
                ] = correlations_by_field_pairs[(field2, field1)]
            else:
                common_samples = set(stats[field1].keys()).intersection(
                    stats[field2].keys(),
                )
                num_common_samples = len(common_samples)
                processed_field_pair_count += 1
                if processed_field_pair_count % 50000 == 0:
                    logging.debug(
                        f"Processed {processed_field_pair_count} field pairs.",
                    )
                if num_common_samples >= min_samples:
                    field1_values = reduce(
                        operator.concat,
                        [stats[field1][sample] for sample in common_samples],
                    )
                    field2_values = reduce(
                        operator.concat,
                        [stats[field2][sample] for sample in common_samples],
                    )

                    num_field1_nans = len(list(filter(math.isnan, field1_values)))
                    num_field2_nans = len(list(filter(math.isnan, field2_values)))
                    at_least_one_field_has_nans = False
                    if num_field1_nans != 0:
                        nan_counter[field1] = True
                        at_least_one_field_has_nans = True
                    # Also add field2 to the counter if it has NaNs before 'continue'ing so we can skip it sooner in next iterations
                    if num_field2_nans != 0:
                        nan_counter[field2] = True
                        at_least_one_field_has_nans = True
                    if at_least_one_field_has_nans:
                        continue

                    if len(field1_values) == len(field2_values):
                        if len(set(field1_values)) == 1 or len(set(field2_values)) == 1:
                            logging.debug(
                                f"Not calculating correlation for fields {field1} and {field2} because at least one of "
                                f"the fields has all the same values for the {num_common_samples} common samples.",
                            )
                            continue
                        corr = np.corrcoef(field1_values, field2_values)[1, 0]
                        if not math.isnan(corr):
                            correlations_by_field_pairs[(field1, field2)] = corr
                        else:
                            logging.warning(
                                f"Pearson correlation for fields {field1} and {field2} is NaN.",
                            )
                    else:
                        logging.debug(
                            f"Not calculating correlation for fields '{field1}' and '{field2}' "
                            f"because they have different number of values ({len(field1_values)} vs. {len(field2_values)}).",
                        )
        else:
            continue

    logging.info(f"Total number of correlations: {len(correlations_by_field_pairs)}")

    fields_with_nans = nan_counter.keys()
    if len(fields_with_nans) != 0:
        logging.warning(
            f"The {len(fields_with_nans)} fields containing NaNs are: {', '.join(fields_with_nans)}.",
        )

    # correlations_by_field_pairs = dict(random.sample(correlations_by_field_pairs.items(), 15))
    ser = pd.Series(
        list(correlations_by_field_pairs.values()),
        index=pd.MultiIndex.from_tuples(correlations_by_field_pairs.keys()),
    )
    df = ser.unstack()

    # Scale the figure size with the number of fields
    width = height = int(0.23 * num_fields)
    plt.subplots(figsize=(width, height))

    # Generate a custom diverging colourmap
    cmap = sns.diverging_palette(240, 10, n=9, as_cmap=True)

    # Plot
    ax = sns.heatmap(
        df, cmap=cmap, square=True, vmin=-1, vmax=1, cbar_kws={"shrink": 0.8},
    )

    # Custom-colour missing values
    ax.set_facecolor("xkcd:light grey")

    # Scale the colourbar label font size with the number of fields
    cbar = ax.collections[0].colorbar
    cbar.ax.tick_params(labelsize=int(0.27 * num_fields))

    # Save figure
    fig = ax.get_figure()
    heatmap_path = os.path.join(output_folder_path, output_file_name + IMAGE_EXT)
    fig.savefig(heatmap_path)
    logging.info(f"Plotted heatmap ({df.shape[0]}x{df.shape[1]}) at: {heatmap_path}")


def plot_ecg(data, label, prefix="./figures/"):
    lw = 3
    matplotlib.rcParams.update({"font.size": 36})

    rows = int(math.ceil(math.sqrt(len(data))))
    cols = math.ceil(len(data) / rows)
    fig, axes = plt.subplots(rows, cols, figsize=(28, 24))
    for i, k in enumerate(data):
        color = _hash_string_to_color(k)
        ax = plt.subplot(rows, cols, i + 1)
        ax.set_title(k)
        ax.plot(data[k], color=color, lw=lw, label=str(k))
    plt.tight_layout()

    figure_path = os.path.join(prefix, label + "_ecg" + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved ECG plot at: {figure_path}")


def _plot_partners_text(
    data: Dict[str, Union[np.ndarray, str, Dict]], fig: plt.Figure, w: float, h: float,
) -> None:
    # top text
    dt = datetime.strptime(data["datetime"], PARTNERS_DATETIME_FORMAT)
    dob = data["dob"]
    if dob != "":
        dob = datetime.strptime(dob, PARTNERS_DATE_FORMAT)
        dob = f"{dob:%d-%b-%Y}".upper()
    age = -1
    if not np.isnan(data["age"]):
        age = int(data["age"])
    sex = {value: key for key, value in data["sex"].items()}

    fig.text(
        0.17 / w, 8.04 / h, f"{data['lastname']}, {data['firstname']}", weight="bold",
    )
    fig.text(3.05 / w, 8.04 / h, f"ID:{data['patientid']}", weight="bold")
    fig.text(4.56 / w, 8.04 / h, f"{dt:%d-%b-%Y %H:%M:%S}".upper(), weight="bold")
    fig.text(6.05 / w, 8.04 / h, f"{data['sitename']}", weight="bold")

    fig.text(0.17 / w, 7.77 / h, f"{dob} ({age} yr)", weight="bold")  # TODO age units
    fig.text(0.17 / w, 7.63 / h, f"{sex[1]}".title(), weight="bold")
    fig.text(0.17 / w, 7.35 / h, f"Room: ", weight="bold")  # TODO room?
    fig.text(0.17 / w, 7.21 / h, f"Loc: {data['location']}", weight="bold")

    fig.text(2.15 / w, 7.77 / h, f"Vent. rate", weight="bold")
    fig.text(2.15 / w, 7.63 / h, f"PR interval", weight="bold")
    fig.text(2.15 / w, 7.49 / h, f"QRS duration", weight="bold")
    fig.text(2.15 / w, 7.35 / h, f"QT/QTc", weight="bold")
    fig.text(2.15 / w, 7.21 / h, f"P-R-T axes", weight="bold")

    fig.text(3.91 / w, 7.77 / h, f"{int(data['rate_md'])}", weight="bold", ha="right")
    fig.text(3.91 / w, 7.63 / h, f"{int(data['pr_md'])}", weight="bold", ha="right")
    fig.text(3.91 / w, 7.49 / h, f"{int(data['qrs_md'])}", weight="bold", ha="right")
    fig.text(
        3.91 / w,
        7.35 / h,
        f"{int(data['qt_md'])}/{int(data['qtc_md'])}",
        weight="bold",
        ha="right",
    )
    fig.text(
        3.91 / w,
        7.21 / h,
        f"{int(data['paxis_md'])}   {int(data['raxis_md'])}",
        weight="bold",
        ha="right",
    )

    fig.text(4.30 / w, 7.77 / h, f"BPM", weight="bold", ha="right")
    fig.text(4.30 / w, 7.63 / h, f"ms", weight="bold", ha="right")
    fig.text(4.30 / w, 7.49 / h, f"ms", weight="bold", ha="right")
    fig.text(4.30 / w, 7.35 / h, f"ms", weight="bold", ha="right")
    fig.text(4.30 / w, 7.21 / h, f"{int(data['taxis_md'])}", weight="bold", ha="right")

    fig.text(4.75 / w, 7.21 / h, f"{data['read_md']}", wrap=True, weight="bold")

    # TODO tensorize these values from XML
    fig.text(1.28 / w, 6.65 / h, f"Technician: {''}", weight="bold")
    fig.text(1.28 / w, 6.51 / h, f"Test ind: {''}", weight="bold")
    fig.text(4.75 / w, 6.25 / h, f"Referred by: {''}", weight="bold")
    fig.text(7.63 / w, 6.25 / h, f"Electronically Signed By: {''}", weight="bold")


def _plot_partners_full(voltage: Dict[str, np.ndarray], ax: plt.Axes) -> None:
    full_voltage = np.full((12, 2500), np.nan)
    for i, lead in enumerate(voltage):
        full_voltage[i] = voltage[lead]

    # convert voltage to millivolts
    full_voltage /= 1000

    # calculate space between leads
    min_y, max_y = ax.get_ylim()
    y_offset = (max_y - min_y) / len(voltage)

    text_xoffset = 5
    text_yoffset = -0.01

    # plot signal and add labels
    for i, lead in enumerate(voltage):
        this_offset = (len(voltage) - i - 0.5) * y_offset
        ax.plot(full_voltage[i] + this_offset, color="black", linewidth=0.375)
        ax.text(
            0 + text_xoffset,
            this_offset + text_yoffset,
            lead,
            ha="left",
            va="top",
            weight="bold",
        )


def _plot_partners_clinical(voltage: Dict[str, np.ndarray], ax: plt.Axes) -> None:
    # get voltage in clinical chunks
    clinical_voltage = np.full((6, 2500), np.nan)
    halfgap = 5

    clinical_voltage[0][0 : 625 - halfgap] = voltage["I"][0 : 625 - halfgap]
    clinical_voltage[0][625 + halfgap : 1250 - halfgap] = voltage["aVR"][
        625 + halfgap : 1250 - halfgap
    ]
    clinical_voltage[0][1250 + halfgap : 1875 - halfgap] = voltage["V1"][
        1250 + halfgap : 1875 - halfgap
    ]
    clinical_voltage[0][1875 + halfgap : 2500] = voltage["V4"][1875 + halfgap : 2500]

    clinical_voltage[1][0 : 625 - halfgap] = voltage["II"][0 : 625 - halfgap]
    clinical_voltage[1][625 + halfgap : 1250 - halfgap] = voltage["aVL"][
        625 + halfgap : 1250 - halfgap
    ]
    clinical_voltage[1][1250 + halfgap : 1875 - halfgap] = voltage["V2"][
        1250 + halfgap : 1875 - halfgap
    ]
    clinical_voltage[1][1875 + halfgap : 2500] = voltage["V5"][1875 + halfgap : 2500]

    clinical_voltage[2][0 : 625 - halfgap] = voltage["III"][0 : 625 - halfgap]
    clinical_voltage[2][625 + halfgap : 1250 - halfgap] = voltage["aVF"][
        625 + halfgap : 1250 - halfgap
    ]
    clinical_voltage[2][1250 + halfgap : 1875 - halfgap] = voltage["V3"][
        1250 + halfgap : 1875 - halfgap
    ]
    clinical_voltage[2][1875 + halfgap : 2500] = voltage["V6"][1875 + halfgap : 2500]

    clinical_voltage[3] = voltage["V1"]
    clinical_voltage[4] = voltage["II"]
    clinical_voltage[5] = voltage["V5"]

    voltage = clinical_voltage

    # convert voltage to millivolts
    voltage /= 1000

    # calculate space between leads
    min_y, max_y = ax.get_ylim()
    y_offset = (max_y - min_y) / len(voltage)

    text_xoffset = 5
    text_yoffset = -0.1

    # plot signal and add labels
    for i in range(len(voltage)):
        this_offset = (len(voltage) - i - 0.5) * y_offset
        ax.plot(voltage[i] + this_offset, color="black", linewidth=0.375)
        if i == 0:
            ax.text(
                0 + text_xoffset,
                this_offset + text_yoffset,
                "I",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                625 + text_xoffset,
                this_offset + text_yoffset,
                "aVR",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                1250 + text_xoffset,
                this_offset + text_yoffset,
                "V1",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                1875 + text_xoffset,
                this_offset + text_yoffset,
                "V4",
                ha="left",
                va="top",
                weight="bold",
            )
        elif i == 1:
            ax.text(
                0 + text_xoffset,
                this_offset + text_yoffset,
                "II",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                625 + text_xoffset,
                this_offset + text_yoffset,
                "aVL",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                1250 + text_xoffset,
                this_offset + text_yoffset,
                "V2",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                1875 + text_xoffset,
                this_offset + text_yoffset,
                "V5",
                ha="left",
                va="top",
                weight="bold",
            )
        elif i == 2:
            ax.text(
                0 + text_xoffset,
                this_offset + text_yoffset,
                "III",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                625 + text_xoffset,
                this_offset + text_yoffset,
                "aVF",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                1250 + text_xoffset,
                this_offset + text_yoffset,
                "V3",
                ha="left",
                va="top",
                weight="bold",
            )
            ax.text(
                1875 + text_xoffset,
                this_offset + text_yoffset,
                "V6",
                ha="left",
                va="top",
                weight="bold",
            )
        elif i == 3:
            ax.text(
                0 + text_xoffset,
                this_offset + text_yoffset,
                "V1",
                ha="left",
                va="top",
                weight="bold",
            )
        elif i == 4:
            ax.text(
                0 + text_xoffset,
                this_offset + text_yoffset,
                "II",
                ha="left",
                va="top",
                weight="bold",
            )
        elif i == 5:
            ax.text(
                0 + text_xoffset,
                this_offset + text_yoffset,
                "V5",
                ha="left",
                va="top",
                weight="bold",
            )


def _plot_partners_figure(
    data: Dict[str, Union[np.ndarray, str, Dict]],
    plot_signal_function: Callable[[Dict[str, np.ndarray], plt.Axes], None],
    plot_mode: str,
    output_folder: str,
    run_id: str,
) -> None:
    plt.rcParams["font.family"] = "Times New Roman"
    plt.rcParams["font.size"] = 9.5

    w, h = 11, 8.5
    fig = plt.figure(
        figsize=(w, h),
        dpi=100,
    )

    # patient info and ecg text
    _plot_partners_text(data, fig, w, h)

    # define plot area in inches
    left = 0.17
    bottom = h - 7.85
    width = w - 2 * left
    height = h - bottom - 2.3

    # ecg plot area
    ax = fig.add_axes([left / w, bottom / h, width / w, height / h])

    # voltage is in microvolts
    # the entire plot area is 5.55 inches tall, 10.66 inches wide (141 mm, 271 mm)
    # the resolution on the y-axis is 10 mm/mV
    # the resolution on the x-axis is 25 mm/s
    inch2mm = lambda inches: inches * 25.4

    # 1. set y-limit to max 14.1 mV
    y_res = 10  # mm/mV
    max_y = inch2mm(height) / y_res
    min_y = 0
    ax.set_ylim(min_y, max_y)

    # 2. set x-limit to max 10.8 s, center 10 s leads
    sampling_frequency = 250  # Hz
    x_res = 25  # mm/s
    max_x = inch2mm(width) / x_res
    x_buffer = (max_x - 10) / 2
    max_x -= x_buffer
    min_x = -x_buffer
    max_x *= sampling_frequency
    min_x *= sampling_frequency
    ax.set_xlim(min_x, max_x)

    # 3. set ticks for every 0.1 mV or every 1/25 s
    y_tick = 1 / y_res
    x_tick = 1 / x_res * sampling_frequency
    x_major_ticks = np.arange(min_x, max_x, x_tick * 5)
    x_minor_ticks = np.arange(min_x, max_x, x_tick)
    y_major_ticks = np.arange(min_y, max_y, y_tick * 5)
    y_minor_ticks = np.arange(min_y, max_y, y_tick)

    ax.set_xticks(x_major_ticks)
    ax.set_xticks(x_minor_ticks, minor=True)
    ax.set_yticks(y_major_ticks)
    ax.set_yticks(y_minor_ticks, minor=True)

    ax.tick_params(
        which="both", left=False, bottom=False, labelleft=False, labelbottom=False,
    )
    ax.grid(b=True, color="r", which="major", lw=0.5)
    ax.grid(b=True, color="r", which="minor", lw=0.2)

    # signal plot
    voltage = data["2500_raw"]
    plot_signal_function(voltage, ax)

    # bottom text
    fig.text(
        0.17 / w,
        0.46 / h,
        f"{x_res}mm/s    {y_res}mm/mV    {sampling_frequency}Hz",
        ha="left",
        va="center",
        weight="bold",
    )

    # save both pdf and png
    title = re.sub(r"[:/. ]", "", f'{plot_mode}_{data["patientid"]}_{data["datetime"]}')
    plt.savefig(os.path.join(output_folder, run_id, f"{title}{PDF_EXT}"))
    plt.savefig(os.path.join(output_folder, run_id, f"{title}{IMAGE_EXT}"))
    plt.close(fig)


def plot_partners_ecgs(args):
    plot_tensors = [
        ml4h.tensormap.mgb.ecg.partners_ecg_patientid,
        ml4h.tensormap.mgb.ecg.partners_ecg_firstname,
        ml4h.tensormap.mgb.ecg.partners_ecg_lastname,
        ml4h.tensormap.mgb.ecg.partners_ecg_sex,
        ml4h.tensormap.mgb.ecg.partners_ecg_dob,
        ml4h.tensormap.mgb.ecg.partners_ecg_age,
        ml4h.tensormap.mgb.ecg.partners_ecg_datetime,
        ml4h.tensormap.mgb.ecg.partners_ecg_sitename,
        ml4h.tensormap.mgb.ecg.partners_ecg_location,
        ml4h.tensormap.mgb.ecg.partners_ecg_read_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_taxis_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_rate_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_pr_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_qrs_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_qt_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_paxis_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_raxis_md,
        ml4h.tensormap.mgb.ecg.partners_ecg_qtc_md,
    ]
    voltage_tensor = make_waveform_maps("partners_ecg_2500_raw")
    tensor_maps_in = plot_tensors + [voltage_tensor]
    tensor_paths = [
        os.path.join(args.tensors, tp)
        for tp in os.listdir(args.tensors)
        if os.path.splitext(tp)[-1].lower() == TENSOR_EXT
    ]

    if "clinical" == args.plot_mode:
        plot_signal_function = _plot_partners_clinical
    elif "full" == args.plot_mode:
        plot_signal_function = _plot_partners_full
    else:
        raise ValueError(f"Unsupported plot mode: {args.plot_mode}")

    # TODO use TensorGenerator here
    # Get tensors for all hd5
    for tp in tensor_paths:
        try:
            with h5py.File(tp, "r") as hd5:
                skip_hd5 = False
                tdict = defaultdict(dict)
                for tm in tensor_maps_in:
                    key = tm.name.split("partners_ecg_")[1]
                    try:
                        tensors = tm.tensor_from_file(tm, hd5)

                        if tm.shape[0] is not None:
                            # If not a multi-tensor tensor, wrap in array to loop through
                            tensors = np.array([tensors])
                        for i, tensor in enumerate(tensors):
                            if tm.channel_map:
                                tdict[i][key] = dict()
                                for cm in tm.channel_map:
                                    tdict[i][key][cm] = (
                                        tensor[:, tm.channel_map[cm]]
                                        if tm.name == voltage_tensor
                                        else tensor[tm.channel_map[cm]]
                                    )
                            else:
                                if 1 == (
                                    tm.shape[0]
                                    if tm.shape[0] is not None
                                    else tm.shape[1]
                                ):
                                    tensor = tensor.item()
                                tdict[i][key] = tensor
                    except (
                        IndexError,
                        KeyError,
                        ValueError,
                        OSError,
                        RuntimeError,
                    ) as e:
                        logging.warning(
                            f"Could not obtain {tm.name}. Skipping plotting for all ECGs at {tp}",
                        )
                        skip_hd5 = True
                    if skip_hd5:
                        break
                if skip_hd5:
                    continue

                # plot each ecg
                for i in tdict:
                    _plot_partners_figure(
                        data=tdict[i],
                        plot_signal_function=plot_signal_function,
                        plot_mode=args.plot_mode,
                        output_folder=args.output_folder,
                        run_id=args.id,
                    )
        except:
            logging.exception(f"Broken tensor at: {tp}")


def plot_cross_reference(
    args, xref_df, title, time_description, window_start, window_end,
):
    # TODO make this work with multiple time windows
    if xref_df.empty:
        logging.info(f'No cross reference found for "{title}"')
        return

    title = title.replace(" ", "_")

    # compute day diffs
    day_diffs = np.array(
        xref_df.apply(
            lambda row: (row[args.time_tensor] - row[window_end]).days, axis=1,
        ),
    )

    plt.rcParams["font.size"] = 18
    fig = plt.figure(figsize=(15, 9))
    ax = fig.add_subplot(111)
    binwidth = 5
    ax.hist(
        day_diffs, bins=range(day_diffs.min(), day_diffs.max() + binwidth, binwidth),
    )
    ax.set_xlabel("Days relative to event")
    ax.set_ylabel("Number of patients")
    ax.set_title(
        f"Distribution of {args.tensors_name} {time_description}: N={len(day_diffs)}",
    )

    ax.text(0.05, 0.90, f"Min: {day_diffs.min()}", transform=ax.transAxes)
    ax.text(0.05, 0.85, f"Max: {day_diffs.max()}", transform=ax.transAxes)
    ax.text(0.05, 0.80, f"Median: {np.median(day_diffs):.0f}", transform=ax.transAxes)
    plt.tight_layout()

    fpath = os.path.join(
        args.output_folder, args.id, f"distribution_{title}{IMAGE_EXT}",
    )
    fig.savefig(fpath)
    logging.info(f"Saved histogram of days relative to {window_end} to {fpath}")


def plot_categorical_tmap_over_time(counts, tmap_name, dates, fpath):
    f, ax = plt.subplots()
    f.set_size_inches(12, 6.75)
    ax.set_title(tmap_name)
    colors = [[c, c, c] for c in np.linspace(1.0, 0.2, len(counts))]
    bottom = np.zeros(len(dates) - 1)
    for i, (cm, counts_cm) in enumerate(
        sorted(counts.items(), key=lambda kv: np.sum(kv[1]), reverse=True),
    ):
        label = cm + f", n={int(np.sum(counts_cm))}" if i < 20 else None
        ax.bar(
            range(1, len(dates)),
            counts_cm,
            width=1,
            bottom=bottom,
            edgecolor="black",
            linewidth=0.5,
            label=label,
            color=colors[i],
        )
        bottom += np.array(counts_cm)
        ax.legend()
    ax.set_xlim(1, len(dates))
    ax.set_xticks(range(1, len(dates), 12))
    ax.set_xticklabels([date.year for date in dates[1::12]], rotation=45)
    ax.set_ylabel("counts per increment of time frequency")
    ax.legend(loc="upper left")
    f.savefig(fpath, dpi=500)


def _ecg_rest_traces_and_text(hd5):
    """Extracts ECG resting traces from HD5 and returns a dictionary based on biosppy template"""
    path_prefix = "ukb_ecg_rest"
    ecg_text_group = "ecg_rest_text"
    if path_prefix not in hd5:
        raise ValueError("Tensor does not contain resting ECGs")
    twelve_leads = defaultdict(dict)
    leads = list(ECG_REST_LEADS.keys()) + list(ECG_REST_MEDIAN_LEADS.keys())
    oldest_instance = sorted([instance for instance in hd5[path_prefix][leads[0]]])[0]
    for lead in leads:
        twelve_leads[lead]["raw"] = np.array(
            hd5[path_prefix][lead][oldest_instance], dtype=np.float32,
        )
        if len(twelve_leads[lead]["raw"]) == 5000:
            try:
                # Attempt analysis by biosppy, which may fail if not enough beats
                (
                    twelve_leads[lead]["ts_reference"],
                    twelve_leads[lead]["filtered"],
                    twelve_leads[lead]["rpeaks"],
                    twelve_leads[lead]["template_ts"],
                    twelve_leads[lead]["templates"],
                    twelve_leads[lead]["heart_rate_ts"],
                    twelve_leads[lead]["heart_rate"],
                ) = ecg.ecg(
                    signal=twelve_leads[lead]["raw"], sampling_rate=500.0, show=False,
                )
            except:
                twelve_leads[lead]["ts_reference"] = np.linspace(
                    0,
                    len(twelve_leads[lead]["raw"]) / 500.0,
                    len(twelve_leads[lead]["raw"]),
                )
    ecg_rest_text = np.array(hd5[path_prefix][ecg_text_group][oldest_instance]).tolist()
    return twelve_leads, ecg_rest_text


def _ecg_rest_ylims(yrange, yplot):
    """Returns ECG plot y-axis limits based on the range covered by ECG traces"""
    deltas = [-1.0, 1.0]
    extremes = np.array([np.min(yplot), np.max(yplot)])
    delta_ext = extremes[1] - extremes[0]
    yrange = np.max([yrange, delta_ext * 1.10])
    ylim_min = -yrange / 2.0
    ylim_max = yrange / 2.0
    if ((extremes[0] - ylim_min) < yrange * 0.2) or (
        (ylim_max - extremes[1]) < yrange * 0.2
    ):
        ylim_min = extremes[0] - (yrange - delta_ext) / 2.0
        ylim_max = extremes[1] + (yrange - delta_ext) / 2.0
    return ylim_min, ylim_max


def _ecg_rest_yrange(twelve_leads, default_yrange, raw_scale, time_interval):
    """Returns y-range necessary not to cut any of the plotted ECG waveforms"""
    yrange = default_yrange
    for is_median, offset in zip([False, True], [3, 0]):
        for i in range(offset, offset + 3):
            for j in range(0, 4):
                lead_name = ECG_REST_PLOT_LEADS[i - offset][j]
                lead = twelve_leads[lead_name]
                y_plot = np.array([elem_ * raw_scale for elem_ in lead["raw"]])
                if not is_median:
                    y_plot = y_plot[
                        np.logical_and(
                            lead["ts_reference"] > j * time_interval,
                            lead["ts_reference"] < (j + 1) * time_interval,
                        )
                    ]
                ylim_min, ylim_max = _ecg_rest_ylims(yrange, y_plot)
                yrange = ylim_max - ylim_min
    return min(yrange, ECG_REST_PLOT_MAX_YRANGE)


def _subplot_ecg_rest(
    twelve_leads,
    raw_scale,
    time_interval,
    hertz,
    lead_mapping,
    f,
    ax,
    yrange,
    offset,
    pat_df,
    is_median,
    is_blind,
):
    """Fills subplots with either median or raw resting ECG waveforms"""
    # plot will be in seconds vs mV, boxes are
    sec_per_box = 0.04
    mv_per_box = 0.1
    median_interval = 1.2  # 600 samples at 500Hz
    # if available, extract patient metadata and ECG interpretation
    if pat_df is not None:
        avl_yn = "Y" if np.argmax(pat_df["aVL"]) > 0.5 else "N"
        sl_yn = "Y" if np.argmax(pat_df["Sokolow_Lyon"]) > 0.5 else "N"
        cor_yn = "Y" if np.argmax(pat_df["Cornell"]) > 0.5 else "N"
        sex_fm = "F" if ((pat_df["sex"] == "F") or (pat_df["sex"] == "female")) else "M"
        text = f"ID: {pat_df['patient_id']}, sex: {sex_fm}\n"
        if not is_blind:
            text += f"{pat_df['ecg_text']}\n"
            text += f"LVH criteria - aVL: {avl_yn}, Sokolow-Lyon: {sl_yn}, Cornell: {cor_yn}"
        st = f.suptitle(
            text, x=0.0, y=1.05, ha="left", bbox=dict(facecolor="black", alpha=0.1),
        )
    for i in range(offset, offset + 3):
        for j in range(0, 4):
            lead_name = lead_mapping[i - offset][j]
            lead = twelve_leads[lead_name]
            # Convert units to mV
            if isinstance(lead, dict):
                yy = np.array([elem_ * raw_scale for elem_ in lead["raw"]])
            else:
                yy = lead
            if not is_median:
                ax[i, j].set_xlim(j * time_interval, (j + 1) * time_interval)
                # extract portion of waveform that is included in the actual plots
                yplot = yy[
                    int(j * time_interval * hertz) : int(
                        (j + 1) * time_interval * hertz,
                    )
                ]
            else:
                yplot = yy
            ylim_min, ylim_max = _ecg_rest_ylims(yrange, yplot)
            ax[i, j].set_ylim(ylim_min, ylim_max)  # 3.0 mV range
            ax[i, j].xaxis.set_major_locator(
                MultipleLocator(0.2),
            )  # major grids at every .2sec = 5 * 0.04 sec
            ax[i, j].yaxis.set_major_locator(
                MultipleLocator(0.5),
            )  # major grids at every .5mV
            ax[i, j].xaxis.set_minor_locator(AutoMinorLocator(5))
            ax[i, j].yaxis.set_minor_locator(AutoMinorLocator(5))
            ax[i, j].grid(which="major", color="#CCCCCC", linestyle="--")
            ax[i, j].grid(which="minor", color="#CCCCCC", linestyle=":")
            for label in ax[i, j].xaxis.get_ticklabels()[::2]:
                label.set_visible(False)
            if len(ax[i, j].yaxis.get_ticklabels()) > 10:
                for label in ax[i, j].yaxis.get_ticklabels()[::2]:
                    label.set_visible(False)
            # normalize data in muv
            if "ts_reference" in lead:
                ax[i, j].plot(lead["ts_reference"], yy, label="raw")
            else:
                ax[i, j].plot(
                    np.arange(0.0, median_interval, median_interval / len(lead["raw"])),
                    yy,
                    label="raw",
                )
            ax[i, j].set_title(lead_name)
            if is_median and (pat_df is not None):
                # Find where to put the R and S amp text based on ECG baseline position
                dy_ecg = (yy[-1] - ylim_min) / yrange
                if dy_ecg > 0.3:  # Put in bottom right
                    dy_amp = 0.2
                else:  # Put in top right
                    dy_amp = 0.85
                ax[i, j].text(
                    0.9,
                    dy_amp * yrange + ylim_min,
                    f"R: {pat_df['ramp'][ECG_REST_PLOT_AMP_LEADS[i-offset][j]]:.0f}",
                )
                ax[i, j].text(
                    0.9,
                    (dy_amp - 0.15) * yrange + ylim_min,
                    f"S: {pat_df['samp'][ECG_REST_PLOT_AMP_LEADS[i-offset][j]]:.0f}",
                )


def _str_to_list_float(str_list: str) -> List[int]:
    """'[ 3. 4. nan 3 ]' --> [ 3.0, 4.0, nan, 3.0 ]"""
    tmp_str = str_list[1:-1].split()
    return list(map(float, tmp_str))


def _ecg_rest_csv_to_df(csv):
    df = pd.read_csv(csv)
    df["ramp"] = df["ramp"].apply(_str_to_list_float)
    df["samp"] = df["samp"].apply(_str_to_list_float)
    df["patient_id"] = df["patient_id"].apply(str)
    df["Sokolow_Lyon"] = df["Sokolow_Lyon"].apply(float)
    df["Cornell"] = df["Cornell"].apply(float)
    df["aVL"] = df["aVL"].apply(float)
    return df


def _remove_duplicate_rows(df, out_folder):
    arr_list = []
    pdfs = glob.glob(out_folder + "/*.pdf")
    for i, row in df.iterrows():
        if os.path.join(out_folder, row["patient_id"] + ".pdf") not in pdfs:
            arr_list.append(i)
    arr = np.array(arr_list, dtype=np.int)
    return arr


def plot_ecg_rest(
    tensor_paths: List[str],
    rows: List[int],
    out_folder: str,
    is_blind: bool,
) -> None:
    """Plots resting ECGs including annotations and LVH criteria

    :param tensor_paths: list of HDF5 file paths with ECG traces
    :param rows: indices of the subset of tensor_paths to be plotted (used by multiprocessing)
    :param out_folder: destination folder for the plots
    :param is_blind: if True, the plot gets blinded (helpful for review and annotation)
    """
    map_fields_to_tmaps = {
        "ramp": ml4h.tensormap.ukb.ecg.ecg_rest_ramplitude_raw,
        "samp": ml4h.tensormap.ukb.ecg.ecg_rest_samplitude_raw,
        "aVL": ml4h.tensormap.ukb.ecg.ecg_rest_lvh_avl,
        "Sokolow_Lyon": ml4h.tensormap.ukb.ecg.ecg_rest_lvh_sokolow_lyon,
        "Cornell": ml4h.tensormap.ukb.ecg.ecg_rest_lvh_cornell,
    }

    raw_scale = 0.005  # Conversion from raw to mV
    default_yrange = ECG_REST_PLOT_DEFAULT_YRANGE  # mV
    time_interval = 2.5  # time-interval per plot in seconds. ts_Reference data is in s, voltage measurement is 5 uv per lsb
    hertz = 500  # number of samples per second
    for row in rows:
        tensor_path = tensor_paths[row]
        patient_dic = {}
        patient_dic["patient_id"] = os.path.basename(tensor_path).replace(
            TENSOR_EXT, "",
        )
        with h5py.File(tensor_path, "r") as hd5:
            traces, text = _ecg_rest_traces_and_text(hd5)
            for field in map_fields_to_tmaps:
                tm = map_fields_to_tmaps[field]
                patient_dic[field] = np.zeros(tm.shape)
                try:
                    patient_dic[field][:] = tm.tensor_from_file(tm, hd5)
                except ValueError as e:
                    logging.warning(e)
            is_female = "Sex_Female_0_0" in hd5["categorical"]
            patient_dic["sex"] = "F" if is_female else "M"
            patient_dic["ecg_text"] = text
        matplotlib.rcParams.update({"font.size": 20})
        fig, ax = plt.subplots(nrows=6, ncols=4, figsize=(24, 18), tight_layout=True)
        yrange = _ecg_rest_yrange(traces, default_yrange, raw_scale, time_interval)
        _subplot_ecg_rest(
            traces,
            raw_scale,
            time_interval,
            hertz,
            ECG_REST_PLOT_LEADS,
            fig,
            ax,
            yrange,
            offset=3,
            pat_df=None,
            is_median=False,
            is_blind=is_blind,
        )
        _subplot_ecg_rest(
            traces,
            raw_scale,
            time_interval,
            hertz,
            ECG_REST_PLOT_MEDIAN_LEADS,
            fig,
            ax,
            yrange,
            offset=0,
            pat_df=patient_dic,
            is_median=True,
            is_blind=is_blind,
        )
        fig.savefig(
            os.path.join(out_folder, patient_dic["patient_id"] + ".svg"),
            bbox_inches="tight",
        )


def plot_ecg_rest_mp(
    tensors: str,
    min_sample_id: int,
    max_sample_id: int,
    output_folder: str,
    num_workers: int = 1,
    is_blind: bool = False,
) -> None:
    """
    Generates (in parallel) plots for 12-lead resting ECGs given a CSV file pointing to the tensor HDF5s
    :param ecg_csv: name of the CSV file listing the HD5s to plot
    :param min_sample_id: smallest patient id to be plotted
    :param max_sample_id: largest patient id to be plotted
    :param output_folder: directory where output PDFs will be written to
    :param num_workers: number of parallel cores to be used
    :param is_blind: whether ECG interpretation should be included in the plot
    :return: None
    """
    tensor_paths = [
        os.path.join(tensors, tp)
        for tp in os.listdir(tensors)
        if os.path.splitext(tp)[-1].lower() == TENSOR_EXT
    ]
    can_filter = True
    try:
        tensor_sample_ids = np.array(
            [
                os.path.basename(tensor_path).replace(TENSOR_EXT, "")
                for tensor_path in tensor_paths
            ],
            dtype=np.int,
        )
        condition = (tensor_sample_ids > min_sample_id) & (
            tensor_sample_ids < max_sample_id
        )
        tensor_paths = [
            tensor_path for i, tensor_path in enumerate(tensor_paths) if condition[i]
        ]
    except ValueError:
        logging.warning(
            "Cannot select subset of tensors based on sample ids. Discarding min_ and max_sample_id",
        )
    row_split = np.array_split(np.arange(len(tensor_paths)), num_workers)
    pool = Pool(num_workers)
    pool.starmap(
        plot_ecg_rest,
        zip(
            [tensor_paths] * num_workers,
            row_split,
            [output_folder] * num_workers,
            [is_blind] * num_workers,
        ),
    )
    pool.close()
    pool.join()


def plot_counter(counts, title, prefix="./figures/"):
    plt.figure(figsize=(28, 32))
    matplotlib.rcParams.update({"font.size": 12})
    idxs = np.arange(len(counts))

    keyz = []
    vals = []
    for k in sorted(list(counts.keys())):
        keyz.append(k.replace("categorical/", "").replace("continuous/", ""))
        vals.append(counts[k])

    plt.barh(idxs, vals)
    plt.yticks(idxs, keyz)
    plt.tight_layout()
    plt.title(title + "\n")

    figure_path = os.path.join(prefix, "counter_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved counter plot at: {figure_path}")


def _bootstrap_performance(
    truth: np.ndarray,
    prediction: np.ndarray,
    metric: Callable[[np.ndarray, np.ndarray], float],
    n_samples: int = 1000,
) -> List[float]:
    performance = []
    for _ in range(n_samples):
        idx = np.random.randint(0, truth.shape[0], size=truth.shape[0])
        performance.append(metric(truth[idx], prediction[idx]))
    return performance


def _performance_by_index(
    separated_metric_inputs: List[Tuple[np.ndarray, np.ndarray]],
    metric: Callable[[np.ndarray, np.ndarray], float],
) -> List[List[float]]:
    """
    Input dictionary {name: (truth, predicted)} and performance metric
    output DataFrame {performance: performances, name: names}
    """
    return [
        _bootstrap_performance(truth, pred, metric)
        for truth, pred in separated_metric_inputs
    ]


def _protected_subplots(
    prediction: np.ndarray,
    truth: np.ndarray,
    protected: Dict[TensorMap, np.ndarray],
    axes: np.ndarray,
    metric: Callable[[np.ndarray, np.ndarray], float],
    metric_name: str,
    continuous_quantiles: Tuple[float, ...] = (1 / 3, 2 / 3),
):
    group_col = "Group"
    for p, ax in zip(protected, axes.ravel()):
        ax.set_title(f"Performance by protected {p.name}")
        if p.is_categorical():
            group_names = list(p.channel_map.keys())
            protected_indexes = [
                protected[p][:, class_label] == 1
                for class_label in p.channel_map.values()
            ]
        elif p.is_continuous():
            thresholds = (
                [np.min(protected[p])]
                + np.quantile(protected[p], continuous_quantiles).tolist()
                + [np.max(protected[p])]
            )
            group_names = [
                f"{p.rescale(thresholds[i]):.1f} < x < {p.rescale(thresholds[i + 1]):.1f}"
                for i in range(len(thresholds) - 1)
            ]
            protected_indexes = [
                (protected[p] > thresholds[i]) & (protected[p] <= thresholds[i + 1])
                for i in range(len(thresholds) - 1)
            ]
        else:
            raise NotImplementedError(
                f"Cannot plot {p} with interpretation {p.interpretation}",
            )
        logging.info(f"Beginning bootstrap performance evaluation of protected {p}")
        performances = _performance_by_index(
            [
                (truth[protected_idx], prediction[protected_idx])
                for protected_idx in protected_indexes
            ],
            metric,
        )
        group_names = [
            f"{group_name}\nn = {idx.sum()}"
            for group_name, idx in zip(group_names, protected_indexes)
        ]
        performance_dict = {group_col: [], metric_name: []}
        for performance, group_name in zip(performances, group_names):
            performance_dict[group_col] += [group_name] * len(performance)
            performance_dict[metric_name] += performance
        sns.boxplot(
            x=group_col, y=metric_name, data=pd.DataFrame(performance_dict), ax=ax,
        )


def subplot_roc_per_class(
    prediction: np.ndarray,
    truth: np.ndarray,
    labels: Dict[str, int],
    protected: Dict[TensorMap, np.ndarray],
    title: str,
    prefix: str = "./figures/",
    dpi: int = 300,
    width: int = 6,
    height: int = 6,
) -> Dict[str, float]:
    lw = 2
    labels_to_areas = {}
    true_sums = np.sum(truth, axis=0)
    total_plots = len(protected) + 1
    if total_plots == 1:
        return plot_roc(prediction, truth, labels, title, prefix)
    cols = max(2, int(math.ceil(math.sqrt(total_plots))))
    rows = max(2, int(math.ceil(total_plots / cols)))
    fig, axes = plt.subplots(
        rows, cols, figsize=(cols * width, rows * height), dpi=dpi,
    )
    _protected_subplots(
        prediction, truth, protected, axes, metric=roc_auc_score, metric_name="ROC AUC",
    )
    fpr, tpr, roc_auc = get_fpr_tpr_roc_pred(prediction, truth, labels)

    for key in labels:
        if "no_" in key and len(labels) == 2:
            continue
        labels_to_areas[key] = roc_auc[labels[key]]
        color = _hash_string_to_color(key)
        label_text = (
            f"{key} area: {roc_auc[labels[key]]:.3f} n={true_sums[labels[key]]:.0f}"
        )
        axes[-1, -1].plot(
            fpr[labels[key]], tpr[labels[key]], color=color, lw=lw, label=label_text,
        )
        logging.info(
            f"ROC Label {label_text} Truth shape {truth.shape}, true sums {true_sums}",
        )

    axes[-1, -1].set_title(f"ROC {title} n={truth.shape[0]:.0f}")
    axes[-1, -1].legend(loc="lower right")
    figure_path = os.path.join(prefix, "per_class_roc_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path, bbox_inches="tight")
    logging.info(
        f"Saved ROC curve at: {figure_path} with {len(protected)} protected TensorMaps.",
    )
    return labels_to_areas


def plot_roc(prediction, truth, labels, title, prefix="./figures/", dpi=300, width: int = 6, height: int = 6):
    lw = 2
    labels_to_areas = {}
    true_sums = np.sum(truth, axis=0)
    plt.figure(figsize=(width, height), dpi=dpi)

    fpr, tpr, roc_auc = get_fpr_tpr_roc_pred(prediction, truth, labels)
    for key in labels:
        if "no_" in str(key) and len(labels) == 2:
            continue
        color = _hash_string_to_color(str(key))
        labels_to_areas[key] = roc_auc[labels[key]]
        label_text = f"{key} area:{roc_auc[labels[key]]:.3f} n={true_sums[labels[key]]:.0f}"
        plt.plot(
            fpr[labels[key]], tpr[labels[key]], color=color, lw=lw, label=label_text,
        )
        logging.info(f"ROC Label {label_text}")

    plt.xlim([0.0, 1.0])
    plt.ylim([-0.02, 1.03])
    plt.ylabel(RECALL_LABEL)
    plt.xlabel(FALLOUT_LABEL)
    plt.legend(loc="lower right")
    plt.plot([0, 1], [0, 1], "k:", lw=0.5)
    plt.title(f"ROC {title} n={np.sum(true_sums):.0f}")

    figure_path = os.path.join(prefix, "per_class_roc_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved ROC curve at: {figure_path}")
    return labels_to_areas


def _figure_and_subplot_axes_from_total(total_plots: int, dpi: int = 300, width: int = 6, height: int = 6):
    cols = max(2, int(math.ceil(math.sqrt(total_plots))))
    rows = max(2, int(math.ceil(total_plots / cols)))
    return plt.subplots(rows, cols, figsize=(cols * width, rows * height), dpi=dpi)


def subplot_rocs(
    rocs: List[Tuple[np.ndarray, np.ndarray, Dict[str, int]]],
    prefix: str = "./figures/", dpi: int = 300, width: int = 6, height: int = 6,
):
    """Log and tabulate AUCs given as nested dictionaries in the format '{model: {label: auc}}'"""
    lw = 2
    fig, axes = _figure_and_subplot_axes_from_total(len(rocs), dpi=dpi, width=width, height=height)
    for roc_data, ax in zip(rocs, axes.ravel()):
        predicted, truth, labels = roc_data
        true_sums = np.sum(truth, axis=0)
        fpr, tpr, roc_auc = get_fpr_tpr_roc_pred(predicted, truth, labels)
        for key in labels:
            if "no_" in key and len(labels) == 2:
                continue
            color = _hash_string_to_color(key)
            label_text = (
                f"{key} area: {roc_auc[labels[key]]:.3f} n={true_sums[labels[key]]:.0f}"
            )
            ax.plot(
                fpr[labels[key]], tpr[labels[key]], color=color, lw=lw, label=label_text,
            )
            logging.info(f"ROC Label {label_text}")
        ax.set_xlim([0.0, 1.0])
        ax.set_ylim([-0.02, 1.03])
        ax.set_ylabel(RECALL_LABEL)
        ax.set_xlabel(FALLOUT_LABEL)
        ax.legend(loc="lower right")
        ax.plot([0, 1], [0, 1], "k:", lw=0.5)
        ax.set_title(f"ROC n={np.sum(true_sums):.0f}")

    figure_path = prefix + "rocs_together" + IMAGE_EXT
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)


def subplot_comparison_rocs(
    rocs: List[Tuple[Dict[str, np.ndarray], np.ndarray, Dict[str, int]]],
    prefix: str = "./figures/",
    dpi: int = 300,
    width: int = 6,
    height: int = 6,
):
    """Log and tabulate AUCs given as nested dictionaries in the format '{model: {label: auc}}'"""
    lw = 3
    fig, axes = _figure_and_subplot_axes_from_total(len(rocs), dpi=dpi, width=width, height=height)
    for roc_data, ax in zip(rocs, axes.ravel()):
        predictions, truth, labels = roc_data
        true_sums = np.sum(truth, axis=0)
        for p in predictions:
            fpr, tpr, roc_auc = get_fpr_tpr_roc_pred(predictions[p], truth, labels)
            for key in labels:
                if "no_" in key and len(labels) == 2:
                    continue
                color = _hash_string_to_color(p + key)
                label_text = f"{p}_{key} area:{roc_auc[labels[key]]:.3f} n={true_sums[labels[key]]:.0f}"
                ax.plot(
                    fpr[labels[key]],
                    tpr[labels[key]],
                    color=color,
                    lw=lw,
                    label=label_text,
                )
                logging.info(f"ROC Label {label_text}")

        ax.set_xlim([0.0, 1.0])
        ax.set_ylim([-0.02, 1.03])
        ax.set_ylabel(RECALL_LABEL)
        ax.set_xlabel(FALLOUT_LABEL)
        ax.legend(loc="lower right")
        ax.plot([0, 1], [0, 1], "k:", lw=0.5)
        ax.set_title(f"ROC n={np.sum(true_sums):.0f}")

    figure_path = os.path.join(prefix, "rocs_compared_together" + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)


def plot_precision_recall_per_class(
    prediction, truth, labels, title, prefix="./figures/", dpi=300, width: int = 6, height: int = 6,
):
    # Compute Precision-Recall and plot curve
    lw = 2.0
    labels_to_areas = {}
    true_sums = np.sum(truth, axis=0)
    plt.figure(figsize=(width, height), dpi=dpi)

    for k in labels:
        c = _hash_string_to_color(str(k))
        precision, recall, _ = precision_recall_curve(
            truth[:, labels[k]], prediction[:, labels[k]],
        )
        average_precision = average_precision_score(
            truth[:, labels[k]], prediction[:, labels[k]],
        )
        label_text = (
            f"{k} mean precision:{average_precision:.3f} n={true_sums[labels[k]]:.0f}"
        )
        plt.plot(recall, precision, lw=lw, color=c, label=label_text)
        logging.info(f"prAUC Label {label_text}")
        labels_to_areas[k] = average_precision

    plt.xlim([0.0, 1.00])
    plt.ylim([-0.02, 1.03])
    plt.xlabel(RECALL_LABEL)
    plt.ylabel(PRECISION_LABEL)
    plt.legend()
    plt.title(f"{title} n={np.sum(true_sums):.0f}")

    figure_path = os.path.join(prefix, "precision_recall_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved Precision Recall curve at: {figure_path}")
    return labels_to_areas


def plot_precision_recalls(predictions, truth, labels, title, prefix="./figures/", dpi=300, width=7, height=7):
    # Compute Precision-Recall and plot curve for each model
    lw = 2.0
    true_sums = np.sum(truth, axis=0)
    plt.figure(figsize=(width, height), dpi=dpi)

    for p in predictions:
        for k in labels:
            c = _hash_string_to_color(p + k)
            precision, recall, _ = precision_recall_curve(
                truth[:, labels[k]], predictions[p][:, labels[k]],
            )
            average_precision = average_precision_score(
                truth[:, labels[k]], predictions[p][:, labels[k]],
            )
            label_text = f"{p}_{k} mean precision:{average_precision:.3f} n={true_sums[labels[k]]:.0f}"
            plt.plot(recall, precision, lw=lw, color=c, label=label_text)
            logging.info(f"prAUC Label {label_text}")

    plt.xlim([0.0, 1.00])
    plt.ylim([-0.02, 1.03])
    plt.xlabel(RECALL_LABEL)
    plt.ylabel(PRECISION_LABEL)
    plt.legend(loc="lower left")
    plt.title(f"{title} n={np.sum(true_sums):.0f}")

    figure_path = os.path.join(prefix, "precision_recall_" + title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info("Saved Precision Recall curve at: {}".format(figure_path))


def get_fpr_tpr_roc_pred(y_pred, test_truth, labels):
    # Compute ROC curve and ROC area for each class
    fpr = dict()
    tpr = dict()
    roc_auc = dict()

    for k in labels:
        cur_idx = labels[k]
        aser = roc_curve(test_truth[:, cur_idx], y_pred[:, cur_idx])
        fpr[labels[k]], tpr[labels[k]], _ = aser
        roc_auc[labels[k]] = auc(fpr[labels[k]], tpr[labels[k]])

    return fpr, tpr, roc_auc


def plot_waves(predicted_waves, true_waves, title, plot_path, rows=6, cols=6):
    f, axes = plt.subplots(rows, cols, sharex=True, figsize=(36, 36))
    for i, ax in zip(range(true_waves.shape[0]), axes.ravel()):
        ax.plot(true_waves[i, :, 0], color="blue", label="Actual Wave")
        if predicted_waves is not None:
            ax.plot(predicted_waves[i, :, 0], color="green", label="Predicted")
        ax.set_xlabel("time")
    plt.legend(loc="lower left")

    figure_path = os.path.join(plot_path, title + IMAGE_EXT)
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info("Saved waves at: {}".format(figure_path))


def plot_tsne(
    x_embed,
    categorical_labels,
    continuous_labels,
    gene_labels,
    label_dict,
    figure_path,
    alpha,
    dpi: int = 300,
    width: int = 7,
    height: int = 7,
    perplexities = [12, 30],
):
    x_embed = np.array(x_embed)
    if len(x_embed.shape) > 2:
        x_embed = np.reshape(x_embed, (x_embed.shape[0], np.prod(x_embed.shape[1:])))

    n_components = 2
    rows = max(2, len(label_dict))

    (fig, subplots) = plt.subplots(
        rows,
        len(perplexities),
        figsize=(len(perplexities) * width, rows * height),
        dpi=dpi,
    )

    p2y = {}
    for i, p in enumerate(perplexities):
        tsne = manifold.TSNE(
            n_components=n_components,
            init="pca",
            random_state=123,
            perplexity=p,
            learning_rate=20,
            n_iter_without_progress=500,
        )
        p2y[p] = tsne.fit_transform(x_embed)

    j = -1
    for tm in label_dict:
        j += 1
        if j == rows:
            break
        categorical_subsets = {}
        categorical_counts = Counter()
        if tm in categorical_labels + gene_labels:
            for c in tm.channel_map:
                categorical_subsets[tm.channel_map[c]] = (
                    label_dict[tm] == tm.channel_map[c]
                )
                categorical_counts[tm.channel_map[c]] = np.sum(
                    categorical_subsets[tm.channel_map[c]],
                )
        elif tm in continuous_labels:
            colors = label_dict[tm]
        for i, p in enumerate(perplexities):
            ax = subplots[j, i]
            ax.set_title(f"{tm.name} | t-SNE perplexity:{p}")
            if tm in categorical_labels + gene_labels:
                color_labels = []
                for c in tm.channel_map:
                    channel_index = tm.channel_map[c]
                    color = _hash_string_to_color(c)
                    color_labels.append(
                        f"{c} n={categorical_counts[tm.channel_map[c]]}",
                    )
                    ax.scatter(
                        p2y[p][categorical_subsets[channel_index], 0],
                        p2y[p][categorical_subsets[channel_index], 1],
                        c=color,
                        alpha=alpha,
                    )
                ax.legend(color_labels, loc="lower left")
            elif tm in continuous_labels:
                points = ax.scatter(
                    p2y[p][:, 0], p2y[p][:, 1], c=colors, alpha=alpha, cmap="jet",
                )
                if i == len(perplexities) - 1:
                    fig.colorbar(points, ax=ax)

            ax.xaxis.set_major_formatter(NullFormatter())
            ax.yaxis.set_major_formatter(NullFormatter())
            ax.axis("tight")

    figure_path += "tsne_plot" + IMAGE_EXT
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)
    logging.info(f"Saved T-SNE plot at: {figure_path}")


def plot_find_learning_rate(
    learning_rates: List[float],
    losses: List[float],
    smoothed_losses: List[float],
    picked_learning_rate: Optional[float],
    figure_path: str,
):
    plt.figure(figsize=(2 * SUBPLOT_SIZE, SUBPLOT_SIZE))
    plt.title("Learning rate finder")
    cutoff = smoothed_losses[0]
    plt.ylim(min(smoothed_losses), cutoff * 1.05)
    plt.axhline(
        cutoff, linestyle="--", color="k", label=f"Deltas ignored above {cutoff:.2f}",
    )
    learning_rates = np.log(learning_rates) / np.log(10)
    plt.plot(learning_rates, losses, label="Loss", c="r")
    plt.plot(learning_rates, smoothed_losses, label="Smoothed loss", c="b")
    if picked_learning_rate is not None:
        plt.axvline(
            np.log(picked_learning_rate) / np.log(10),
            label=f"Learning rate found {picked_learning_rate:.2E}",
            color="g",
            linestyle="--",
        )
    plt.xlabel("Log_10 learning rate")
    plt.legend()
    plt.savefig(os.path.join(figure_path, f"find_learning_rate{IMAGE_EXT}"))


def plot_saliency_maps(
    data: np.ndarray, gradients: np.ndarray, paths: List, prefix: str,
):
    """Plot saliency maps of a batch of input tensors.

    Saliency maps for each input tensor in the batch will be saved at the file path indicated by prefix.
    Also creates a mean saliency map across the batch
    2D tensors are assumed to be ECGs and 3D tensors are plotted with each slice as an RGB image.
    The red channel indicates negative gradients, and the green channel positive ones.

    :param data: A batch of input tensors
    :param gradients: A corresponding batch of gradients for those inputs, must be the same shape as data
    :param paths: A List of paths corresponding to each input tensor
    :param prefix: file path prefix where saliency maps will be saved
    """
    if data.shape[-1] == 1:
        data = data[..., 0]
        gradients = gradients[..., 0]

    mean_saliency = np.zeros(data.shape[1:4] + (3,))
    for batch_i, path in enumerate(paths):
        sample_id = os.path.basename(path).replace(TENSOR_EXT, "")
        if len(data.shape) == 3:
            ecgs = {f"{sample_id}_raw": data[batch_i], "gradients": gradients[batch_i]}
            _plot_ecgs(ecgs, f"{prefix}_{sample_id}_saliency_{batch_i}{IMAGE_EXT}")
        elif len(data.shape) == 4:
            cols = max(2, int(math.ceil(math.sqrt(data.shape[-1]))))
            rows = max(2, int(math.ceil(data.shape[-1] / cols)))
            title = f"{prefix}_{sample_id}_saliency_{batch_i}{IMAGE_EXT}"
            _plot_3d_tensor_slices_as_rgb(
                _saliency_map_rgb(data[batch_i], gradients[batch_i]), title, cols, rows,
            )
            saliency = _saliency_blurred_and_scaled(
                gradients[batch_i], blur_radius=5.0, max_value=1.0 / data.shape[0],
            )
            mean_saliency[..., 0] -= saliency
            mean_saliency[..., 1] += saliency
        elif len(data.shape) == 5:
            for j in range(data.shape[-1]):
                cols = max(2, int(math.ceil(math.sqrt(data.shape[-2]))))
                rows = max(2, int(math.ceil(data.shape[-2] / cols)))
                name = f"{prefix}_saliency_{batch_i}_channel_{j}{IMAGE_EXT}"
                _plot_3d_tensor_slices_as_rgb(
                    _saliency_map_rgb(
                        data[batch_i, ..., j], gradients[batch_i, ..., j],
                    ),
                    name,
                    cols,
                    rows,
                )
                saliency = _saliency_blurred_and_scaled(
                    gradients[batch_i, ..., j],
                    blur_radius=5.0,
                    max_value=1.0 / data.shape[0],
                )
                mean_saliency[..., 0] -= saliency
                mean_saliency[..., 1] += saliency
        else:
            logging.warning(f"No method to plot saliency for data shape: {data.shape}")

    if len(data.shape) == 4:
        _plot_3d_tensor_slices_as_rgb(
            _scale_tensor_inplace(mean_saliency),
            f"{prefix}_batch_mean_saliency{IMAGE_EXT}",
            cols,
            rows,
        )
    logging.info(f"Saved saliency maps at:{prefix}")


def _scale_tensor_inplace(tensor, min_value=0.0, max_value=1.0):
    tensor -= tensor.min()
    tensor *= (max_value - min_value) / tensor.max()
    tensor += min_value
    return tensor


def _saliency_blurred_and_scaled(gradients, blur_radius, max_value=1.0):
    blurred = gaussian_filter(gradients, sigma=blur_radius)
    _scale_tensor_inplace(blurred, max_value=max_value)
    blurred -= blurred.mean()
    return blurred


def _saliency_map_rgb(image, gradients, blur_radius=0):
    _scale_tensor_inplace(image)
    rgb_map = np.zeros(image.shape + (3,))
    blurred = _saliency_blurred_and_scaled(gradients, blur_radius)
    rgb_map[..., 0] = image - blurred
    rgb_map[..., 1] = image + blurred
    rgb_map[..., 2] = image
    rgb_map = np.clip(rgb_map, 0, 1)
    # _scale_tensor_inplace(rgb_map)
    return rgb_map


def _plot_ecgs(
    ecgs,
    figure_path,
    rows=3,
    cols=4,
    time_interval=2.5,
    raw_scale=0.005,
    hertz=500,
    lead_dictionary=ECG_REST_LEADS,
):
    index2leads = {v: k for k, v in lead_dictionary.items()}
    _, axes = plt.subplots(rows, cols, figsize=(18, 16), sharey=True)
    for i in range(rows):
        for j in range(cols):
            start = int(i * time_interval * hertz)
            stop = int((i + 1) * time_interval * hertz)
            axes[i, j].set_xlim(start, stop)
            for label in ecgs:
                axes[i, j].plot(
                    range(start, stop),
                    ecgs[label][start:stop, j + i * cols] * raw_scale,
                    label=label,
                )
            axes[i, j].legend(loc="lower right")
            axes[i, j].set_xlabel("milliseconds")
            axes[i, j].set_ylabel("mV")
            axes[i, j].set_title(index2leads[j + i * cols])
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)


def _plot_3d_tensor_slices_as_rgb(tensor, figure_path, cols=3, rows=10):
    _, axes = plt.subplots(rows, cols, figsize=(cols * 4, rows * 4))
    for i in range(tensor.shape[-2]):
        axes[i // cols, i % cols].imshow(tensor[:, :, i, :])
        axes[i // cols, i % cols].set_yticklabels([])
        axes[i // cols, i % cols].set_xticklabels([])

    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)


def _hash_string_to_color(string):
    """Hash a string to color (using hashlib and not the built-in hash for consistency between runs)"""
    return COLOR_ARRAY[
        int(hashlib.sha1(string.encode("utf-8")).hexdigest(), 16) % len(COLOR_ARRAY)
    ]


def _text_on_plot(axes, x, y, text, alpha=0.8, background="white"):
    t = axes.text(x, y, text)
    t.set_bbox({"facecolor": background, "alpha": alpha, "edgecolor": background})


def plot_reconstruction(
    tm: TensorMap,
    y_true: np.ndarray,
    y_pred: np.ndarray,
    folder: str,
    paths: List[str],
    num_samples: int = 4,
    dpi: int = 300,
    width: int = 7,
    height: int = 7,
):
    logging.info(f"Plotting {num_samples} reconstructions of {tm}.")
    if None in tm.shape or paths is None or len(paths) == 0:  # can't handle dynamic shapes
        return
    os.makedirs(os.path.dirname(folder), exist_ok=True)
    for i in range(num_samples):
        sample_id = os.path.basename(paths[i]).replace(TENSOR_EXT, "")
        title = f"{tm.name}_{sample_id}_reconstruction"
        y = y_true[i]
        yp = y_pred[i]
        if tm.axes() == 2:
            index2channel = {v: k for k, v in tm.channel_map.items()}
            fig, axes = plt.subplots(
                max(2,tm.shape[1]), 3, figsize=(width, height), dpi=dpi,
            )  # , sharey=True)
            for j in range(tm.shape[1]):
                axes[j, 0].plot(y[:, j], c="k", label="original")
                axes[j, 1].plot(yp[:, j], c="b", label="reconstruction")
                axes[j, 2].plot(y[:, j], c="k", label="original")
                axes[j, 2].plot(yp[:, j], c="b", label="reconstruction")
                axes[j, 0].set_title(f"Lead: {index2channel[j]}")
                axes[j, 0].legend()
                axes[j, 1].legend()
                axes[j, 2].legend()
            plt.tight_layout()
            plt.savefig(os.path.join(folder, title + IMAGE_EXT))
        elif tm.axes() == 3:
            if tm.is_categorical():
                plt.imsave(
                    f"{folder}{sample_id}_{tm.name}_truth_{i:02d}{IMAGE_EXT}",
                    np.argmax(y, axis=-1),
                    cmap="plasma",
                )
                plt.imsave(
                    f"{folder}{sample_id}_{tm.name}_prediction_{i:02d}{IMAGE_EXT}",
                    np.argmax(yp, axis=-1),
                    cmap="plasma",
                )
            else:
                if y.shape[-1] not in [3, 4]:
                    plt.imsave(
                        f"{folder}{sample_id}_{tm.name}_truth_{i:02d}{IMAGE_EXT}",
                        y[:, :, 0],
                        cmap="gray",
                    )
                    plt.imsave(
                        f"{folder}{sample_id}_{tm.name}_prediction_{i:02d}{IMAGE_EXT}",
                        yp[:, :, 0],
                        cmap="gray",
                    )
                else:
                    plt.imsave(
                        f"{folder}{sample_id}_{tm.name}_truth_{i:02d}{IMAGE_EXT}",
                        (y[:, :, :]-y.min())/ (1e-6+ y.max()-y.min()),
                    )
                    plt.imsave(
                        f"{folder}{sample_id}_{tm.name}_prediction_{i:02d}{IMAGE_EXT}",
                        (yp[:, :, :]-yp.min())/ (1e-6+ yp.max()-yp.min()),
                    )
        elif tm.axes() == 4:
            for j in range(y.shape[3]):
                image_path_base = f"{folder}{sample_id}_{tm.name}_{i:03d}_{j:03d}"
                if tm.is_categorical():
                    truth = np.argmax(yp[:, :, j, :], axis=-1)
                    prediction = np.argmax(y[:, :, j, :], axis=-1)
                    plt.imsave(
                        f"{image_path_base}_truth{IMAGE_EXT}", truth, cmap="plasma",
                    )
                    plt.imsave(
                        f"{image_path_base}_prediction{IMAGE_EXT}",
                        prediction,
                        cmap="plasma",
                    )
                else:
                    plt.imsave(
                        f"{image_path_base}_truth{IMAGE_EXT}",
                        y[:, :, j, 0],
                        cmap="gray",
                    )
                    plt.imsave(
                        f"{image_path_base}_prediction{IMAGE_EXT}",
                        yp[:, :, j, 0],
                        cmap="gray",
                    )


def pca_on_matrix(matrix, pca_components, prefix="./figures/"):
    pca = PCA()
    pca.fit(matrix)
    print(
        f"PCA explains {100 * np.sum(pca.explained_variance_ratio_[:pca_components]):0.1f}% of variance with {pca_components} top PCA components.",
    )
    matrix_reduced = pca.transform(matrix)[:, :pca_components]
    print(
        f"PCA reduces matrix shape:{matrix_reduced.shape} from matrix shape: {matrix.shape}",
    )
    plot_scree(pca_components, 100 * pca.explained_variance_ratio_, prefix)
    return pca, matrix_reduced


def plot_scree(pca_components, percent_explained, prefix="./figures/"):
    _ = plt.figure(figsize=(6, 4))
    plt.plot(range(len(percent_explained)), percent_explained, "g.-", linewidth=1)
    plt.axvline(x=pca_components, c="r", linewidth=3)
    label = f"{np.sum(percent_explained[:pca_components]):0.1f}% of variance explained by top {pca_components} of {len(percent_explained)} components"
    plt.text(
        pca_components + 0.02 * len(percent_explained), percent_explained[1], label,
    )
    plt.title("Scree Plot")
    plt.xlabel("Principal Components")
    plt.ylabel("% of Variance Explained by Each Component")
    figure_path = f"{prefix}pca_{pca_components}_of_{len(percent_explained)}.png"
    if not os.path.exists(os.path.dirname(figure_path)):
        os.makedirs(os.path.dirname(figure_path))
    plt.savefig(figure_path)


def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)


def angle_between(v1, v2):
    """Returns the angle in radians between vectors 'v1' and 'v2'::
    angle_between((1, 0, 0), (0, 1, 0))
    90
    angle_between((1, 0, 0), (1, 0, 0))
    0.0
    angle_between((1, 0, 0), (-1, 0, 0))
    180
    """
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)) * 180 / 3.141592


def stratify_latent_space(stratify_column, stratify_thresh, latent_cols, latent_df):
    hit = latent_df.loc[latent_df[stratify_column] >= stratify_thresh][
        latent_cols
    ].to_numpy()
    miss = latent_df.loc[latent_df[stratify_column] < stratify_thresh][
        latent_cols
    ].to_numpy()
    miss_mean_vector = np.mean(miss, axis=0)
    hit_mean_vector = np.mean(hit, axis=0)
    angle = angle_between(miss_mean_vector, hit_mean_vector)
    print(
        f"Angle between {stratify_column} and all others: {angle}, \n"
        f"Hit shape {hit.shape}, miss:{miss.shape} threshold:{stratify_thresh}\n"
        f"Distance: {np.linalg.norm(hit_mean_vector - miss_mean_vector):.3f}, "
        f"Hit std {np.std(hit, axis=1).mean():.3f}, miss std:{np.std(miss, axis=1).mean():.3f}\n",
    )
    return hit_mean_vector, miss_mean_vector


def plot_hit_to_miss_transforms(
    latent_df,
    decoders,
    feature="Sex_Female_0_0",
    prefix="./figures/",
    thresh=1.0,
    latent_dimension=256,
    samples=16,
    scalar=3.0,
    cmap="plasma",
    test_csv=None,
):
    latent_cols = [f"latent_{i}" for i in range(latent_dimension)]
    female, male = stratify_latent_space(feature, thresh, latent_cols, latent_df)
    sex_vector = female - male
    if test_csv is not None:
        sample_ids = [
            int(s)
            for s in _sample_csv_to_set(test_csv)
            if len(s) > 4 and int(s) in latent_df.index
        ]
        latent_df = latent_df.loc[sample_ids]
        latent_df.info()
        logging.info(f"Subset to test set with {len(sample_ids)} samples")

    samples = min(len(latent_df.index), samples)
    embeddings = latent_df.iloc[:samples][latent_cols].to_numpy()
    sexes = latent_df.iloc[:samples][feature].to_numpy()
    logging.info(f"Embedding shape: {embeddings.shape} sexes  shape: {sexes.shape}")

    sex_vectors = np.tile(sex_vector, (samples, 1))
    male_to_female = embeddings + (scalar * sex_vectors)
    female_to_male = embeddings - (scalar * sex_vectors)
    for dtm in decoders:
        predictions = decoders[dtm].predict(embeddings)
        m2f = decoders[dtm].predict(male_to_female)
        f2m = decoders[dtm].predict(female_to_male)
        if dtm.axes() == 3:
            fig, axes = plt.subplots(max(2, samples), 2, figsize=(18, samples * 4))
            for i in range(samples):
                axes[i, 0].set_title(f"{feature}: {sexes[i]} ?>=<? {thresh}")
                axes[i, 0].axis("off")
                axes[i, 1].axis("off")
                if dtm.is_categorical():
                    axes[i, 0].imshow(
                        np.argmax(predictions[i, ...], axis=-1), cmap=cmap,
                    )
                    if sexes[i] >= thresh:
                        axes[i, 1].imshow(np.argmax(f2m[i, ...], axis=-1), cmap=cmap)
                        axes[i, 1].set_title(f"{feature} to less than {thresh}")
                    else:
                        axes[i, 1].imshow(np.argmax(m2f[i, ...], axis=-1), cmap=cmap)
                        axes[i, 1].set_title(
                            f"{feature} to more than or equal to {thresh}",
                        )
                else:
                    axes[i, 0].imshow(predictions[i, ..., 0], cmap="gray")
                    if sexes[i] >= thresh:
                        axes[i, 1].imshow(f2m[i, ..., 0], cmap="gray")
                        axes[i, 1].set_title(f"{feature} to less than {thresh}")
                    else:
                        axes[i, 1].imshow(m2f[i, ..., 0], cmap="gray")
                        axes[i, 1].set_title(
                            f"{feature} to more than or equal to {thresh}",
                        )
            figure_path = f"{prefix}/{dtm.name}_{feature}_transform_scalar_{scalar}.png"
            if not os.path.exists(os.path.dirname(figure_path)):
                os.makedirs(os.path.dirname(figure_path))
            plt.savefig(figure_path)


def regplot(
    x: str,
    y: str,
    dataframe: pd.DataFrame,
    hue: str = None,
    destination: str = None,
    res_hue_limit: float = None,
    sizes=(25, 25),
    markers=["o", "o"],
    bins: int = 50,
    figsize=(20, 8),
) -> sm.regression.linear_model.RegressionResultsWrapper:
    """Plots three panels of pair-wise correlation. The arguments `hue` and `res_hue_limit`
    can be used to color points given either a provided set of categories, or a residual threshold,
    respectively. Arguments `sizes`, and `markers` can be used to control the plot output if more
    than two categories are provided to `hue`.

    Plots are in order from left to right:
    1) Feature 1 vs Feature 2
    2) Residuals of linear fit
    3) Density of residuals

    Args:
        x (str): Feature 1 string name corresponding to a column in the provided DataFrame.
        y (str): Feature 2 string name corresponding to a column in the provided DataFrame.
        dataframe (pd.DataFrame): Pandas DataFrame with columns `x` and `y` and optionally `hue`
        hue (str, optional): Hue string name corresponding to a column in the provided DataFrame.. Defaults to None.
        destination (str optional): Path to store output image. Defaults to None.
        res_hue_limit (float, optional): Threshold of residuals deemed outliers. Defaults to None.
        sizes (tuple, optional): Tuple of integers controlling point sizes for hue categories. Defaults to (25, 25).
        markers (tuple, optional): Tuple of chars controlling point types for hue categories. Defalts to ["o", "o"].
        bins (int, optional): Number of bins in the density plot. Defaults to 50.
        figsize (tuple, optional): Tuple of ints controlling the image size. Defaults to (20, 8).

    Example:

    Highlighting residuals using a herustic cut-off value of 0.05:
    ```python
        dataframe=df,destination='output_image.jpg')
    ```

    Highlighting points using a provided column:
    ```python
        dataframe=df)
    ```

    Returns:
        sm.regression.linear_model.RegressionResultsWrapper: Fitted OLS results
    """
    sns.set_style("whitegrid")
    mod = sm.OLS(dataframe[y], sm.add_constant(dataframe[x]))
    res = mod.fit()

    test = dataframe[x].drop_duplicates().sort_values()
    predictions = res.get_prediction(sm.add_constant(test))
    predictions = pd.concat([test, predictions.summary_frame(alpha=0.05)], axis=1)

    fig, axes = plt.subplots(1, 3, figsize=figsize)

    res_df = pd.DataFrame({x: dataframe[x], y: dataframe[y], "resy": res.resid})
    if hue is not None:
        res_df[hue] = dataframe[hue]

    if res_hue_limit is not None:
        res_df[hue] = (np.abs(res.resid) >= res_hue_limit).astype(np.uint8)

    sns.scatterplot(
        x=x,
        y=y,
        hue=hue,
        data=res_df,
        ax=axes[0],
        sizes=sizes,
        markers=markers,
    )
    axes[0].plot(predictions[x], predictions["mean"], lw=2)
    axes[0].fill_between(
        x=predictions[x],
        y1=predictions["mean_ci_lower"],
        y2=predictions["mean_ci_upper"],
        facecolor="blue",
        alpha=0.2,
    )

    sns.scatterplot(
        x=x,
        y="resy",
        ax=axes[1],
        hue=hue,
        sizes=sizes,
        markers=markers,
        data=res_df,
    )
    axes[1].axhline(0, ls="--", color="k")

    sns.distplot(res.resid, ax=axes[2], bins=bins)

    if destination is not None:
        fig.savefig(destination)

    return res
