Base

def fit(
    self,
    x: Tensor,
    y: Tensor,
    epochs: int = 200,
    learning_rate: float = 1e-3,
    criterion: Callable[[Tensor, Tensor], Tensor] | None = None,
    optimizer: torch.optim.Optimizer | None = None,
    batch_size: int | None = None,
    shuffle: bool = True,
    ur_weight: float = 0.0,
    ur_target: float | None = None,
    verbose: bool | int = False,
    x_val: Tensor | None = None,
    y_val: Tensor | None = None,
    patience: int | None = 20,
    restore_best: bool = True,
    weight_decay: float = 1e-8,
) -> dict[str, Any]:
    """Train the model with optional early stopping.

    When *x_val* and *y_val* are provided the model evaluates a
    task-specific metric (via :meth:`_evaluate_validation`) after every
    epoch and applies early stopping when the metric has not improved for
    *patience* consecutive epochs.
    By default the best model weights from validation are restored when
    ``restore_best=True``.

    Args:
        x: Training features of shape ``(N, n_inputs)``.
        y: Training targets of shape ``(N,)``.
        epochs: Maximum number of training epochs.
        learning_rate: Learning rate for the default AdamW optimizer.
        criterion: Optional loss function.  Defaults to
            :meth:`_default_criterion`.
        optimizer: Optional pre-built optimizer.  When ``None``, AdamW
            is constructed with separate parameter groups for antecedent
            (no weight decay) and consequent (*weight_decay*) layers.
        batch_size: Mini-batch size.  ``None`` uses the full dataset.
        shuffle: If ``True``, reshuffle sample indices each epoch.
        ur_weight: Non-negative weight for the uniform rule
            regularization term.  ``0.0`` disables it.
        ur_target: Target uniform activation for UR.  Must be in
            ``(0, 1]`` when provided.  ``None`` defaults to
            ``1 / n_rules``.
        verbose: Verbosity level. ``0`` = quiet, ``1`` = progress bar,
            ``2`` = per-epoch summary logging, ``3`` = per-epoch detailed
            logging. ``True`` is accepted as an alias for ``1``.
        x_val: Optional validation features of shape
            ``(M, n_inputs)``.
        y_val: Optional validation targets of shape ``(M,)``.
        patience: Number of consecutive epochs without improvement
            before early stopping.  Set to ``None`` to disable early
            stopping.  Only active when *x_val* and *y_val* are given.
        restore_best: If ``True`` (default), restore the model weights
            from the best validation epoch when early stopping is used.
        weight_decay: L2 weight decay applied to consequent parameters
            by the default AdamW optimizer.

    Returns:
        A dictionary with keys ``"train"``, ``"ur"``, and ``"val"``
        containing per-epoch loss lists.

    Raises:
        ValueError: If shapes of *x*, *y*, *x_val*, or *y_val* are
            incompatible, or if *ur_weight* < 0 or *ur_target* is
            outside ``(0, 1]``.
    """
    has_val = self._validate_fit_inputs(x, y, x_val, y_val, ur_weight, ur_target)
    train_criterion = criterion or self._default_criterion()
    train_optimizer = self._build_optimizer(optimizer, learning_rate, weight_decay)

    history: dict[str, Any] = {"train": [], "ur": [], "val": []}
    best_metric = float("-inf")
    epochs_no_improve = 0
    best_state: dict[str, Any] | None = None
    verbose_level = self._resolve_verbose(verbose)

    self.train()
    pbar = None
    if verbose_level == 1:
        pbar = trange(epochs, desc="Training", leave=False)
        epoch_iterator = pbar
    else:
        epoch_iterator = range(epochs)

    stopped_epoch = 0
    for epoch in epoch_iterator:
        stopped_epoch = epoch + 1
        epoch_train_loss, epoch_ur_loss = self._run_minibatch_epoch(
            x, y, train_criterion, train_optimizer, batch_size, shuffle, ur_weight, ur_target
        )
        history["train"].append(epoch_train_loss)
        history["ur"].append(epoch_ur_loss)

        if has_val and x_val is not None and y_val is not None:
            self.eval()
            val_info = self._evaluate_validation(train_criterion, x_val, y_val)
            history["val"].append(val_info.get("val_loss", 0.0))
            for k, v in val_info.items():
                if k not in ("val_loss", "metric"):
                    history.setdefault(k, []).append(v)
            self.train()

            metric = val_info["metric"]
            if metric > best_metric:
                best_metric = metric
                epochs_no_improve = 0
                best_state = copy.deepcopy(self.state_dict())
            else:
                epochs_no_improve += 1

            self._log_epoch_with_val(epoch, epochs, epoch_train_loss, val_info, verbose_level, pbar)

            if patience is not None and epochs_no_improve >= patience:
                if verbose_level >= 2:
                    self._log(
                        "early stopping at epoch %s (patience=%s)",
                        epoch + 1,
                        patience,
                        verbose=verbose_level,
                    )
                break
        else:
            self._log_epoch_no_val(epoch, epochs, epoch_train_loss, verbose_level, pbar)

    if pbar is not None:
        pbar.close()

    if restore_best and best_state is not None:
        self.load_state_dict(best_state)

    history["stopped_epoch"] = stopped_epoch

    return history

def forward(self, x: Tensor) -> Tensor:
    """Full forward pass through the TSK pipeline."""
    output, _ = self._forward_train(x)
    return output

def forward_antecedents(self, x: Tensor) -> Tensor:
    """Compute normalized rule strengths from model antecedents."""
    mu = self.membership_layer(x)
    w = self.rule_layer(mu)
    return self.defuzzifier(w)

def get_consequent_weights(self) -> Tensor | None:
    """Return the consequent layer weights or ``None`` when unavailable."""
    weight = getattr(self.consequent_layer, "weight", None)
    if isinstance(weight, Tensor):
        return weight.detach()
    return None

def get_mf_params(self) -> dict[str, list[dict[str, Any]]]:
    """Return a serializable description of the model's membership functions."""
    return {
        name: [{"type": type(mf).__name__, **mf.inspect_params()} for mf in cast(Sequence[MembershipFunction], mfs)]
        for name, mfs in self.membership_layer.input_mfs.items()
    }

def get_rule_table(self) -> list[dict[str, Any]]:
    """Return the rule base as a table of feature-to-MF indices."""
    return [
        dict(
            rule_id=rule_index,
            **dict(zip(self.rule_layer.input_names, rule, strict=False)),
        )
        for rule_index, rule in enumerate(self.rule_layer.rules)
    ]