#!/usr/bin/env python
#
# Copyright (c) 2019 10X Genomics, Inc. All rights reserved.
#
from __future__ import absolute_import, division

import collections
import copy
import csv
import math
import random
import re
import sys
import numpy as np
from six import iteritems, itervalues, viewvalues, PY3
from six.moves import xrange as range, zip
import bmkcc.tenkit.safe_json as tk_safe_json
import bmkcc.cellranger.analysis.clustering as cr_clustering
#from cellranger.analysis.singlegenome import SingleGenomeAnalysis
#from cellranger.analysis.multigenome import MultiGenomeAnalysis
import bmkcc.cellranger.constants as cr_constants
#import cellranger.reference as cr_reference
#import cellranger.report as cr_report  # pylint: disable=no-name-in-module
import bmkcc.cellranger.rna.library as rna_library
import bmkcc.cellranger.utils as cr_utils
#import cellranger.targeted.utils as cr_tgt_utils
#import cellranger.vdj.constants as vdj_constants
#import cellranger.vdj.report as vdj_report
import bmkcc.cellranger.webshim.constants.shared as shared_constants
import bmkcc.cellranger.webshim.constants.gex as ws_gex_constants
#import cellranger.webshim.constants.vdj as ws_vdj_constants
from bmkcc.cellranger.webshim.data import SampleData, generate_counter_barcode_rank_plot_data
'''
from cellranger.websummary.sample_properties import (
    CountSampleProperties,
    SampleProperties,
)
'''
#from cellranger.websummary import react_components as rc
#from cellranger.analysis import jibes
#from cellranger.feature.jibes_tag_assigner import JibesTagAssigner
from six import ensure_str

if PY3:
    from typing import Any, Dict, List  # pylint: disable=unused-import

JIBES_BIPLOT_N_POINTS = 100000
MIN_QV_VAL = 1.0
MAX_QV_VAL = 100.0
# sample 100k points for jibes biplot

# https://yutannihilation.github.io/allYourFigureAreBelongToUs/ggthemes/economist_pal/
JIBES_COLOR_PALETTE = [
    "#3E647D",
    "#7B92AB",
    "#82C0E9",
    "#2D6D66",
    "#BFA19C",
    "#008BBC",
    "#97B6B0",
    "#D7D29E",
    "#1A476F",
    "#90353B",
    "#9C8847",
    "#938DD2",
    "#C10534",
    "#6E8E84",
    "#CAC27E",
]


def add_prefix(prefix, name):
    if "%s" in name:
        return name % prefix
    else:
        return "_".join([prefix, name]) if prefix else name


def format_name(display_name, prefix, prefixes, prefix_format_func=None):
    if prefix is not None and prefix_format_func is not None:
        prefix = prefix_format_func(prefix)

    # Default multi -> '' if no format func given
    if prefix_format_func is None and prefix == rna_library.MULTI_REFS_PREFIX:
        prefix = ""

    if len(prefixes) > 1 or "%s" in display_name:
        display_name = add_prefix(prefix, display_name)

    # Replace underscores w/ spaces
    display_name = display_name.replace("_", " ")

    # Collapse whitespace
    display_name = re.sub(r"\s+", " ", display_name)

    return display_name


def format_description(description, prefix, prefixes, prefix_format_func=None):
    if prefix is not None and prefix_format_func is not None:
        prefix = prefix_format_func(prefix)

    # Default multi -> '' if no format func given
    if prefix_format_func is None and prefix == rna_library.MULTI_REFS_PREFIX:
        prefix = ""

    if "%s" in description:
        # Escape stray percents
        description = re.sub("%([^s])", "%%\\1", description)

        # Only add the prefix if there are multiple possibilities
        s = str(prefix) if len(prefixes) > 1 and prefix else ""
        description = description % s

    # Collapse whitespace
    description = re.sub(r"\s+", " ", description)

    return description


def format_value(value, format_type):
    if format_type == "string":
        return str(value)

    value = value if not math.isnan(float(value)) else 0

    if format_type == "percent":
        return "{:.1%}".format(float(value))
    elif format_type == "integer":
        return "{:,.0f}".format(value)
    elif format_type[0] == "%":
        return format_type % float(value)
    raise Exception("Invalid format type: %s" % format_type)


def lookup_name(data, name):
    name_parts = name.split("/", 1)
    data = data.get(name_parts[0])
    if len(name_parts) == 1 or data is None:
        return data
    else:
        return lookup_name(data, name_parts[1])


def add_alarm(value, formatted_value, level, alarm_dict, alarms):
    # The numeric value may have become a string via JSON ('NaN', 'Infinity', '-Infinity')
    # Note that by default python's json encoder outputs a bare NaN symbol which is even worse
    test = 'float("%s") %s' % (str(value), alarm_dict["test"])
    raised = eval(test, {}, {})
    if raised:
        title = alarm_dict["title"]
        message = alarm_dict["message"]
        alarms[title] = {
            "title": title,
            "message": message,
            "value": formatted_value,
            "level": level,
        }
    return raised


def add_alarms(data, alarms_dict, alarms):
    # Prefixes on alarm metrics not implemented
    full_name = alarms_dict["name"]
    value = lookup_name(data, full_name)
    if value is None:
        return

    formatted_value = format_value(value, alarms_dict["format"])
    error_dict = alarms_dict.get(shared_constants.ALARM_ERROR)
    if error_dict is not None and add_alarm(
        value, formatted_value, shared_constants.ALARM_ERROR, error_dict, alarms
    ):
        return
    warn_dict = alarms_dict.get(shared_constants.ALARM_WARN)
    if warn_dict is not None:
        add_alarm(value, formatted_value, shared_constants.ALARM_WARN, warn_dict, alarms)


