Plotting

def plot(estimator: Any, kind: str = "memberships", **kwargs: Any) -> Any:
    """Draw the plot named *kind* for a fitted estimator.

    Args:
        estimator: A fitted estimator.
        kind: One of ``"memberships"`` (default), ``"history"``, ``"rule_activation"``,
            ``"diagnostics"``, ``"feature_importance"`` or ``"gates"``.
        **kwargs: Arguments of the corresponding ``plot_<kind>`` function.

    Returns:
        The ``Axes`` or ``Figure`` returned by that function.
    """
    if kind not in _PLOTTERS:
        raise ValueError(f"kind must be one of {list(PLOT_KINDS)}; got {kind!r}")
    return _PLOTTERS[kind](estimator, **kwargs)

def plot_diagnostics(estimator: Any, X: Any, y: Any, *, ax: Any = None) -> Any:
    """Plot two views of the predictions on ``(X, y)``.

    For a classifier: the confusion matrix, and the decision margin of every sample
    grouped by its true class. The margin is the probability given to the true class minus
    the highest probability given to another class, so it is negative for the samples that
    are misclassified.

    For a regressor: predicted against observed values, and the residuals against the
    observed values.

    Args:
        estimator: A fitted estimator.
        X: Data to predict.
        y: True labels or targets for ``X``.
        ax: Two axes to draw on. ``None`` creates a new figure.

    Returns:
        The ``Figure`` holding the two panels.
    """
    check_is_fitted(estimator, "model_")
    fig, axes = _axes(ax, 2, ncols=2, panel_size=(3.6, 3.2))
    draw = _classification_diagnostics if is_classifier(estimator) else _regression_diagnostics
    draw(estimator, X, np.asarray(y), axes)
    return fig

def plot_feature_importance(estimator: Any, *, top: int | None = 20, ax: Any = None) -> Any:
    """Plot the features that weigh most in the consequents of the fitted model.

    The importance is the one of ``estimator.feature_importance()``: the mean magnitude of
    the consequent coefficients of each feature, normalized to sum to one. Features are
    named as in ``fit`` when the data had names, and by their column otherwise.

    Args:
        estimator: A fitted estimator.
        top: Number of features to draw, from the most important. ``None`` draws all.
        ax: Axes to draw on. ``None`` creates a new figure.

    Returns:
        The ``Axes`` with the bar chart.
    """
    check_is_fitted(estimator, "model_")
    importance = estimator.feature_importance()
    if importance is None:
        raise ValueError("this model has no consequent coefficients to compute a feature importance from")
    importance = np.asarray(importance, dtype=float)
    labels = _feature_labels(estimator)
    order = np.argsort(-importance)
    if top is not None:
        order = order[: int(top)]

    _, (panel,) = _axes(ax, 1, ncols=1, panel_size=(4.8, max(2.4, 0.22 * len(order) + 1.0)))
    panel.barh(np.arange(len(order)), importance[order])
    panel.set_yticks(np.arange(len(order)), [labels[position] for position in order])
    panel.invert_yaxis()
    panel.set_xlabel("Feature importance")
    return panel

