Estimators

def evaluate(
    self,
    X: npt.ArrayLike,
    y: npt.ArrayLike,
    metrics: list[str] | None = None,
    sample_weight: npt.ArrayLike | None = None,
) -> dict[str, Any]:
    """Compute classification evaluation metrics for the provided dataset."""
    y_true = np.asarray(y)
    y_pred = self.predict(X)
    return compute_metrics(
        task="classification",
        y_true=y_true,
        y_pred=y_pred,
        sample_weight=sample_weight,
        metrics=metrics,
    )

def feature_importance(self) -> np.ndarray | None:
    """Compute a normalized feature importance vector from consequent weights."""
    check_is_fitted(self, "model_")
    weights = self.model_.get_consequent_weights()
    if weights is None:
        return None

    consequent_layer = self.model_.consequent_layer
    if hasattr(consequent_layer, "rule_feature_mask"):
        rule_feature_mask = cast(Tensor, consequent_layer.rule_feature_mask)
        weights = weights * rule_feature_mask.unsqueeze(1) if weights.ndim == 3 else weights * rule_feature_mask

    abs_weights = weights.abs()
    if abs_weights.ndim == 3:
        importance = abs_weights.mean(dim=(0, 1))
    elif abs_weights.ndim == 2:
        importance = abs_weights.mean(dim=0)
    else:
        raise ValueError("unsupported consequent weight shape for feature importance")

    return _normalize_importance(importance)

def fit(
    self,
    x: npt.ArrayLike,
    y: npt.ArrayLike,
    *,
    x_val: npt.ArrayLike | None = None,
    y_val: npt.ArrayLike | None = None,
    metrics: list[str] | None = None,
) -> ADATSKClassifier:
    """Fit the ADATSK classifier estimator."""
    return cast(ADATSKClassifier, super().fit(x, y, x_val=x_val, y_val=y_val, metrics=metrics))

def get_mf_params(self) -> dict[str, list[dict[str, Any]]]:
    """Return model membership function metadata after fitting."""
    check_is_fitted(self, "model_")
    return self.model_.get_mf_params()

def inspect(self) -> dict[str, Any]:
    """Return a structured summary of fitted model state and rule metadata."""
    check_is_fitted(self, "model_")
    return {
        "n_rules": int(self.model_.n_rules),
        "n_inputs": int(self.model_.n_inputs),
        "feature_names": list(self.model_.input_names),
        "rule_base": self.rule_base_,
        "defuzzifier_type": type(self.model_.defuzzifier).__name__,
        "mf_params": self.get_mf_params(),
        "rule_table": self.model_.get_rule_table(),
    }

@classmethod
def load(cls, path: str) -> Self:
    """Load a persisted estimator created by save."""
    checkpoint = load_checkpoint(path)
    validate_checkpoint_payload(checkpoint, expected_estimator_class=cls.__name__)

    params: dict[str, Any] = dict(checkpoint["estimator_params"])
    if params.get("input_configs") is not None:
        params["input_configs"] = [InputConfig(**c) for c in params["input_configs"]]
    estimator = cls(**params)
    model_init = checkpoint["model_init"]
    estimator.rule_base_ = model_init["rule_base"]

    import inspect

    sig = inspect.signature(estimator._build_model)
    if "rules" in sig.parameters:
        estimator.model_ = estimator._build_model(
            deserialize_input_mfs(model_init["input_mfs_config"]),
            int(model_init["n_classes"]),
            str(model_init["rule_base"]),
            rules=model_init.get("rules"),
        )
    else:
        estimator.model_ = estimator._build_model(
            deserialize_input_mfs(model_init["input_mfs_config"]),
            int(model_init["n_classes"]),
            str(model_init["rule_base"]),
        )

    consequent_mode = model_init.get("consequent_mode")
    if consequent_mode is not None:
        if hasattr(estimator.model_, "set_consequent_mode"):
            cast(Any, estimator.model_).set_consequent_mode(consequent_mode)
        elif hasattr(estimator.model_.consequent_layer, "mode"):
            estimator.model_.consequent_layer.mode = consequent_mode

    estimator.model_.load_state_dict(checkpoint["model_state_dict"])
    estimator.model_.to(torch.device(str(estimator.device)))

    fitted = checkpoint["fitted_attrs"]
    estimator.n_features_in_ = int(fitted["n_features_in"])
    if fitted.get("feature_names_in") is not None:
        estimator.feature_names_in_ = np.asarray(fitted["feature_names_in"], dtype=object)
    elif hasattr(estimator, "feature_names_in_"):
        delattr(estimator, "feature_names_in_")
    estimator.classes_ = np.asarray(fitted["classes"], dtype=object)
    label_encoder = LabelEncoder()
    label_encoder.classes_ = estimator.classes_
    estimator._label_encoder_ = label_encoder
    history = checkpoint.get("history")
    estimator.history_ = history if history is not None else {}
    return estimator