def add_table_rows(data, name, metric_dict, rows, style_func, target_func, prefixes=None):
    """Args:
    data (dict) - Dict containing summary metrics
    name (str) - Metric name (key in data dict)
    metric_dict (dict) - Metric display definition dict
    """
    format_type = metric_dict["format"]
    display_name = metric_dict.get("display_name", name)
    description = metric_dict.get("description", display_name)

    if not prefixes:
        prefixes = [None]

    values = []
    for prefix in prefixes:
        # Find the metric and construct a table row
        full_name = add_prefix(prefix, name)
        value = lookup_name(data, full_name)

        if value is not None:
            prefix_format_func = metric_dict.get("prefix_format")
            formatted_name = format_name(
                display_name, prefix, prefixes, prefix_format_func=prefix_format_func
            )
            formatted_description = format_description(
                description, prefix, prefixes, prefix_format_func=prefix_format_func
            )
            formatted_value = format_value(value, format_type)

            style = style_func(metric_dict, value)

            rows.append(
                [
                    {"v": formatted_name, "f": formatted_description, "s": style,},
                    {"v": target_func(metric_dict), "s": style,},
                    {"v": formatted_value, "s": style, "r": value,},
                ]
            )
            values.append(value)


def _is_metric_hidden(sample_properties, metric_dict):
    if "hidden" in metric_dict:
        return eval(metric_dict["hidden"], sample_properties, sample_properties)
    return False


def build_tables(
    sample_properties,
    table_dicts,
    alarm_table_dicts,
    sample_data,
    style_func=lambda *args: "",
    target_func=lambda *args: "",
    metric_headers=[],
    all_prefixes={},
):
    tables = []
    for table_dict in table_dicts:
        table_name = table_dict["name"]
        table_metrics = table_dict["metrics"]

        rows = []
        for metric_dict in table_metrics:
            name = metric_dict["name"]
            prefix = metric_dict.get("prefix")

            if _is_metric_hidden(sample_properties, metric_dict):
                continue

            if prefix is not None:
                # Search sample properties for prefix values
                prefixes = sample_properties.get(prefix)

                # Otherwise search the given prefix values
                if not prefixes:
                    prefixes = all_prefixes.get(prefix)

                # Apply the prefix filter for this metric
                if "prefix_filter" in metric_dict:
                    if prefixes is not None:
                        prefixes = [x for x in prefixes if x in metric_dict["prefix_filter"]]
                    else:
                        sys.stderr.write(
                            "Warning: no metric prefix values to filter for metric %s. The prefix name %s is probably wrong.\n"
                            % (name, prefix)
                        )

                add_table_rows(
                    sample_data.summary,
                    name,
                    metric_dict,
                    rows,
                    style_func,
                    target_func,
                    prefixes=prefixes,
                )
            else:
                add_table_rows(
                    sample_data.summary, name, metric_dict, rows, style_func, target_func
                )

        if len(rows) > 0:
            tables.append(
                {"name": table_name, "rows": rows, "headers": metric_headers,}
            )

        # TODO: Temporary hack to view antibody rank plot, until the new websummary arrives
        if (len(rows) == 0 and table_name == "Mapping") or (
            len(rows) == 0 and table_name == "Cells"
        ):
            tables.append(
                {"name": table_name, "rows": rows, "headers": metric_headers,}
            )

    # Build alarm table
    alarms = collections.OrderedDict()
    for alarm_dict in alarm_table_dicts:
        add_alarms(sample_data.summary, alarm_dict, alarms)

    return tables, list(viewvalues(alarms))


def convert_numpy_array_to_line_chart(array, ntype):
    """
    Given an array of values, sort them in descending order, then give
    the x,y coordinates for a line that could represent this sorted order by collapsing
    identical Y values into a single range of ranks.

    Example:
        array = np.array([1,1,1,2,3,3,3,3,3,4,5,9,20])
        convert_numpy_array_to_line_chart(array, float)
        # Result
        [[0, 20.0],
         [1, 9.0],
         [2, 5.0],
         [3, 4.0],
         [4, 3.0],
         [8, 3.0],
         [9, 2.0],
         [10, 1.0],
         [12, 1.0]]
    :param array: a numpy array
    :param ntype: a datatype to return the array values as
    :return:
    """
    array = np.sort(array)[::-1]
    rows = []
    previous_count = None
    for (index,), count in np.ndenumerate(array):
        if index == 0 or index == len(array) - 1:
            rows.append([index, ntype(count)])
        elif previous_count != count:
            assert previous_count is not None
            previous_index = rows[-1][0]
            if previous_index != index - 1:
                rows.append([index - 1, ntype(previous_count)])
            rows.append([index, ntype(count)])
        previous_count = count
    return rows


def _plot_barcode_rank(chart, counts, num_cells):
    """ Generate a generic barcode rank plot """
    rows = convert_numpy_array_to_line_chart(counts, int)

    for row in rows:
        index, count = row[0], row[1]
        if index < num_cells:
            series_list = [chart["data"][0]]
        elif index == num_cells:
            # Connect the two lines
            series_list = [chart["data"][0], chart["data"][1]]
        else:
            series_list = [chart["data"][1]]

        for series in series_list:
            series["x"].append(index)
            series["y"].append(count)

    # Handle case where there is no background
    bg_series = chart["data"][1]
    if len(bg_series["x"]) == 0:
        bg_series["x"].append(0)
        bg_series["y"].append(0)

    return chart


def build_plot_data_dict(plot_segment, counts, show_name_and_hover=True, color=None):
    """
    Construct the data for a plot segment by appropriately slicing the
    counts
    Inputs:
    - plot_segment: BarcodeRankPlotSegment containing [start, end)
        of the segment, the cell density and legend visibility option
    - counts: Reverse sorted UMI counts for all barcodes.
    - show_name_and_hover: boolean whether to add hover and name text specific to the classic
    CR rank plot (the user might want to set these themselves for alternate ranke plots)
    """

    # -1 for continuity between two charts
    start = max(0, plot_segment.start - 1)
    end = plot_segment.end
    plot_rows = convert_numpy_array_to_line_chart(counts[start:end], int)

    if color is None:
        color = shared_constants.BC_PLOT_CMAP(plot_segment.cell_density)

    data_dict = {
        "x": [],
        "y": [],
        "type": "scattergl",
        "mode": "lines",
        "line": {"color": color, "width": shared_constants.BC_RANK_PLOT_LINE_WIDTH,},
        "showlegend": plot_segment.legend,
    }
    if show_name_and_hover:
        name = "Cells" if plot_segment.cell_density > 0 else "Background"
        data_dict["name"] = name
        # Setup the tooltip
        if plot_segment.cell_density > 0.0:
            n_barcodes = plot_segment.end - plot_segment.start
            n_cells = int(
                round(plot_segment.cell_density * n_barcodes)  # pylint: disable=round-builtin
            )
            hover = "{:.0f}% Cells<br>({}/{})".format(
                100 * plot_segment.cell_density, n_cells, n_barcodes
            )
        else:
            hover = "Background"
        data_dict["hoverinfo"] = "text"
        data_dict["text"] = hover
    offset = 1 + start  # it's a log-log plot, hence the 1
    for index, count in plot_rows:
        data_dict["x"].append(index + offset)
        data_dict["y"].append(count)

    # Handle case where the data is empty
    if len(data_dict["x"]) == 0:
        data_dict["x"].append(0)
        data_dict["y"].append(0)

    return data_dict


