import io
import itertools

import pytest
from pandas.util.testing import (
    assert_series_equal, assert_frame_equal, assert_index_equal)
from numpy.testing import assert_array_equal
import pandas as pd
import numpy as np
import matplotlib.figure
import matplotlib.pyplot as plt
from matplotlib.text import Text
from matplotlib.colors import to_hex
from matplotlib import cm

from upsetplot import plot
from upsetplot import UpSet
from upsetplot import generate_counts, generate_samples
from upsetplot.plotting import _process_data

# TODO: warnings should raise errors


def is_ascending(seq):
    # return np.all(np.diff(seq) >= 0)
    return sorted(seq) == list(seq)


def get_all_texts(mpl_artist):
    out = [text.get_text() for text in mpl_artist.findobj(Text)]
    return [text for text in out if text]


@pytest.mark.parametrize('x', [
    generate_counts(),
    generate_counts().iloc[1:-2],
])
@pytest.mark.parametrize('sort_by', ['cardinality', 'degree', None])
@pytest.mark.parametrize('sort_categories_by', [None, 'cardinality'])
def test_process_data_series(x, sort_by, sort_categories_by):
    assert x.name == 'value'
    for subset_size in ['auto', 'sum', 'count']:
        for sum_over in ['abc', False]:
            with pytest.raises(ValueError, match='sum_over is not applicable'):
                _process_data(x, sort_by=sort_by,
                              sort_categories_by=sort_categories_by,
                              subset_size=subset_size, sum_over=sum_over)

    # shuffle input to test sorting
    x = x.sample(frac=1., replace=False, random_state=0)

    total, df, intersections, totals = _process_data(
        x, subset_size='auto', sort_by=sort_by,
        sort_categories_by=sort_categories_by, sum_over=None)

    assert total == x.sum()

    assert intersections.name == 'value'
    x_reordered = (x
                   .reorder_levels(intersections.index.names)
                   .reindex(index=intersections.index))
    assert len(x) == len(x_reordered)
    assert x_reordered.index.is_unique
    assert_series_equal(x_reordered, intersections,
                        check_dtype=False)

    if sort_by == 'cardinality':
        assert is_ascending(intersections.values[::-1])
    elif sort_by == 'degree':
        # check degree order
        assert is_ascending(intersections.index.to_frame().sum(axis=1))
        # TODO: within a same-degree group, the tuple of active names should
        #       be in sort-order
    else:
        find_first_in_orig = x_reordered.index.tolist().index
        orig_order = list(map(find_first_in_orig,
                              intersections.index.tolist()))
        assert orig_order == sorted(orig_order)
    if sort_categories_by:
        assert is_ascending(totals.values[::-1])

    assert np.all(totals.index.values == intersections.index.names)

    assert np.all(df.index.names == intersections.index.names)
    assert set(df.columns) == {'_value', '_bin'}
    assert_index_equal(df['_value'].reorder_levels(x.index.names).index,
                       x.index)
    assert_array_equal(df['_value'], x)
    assert_index_equal(intersections.iloc[df['_bin']].index,
                       df.index)
    assert len(df) == len(x)


@pytest.mark.parametrize('x', [
    generate_samples()['value'],
    generate_counts(),
])
def test_subset_size_series(x):
    kw = {'sort_by': 'cardinality',
          'sort_categories_by': 'cardinality',
          'sum_over': None}
    total, df_sum, intersections_sum, totals_sum = _process_data(
        x, subset_size='sum', **kw)
    assert total == intersections_sum.sum()

    if x.index.is_unique:
        total, df, intersections, totals = _process_data(
            x, subset_size='auto', **kw)
        assert total == intersections.sum()
        assert_frame_equal(df, df_sum)
        assert_series_equal(intersections, intersections_sum)
        assert_series_equal(totals, totals_sum)
    else:
        with pytest.raises(ValueError):
            _process_data(x, subset_size='auto', **kw)

    total, df_count, intersections_count, totals_count = _process_data(
        x, subset_size='count', **kw)
    assert total == intersections_count.sum()
    total, df, intersections, totals = _process_data(
        x.groupby(level=list(range(len(x.index.levels)))).count(),
        subset_size='sum', **kw)
    assert total == intersections.sum()
    assert_series_equal(intersections, intersections_count, check_names=False)
    assert_series_equal(totals, totals_count)