def predict(self, x: npt.ArrayLike) -> np.ndarray:
    """Predict class labels for input samples."""
    proba = self.predict_proba(x)
    y_idx = np.argmax(proba, axis=1)
    return np.asarray(self._label_encoder_.inverse_transform(y_idx))

def predict_proba(self, x: npt.ArrayLike) -> np.ndarray:
    """Predict class probabilities for input samples."""
    check_is_fitted(self, "model_")
    x_arr = validate_data(self, x, reset=False)
    device_str = str(self.device).lower()
    x_tensor = torch.as_tensor(x_arr, dtype=torch.float32, device=torch.device(device_str))
    probs = cast(Any, self.model_).predict_proba(x_tensor)
    return probs.detach().cpu().numpy()

def rule_activation(self, X: npt.ArrayLike) -> np.ndarray:
    """Return normalized rule activations for the provided inputs."""
    check_is_fitted(self, "model_")
    x_arr = validate_data(self, X, reset=False)

    was_training = self.model_.training
    try:
        self.model_.eval()
        with torch.no_grad():
            norm_w = self.model_.forward_antecedents(self._as_tensor_x(x_arr, torch.device(str(self.device))))
    finally:
        self.model_.train(was_training)

    return _to_numpy(norm_w)

def save(self, path: str) -> None:
    """Persist estimator configuration, model weights and fitted metadata."""
    _fnames: np.ndarray | None = getattr(self, "feature_names_in_", None)
    rules = None
    if hasattr(self.model_, "rule_layer") and hasattr(self.model_.rule_layer, "rules"):
        rules = self.model_.rule_layer.rules
    consequent_mode = getattr(self.model_.consequent_layer, "mode", None)
    checkpoint = self._build_checkpoint_base(
        model_init={
            "input_mfs_config": serialize_input_mfs(self.model_.input_mfs),
            "n_classes": len(self.classes_),
            "rule_base": self.rule_base_,
            "rules": rules,
            "consequent_mode": consequent_mode,
        },
        fitted_attrs={
            "n_features_in": int(self.n_features_in_),
            "feature_names_in": _fnames.tolist() if _fnames is not None else None,
            "classes": self.classes_.tolist(),
        },
    )
    save_checkpoint(path, checkpoint)

def score(self, X: npt.ArrayLike, y: npt.ArrayLike, sample_weight: npt.ArrayLike | None = None) -> float:
    """Return classification accuracy on the provided dataset."""
    y_true = np.asarray(y)
    y_pred = self.predict(X)
    return float(accuracy_score(y_true, y_pred, sample_weight=sample_weight))

def evaluate(
    self,
    X: npt.ArrayLike,
    y: npt.ArrayLike,
    metrics: list[str] | None = None,
    sample_weight: npt.ArrayLike | None = None,
) -> dict[str, Any]:
    """Compute regression evaluation metrics for the provided dataset."""
    y_true = np.asarray(y)
    y_pred = self.predict(X)
    return compute_metrics(
        task="regression",
        y_true=y_true,
        y_pred=y_pred,
        sample_weight=sample_weight,
        metrics=metrics,
    )

def feature_importance(self) -> np.ndarray | None:
    """Compute a normalized feature importance vector from consequent weights."""
    check_is_fitted(self, "model_")
    weights = self.model_.get_consequent_weights()
    if weights is None:
        return None

    consequent_layer = self.model_.consequent_layer
    if hasattr(consequent_layer, "rule_feature_mask"):
        rule_feature_mask = cast(Tensor, consequent_layer.rule_feature_mask)
        weights = weights * rule_feature_mask.unsqueeze(1) if weights.ndim == 3 else weights * rule_feature_mask

    abs_weights = weights.abs()
    if abs_weights.ndim == 3:
        importance = abs_weights.mean(dim=(0, 1))
    elif abs_weights.ndim == 2:
        importance = abs_weights.mean(dim=0)
    else:
        raise ValueError("unsupported consequent weight shape for feature importance")

    return _normalize_importance(importance)