def _plot_segmented_barcode_rank(chart, counts, plot_segments):
    """
    Generate the RNA counter barcode rank plot
    Inputs:
        - chart: chart element to populate data
        - counts: UMI counts reverse sorted
        - plot_segments: A list of BarcodeRankPlotSegments
    """

    for segment in plot_segments:
        chart["data"].append(build_plot_data_dict(segment, counts))

    return chart


def plot_basic_barcode_rank(chart, cell_barcodes, barcode_summary, genomes, lib_prefix):
    """
     Generate a basic RNA counter barcode rank plot without depending on SampleData/SampleProperties
    Inputs:
        - chart: chart element to populate data
        - cell_barcodes: set of cell barcodes as bytes
        - barcode_summary: barcode summary from the barcode_summary.h5
        - lib_prefix: The library prefix to create the plot for
    """
    genome = rna_library.MULTI_REFS_PREFIX

    if (
        lib_prefix
        == rna_library.get_library_type_metric_prefix(rna_library.GENE_EXPRESSION_LIBRARY_TYPE)
        and len(genomes) == 1
    ):
        genome = genomes[0]

    key = cr_utils.format_barcode_summary_h5_key(
        lib_prefix,
        genome,
        cr_constants.TRANSCRIPTOME_REGION,
        cr_constants.CONF_MAPPED_DEDUPED_READ_TYPE,
    )

    if key not in barcode_summary:
        return None

    counts_per_bc, plot_segments = generate_counter_barcode_rank_plot_data(
        cell_barcodes, barcode_summary, key
    )
    return _plot_segmented_barcode_rank(chart, counts_per_bc, plot_segments)


def plot_barcode_rank(chart, sample_properties, sample_data):
    """ Generate the RNA counter barcode rank plot """

    # TODO: The old PD plots use this function and pass things around as dictionaries,
    # so we convert here if this function was called from the PD code
    if isinstance(sample_properties, dict):
        sample_properties = CountSampleProperties(
            sample_id=sample_properties.get("sample_id", None),
            sample_desc=sample_properties.get("sample_desc", None),
            genomes=sample_properties.get("genomes", None),
        )

    assert isinstance(sample_data, SampleData)
    assert isinstance(sample_properties, SampleProperties)
    if (
        not isinstance(sample_properties, CountSampleProperties)
        or sample_data.barcode_summary is None
        or sample_data.cell_barcodes is None
    ):
        return None

    if len(sample_properties.genomes) == 0:
        return None

    # UMI counts per BC across all genomes present
    if len(sample_properties.genomes) > 1:
        genome = rna_library.MULTI_REFS_PREFIX
    else:
        genome = sample_properties.genomes[0]

    gex_prefix = rna_library.get_library_type_metric_prefix(
        rna_library.GENE_EXPRESSION_LIBRARY_TYPE
    )
    ab_prefix = rna_library.get_library_type_metric_prefix(rna_library.ANTIBODY_LIBRARY_TYPE)
    gex_key = cr_utils.format_barcode_summary_h5_key(
        gex_prefix,
        genome,
        cr_constants.TRANSCRIPTOME_REGION,
        cr_constants.CONF_MAPPED_DEDUPED_READ_TYPE,
    )
    ab_key = cr_utils.format_barcode_summary_h5_key(
        ab_prefix,
        rna_library.MULTI_REFS_PREFIX,
        cr_constants.TRANSCRIPTOME_REGION,
        cr_constants.CONF_MAPPED_DEDUPED_READ_TYPE,
    )

    if gex_key in sample_data.barcode_summary:
        counts_per_bc, plot_segments = sample_data.counter_barcode_rank_plot_data(gex_key)
        return _plot_segmented_barcode_rank(chart, counts_per_bc, plot_segments)
    elif ab_key in sample_data.barcode_summary:  # in case there was only Antibody library
        counts_per_bc, plot_segments = sample_data.counter_barcode_rank_plot_data(ab_key)
        return _plot_segmented_barcode_rank(chart, counts_per_bc, plot_segments)
    else:
        # Not guaranteed to exist, depending on pipeline
        return None


# TODO: This unneccesary argument is needed because of how the PD code is structured.
def plot_vdj_barcode_rank(chart, sample_properties, sample_data):
    """ Generate the VDJ barcode rank plot """
    if not sample_data.cell_barcodes:  # No cells
        return None

    counts_per_bc, plot_segments = sample_data.vdj_barcode_rank_plot_data()
    return _plot_segmented_barcode_rank(chart, counts_per_bc, plot_segments)


def plot_clonotype_table(chart, sample_properties, sample_data):
    if sample_data.vdj_clonotype_summary is None:
        return None

    clonotypes = sample_data.vdj_clonotype_summary.iloc[0:10]

    # This column used to be called 'cdr3s'; allow the webshim to work on older data
    cdr3_aa_col = "cdr3s_aa"
    if cdr3_aa_col not in clonotypes:
        cdr3_aa_col = "cdr3s"

    col_defs = collections.OrderedDict(
        [
            (
                "clonotype_id",
                {
                    "label": "Clonotype ID",
                    "format": "string",
                    "title": "Clonotype ID",
                    "style": "text-align: left",
                },
            ),
            (
                cdr3_aa_col,
                {
                    "label": "CDR3s",
                    "format": "string",
                    "title": "CDR3s in clonotype",
                    "style": "text-align: left",
                },
            ),
            (
                "frequency",
                {
                    "label": "Frequency",
                    "format": "integer",
                    "title": "Number of cells with clonotype",
                    "style": "text-align: right",
                },
            ),
            (
                "proportion",
                {
                    "label": "Proportion",
                    "format": "%0.4f",
                    "title": "Fraction of cell with clonotype",
                    "style": "text-align: right",
                },
            ),
        ]
    )

    cols = []
    for name, col_def in iteritems(col_defs):
        if name not in clonotypes:
            raise ValueError("Column not found in clonotype summary: %s" % name)
        cols.append(
            {"label": col_defs[name]["label"], "title": col_defs[name]["title"],}
        )

    rows = []
    for _, cl_row in clonotypes.iterrows():
        row = []
        for col_name, col_def in iteritems(col_defs):
            value = cl_row[col_name]
            formatted_value = format_value(value, col_def["format"])

            # Make the CDR3 list bit more readable
            formatted_value = formatted_value.replace(";", "; ")

            row.append(
                {
                    "v": tk_safe_json.json_sanitize(value),
                    "f": formatted_value,
                    "s": col_def["style"],
                }
            )
        rows.append(row)

    chart["table"].update({"rows": rows, "cols": cols})

    return chart