@pytest.mark.parametrize('x', [
    generate_samples()['value'],
])
@pytest.mark.parametrize('sort_by', ['cardinality', 'degree', None])
@pytest.mark.parametrize('sort_categories_by', [None, 'cardinality'])
def test_process_data_frame(x, sort_by, sort_categories_by):
    # shuffle input to test sorting
    x = x.sample(frac=1., replace=False, random_state=0)

    X = pd.DataFrame({'a': x})

    with pytest.warns(None):
        total, df, intersections, totals = _process_data(
            X, sort_by=sort_by, sort_categories_by=sort_categories_by,
            sum_over='a', subset_size='auto')
    assert df is not X
    assert total == pytest.approx(intersections.sum())

    # check equivalence to Series
    total1, df1, intersections1, totals1 = _process_data(
        x, sort_by=sort_by, sort_categories_by=sort_categories_by,
        subset_size='sum', sum_over=None)

    assert intersections.name == 'a'
    assert_frame_equal(df, df1.rename(columns={'_value': 'a'}))
    assert_series_equal(intersections, intersections1, check_names=False)
    assert_series_equal(totals, totals1)

    # check effect of extra column
    X = pd.DataFrame({'a': x, 'b': np.arange(len(x))})
    total2, df2, intersections2, totals2 = _process_data(
        X, sort_by=sort_by, sort_categories_by=sort_categories_by,
        sum_over='a', subset_size='auto')
    assert total2 == pytest.approx(intersections2.sum())
    assert_series_equal(intersections, intersections2)
    assert_series_equal(totals, totals2)
    assert_frame_equal(df, df2.drop('b', axis=1))
    assert_array_equal(df2['b'], X['b'])  # disregard levels, tested above

    # check effect not dependent on order/name
    X = pd.DataFrame({'b': np.arange(len(x)), 'c': x})
    total3, df3, intersections3, totals3 = _process_data(
        X, sort_by=sort_by, sort_categories_by=sort_categories_by,
        sum_over='c', subset_size='auto')
    assert total3 == pytest.approx(intersections3.sum())
    assert_series_equal(intersections, intersections3, check_names=False)
    assert intersections.name == 'a'
    assert intersections3.name == 'c'
    assert_series_equal(totals, totals3)
    assert_frame_equal(df.rename(columns={'a': 'c'}), df3.drop('b', axis=1))
    assert_array_equal(df3['b'], X['b'])

    # check subset_size='count'
    X = pd.DataFrame({'b': np.ones(len(x), dtype=int), 'c': x})
    total4, df4, intersections4, totals4 = _process_data(
        X, sort_by=sort_by, sort_categories_by=sort_categories_by,
        sum_over='b', subset_size='auto')
    total5, df5, intersections5, totals5 = _process_data(
        X, sort_by=sort_by, sort_categories_by=sort_categories_by,
        subset_size='count', sum_over=None)
    assert total5 == pytest.approx(intersections5.sum())
    assert_series_equal(intersections4, intersections5, check_names=False)
    assert intersections4.name == 'b'
    assert intersections5.name == 'size'
    assert_series_equal(totals4, totals5)
    assert_frame_equal(df4, df5)


@pytest.mark.parametrize('x', [
    generate_samples()['value'],
    generate_counts(),
])
def test_subset_size_frame(x):
    kw = {'sort_by': 'cardinality',
          'sort_categories_by': 'cardinality'}
    X = pd.DataFrame({'x': x})
    total_sum, df_sum, intersections_sum, totals_sum = _process_data(
        X, subset_size='sum', sum_over='x', **kw)
    total_count, df_count, intersections_count, totals_count = _process_data(
        X, subset_size='count', sum_over=None, **kw)

    # error cases: sum_over=False
    for subset_size in ['auto', 'sum', 'count']:
        with pytest.raises(ValueError, match='sum_over'):
            _process_data(
                X, subset_size=subset_size, sum_over=False, **kw)

    with pytest.raises(ValueError, match='sum_over'):
        _process_data(
            X, subset_size=subset_size, sum_over=False, **kw)

    # error cases: sum_over incompatible with subset_size
    with pytest.raises(ValueError, match='sum_over should be a field'):
        _process_data(
            X, subset_size='sum', sum_over=None, **kw)
    with pytest.raises(ValueError, match='sum_over cannot be set'):
        _process_data(
            X, subset_size='count', sum_over='x', **kw)

    # check subset_size='auto' with sum_over=str => sum
    total, df, intersections, totals = _process_data(
        X, subset_size='auto', sum_over='x', **kw)
    assert total == intersections.sum()
    assert_frame_equal(df, df_sum)
    assert_series_equal(intersections, intersections_sum)
    assert_series_equal(totals, totals_sum)

    # check subset_size='auto' with sum_over=None => count
    total, df, intersections, totals = _process_data(
        X, subset_size='auto', sum_over=None, **kw)
    assert total == intersections.sum()
    assert_frame_equal(df, df_count)
    assert_series_equal(intersections, intersections_count)
    assert_series_equal(totals, totals_count)