def plot_gates(estimator: Any, *, ax: Any = None) -> Any:
    """Plot the gates of the features and of the rules that a gated family kept.

    Only DG-TSK, DG-ALETSK and FSRE-ADATSK have gates. The bars are the gate values at the
    end of the phase that trained them, for the features in ``selected_features_`` and for
    the rules of the fitted model, and the dashed line is the threshold they passed. The
    gate of FSRE-ADATSK is an odd function, so its magnitude is drawn.

    Args:
        estimator: A fitted estimator of a gated family.
        ax: Two axes to draw on, for the features and for the rules. ``None`` creates a
            new figure.

    Returns:
        The ``Figure`` with the two panels.
    """
    check_is_fitted(estimator, "model_")
    feature_gates, rule_gates = estimator.get_feature_gates(), estimator.get_rule_gates()
    if feature_gates is None or rule_gates is None:
        raise ValueError("only DG-TSK, DG-ALETSK and FSRE-ADATSK have gates to plot")
    thresholds = _gate_thresholds(estimator.history_)
    panels = (
        ("Feature gates", np.abs(np.asarray(feature_gates, dtype=float)).ravel(), _feature_labels(estimator)),
        ("Rule gates", np.abs(np.asarray(rule_gates, dtype=float)).ravel(), None),
    )

    fig, axes = _axes(ax, 2, ncols=2, panel_size=(4.4, 3.0))
    for panel, (title, values, labels), threshold in zip(axes, panels, thresholds, strict=True):
        positions = np.arange(len(values))
        panel.bar(positions, values)
        if threshold is not None:
            panel.axhline(threshold, color="0.3", linestyle="--", label=f"threshold ({threshold:.2f})")
            panel.legend(loc="lower right")
        if labels is None:
            labels = [f"Rule {rule}" for rule in positions]
        if len(values) <= 30:
            slanted = len(values) > 4
            panel.set_xticks(positions, labels, rotation=60 if slanted else 0, ha="right" if slanted else "center")
        else:
            panel.set_xticks([])
        panel.set_ylim(0.0, max(1.05, float(values.max()) * 1.05))
        panel.set_ylabel("Gate value")
        panel.set_title(title)
    return fig

def plot_history(estimator: Any, metric: str | None = None, *, ax: Any = None) -> Any:
    """Plot the training loss, or a metric, per epoch.

    The validation curve is drawn when a validation set was passed to ``fit``, and the
    best epoch is marked when early stopping recorded one. Families that train in several
    phases (DG-TSK, DG-ALETSK, FSRE-ADATSK) show the phases one after the other.

    Args:
        estimator: A fitted estimator.
        metric: Name of a metric recorded during ``fit`` (for example ``"accuracy"``).
            ``None`` plots the loss.
        ax: Axes to draw on. ``None`` creates a new figure.

    Returns:
        The ``Axes`` with the curves.
    """
    check_is_fitted(estimator, "model_")
    phases = _history_phases(estimator.history_)
    name = "loss" if metric is None else metric
    if metric is not None and metric not in _available_metrics(phases):
        raise ValueError(f"metric {metric!r} was not recorded during fit; available: {_available_metrics(phases)}")

    _, (panel,) = _axes(ax, 1, ncols=1, panel_size=(4.8, 3.2))
    offset = 0
    for index, (label, phase) in enumerate(phases):
        n_epochs = len(phase["train_loss"])
        for prefix, colour in (("train", "C0"), ("val", "C1")):
            series = phase.get(f"{prefix}_{name}")
            if not series:
                continue
            # Metrics may be evaluated every few epochs; spread them over the phase.
            epochs = offset + (np.arange(len(series)) + 1) * n_epochs / len(series)
            legend = {"train": "training", "val": "validation"}[prefix] if index == 0 else None
            panel.plot(epochs, series, color=colour, label=legend)
        best = phase.get("best_epoch")
        if best is not None:
            panel.axvline(offset + best + 1, color="0.4", linestyle="--", label="best epoch" if index == 0 else None)
        if label:
            panel.axvline(offset + 0.5, color="0.8", linewidth=0.8)
            panel.text(offset + 0.5, 1.01, f" {label}", transform=panel.get_xaxis_transform(), va="bottom")
        offset += n_epochs
    panel.set_xlabel("Epoch")
    panel.set_ylabel(name.replace("_", " ").capitalize())
    panel.legend()
    return panel