def plot_histogram_metric(chart, sample_properties, sample_data, **kwargs):
    """ Plot a HistogramMetric from the summary json """
    summary_data = sample_data.summary
    items = list(summary_data.get(kwargs["metric_name"], {}).items())
    if len(items) < 1:
        return None

    ordering = kwargs.get("order_by", shared_constants.HISTOGRAM_METRIC_DEFAULT_ORDERING)

    if ordering == shared_constants.HISTOGRAM_METRIC_ORDER_INTEGER_BIN:
        items.sort(key=lambda x: nullable_int_sort_key(x[0]))

    elif ordering == shared_constants.HISTOGRAM_METRIC_ORDER_DECREASING_FREQUENCY:
        items.sort(key=lambda x: nullable_int_sort_key(x[1], negate=True))

    elif ordering == shared_constants.HISTOGRAM_METRIC_ORDER_DECREASING_PROPORTION:
        items.sort(key=lambda x: -convert_to_float_gracefully(x[1]))

    x, y = zip(*items)
    chart["data"][0].update({"x": x, "y": y})

    return chart


def plot_barnyard_barcode_counts(chart, sample_properties, sample_data):
    analysis = sample_data.get_analysis(MultiGenomeAnalysis)
    if analysis is None:
        return None

    chart["data"] = []
    for label_info in ws_gex_constants.GEM_CALL_LABELS:
        name = label_info["label"]
        name = name.replace("genome0", analysis.result["genome0"])
        name = name.replace("genome1", analysis.result["genome1"])
        chart["data"].append(
            {
                "x": [],
                "y": [],
                "name": name,
                "mode": "markers",
                "marker": {"color": label_info["color"],},
            }
        )

    call_to_series = {v["key"]: i for i, v in enumerate(ws_gex_constants.GEM_CALL_LABELS)}
    for count0, count1, call in zip(
        analysis.result["count0"], analysis.result["count1"], analysis.result["call"]
    ):
        series = chart["data"][call_to_series[call]]
        series["x"].append(int(count0))
        series["y"].append(int(count1))

    chart["layout"]["xaxis"] = {
        "title": "%s UMI counts" % analysis.result["genome0"],
        "rangemode": "tozero",
        "autorange": True,
        "fixedrange": False,
    }
    chart["layout"]["yaxis"] = {
        "title": "%s UMI counts" % analysis.result["genome1"],
        "rangemode": "tozero",
        "autorange": True,
        "fixedrange": False,
    }

    return chart


def plot_preprocess(analyses):
    if analyses is None or len(analyses) == 0:
        return None

    sg_analyses = [an for an in analyses if isinstance(an, SingleGenomeAnalysis)]
    sg_analysis = sg_analyses[0] if len(sg_analyses) > 0 else None
    if sg_analysis is None or sg_analysis.is_zero_matrix():
        return None

    # Limit the number of K-means clusterings displayed to limit the HTML filesize
    new_clusterings = {}
    for key, clu in iteritems(sg_analysis.clusterings):
        if not (
            clu.clustering_type == cr_clustering.CLUSTER_TYPE_KMEANS
            and clu.num_clusters > ws_gex_constants.MAX_WEBSHIM_KMEANS_K
        ):
            new_clusterings[key] = clu
    sg_analysis.clusterings = new_clusterings

    return analyses


def load_sample_data(sample_properties, sample_data_paths):
    return SampleData(sample_properties, sample_data_paths, plot_preprocess)


def plot_tsne(chart, sample_properties, sample_data):
    """ Plot cells in t-SNE space, colored by clustering label """
    analysis = sample_data.get_analysis(SingleGenomeAnalysis)
    if analysis is None:
        return None

    args = [
        analysis.get_tsne().transformed_tsne_matrix,
        ws_gex_constants.TSNE_CLUSTER_DESCRIPTION,
        1,
        2,
    ]
    return clustering_plot_func(chart, sample_properties, sample_data, plot_dimensions, args)


def plot_dimensions(
    chart,
    transformed_matrix,
    description,
    pc1,
    pc2,
    clip=None,
    clustering=None,
    diff_expr=None,
    values=None,
    original_cluster_sizes=None,
):
    """ Plot cells in a 2-d space, colored by clustering label """
    assert (values is not None and clustering is None) or (
        clustering is not None and values is None and original_cluster_sizes is not None
    )

    if clustering is not None:
        values = clustering.clusters

    if original_cluster_sizes is None:
        value_freqs = np.bincount(values)[1:].tolist()
    else:
        # Cluster size array starts w/ cluster1 at index0; make it equivalent to np.bincount on arbitrary values
        value_freqs = [0] + original_cluster_sizes.tolist()

    n, m = transformed_matrix.shape
    if m < max(pc1, pc2):
        return None

    max_value = values.max()

    chart["data"] = [None] * max_value
    for value in np.unique(values):
        chart["data"][value - 1] = {
            "x": [],
            "y": [],
            "name": "%s %s - %s cells"
            % (
                description,
                format_value(value, "integer"),
                format_value(value_freqs[value], "integer"),
            ),
            "hoverinfo": "name",
            "mode": "markers",
            "type": "scattergl",
            "marker": {"size": 4,},
        }

    for i in range(n):
        r1 = shared_constants.DATA_VALUE_FORMAT % transformed_matrix[i, pc1 - 1]
        r2 = shared_constants.DATA_VALUE_FORMAT % transformed_matrix[i, pc2 - 1]
        value = values[i]

        series = chart["data"][value - 1]
        series["x"].append(r1)
        series["y"].append(r2)

    if clip is not None:
        xmin, xmax = np.percentile(transformed_matrix[:, pc1 - 1], clip)
        ymin, ymax = np.percentile(transformed_matrix[:, pc2 - 1], clip)
        chart["layout"]["xaxis"] = {
            "range": [xmin, xmax],
        }
        chart["layout"]["yaxis"] = {
            "range": [ymin, ymax],
        }

    chart["config"] = shared_constants.CHARTS_PLOTLY_MOVABLE_CONFIG
    return chart