@pytest.mark.parametrize('sort_by', ['cardinality', 'degree'])
@pytest.mark.parametrize('sort_categories_by', [None, 'cardinality'])
def test_not_unique(sort_by, sort_categories_by):
    kw = {'sort_by': sort_by,
          'sort_categories_by': sort_categories_by,
          'subset_size': 'sum',
          'sum_over': None}
    Xagg = generate_counts()
    total1, df1, intersections1, totals1 = _process_data(Xagg, **kw)
    Xunagg = generate_samples()['value']
    Xunagg.loc[:] = 1
    total2, df2, intersections2, totals2 = _process_data(Xunagg, **kw)
    assert_series_equal(intersections1, intersections2,
                        check_dtype=False)
    assert total2 == intersections2.sum()
    assert_series_equal(totals1, totals2, check_dtype=False)
    assert set(df1.columns) == {'_value', '_bin'}
    assert set(df2.columns) == {'_value', '_bin'}
    assert len(df2) == len(Xunagg)
    assert df2['_bin'].nunique() == len(intersections2)


@pytest.mark.parametrize('kw', [{'sort_by': 'blah'},
                                {'sort_by': True},
                                {'sort_categories_by': 'blah'},
                                {'sort_categories_by': True}])
def test_param_validation(kw):
    X = generate_counts(n_samples=100)
    with pytest.raises(ValueError):
        UpSet(X, **kw)


@pytest.mark.parametrize('kw', [{},
                                {'element_size': None},
                                {'orientation': 'vertical'},
                                {'intersection_plot_elements': 0},
                                {'facecolor': 'red'},
                                {'shading_color': 'lightgrey',
                                 'other_dots_color': 'pink'}])
def test_plot_smoke_test(kw):
    fig = matplotlib.figure.Figure()
    X = generate_counts(n_samples=100)
    axes = plot(X, fig, **kw)
    fig.savefig(io.BytesIO(), format='png')

    attr = ('get_xlim'
            if kw.get('orientation', 'horizontal') == 'horizontal'
            else 'get_ylim')
    lim = getattr(axes['matrix'], attr)()
    expected_width = len(X)
    assert expected_width == lim[1] - lim[0]

    # Also check fig is optional
    n_nums = len(plt.get_fignums())
    plot(X, **kw)
    assert len(plt.get_fignums()) - n_nums == 1
    assert plt.gcf().axes


@pytest.mark.parametrize('set1',
                         itertools.product([False, True], repeat=2))
@pytest.mark.parametrize('set2',
                         itertools.product([False, True], repeat=2))
def test_two_sets(set1, set2):
    # we had a bug where processing failed if no items were in some set
    fig = matplotlib.figure.Figure()
    plot(pd.DataFrame({'val': [5, 7],
                       'set1': set1,
                       'set2': set2}).set_index(['set1', 'set2'])['val'],
         fig, subset_size='sum')


def test_vertical():
    X = generate_counts(n_samples=100)

    fig = matplotlib.figure.Figure()
    UpSet(X, orientation='horizontal').make_grid(fig)
    horz_height = fig.get_figheight()
    horz_width = fig.get_figwidth()
    assert horz_height < horz_width

    fig = matplotlib.figure.Figure()
    UpSet(X, orientation='vertical').make_grid(fig)
    vert_height = fig.get_figheight()
    vert_width = fig.get_figwidth()
    assert horz_width / horz_height > vert_width / vert_height

    # TODO: test axes positions, plot order, bar orientation
    pass


def test_element_size():
    X = generate_counts(n_samples=100)
    figsizes = []
    for element_size in range(10, 50, 5):
        fig = matplotlib.figure.Figure()
        UpSet(X, element_size=element_size).make_grid(fig)
        figsizes.append((fig.get_figwidth(), fig.get_figheight()))

    figwidths, figheights = zip(*figsizes)
    # Absolute width increases
    assert np.all(np.diff(figwidths) > 0)
    aspect = np.divide(figwidths, figheights)
    # Font size stays constant, so aspect ratio decreases
    assert np.all(np.diff(aspect) <= 1e-8)  # allow for near-equality
    assert np.any(np.diff(aspect) < 1e-4)  # require some significant decrease
    # But doesn't decrease by much
    assert np.all(aspect[:-1] / aspect[1:] < 1.1)

    fig = matplotlib.figure.Figure()
    figsize_before = fig.get_figwidth(), fig.get_figheight()
    UpSet(X, element_size=None).make_grid(fig)
    figsize_after = fig.get_figwidth(), fig.get_figheight()
    assert figsize_before == figsize_after

    # TODO: make sure axes are all within figure
    # TODO: make sure text does not overlap axes, even with element_size=None