def fit(
    self,
    x: npt.ArrayLike,
    y: npt.ArrayLike,
    *,
    x_val: npt.ArrayLike | None = None,
    y_val: npt.ArrayLike | None = None,
    metrics: list[str] | None = None,
) -> Self:
    """Train the TSK regressor on labeled samples.

    Validation data should be supplied using ``x_val`` and ``y_val``
    when available.
    """
    x_arr, y_arr = validate_data(self, x, y, reset=True)
    n_samples = x_arr.shape[0]
    n_mfs_val = getattr(self, "n_mfs", 3)
    n_mfs_val = 3 if n_mfs_val is None else n_mfs_val
    mf_init_val = getattr(self, "mf_init", "kmeans")
    required_samples = 2 if mf_init_val == "grid" else max(2, n_mfs_val)
    if n_samples < required_samples:
        raise ValueError(
            f"Found array with {n_samples} sample(s). Estimator requires at least {required_samples} samples."
        )

    if self.random_state is not None:
        torch.manual_seed(int(self.random_state))

    input_mfs, _, effective_rule_base = self._build_input_mfs(x_arr)

    self.n_features_in_ = x_arr.shape[1]

    _device = torch.device(str(self.device))
    self.model_ = self._build_regressor_model(input_mfs, effective_rule_base).to(_device)

    y_t = torch.as_tensor(np.asarray(y_arr, dtype=np.float32), dtype=torch.float32, device=_device)

    # Prepare validation tensors if provided via fit.
    x_val_t: torch.Tensor | None = None
    y_val_t: torch.Tensor | None = None
    if (x_val is None) != (y_val is None):
        raise ValueError("x_val and y_val must be provided together")
    if x_val is not None and y_val is not None:
        x_v_arr, y_v_arr = validate_data(self, x_val, y_val, reset=False)
        x_val_t = self._as_tensor_x(x_v_arr, _device)
        y_val_t = torch.as_tensor(np.asarray(y_v_arr, dtype=np.float32), dtype=torch.float32, device=_device)

    x_t = self._as_tensor_x(x_arr, _device)
    self.rule_base_ = effective_rule_base
    self._pre_train_hook(self.model_, x_t, y_t)
    _trainer = self.trainer if self.trainer is not None else self._get_trainer()
    self.history_ = _trainer.fit(self.model_, x_t, y_t, x_val=x_val_t, y_val=y_val_t, metrics=metrics)
    return self

def get_mf_params(self) -> dict[str, list[dict[str, Any]]]:
    """Return model membership function metadata after fitting."""
    check_is_fitted(self, "model_")
    return self.model_.get_mf_params()

def inspect(self) -> dict[str, Any]:
    """Return a structured summary of fitted model state and rule metadata."""
    check_is_fitted(self, "model_")
    return {
        "n_rules": int(self.model_.n_rules),
        "n_inputs": int(self.model_.n_inputs),
        "feature_names": list(self.model_.input_names),
        "rule_base": self.rule_base_,
        "defuzzifier_type": type(self.model_.defuzzifier).__name__,
        "mf_params": self.get_mf_params(),
        "rule_table": self.model_.get_rule_table(),
    }

@classmethod
def load(cls, path: str) -> Self:
    """Load a persisted estimator created by save."""
    checkpoint = load_checkpoint(path)
    validate_checkpoint_payload(checkpoint, expected_estimator_class=cls.__name__)

    params: dict[str, Any] = dict(checkpoint["estimator_params"])
    if params.get("input_configs") is not None:
        params["input_configs"] = [InputConfig(**c) for c in params["input_configs"]]
    estimator = cls(**params)
    model_init = checkpoint["model_init"]
    estimator.rule_base_ = model_init["rule_base"]

    import inspect

    sig = inspect.signature(estimator._build_regressor_model)
    if "rules" in sig.parameters:
        estimator.model_ = estimator._build_regressor_model(
            deserialize_input_mfs(model_init["input_mfs_config"]),
            str(model_init["rule_base"]),
            rules=model_init.get("rules"),
        )
    else:
        estimator.model_ = estimator._build_regressor_model(
            deserialize_input_mfs(model_init["input_mfs_config"]),
            str(model_init["rule_base"]),
        )

    consequent_mode = model_init.get("consequent_mode")
    if consequent_mode is not None:
        if hasattr(estimator.model_, "set_consequent_mode"):
            cast(Any, estimator.model_).set_consequent_mode(consequent_mode)
        elif hasattr(estimator.model_.consequent_layer, "mode"):
            estimator.model_.consequent_layer.mode = consequent_mode

    estimator.model_.load_state_dict(checkpoint["model_state_dict"])
    estimator.model_.to(torch.device(str(estimator.device)))

    fitted = checkpoint["fitted_attrs"]
    estimator.n_features_in_ = int(fitted["n_features_in"])
    if fitted.get("feature_names_in") is not None:
        estimator.feature_names_in_ = np.asarray(fitted["feature_names_in"], dtype=object)
    elif hasattr(estimator, "feature_names_in_"):
        delattr(estimator, "feature_names_in_")
    history = checkpoint.get("history")
    estimator.history_ = history if history is not None else {}
    return estimator