def plot_dimensions_color(chart, transformed_matrix, values, description, vmin, vmax, pc1, pc2):
    _, m = transformed_matrix.shape
    if m < max(pc1, pc2):
        return None

    series = chart["data"][0]

    index_order = list(range(transformed_matrix.shape[0]))
    random.shuffle(index_order)
    for i in index_order:
        r1 = shared_constants.DATA_VALUE_FORMAT % transformed_matrix[i, pc1 - 1]
        r2 = shared_constants.DATA_VALUE_FORMAT % transformed_matrix[i, pc2 - 1]
        value = int(values[i])
        text = "%s: %s" % (description, format_value(value, "integer"))

        if value < vmin:
            value = vmin
        if value > vmax:
            value = vmax

        series["x"].append(r1)
        series["y"].append(r2)
        series["marker"]["color"].append(value)
        series["text"].append(text)

    return chart


def plot_tsne_totalcounts(chart, sample_properties, sample_data):
    """ Plot cells colored by total counts """
    analysis = sample_data.get_analysis(SingleGenomeAnalysis)
    if not analysis:
        return None

    reads_per_bc = analysis.matrix.get_counts_per_bc()
    vmin, vmax = np.percentile(reads_per_bc, ws_gex_constants.TSNE_TOTALCOUNTS_PRCT_CLIP)

    return plot_dimensions_color(
        chart,
        analysis.get_tsne().transformed_tsne_matrix,
        reads_per_bc,
        ws_gex_constants.TSNE_TOTALCOUNTS_DESCRIPTION,
        vmin,
        vmax,
        1,
        2,
    )


def _plot_differential_expression(
    chart, analysis, clustering=None, diff_expr=None, original_cluster_sizes=None
):
    n_clusters = clustering.clusters.max()

    # Get the union of top DE genes
    top_genes = set()

    # Limit the number of entries in the DE table
    n_genes = int(np.floor(float(ws_gex_constants.MAX_DE_TABLE_ENTRIES) / (n_clusters ** 2)))
    if n_genes < 1:
        n_genes = 1
    elif n_genes > ws_gex_constants.MAX_TOP_N_GENES:
        n_genes = ws_gex_constants.MAX_TOP_N_GENES

    cols = [
        {"type": "string", "label": "Gene ID"},
        {"type": "string", "label": "Gene name"},
    ]

    for i in range(n_clusters):
        # Filter genes by mean count and sort by log2 fold-change, descending
        means = diff_expr.data[:, 0 + 3 * i]
        log2fcs = diff_expr.data[:, 1 + 3 * i]

        keep_indices = np.flatnonzero(means >= ws_gex_constants.TOP_DE_GENES_MIN_MEAN)
        top_gene_indices = keep_indices[log2fcs[keep_indices].argsort()[::-1]][:n_genes]

        for j in top_gene_indices:
            top_genes.add(analysis.matrix.int_to_feature_id(j))

        cols.append(
            {
                "type": "number",
                "label": "L2FC",
                "title": "Log2 fold-change in cluster %d vs other cells" % (i + 1),
            }
        )
        cols.append(
            {
                "type": "number",
                "label": "p-value",
                "title": "Adjusted p-value of differential expression in cluster %d" % (i + 1),
            }
        )

    rows = []
    for gene_id in top_genes:
        i = analysis.matrix.feature_id_to_int(gene_id)
        gene_name = analysis.matrix.feature_id_to_name(gene_id)

        row = [gene_id, gene_name]
        for j in range(n_clusters):
            log2fc = diff_expr.data[i, 1 + (3 * j)]
            adj_p_value = diff_expr.data[i, 2 + (3 * j)]

            if log2fc <= 0 or adj_p_value >= ws_gex_constants.PVALUE_DEEMPHASIS_CUTOFF:
                style = "#DDD"
            else:
                style = "#000"

            row.append(
                {
                    "v": tk_safe_json.json_sanitize(log2fc),
                    "f": format_value(log2fc, "%.2f"),
                    "s": style,
                }
            )
            row.append(
                {
                    "v": tk_safe_json.json_sanitize(adj_p_value),
                    "f": format_value(adj_p_value, "%.0e"),
                    "s": style,
                }
            )

        rows.append(row)

    # Sort by log2fc, descending, in first cluster
    if n_clusters > 0:
        rows = sorted(rows, key=lambda row: row[2]["v"], reverse=True)

    chart["table"].update({"rows": rows, "cols": cols})
    return chart


def plot_differential_expression(chart, sample_properties, sample_data):
    sg_analysis = sample_data.get_analysis(SingleGenomeAnalysis)
    return clustering_plot_func(
        chart, sample_properties, sample_data, _plot_differential_expression, [sg_analysis]
    )


def clustering_plot_func(chart, sample_properties, sample_data, plot_func, args=[], kwargs={}):
    analysis = sample_data.get_analysis(SingleGenomeAnalysis)
    if analysis is None:
        return None

    new_charts = []
    for clustering_key, clustering in iteritems(analysis.clusterings):
        kwargs["clustering"] = clustering
        kwargs["original_cluster_sizes"] = sample_data.original_cluster_sizes[clustering_key]
        kwargs["diff_expr"] = analysis.differential_expression[clustering_key]
        new_chart = plot_func(copy.deepcopy(chart), *args, **kwargs)
        if new_chart is not None:
            new_chart["filters"] = {ws_gex_constants.CLUSTERS_FILTER_TITLE: clustering.description}
            new_charts.append(new_chart)
    return new_charts