def _walk_artists(el):
    children = el.get_children()
    yield el, children
    for ch in children:
        for x in _walk_artists(ch):
            yield x


def _count_descendants(el):
    return sum(len(children) for x, children in _walk_artists(el))


@pytest.mark.parametrize('orientation', ['horizontal', 'vertical'])
def test_show_counts(orientation):
    fig = matplotlib.figure.Figure()
    X = generate_counts(n_samples=10000)
    plot(X, fig, orientation=orientation)
    n_artists_no_sizes = _count_descendants(fig)

    fig = matplotlib.figure.Figure()
    plot(X, fig, orientation=orientation, show_counts=True)
    n_artists_yes_sizes = _count_descendants(fig)
    assert n_artists_yes_sizes - n_artists_no_sizes > 6
    assert '9547' in get_all_texts(fig)  # set size
    assert '283' in get_all_texts(fig)   # intersection size

    fig = matplotlib.figure.Figure()
    plot(X, fig, orientation=orientation, show_counts='%0.2g')
    assert n_artists_yes_sizes == _count_descendants(fig)
    assert '9.5e+03' in get_all_texts(fig)
    assert '2.8e+02' in get_all_texts(fig)

    fig = matplotlib.figure.Figure()
    plot(X, fig, orientation=orientation, show_percentages=True)
    assert n_artists_yes_sizes == _count_descendants(fig)
    assert '95.5%' in get_all_texts(fig)
    assert '2.8%' in get_all_texts(fig)

    fig = matplotlib.figure.Figure()
    plot(X, fig, orientation=orientation, show_counts=True,
         show_percentages=True)
    assert n_artists_yes_sizes == _count_descendants(fig)
    if orientation == 'vertical':
        assert '9547\n(95.5%)' in get_all_texts(fig)
        assert '283 (2.8%)' in get_all_texts(fig)
    else:
        assert '9547 (95.5%)' in get_all_texts(fig)
        assert '283\n(2.8%)' in get_all_texts(fig)

    with pytest.raises(ValueError):
        fig = matplotlib.figure.Figure()
        plot(X, fig, orientation=orientation, show_counts='%0.2h')


def test_add_catplot():
    pytest.importorskip('seaborn')
    X = generate_counts(n_samples=100)
    upset = UpSet(X)
    # smoke test
    upset.add_catplot('violin')
    fig = matplotlib.figure.Figure()
    upset.plot(fig)

    # can't provide value with Series
    with pytest.raises(ValueError):
        upset.add_catplot('violin', value='foo')

    # check the above add_catplot did not break the state
    upset.plot(fig)

    X = generate_counts(n_samples=100)
    X.name = 'foo'
    X = X.to_frame()
    upset = UpSet(X, subset_size='count')
    # must provide value with DataFrame
    with pytest.raises(ValueError):
        upset.add_catplot('violin')
    upset.add_catplot('violin', value='foo')
    with pytest.raises(ValueError):
        # not a known column
        upset.add_catplot('violin', value='bar')
    upset.plot(fig)

    # invalid plot kind raises error when plotting
    upset.add_catplot('foobar', value='foo')
    with pytest.raises(AttributeError):
        upset.plot(fig)


def _get_patch_data(axes, is_vertical):
    out = [{"y": patch.get_y(), "x": patch.get_x(),
            "h": patch.get_height(), "w": patch.get_width(),
            "fc": patch.get_facecolor(),
            "ec": patch.get_edgecolor(),
            "lw": patch.get_linewidth(),
            "ls": patch.get_linestyle(),
            "hatch": patch.get_hatch(),
            }
           for patch in axes.patches]
    if is_vertical:
        out = [{"y": patch["x"], "x": 6.5 - patch["y"],
                "h": patch["w"], "w": patch["h"],
                "fc": patch["fc"],
                "ec": patch["ec"],
                "lw": patch["lw"],
                "ls": patch["ls"],
                "hatch": patch["hatch"],
                }
               for patch in out]
    return pd.DataFrame(out).sort_values("x").reset_index(drop=True)


def _get_color_to_label_from_legend(ax):
    handles, labels = ax.get_legend_handles_labels()
    color_to_label = {
        patches[0].get_facecolor(): label
        for patches, label in zip(handles, labels)
    }
    return color_to_label


