import io
import itertools

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

from upsetplot import UpSet, generate_counts, generate_samples, plot
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", "-cardinality", "-degree", None, "input", "-input"],
)
@pytest.mark.parametrize(
    "sort_categories_by", [None, "input", "-input", "cardinality", "-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.0, 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_levels = x.reorder_levels(intersections.index.names)
    x_reordered = x_reordered_levels.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 == "-cardinality":
        assert is_ascending(intersections.values)
    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
    elif sort_by == "-degree":
        # check degree order
        assert is_ascending(intersections.index.to_frame().sum(axis=1)[::-1])
    else:
        find_first_in_orig = x_reordered_levels.index.tolist().index
        orig_order = [find_first_in_orig(key) for key in intersections.index.tolist()]
        assert orig_order == sorted(
            orig_order, reverse=sort_by is not None and sort_by.startswith("-")
        )

    if sort_categories_by == "cardinality":
        assert is_ascending(totals.values[::-1])
    elif sort_categories_by == "-cardinality":
        assert is_ascending(totals.values)

    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.0, 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="int64"), "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)


def test_include_empty_subsets():
    X = generate_counts(n_samples=2, n_categories=3)

    no_empty_upset = UpSet(X, include_empty_subsets=False)
    assert len(no_empty_upset.intersections) <= 2

    include_empty_upset = UpSet(X, include_empty_subsets=True)
    assert len(include_empty_upset.intersections) == 2**3
    common_intersections = include_empty_upset.intersections.loc[
        no_empty_upset.intersections.index
    ]
    assert_series_equal(no_empty_upset.intersections, common_intersections)
    include_empty_upset.plot()  # smoke test


@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"},
        {"totals_plot_elements": 0},
    ],
)
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:
        yield from _walk_artists(ch)


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_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_percentages="!{:0.2f}!")
    assert n_artists_yes_sizes == _count_descendants(fig)
    assert "!0.95!" in get_all_texts(fig)
    assert "!0.03!" 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)

    fig = matplotlib.figure.Figure()
    with pytest.raises(ValueError):
        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,
            },
        ),
        (
            {"max_subset_rank": 3},
            {
                (True, False, False): 884,
                (True, True, False): 1547,
                (True, True, True): 990,
            },
        ),
        (
            {"min_subset_size": 800, "max_subset_size": 990},
            {
                (True, False, False): 884,
                (True, True, True): 990,
            },
        ),
        (
            {"min_subset_size": "15%", "max_subset_size": "30.1%"},
            {
                (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)
    assert upset_full.total == pytest.approx(upset_filtered.total)


def test_filter_subsets_max_subset_rank_tie():
    data = generate_samples(seed=0, n_samples=5, n_categories=3)
    tested_non_tie = False
    tested_tie = True
    full = UpSet(data, subset_size="count").intersections
    prev = None
    for max_rank in range(1, 5):
        cur = UpSet(data, subset_size="count", max_subset_rank=max_rank).intersections
        if prev is not None:
            if cur.shape[0] > prev.shape[0]:
                # check we add rows only when they are new
                assert cur.min() < prev.min()
                tested_non_tie = True
            elif cur.shape[0] != full.shape[0]:
                assert (cur == cur.min()).sum() > 1
                tested_tie = True

        prev = cur
    assert tested_non_tie
    assert tested_tie
    assert cur.shape[0] == full.shape[0]


@pytest.mark.parametrize(
    "value",
    [
        "1",
        "-1%",
        "1%%",
        "%1",
        "hello",
    ],
)
def test_bad_percentages(value):
    data = generate_samples(seed=0, n_samples=5, n_categories=3)
    with pytest.raises(ValueError, match="percentage"):
        UpSet(data, min_subset_size=value)


@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]


@pytest.mark.parametrize(
    (
        "kwarg_list",
        "expected_category_styles",
    ),
    [
        # Different forms of including two categories
        (
            [{"categories": ["cat1", "cat2"], "shading_facecolor": "red"}],
            {
                "cat1": {"shading_facecolor": "red"},
                "cat2": {"shading_facecolor": "red"},
            },
        ),
        (
            [
                {"categories": ["cat1", "cat2"], "shading_facecolor": "red"},
                {"categories": "cat1", "shading_facecolor": "green"},
            ],
            {
                "cat1": {"shading_facecolor": "green"},
                "cat2": {"shading_facecolor": "red"},
            },
        ),
        (
            [
                {"categories": ["cat1", "cat2"], "shading_facecolor": "red"},
                {"categories": "cat1", "shading_edgecolor": "green"},
            ],
            {
                "cat1": {"shading_facecolor": "red", "shading_edgecolor": "green"},
                "cat2": {"shading_facecolor": "red"},
            },
        ),
    ],
)
def test_categories(kwarg_list, expected_category_styles):
    data = generate_counts()
    upset = UpSet(data, facecolor="blue")
    for kw in kwarg_list:
        upset.style_categories(**kw)
    actual_category_styles = upset.category_styles
    assert actual_category_styles == expected_category_styles


def test_many_categories():
    # Tests regressions against GH#193
    n_cats = 250
    index1 = [True, False] + [False] * (n_cats - 2)
    index2 = [False, True] + [False] * (n_cats - 2)
    columns = [chr(i + 33) for i in range(n_cats)]
    data = pd.DataFrame([index1, index2], columns=columns)
    data["value"] = 1
    data = data.set_index(columns)["value"]
    UpSet(data)