def make_chart_filters(analyses):
    if analyses is None:
        return {}
    sg_analyses = [an for an in analyses if isinstance(an, SingleGenomeAnalysis)]
    if len(sg_analyses) == 0 or analyses[0] is None:
        return {}
    assert len(sg_analyses) == 1
    analysis = sg_analyses[0]

    filter_values = [
        x.description for x in cr_clustering.sort_clusterings(analysis.clusterings.values())
    ]

    return {
        ws_gex_constants.CLUSTERS_FILTER_TITLE: {
            "values": filter_values,
            "selected": filter_values[0],
        },
    }


def plot_subsampled_scatterplot_metric(chart, sample_properties, sample_data, **kwargs):
    """
    Modifies chart data entry to add traces for subsampled metrics with specified suffixes.
    Metrics must take the form <reference>_<subsample_type>_<subsample_depth>_<metric_suffix>, where metric suffix is specified as a kwarg

    KWARGS:
        metric_suffix (str): suffix for the subsampled metric given a metric of the form <reference>_<subsample_type>_<subsample_depth>_<metric_suffix>

    """
    summary_data = sample_data.summary or {}

    subsample_type = kwargs.get("subsample_type", "raw_rpc")

    # Just use ref_prefix if specified; otherwise try to construct it out of the set of references
    if "ref_prefix" in kwargs:
        ref_prefix = kwargs["ref_prefix"]
    elif "references" in kwargs:
        references = kwargs["references"]
        ref_prefix = "(" + "|".join(references) + ")_"
    else:
        ref_prefix = "(.+)_"

    # is the sample targeted, which lines to show
    is_targeted = kwargs.get("is_targeted", False)
    if kwargs.get("show_targeted_only", False):
        targeting_groups_to_plot = [cr_constants.ON_TARGET_SUBSAMPLE]
    else:
        targeting_groups_to_plot = [
            cr_constants.ON_TARGET_SUBSAMPLE,
            cr_constants.OFF_TARGET_SUBSAMPLE,
        ]

    # Regular expression to match <reference>_<subsample_type>_<subsample_depth>_<metric_suffix> pattern
    metric_pattern = r"^%s(%s)_([0-9]+)_%s_*(\w*)" % (
        ref_prefix,
        subsample_type,
        kwargs.get("metric_suffix", None),
    )

    metric_search_results = [re.search(metric_pattern, key) for key in summary_data.keys()]

    # Only return chart when multiple references are present for multi-genome only metrics
    if kwargs.get("multi_genome_only", False):
        references = {result.group(1) for result in metric_search_results if result is not None}
        if len(references) <= 2:
            return None

    # Find relevant metrics and extract data
    points = []

    for val, search_result in zip(summary_data.values(), metric_search_results):
        if search_result and not isinstance(val, float):
            val = float(val)
        # Only count corresponding keys and non-zero values
        if search_result and val > 0:
            genome = search_result.group(1)
            subsample_type = search_result.group(2).replace("_", " ")
            targeting_group = search_result.group(4)

            if kwargs.get("show_multi_genome_only", False) and (genome != "multi"):
                continue
            subsample_depth = int(search_result.group(3))

            if is_targeted and targeting_group not in targeting_groups_to_plot:
                continue
            elif is_targeted:
                subsample_type = "{}_{}".format(subsample_type, targeting_group)

            if genome == "multi":
                if is_targeted:
                    trace_label = cr_tgt_utils.reformat_targeted_label(targeting_group)
                else:
                    trace_label = ""
            else:
                if is_targeted:
                    trace_label = "{}_{}".format(
                        genome, cr_tgt_utils.reformat_targeted_label(targeting_group)
                    )
                else:
                    trace_label = genome

            points.append((subsample_depth, val, trace_label, genome, subsample_type,))

    # Don't produce chart if there's no data
    if len(points) == 0:
        return None

    # Sort extracted points by subsample depth (makes line plot connect properly)
    sorted_points = sorted(points, key=lambda x: (x[2], x[0]))

    # Build up a set of traces for <reference>_<subsample_type> pairs
    traces = {}

    for point in sorted_points:
        depth, value, trace_label, genome, subsample_type = point
        trace_key = (genome, subsample_type)

        if trace_key not in traces:
            # add (0,0)
            traces[trace_key] = {
                "x": [0, depth],
                "y": [0, value],
                "name": trace_label,
                "mode": "lines",
                "line": {"width": 3,},
            }
        else:
            traces[trace_key]["x"].append(depth)
            traces[trace_key]["y"].append(value)

    # Set extents of saturation line
    if (
        chart.get("layout")
        and chart["layout"].get("shapes")
        and chart["layout"]["shapes"][0]["type"] == "line"
    ):
        chart["layout"]["shapes"][0]["x1"] = max([max(trace["x"]) for trace in itervalues(traces)])

    # Add new data into chart
    chart["data"] = [v for _, v in sorted(traces.items())]
    return chart


def build_charts(sample_properties, chart_dicts, sample_data, all_prefixes={}, module=None):
    modules = [module, globals()] if module else [globals()]

    filters = make_chart_filters(sample_data.analyses)

    charts = []
    for chart_dict in chart_dicts:
        chart_dict = copy.deepcopy(chart_dict)
        function = chart_dict.pop("function")
        f = None
        for mod in modules:
            f = mod.get(function)
            if f is not None:
                break

        if function is not None and f is None:
            raise ValueError('Could not find webshim chart function "%s"' % function)

        kwargs = chart_dict.pop("kwargs", {})  # type: Dict[str, Any]

        kwargs_prefixes = chart_dict.pop("kwargs_prefixes", [])  # type: List[str]
        for prefix in kwargs_prefixes:
            kwargs[prefix] = all_prefixes[prefix]

        new_chart_obj = f(chart_dict, sample_properties, sample_data, **kwargs)
        if new_chart_obj is None:
            continue

        new_charts = new_chart_obj if isinstance(new_chart_obj, list) else [new_chart_obj]
        charts.extend(new_charts)

    return charts, filters