@pytest.mark.parametrize('orientation', ['horizontal', 'vertical'])
@pytest.mark.parametrize('show_counts', [False, True])
def test_add_stacked_bars(orientation, show_counts):
    df = generate_samples()
    df["label"] = (pd.cut(generate_samples().value + np.random.rand() / 2, 3)
                   .cat.codes
                   .map({0: "foo", 1: "bar", 2: "baz"}))

    upset = UpSet(df, show_counts=show_counts, orientation=orientation)
    upset.add_stacked_bars(by="label")
    upset_axes = upset.plot()

    int_axes = upset_axes["intersections"]
    stacked_axes = upset_axes["extra1"]

    is_vertical = orientation == 'vertical'
    int_rects = _get_patch_data(int_axes, is_vertical)
    stacked_rects = _get_patch_data(stacked_axes, is_vertical)

    # check bar heights match between int_rects and stacked_rects
    assert_series_equal(int_rects.groupby("x")["h"].sum(),
                        stacked_rects.groupby("x")["h"].sum(),
                        check_dtype=False)
    # check count labels match (TODO: check coordinate)
    assert ([elem.get_text() for elem in int_axes.texts] ==
            [elem.get_text() for elem in stacked_axes.texts])

    color_to_label = _get_color_to_label_from_legend(stacked_axes)
    stacked_rects["label"] = stacked_rects["fc"].map(color_to_label)
    # check totals for each label
    assert_series_equal(stacked_rects.groupby("label")["h"].sum(),
                        df.groupby("label").size(),
                        check_dtype=False, check_names=False)

    label_order = [text_obj.get_text()
                   for text_obj in stacked_axes.get_legend().get_texts()]
    # label order should be lexicographic
    assert label_order == sorted(label_order)

    if orientation == "horizontal":
        # order of labels in legend should match stack, top to bottom
        for prev, curr in zip(label_order, label_order[1:]):
            assert (stacked_rects.query("label == @prev")
                    .sort_values("x")["y"].values >=
                    stacked_rects.query("label == @curr")
                    .sort_values("x")["y"].values).all()
    else:
        # order of labels in legend should match stack, left to right
        for prev, curr in zip(label_order, label_order[1:]):
            assert (stacked_rects.query("label == @prev")
                    .sort_values("x")["y"].values <=
                    stacked_rects.query("label == @curr")
                    .sort_values("x")["y"].values).all()


@pytest.mark.parametrize("colors, expected", [
    (["blue", "red", "green"], ["blue", "red", "green"]),
    ({"bar": "blue", "baz": "red", "foo": "green"}, ["blue", "red", "green"]),
    ("Pastel1", ["#fbb4ae", "#b3cde3", "#ccebc5"]),
    (cm.viridis, ["#440154", "#440256", "#450457"]),
    (lambda x: cm.Pastel1(x), ["#fbb4ae", "#b3cde3", "#ccebc5"]),
])
def test_add_stacked_bars_colors(colors, expected):
    df = generate_samples()
    df["label"] = (pd.cut(generate_samples().value + np.random.rand() / 2, 3)
                   .cat.codes
                   .map({0: "foo", 1: "bar", 2: "baz"}))

    upset = UpSet(df)
    upset.add_stacked_bars(by="label", colors=colors,
                           title="Count by gender")
    upset_axes = upset.plot()
    stacked_axes = upset_axes["extra1"]
    color_to_label = _get_color_to_label_from_legend(stacked_axes)
    label_to_color = {v: k for k, v in color_to_label.items()}
    actual = [to_hex(label_to_color[label]) for label in ["bar", "baz", "foo"]]
    expected = [to_hex(color) for color in expected]
    assert actual == expected


@pytest.mark.parametrize('int_sum_over', [False, True])
@pytest.mark.parametrize('stack_sum_over', [False, True])
@pytest.mark.parametrize('show_counts', [False, True])
def test_add_stacked_bars_sum_over(int_sum_over, stack_sum_over, show_counts):
    # A rough test of sum_over
    df = generate_samples()
    df["label"] = (pd.cut(generate_samples().value + np.random.rand() / 2, 3)
                   .cat.codes
                   .map({0: "foo", 1: "bar", 2: "baz"}))

    upset = UpSet(df, sum_over="value" if int_sum_over else None,
                  show_counts=show_counts)
    upset.add_stacked_bars(by="label",
                           sum_over="value" if stack_sum_over else None,
                           colors='Pastel1')
    upset_axes = upset.plot()

    int_axes = upset_axes["intersections"]
    stacked_axes = upset_axes["extra1"]

    int_rects = _get_patch_data(int_axes, is_vertical=False)
    stacked_rects = _get_patch_data(stacked_axes, is_vertical=False)

    if int_sum_over == stack_sum_over:
        # check bar heights match between int_rects and stacked_rects
        assert_series_equal(int_rects.groupby("x")["h"].sum(),
                            stacked_rects.groupby("x")["h"].sum(),
                            check_dtype=False)
        # and check labels match with show_counts
        assert ([elem.get_text() for elem in int_axes.texts] ==
                [elem.get_text() for elem in stacked_axes.texts])
    else:
        assert (int_rects.groupby("x")["h"].sum() !=
                stacked_rects.groupby("x")["h"].sum()).all()
        if show_counts:
            assert ([elem.get_text() for elem in int_axes.texts] !=
                    [elem.get_text() for elem in stacked_axes.texts])


