# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/08_tree.ipynb.

# %% auto 0
__all__ = ['_TreeRegInspector', '_TreeBinInspector', '_TreeMultiInspector']

# %% ../../nbs/08_tree.ipynb 3
from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from ..delegate import delegates
from .any_model import _Inspector
from .classifier import _BinInspector, _MultiInspector
from .regressor import _RegInspector
from sklearn.tree import plot_tree

# %% ../../nbs/08_tree.ipynb 4
class _TreeInspector(_Inspector):
    """Decision tree model inspector"""

    @delegates(plot_tree)
    def plot_tree(self, ax: Optional[Axes] = None, **kwargs):
        """Show decision tree."""
        if ax is None:
            # these dimensions seem to work well empirically
            max_size = 50
            depth = (
                kwargs["max_depth"] + 1
                if "max_depth" in kwargs
                else self.model.get_depth()
            )
            fig_height = min(depth * 2.2, max_size)
            width = (
                2 * kwargs["max_depth"]
                if "max_depth" in kwargs
                else self.model.get_n_leaves()
            )
            fig_width = min(width * 3.5, max_size)
            _, ax = plt.subplots(figsize=(fig_width, fig_height))
        kwargs = {"filled": True, "fontsize": 12, **kwargs}
        return plot_tree(
            self.model,
            feature_names=self.X.columns,
            class_names=np.unique(self.y).astype(str),
            ax=ax,
            **kwargs,
        )[0].axes

# %% ../../nbs/08_tree.ipynb 5
class _TreeRegInspector(_TreeInspector, _RegInspector):
    """Decision tree regression model inspector"""

    pass

# %% ../../nbs/08_tree.ipynb 6
class _TreeBinInspector(_TreeInspector, _BinInspector):
    """Decision tree binary classification model inspector"""

    pass

# %% ../../nbs/08_tree.ipynb 7
class _TreeMultiInspector(_TreeInspector, _MultiInspector):
    """Decision tree multiclass model inspector"""

    pass

# %% ../../nbs/08_tree.ipynb 33
_all_ = ["_TreeRegInspector", "_TreeBinInspector", "_TreeMultiInspector"]