def filter_vdj_prefixes(all_prefixes, sample_properties):
    """ Only get subset of metric prefix values """
    chain_filter = sample_properties.chain_type
    if chain_filter is None:
        return all_prefixes

    # NOTE: Assumes chains (TRA, TRB, etc) are prefixed with the chain_type (TR or IG)
    result = {}
    for key, values in iteritems(all_prefixes):
        # Only filter any prefix that is a candidate for filtering
        # (i.e., contains some values prefixed by the selected chain type)
        if values is not None and any(v.startswith(chain_filter) for v in values):
            result[key] = [
                v
                for v in values
                if v.startswith(chain_filter) or v == rna_library.MULTI_REFS_PREFIX
            ]
        else:
            result[key] = values
    return result


def filter_vdj_alarms(all_alarms, sample_properties):
    """ Only get subset of metric alarms """
    chain_filter = sample_properties.chain_type
    if chain_filter is None:
        return all_alarms

    result = []
    for alarm in all_alarms:
        # No filters specified; don't filter
        if "filters" not in alarm:
            result.append(alarm)
            continue

        for f in alarm["filters"]:
            if f.get("chain_type") == chain_filter:
                result.append(alarm)

    return result


def get_constants_for_pipeline(pipeline, sample_properties):
    """ Get the appropriate metrics/alarms/charts for a pipeline """
    if pipeline == shared_constants.PIPELINE_VDJ:
        metrics, alarms, charts = (
            ws_vdj_constants.METRICS,
            ws_vdj_constants.METRIC_ALARMS,
            ws_vdj_constants.CHARTS,
        )

        metric_prefixes = filter_vdj_prefixes(
            vdj_report.VdjReporter().get_all_prefixes(), sample_properties
        )

        alarms = filter_vdj_alarms(alarms, sample_properties)

    else:
        metrics, alarms, charts = (
            ws_gex_constants.METRICS,
            ws_gex_constants.METRIC_ALARMS,
            ws_gex_constants.CHARTS,
        )

        metric_prefixes = cr_report.Reporter().get_all_prefixes()

    return metrics, alarms, charts, metric_prefixes


def get_custom_features(sample_data):
    """Infer the set of distinct custom feature types present in a dataset"""
    # Add the "custom feature type" prefix using the sample data
    analysis = sample_data.get_analysis(SingleGenomeAnalysis)

    # Adding a dummy value here suppresses extraneous web summary dashboards
    custom_features = ["dummy"]

    if (
        analysis
        and analysis.matrix.feature_ref.get_count_of_feature_type(rna_library.CUSTOM_LIBRARY_TYPE)
        > 0
    ):
        custom_features.append(rna_library.CUSTOM_LIBRARY_TYPE)
    return custom_features


def get_genomes(sample_data):
    """Infer the set of genomes present in a dataset"""

    analysis = sample_data.get_analysis(SingleGenomeAnalysis)
    if analysis:
        return [x for x in analysis.matrix.get_genomes() if x != ""]
    else:
        return ["dummy"]


def build_web_summary_json(sample_properties, sample_data, pipeline):
    """sample_properties - SampleProperties object
    sample_data - *SampleData class
    pipeline - string"""
    view = copy.deepcopy(sample_properties.__dict__)
    sample_properties_as_dict = copy.deepcopy(view)

    metrics, alarms, charts, all_prefixes = get_constants_for_pipeline(pipeline, sample_properties)

    all_prefixes["custom_features"] = get_custom_features(sample_data)

    tables, alarms = build_tables(
        sample_properties_as_dict, metrics, alarms, sample_data, all_prefixes=all_prefixes
    )
    if tables:
        view["tables"] = tables
    if alarms:
        view["alarms"] = alarms

    all_prefixes["references"] = list(set(all_prefixes["references"] + get_genomes(sample_data)))

    charts, filters = build_charts(
        sample_properties_as_dict, charts, sample_data, all_prefixes=all_prefixes
    )
    if charts:
        view["charts"] = charts
    if filters:
        view["filters"] = filters

    # Selected metrics that the web summary template needs
    info = build_info_dict(sample_properties, sample_data, pipeline)
    if info:
        view["info"] = info

    return view


def build_info_dict(sample_properties, sample_data, pipeline):
    """ Add miscellaneous metrics required by the web summary template """
    if sample_data.summary is None:
        return None
    info = {}

    info["chemistry_description"] = sample_data.summary.get("chemistry_description")

    info["references"] = []

    if pipeline in [shared_constants.PIPELINE_AGGR, shared_constants.PIPELINE_REANALYZE]:
        genomes = sample_properties.genomes
        if genomes is not None:
            info["references"].append(
                {
                    "type": cr_constants.REFERENCE_TYPE,
                    "name": cr_reference.get_ref_name_from_genomes(genomes),
                }
            )

    else:
        # Find all references in the summary
        reference_metric_prefixes = [
            cr_constants.REFERENCE_METRIC_PREFIX,
            vdj_constants.REFERENCE_METRIC_PREFIX,
        ]

        for prefix in reference_metric_prefixes:
            type_metric = "%s%s" % (prefix, cr_constants.REFERENCE_TYPE_KEY)

            if type_metric not in sample_data.summary:
                continue

            ref_type = sample_data.summary[type_metric]

            name_metric = "%s%s" % (prefix, cr_constants.REFERENCE_GENOMES_KEY)
            if name_metric not in sample_data.summary:
                raise ValueError(
                    "Reference metadata metric %s not found in metrics summary." % name_metric
                )
            ref_name = sample_data.summary.get(name_metric)

            if isinstance(ref_name, list):
                ref_name = cr_reference.get_ref_name_from_genomes(ref_name)

            info["references"].append({"type": ref_type, "name": ref_name})

    return info


def build_metrics_summary_csv(filename, sample_properties, sample_data, pipeline):
    metrics, alarms, _, all_prefixes = get_constants_for_pipeline(pipeline, sample_properties)

    sample_properties_as_dict = copy.deepcopy(sample_properties.__dict__)

    all_prefixes["custom_features"] = get_custom_features(sample_data)

    tables, _ = build_tables(
        sample_properties_as_dict, metrics, alarms, sample_data, all_prefixes=all_prefixes
    )
    if not tables:
        sys.stderr.write("No metrics tables were generated, skipping CSV generation.\n")
        return

    csv_metrics = collections.OrderedDict()
    for table in tables:
        if not table:
            continue
        for metric, _, value in table["rows"]:
            if isinstance(metric, dict):
                metric = metric["v"]
            if isinstance(value, dict):
                value = value["v"]
            if metric not in csv_metrics:
                csv_metrics[metric] = value

    with open(filename, "w") as f:
        writer = csv.writer(f, lineterminator="\n")
        writer.writerow(csv_metrics.keys())  # pylint: disable=dict-keys-not-iterating
        writer.writerow(csv_metrics.values())  # pylint: disable=dict-values-not-iterating