@pytest.mark.parametrize('x', [
    generate_counts(),
])
def test_index_must_be_bool(x):
    # Truthy ints are okay
    x = x.reset_index()
    x[['cat0', 'cat2', 'cat2']] = x[['cat0', 'cat1', 'cat2']].astype(int)
    x = x.set_index(['cat0', 'cat1', 'cat2']).iloc[:, 0]

    UpSet(x)

    # other ints are not
    x = x.reset_index()
    x[['cat0', 'cat2', 'cat2']] = x[['cat0', 'cat1', 'cat2']] + 1
    x = x.set_index(['cat0', 'cat1', 'cat2']).iloc[:, 0]
    with pytest.raises(ValueError, match='not boolean'):
        UpSet(x)


@pytest.mark.parametrize(
    "filter_params, expected",
    [
        ({"min_subset_size": 623},
         {(True, False, False): 884,
          (True, True, False): 1547,
          (True, False, True): 623,
          (True, True, True): 990,
          }),
        ({"min_subset_size": 800, "max_subset_size": 990},
         {(True, False, False): 884,
          (True, True, True): 990,
          }),
        ({"min_degree": 2},
         {(True, True, False): 1547,
          (True, False, True): 623,
          (False, True, True): 258,
          (True, True, True): 990,
          }),
        ({"min_degree": 2, "max_degree": 2},
         {(True, True, False): 1547,
          (True, False, True): 623,
          (False, True, True): 258,
          }),
        ({"max_subset_size": 500, "max_degree": 2},
         {(False, False, False): 220,
          (False, True, False): 335,
          (False, False, True): 143,
          (False, True, True): 258,
          }),
    ]
)
@pytest.mark.parametrize('sort_by', ['cardinality', 'degree'])
def test_filter_subsets(filter_params, expected, sort_by):
    data = generate_samples(seed=0, n_samples=5000, n_categories=3)
    # data =
    #   cat1   cat0   cat2
    #   False  False  False     220
    #   True   False  False     884
    #   False  True   False     335
    #          False  True      143
    #   True   True   False    1547
    #          False  True      623
    #   False  True   True      258
    #   True   True   True      990
    upset_full = UpSet(data, subset_size='auto', sort_by=sort_by)
    upset_filtered = UpSet(data, subset_size='auto',
                           sort_by=sort_by,
                           **filter_params)
    intersections = upset_full.intersections
    df = upset_full._df
    # check integrity of expected, just to be sure
    for key, value in expected.items():
        assert intersections.loc[key] == value
    subset_intersections = intersections[
        intersections.index.isin(list(expected.keys()))]
    subset_df = df[df.index.isin(list(expected.keys()))]
    assert len(subset_intersections) < len(intersections)
    assert_series_equal(upset_filtered.intersections, subset_intersections)
    assert_frame_equal(upset_filtered._df.drop("_bin", axis=1),
                       subset_df.drop("_bin", axis=1))
    # category totals should not be affected
    assert_series_equal(upset_full.totals, upset_filtered.totals)


@pytest.mark.parametrize('x', [
    generate_counts(n_categories=3),
    generate_counts(n_categories=8),
    generate_counts(n_categories=15),
])
@pytest.mark.parametrize('orientation', [
    'horizontal',
    'vertical',
])
def test_matrix_plot_margins(x, orientation):
    """Non-regression test addressing a bug where there is are large whitespace
       margins around the matrix when the number of intersections is large"""
    axes = plot(x, orientation=orientation)

    # Expected behavior is that each matrix column takes up one unit on x-axis
    expected_width = len(x)
    attr = 'get_xlim' if orientation == 'horizontal' else 'get_ylim'
    lim = getattr(axes['matrix'], attr)()
    assert expected_width == lim[1] - lim[0]


def _make_facecolor_list(colors):
    return [{"facecolor": c} for c in colors]


CAT1_2_RED_STYLES = _make_facecolor_list(["blue", "blue", "blue", "blue",
                                          "red", "blue", "blue", "red"])
CAT1_RED_STYLES = _make_facecolor_list(["blue", "red", "blue", "blue",
                                        "red", "red", "blue", "red"])
CAT_NOT1_RED_STYLES = _make_facecolor_list(["red", "blue", "red", "red",
                                            "blue", "blue", "red", "blue"])