def plot_memberships(
    estimator: Any,
    features: Sequence[int | str] | None = None,
    *,
    X: Any = None,
    max_features: int = 6,
    n_points: int = 200,
    ax: Any = None,
) -> Any:
    """Plot the learned fuzzy sets, one panel per feature.

    Args:
        estimator: A fitted estimator.
        features: Features to show, as names or as column indices of the data seen in
            ``fit``. ``None`` shows every feature, or the ``max_features`` most important
            ones when the model has more.
        X: Optional data. When given, each panel spans the range of its column; otherwise
            the range is taken from the parameters of the fuzzy sets.
        max_features: Number of features shown when ``features`` is ``None``.
        n_points: Number of points used to draw each curve.
        ax: Axes to draw on, one per feature. ``None`` creates a new figure.

    Returns:
        The ``Figure`` holding the panels.
    """
    check_is_fitted(estimator, "model_")
    positions = _select_features(estimator, features, max_features)
    labels = _feature_labels(estimator)
    model_names = [str(name) for name in estimator.model_.input_names]
    mf_params = estimator.get_mf_params()
    input_mfs = estimator.model_.membership_layer.input_mfs
    data = None if X is None else estimator._select_model_features(np.asarray(X))

    fig, axes = _axes(ax, len(positions), ncols=min(3, len(positions)), panel_size=(3.2, 2.4))
    for panel, position in zip(axes, positions, strict=True):
        name = model_names[position]
        lo, hi = _feature_range(mf_params[name], None if data is None else data[:, position])
        # The centres are added to the grid so that the peak of a narrow set is not missed.
        centres = [p.get("mean", p.get("center")) for p in mf_params[name]]
        centres = [c for c in centres if c is not None and lo <= c <= hi]
        points = np.unique(np.concatenate([np.linspace(lo, hi, n_points), centres]))
        grid = torch.as_tensor(points, dtype=estimator._model_dtype(), device=str(estimator.device))
        sets = [(j, mf) for j, mf in enumerate(input_mfs[name]) if type(mf).__name__ != "ConstantMF"]
        with torch.no_grad():
            for j, mf in sets:
                panel.plot(grid.cpu().numpy(), mf(grid).cpu().numpy(), label=f"MF{j}")
        panel.set_title(labels[position])
        panel.set_ylim(0.0, 1.05)
        panel.set_xlabel("Feature value")
        panel.set_ylabel("Membership degree")
        if 0 < len(sets) <= _MAX_LEGEND_ENTRIES:
            panel.legend()
    return fig

def plot_rule_activation(estimator: Any, X: Any, y: Any = None, *, ax: Any = None) -> Any:
    """Plot the mean normalized firing strength of each rule.

    Args:
        estimator: A fitted estimator.
        X: Data on which the rules are evaluated.
        y: Optional class labels, for classifiers only. When given, the mean activation is
            computed separately for the samples of each class.
        ax: Axes to draw on. ``None`` creates a new figure.

    Returns:
        The ``Axes`` with the heatmap (labels given) or the bar chart (no labels).
    """
    check_is_fitted(estimator, "model_")
    activation = estimator.rule_activation(X)
    n_rules = activation.shape[1]
    rule_labels = [f"Rule {r}" for r in range(n_rules)]
    _, (panel,) = _axes(ax, 1, ncols=1, panel_size=(4.8, max(2.4, 0.3 * n_rules + 1.2)))

    if y is None:
        panel.barh(np.arange(n_rules), activation.mean(axis=0))
        panel.set_xlabel("Mean rule activation")
        panel.invert_yaxis()
    else:
        if not is_classifier(estimator):
            raise ValueError("y groups the activations by class, so it only applies to classifiers")
        y = np.asarray(y)
        classes = estimator.classes_
        by_class = np.stack([activation[y == c].mean(axis=0) for c in classes], axis=1)
        _annotated_heatmap(panel, by_class, "Greens", lambda v: f"{v:.2f}")
        panel.set_xticks(np.arange(len(classes)), [str(c) for c in classes])
        panel.set_xlabel("True class")
        panel.set_title("Mean rule activation")
    if n_rules <= 30:
        panel.set_yticks(np.arange(n_rules), rule_labels)
    else:
        panel.set_ylabel("Rule")
    return panel