import numpy as np
import pytest

import xgboost as xgb
from xgboost import testing as tm
from xgboost.testing.plotting import run_categorical

try:
    import matplotlib

    matplotlib.use("Agg")
    from graphviz import Source
    from matplotlib.axes import Axes
except ImportError:
    pass

pytestmark = pytest.mark.skipif(**tm.no_multiple(tm.no_matplotlib(), tm.no_graphviz()))


class TestPlotting:
    def test_plotting(self):
        m, _ = tm.load_agaricus(__file__)
        booster = xgb.train(
            {"max_depth": 2, "eta": 1, "objective": "binary:logistic"},
            m,
            num_boost_round=2,
        )

        ax = xgb.plot_importance(booster)
        assert isinstance(ax, Axes)
        assert ax.get_title() == "Feature importance"
        assert ax.get_xlabel() == "Importance score"
        assert ax.get_ylabel() == "Features"
        assert len(ax.patches) == 4

        ax = xgb.plot_importance(booster, color="r", title="t", xlabel="x", ylabel="y")
        assert isinstance(ax, Axes)
        assert ax.get_title() == "t"
        assert ax.get_xlabel() == "x"
        assert ax.get_ylabel() == "y"
        assert len(ax.patches) == 4
        for p in ax.patches:
            assert p.get_facecolor() == (1.0, 0, 0, 1.0)  # red

        ax = xgb.plot_importance(
            booster, color=["r", "r", "b", "b"], title=None, xlabel=None, ylabel=None
        )
        assert isinstance(ax, Axes)
        assert ax.get_title() == ""
        assert ax.get_xlabel() == ""
        assert ax.get_ylabel() == ""
        assert len(ax.patches) == 4
        assert ax.patches[0].get_facecolor() == (1.0, 0, 0, 1.0)  # red
        assert ax.patches[1].get_facecolor() == (1.0, 0, 0, 1.0)  # red
        assert ax.patches[2].get_facecolor() == (0, 0, 1.0, 1.0)  # blue
        assert ax.patches[3].get_facecolor() == (0, 0, 1.0, 1.0)  # blue

        g = xgb.to_graphviz(booster, tree_idx=0)
        assert isinstance(g, Source)

        ax = xgb.plot_tree(booster, tree_idx=0)
        assert isinstance(ax, Axes)

    def test_importance_plot_lim(self):
        np.random.seed(1)
        dm = xgb.DMatrix(np.random.randn(100, 100), label=[0, 1] * 50)
        bst = xgb.train({}, dm)
        assert len(bst.get_fscore()) == 71
        ax = xgb.plot_importance(bst)
        assert ax.get_xlim() == (0.0, 11.0)
        assert ax.get_ylim() == (-1.0, 71.0)

        ax = xgb.plot_importance(bst, xlim=(0, 5), ylim=(10, 71))
        assert ax.get_xlim() == (0.0, 5.0)
        assert ax.get_ylim() == (10.0, 71.0)

    @pytest.mark.skipif(**tm.no_pandas())
    def test_categorical(self) -> None:
        run_categorical("approx", "cpu")