CAT1_NOT2_RED_STYLES = _make_facecolor_list(["blue", "red", "blue", "blue",
                                             "blue", "red", "blue", "blue"])
CAT_NOT1_2_RED_STYLES = _make_facecolor_list(["red", "blue", "blue", "red",
                                              "blue", "blue", "blue", "blue"])


@pytest.mark.parametrize(
    "kwarg_list,expected_subset_styles,expected_legend",
    [
        # Different forms of including two categories
        ([{"present": ["cat1", "cat2"], "facecolor": "red"}],
         CAT1_2_RED_STYLES, []),
        ([{"present": {"cat1", "cat2"}, "facecolor": "red"}],
         CAT1_2_RED_STYLES, []),
        ([{"present": ("cat1", "cat2"), "facecolor": "red"}],
         CAT1_2_RED_STYLES, []),
        # with legend
        ([{"present": ("cat1", "cat2"), "facecolor": "red", "label": "foo"}],
         CAT1_2_RED_STYLES, [({"facecolor": "red"}, "foo")]),
        # present only cat1
        ([{"present": ("cat1",), "facecolor": "red"}],
         CAT1_RED_STYLES, []),
        ([{"present": "cat1", "facecolor": "red"}],
         CAT1_RED_STYLES, []),
        # Some uses of absent
        ([{"absent": "cat1", "facecolor": "red"}],
         CAT_NOT1_RED_STYLES, []),
        ([{"present": "cat1", "absent": ["cat2"], "facecolor": "red"}],
         CAT1_NOT2_RED_STYLES, []),
        ([{"absent": ["cat2", "cat1"], "facecolor": "red"}],
         CAT_NOT1_2_RED_STYLES, []),
        # min/max args
        ([{"present": ["cat1", "cat2"], "min_degree": 3, "facecolor": "red"}],
         _make_facecolor_list(["blue"] * 7 + ["red"]), []),
        ([{"present": ["cat1", "cat2"], "max_subset_size": 3000,
           "facecolor": "red"}],
         _make_facecolor_list(["blue"] * 7 + ["red"]), []),
        ([{"present": ["cat1", "cat2"], "max_degree": 2, "facecolor": "red"}],
         _make_facecolor_list(["blue"] * 4 + ["red"] + ["blue"] * 3), []),
        ([{"present": ["cat1", "cat2"], "min_subset_size": 3000,
           "facecolor": "red"}],
         _make_facecolor_list(["blue"] * 4 + ["red"] + ["blue"] * 3), []),
        # cat1 _or_ cat2
        ([{"present": "cat1", "facecolor": "red"},
          {"present": "cat2", "facecolor": "red"}],
         _make_facecolor_list(["blue", "red", "red", "blue",
                               "red", "red", "red", "red"]), []),
        # With multiple uses of label
        ([{"present": "cat1", "facecolor": "red", "label": "foo"},
          {"present": "cat2", "facecolor": "red", "label": "bar"}],
         _make_facecolor_list(["blue", "red", "red", "blue",
                               "red", "red", "red", "red"]),
         [({"facecolor": "red"}, "foo; bar")]),
        ([{"present": "cat1", "facecolor": "red", "label": "foo"},
          {"present": "cat2", "facecolor": "red", "label": "foo"}],
         _make_facecolor_list(["blue", "red", "red", "blue",
                               "red", "red", "red", "red"]),
         [({"facecolor": "red"}, "foo")]),
        # With multiple colours, the latest overrides
        ([{"present": "cat1", "facecolor": "red", "label": "foo"},
          {"present": "cat2", "facecolor": "green", "label": "bar"}],
         _make_facecolor_list(["blue", "red", "green", "blue",
                               "green", "red", "green", "green"]),
         [({"facecolor": "red"}, "foo"),
          ({"facecolor": "green"}, "bar")]),
        # Combining multiple style properties
        ([{"present": "cat1", "facecolor": "red", "hatch": "//"},
          {"present": "cat2", "edgecolor": "green", "linestyle": "dotted"}],
         [{"facecolor": "blue"},
          {"facecolor": "red", "hatch": "//"},
          {"facecolor": "blue", "edgecolor": "green", "linestyle": "dotted"},
          {"facecolor": "blue"},
          {"facecolor": "red", "hatch": "//", "edgecolor": "green",
           "linestyle": "dotted"},
          {"facecolor": "red", "hatch": "//"},
          {"facecolor": "blue", "edgecolor": "green",
           "linestyle": "dotted"},
          {"facecolor": "red", "hatch": "//", "edgecolor": "green",
           "linestyle": "dotted"},
          ],
         []),
    ])