def predict(self, x: npt.ArrayLike) -> np.ndarray:
    """Predict continuous target values for input samples."""
    check_is_fitted(self, "model_")
    x_arr = validate_data(self, x, reset=False)
    device_str = str(self.device).lower()
    x_tensor = torch.as_tensor(x_arr, dtype=torch.float32, device=torch.device(device_str))
    preds = cast(Any, self.model_).predict(x_tensor)
    return preds.detach().cpu().numpy()

def rule_activation(self, X: npt.ArrayLike) -> np.ndarray:
    """Return normalized rule activations for the provided inputs."""
    check_is_fitted(self, "model_")
    x_arr = validate_data(self, X, reset=False)

    was_training = self.model_.training
    try:
        self.model_.eval()
        with torch.no_grad():
            norm_w = self.model_.forward_antecedents(self._as_tensor_x(x_arr, torch.device(str(self.device))))
    finally:
        self.model_.train(was_training)

    return _to_numpy(norm_w)

def save(self, path: str) -> None:
    """Persist estimator configuration, model weights and fitted metadata."""
    _fnames: np.ndarray | None = getattr(self, "feature_names_in_", None)
    rules = None
    if hasattr(self.model_, "rule_layer") and hasattr(self.model_.rule_layer, "rules"):
        rules = self.model_.rule_layer.rules
    consequent_mode = getattr(self.model_.consequent_layer, "mode", None)
    checkpoint = self._build_checkpoint_base(
        model_init={
            "input_mfs_config": serialize_input_mfs(self.model_.input_mfs),
            "rule_base": self.rule_base_,
            "rules": rules,
            "consequent_mode": consequent_mode,
        },
        fitted_attrs={
            "n_features_in": int(self.n_features_in_),
            "feature_names_in": _fnames.tolist() if _fnames is not None else None,
        },
    )
    save_checkpoint(path, checkpoint)

def __sklearn_is_fitted__(self) -> bool:
    """Check fitted status via model_ only, ignoring lambda_ (trailing underscore)."""
    return hasattr(self, "model_")

def evaluate(
    self,
    X: npt.ArrayLike,
    y: npt.ArrayLike,
    metrics: list[str] | None = None,
    sample_weight: npt.ArrayLike | None = None,
) -> dict[str, Any]:
    """Compute classification evaluation metrics for the provided dataset."""
    y_true = np.asarray(y)
    y_pred = self.predict(X)
    return compute_metrics(
        task="classification",
        y_true=y_true,
        y_pred=y_pred,
        sample_weight=sample_weight,
        metrics=metrics,
    )

def feature_importance(self) -> np.ndarray | None:
    """Compute a normalized feature importance vector from consequent weights."""
    check_is_fitted(self, "model_")
    weights = self.model_.get_consequent_weights()
    if weights is None:
        return None

    consequent_layer = self.model_.consequent_layer
    if hasattr(consequent_layer, "rule_feature_mask"):
        rule_feature_mask = cast(Tensor, consequent_layer.rule_feature_mask)
        weights = weights * rule_feature_mask.unsqueeze(1) if weights.ndim == 3 else weights * rule_feature_mask

    abs_weights = weights.abs()
    if abs_weights.ndim == 3:
        importance = abs_weights.mean(dim=(0, 1))
    elif abs_weights.ndim == 2:
        importance = abs_weights.mean(dim=0)
    else:
        raise ValueError("unsupported consequent weight shape for feature importance")

    return _normalize_importance(importance)