def nullable_int_sort_key(s, negate=False):
    """Return a sort key for s that will sort
    by int(s). Values of s that cannot be
    converted to an integer will be sorted
    to the begining."""

    try:
        if not negate:
            return (True, int(s))
        else:
            return (True, -int(s))
    except ValueError:
        return (False, 0)


def convert_to_float_gracefully(s):
    try:
        return float(s)
    except ValueError:
        return sys.float_info.max


def _make_jibes_color_map(tags):

    """
    Make an economist style color map for the websummary
    :param tags: list of tags
    :return: dict of {tag:color}
    """
    color_map = {}
    for i, tag in enumerate(sorted(tags)):
        color_map[ensure_str(tag)] = JIBES_COLOR_PALETTE[i]
    color_map[ensure_str(jibes.MULTIPLETS_FACTOR_NAME)] = JIBES_COLOR_PALETTE[12]
    color_map[ensure_str(jibes.BLANK_FACTOR_NAME)] = JIBES_COLOR_PALETTE[13]
    color_map[ensure_str(jibes.UNASSIGNED_FACTOR_NAME)] = JIBES_COLOR_PALETTE[14]
    return color_map


def make_jibes_histogram_plot(tag_names, counts, color_map):
    """
    :param tag names: labels for tags
     counts: the data to plot
     color_map: the colors to plot
    :return:
    jibes histogram data for websummary JSON
    """
    traces = []
    for i, tag in enumerate(tag_names):
        trace = {
            "x": counts[i],
            "type": "histogram",
            "opacity": 0.6,
            "name": ensure_str(tag),
            "marker": {"color": color_map[ensure_str(tag)],},
        }
        traces.append(trace)
    plot = {
        "data": traces,
        "layout": {"barmode": "overlay", "xaxis": {"title": "Log 10 (1 + Count)"}},
    }
    return plot


# the shared_resources implementation in python doesn't interop with Rust websummary at all
# need to just define a key to put the data under and use that
JIBES_SHARED_VECTORS_RESOURCE_KEY_PREFIX = "shared_jibes_data_vector"
RESOURCE_PREFIX = "_resources"


def make_jibes_biplot_histogram(assigner):
    """
    :param fitter: instance of JibesTagAssigner
    :return:
    dict containing jibes biplot and histogram data for websummary JSON
    """

    # pylint: disable=too-many-locals
    isinstance(assigner, JibesTagAssigner)
    fitter = assigner.fitter
    if fitter is None:
        return None
    assignments = assigner.jibes_assignments

    # we want roughly JIBES_BIPLOT_N_POINTS so find out how many rows to subsample
    n_rows_to_downsample = int(JIBES_BIPLOT_N_POINTS)
    # create a random mask for the rows with that many points
    downsample_indices = np.full(fitter.data.counts.shape[0], False)
    downsample_indices[:n_rows_to_downsample] = True
    np.random.shuffle(downsample_indices)

    # downsample the assignments
    assignments = assignments[downsample_indices]

    # downsample the data vectors
    original_vectors = [
        np.asarray(fitter.data.counts[downsample_indices, i]).ravel().tolist()
        for i in range(len(fitter.data.column_names))
    ]

    for vector in original_vectors:
        assert len(vector) == len(assignments)

    shared_resources = {}
    # this is sort of duplicating functionality that should be encapsulated by ResourceCollection in react_components.py
    # but we are sort of doing this in an ad-hoc way and passing it off to the Rust websummary code to
    # put it in the right place
    vectors = []
    for i, v in enumerate(original_vectors):
        store_key = "{}_{}".format(JIBES_SHARED_VECTORS_RESOURCE_KEY_PREFIX, i)
        use_key = "{}_{}".format(RESOURCE_PREFIX, store_key)
        shared_resources[store_key] = rc.array_to_float32_base64(v)
        vectors.append(use_key)

    vector_names = fitter.data.column_names
    color_map = _make_jibes_color_map(vector_names)
    histogram_data = make_jibes_histogram_plot(vector_names, vectors, color_map)

    # Clamp confidence values to avoid numeric issues
    assn_probs = assignments[jibes.ASSIGNMENT_PROB_COL_NAME].values
    ### Numeric issues values close to 1 by a small delta, sanity check we're close to
    # the expected range, then clamp the values to the QV min/max.
    over_1 = np.sum(assn_probs > 1.0001).sum() > 0
    under_0 = np.sum(assn_probs < -0.0001).sum() > 0
    if over_1 or under_0:
        # Should never ever happen
        raise Exception(
            "Normalized probabilities were outside of tolerance ranges.  This likely indicates a bug, please contact support@10xgenomics.com."
        )
    max_confidence = 1 - np.power(10, -(MAX_QV_VAL / 10.0))
    min_confidence = 1 - np.power(10, -(MIN_QV_VAL / 10.0))
    assn_probs = assn_probs.clip(min_confidence, max_confidence)
    phred_values = -10.0 * np.log10(1.0 - assn_probs)
    phred_values = phred_values.tolist()
    slider_values = [-10.0 * np.log10(1.0 - jibes.JIBES_MIN_CONFIDENCE)]

    biplot_data = rc.BiplotsWithThresholding(
        vectors,
        vector_names,
        True,
        False,
        phred_values,
        slider_values,
        "Assignment Phred Score Threshold",
        "Singlets",
        "Filtered Singlets",
        "% Filtered",
        [
            ensure_str(jibes.BLANK_FACTOR_NAME),
            ensure_str(jibes.MULTIPLETS_FACTOR_NAME),
            ensure_str(jibes.UNASSIGNED_FACTOR_NAME),
        ],
        [ensure_str(x) for x in assignments[jibes.ASSIGNMENT_COL_NAME].values.tolist()],
        color_map,
    )

    jibes_biplot_histogram = {
        "biplot": biplot_data,
        "histogram": histogram_data,
        "_resources": shared_resources,
    }

    return jibes_biplot_histogram