def test_style_subsets(kwarg_list, expected_subset_styles, expected_legend):
    data = generate_counts()
    upset = UpSet(data, facecolor="blue")
    for kw in kwarg_list:
        upset.style_subsets(**kw)
    actual_subset_styles = upset.subset_styles
    assert actual_subset_styles == expected_subset_styles
    assert upset.subset_legend == expected_legend


def _dots_to_dataframe(ax, is_vertical):
    matrix_path_collection = ax.collections[0]
    matrix_dots = pd.DataFrame(
        matrix_path_collection.get_offsets(), columns=["x", "y"]
    ).join(
        pd.DataFrame(matrix_path_collection.get_facecolors(),
                     columns=["fc_r", "fc_g", "fc_b", "fc_a"]),
    ).join(
        pd.DataFrame(matrix_path_collection.get_edgecolors(),
                     columns=["ec_r", "ec_g", "ec_b", "ec_a"]),
    ).assign(
        lw=matrix_path_collection.get_linewidths(),
        ls=matrix_path_collection.get_linestyles(),
        hatch=matrix_path_collection.get_hatch(),
    )

    matrix_dots["ls_offset"] = matrix_dots["ls"].map(
        lambda tup: tup[0]).astype(float)
    matrix_dots["ls_seq"] = matrix_dots["ls"].map(
        lambda tup: None if tup[1] is None else tuple(tup[1]))
    del matrix_dots["ls"]

    if is_vertical:
        matrix_dots[["x", "y"]] = matrix_dots[["y", "x"]]
        matrix_dots["x"] = 7 - matrix_dots["x"]
    return matrix_dots


@pytest.mark.parametrize('orientation', ['horizontal', 'vertical'])
def test_style_subsets_artists(orientation):
    # Check that subset_styles are all appropriately reflected in matplotlib
    # artists.
    # This may be a bit overkill, and too coupled with implementation details.
    is_vertical = orientation == 'vertical'
    data = generate_counts()
    upset = UpSet(data, orientation=orientation)
    subset_styles = [
        {"facecolor": "black"},
        {"facecolor": "red"},
        {"edgecolor": "red"},
        {"edgecolor": "red", "linewidth": 4},
        {"linestyle": "dotted"},
        {"edgecolor": "red", "facecolor": "blue", "hatch": "//"},
        {"facecolor": "blue"},
        {},
    ]

    if is_vertical:
        upset.subset_styles = subset_styles[::-1]
    else:
        upset.subset_styles = subset_styles

    upset_axes = upset.plot()

    int_rects = _get_patch_data(upset_axes["intersections"], is_vertical)
    int_rects[["fc_r", "fc_g", "fc_b", "fc_a"]] = (
        int_rects.pop("fc").apply(lambda x: pd.Series(x)))
    int_rects[["ec_r", "ec_g", "ec_b", "ec_a"]] = (
        int_rects.pop("ec").apply(lambda x: pd.Series(x)))
    int_rects["ls_is_solid"] = int_rects.pop("ls").map(
        lambda x: x == "solid" or pd.isna(x))
    expected = pd.DataFrame({
        "fc_r": [0, 1, 0, 0, 0, 0, 0, 0],
        "fc_g": [0, 0, 0, 0, 0, 0, 0, 0],
        "fc_b": [0, 0, 0, 0, 0, 1, 1, 0],
        "ec_r": [0, 1, 1, 1, 0, 1, 0, 0],
        "ec_g": [0, 0, 0, 0, 0, 0, 0, 0],
        "ec_b": [0, 0, 0, 0, 0, 0, 1, 0],
        "lw": [1, 1, 1, 4, 1, 1, 1, 1],
        "ls_is_solid": [True, True, True, True, False, True, True, True],
    })

    assert_frame_equal(expected, int_rects[expected.columns],
                       check_dtype=False)

    styled_dots = _dots_to_dataframe(upset_axes["matrix"], is_vertical)
    baseline_dots = _dots_to_dataframe(
        UpSet(data, orientation=orientation).plot()["matrix"],
        is_vertical
    )
    inactive_dot_mask = (baseline_dots[["fc_a"]] < 1).values.ravel()
    assert_frame_equal(baseline_dots.loc[inactive_dot_mask],
                       styled_dots.loc[inactive_dot_mask])

    styled_dots = styled_dots.loc[~inactive_dot_mask]

    styled_dots = styled_dots.drop(columns="y").groupby("x").apply(
        lambda df: df.drop_duplicates())
    styled_dots["ls_is_solid"] = styled_dots.pop("ls_seq").isna()
    assert_frame_equal(expected.iloc[1:].reset_index(drop=True),
                       styled_dots[expected.columns].reset_index(drop=True),
                       check_dtype=False)

    # TODO: check lines between dots
    # matrix_line_collection = upset_axes["matrix"].collections[1]
