Extract

class ExtractFeatures:
    """Port of ``FantailConnector``. Construct with a ``characterizer_set``
    of ``"simple"`` or ``"all"``; call ``process(did)`` for one dataset or
    ``poll()`` to loop the qualities-unprocessed endpoint."""

    def __init__(
        self,
        client: OpenmlClient,
        mode: str = "normal",
        characterizer_set: str = "simple",
        priority_tag: Optional[str] = None,
        dataset_format: DataFormat = "arff",
    ) -> None:
        if characterizer_set not in ("simple", "all"):
            raise ValueError(
                f"characterizer_set must be 'simple' or 'all', got {characterizer_set!r}"
            )
        self.client = client
        self.mode = mode
        self.characterizer_set = characterizer_set
        self.priority_tag = priority_tag
        self.dataset_format: DataFormat = dataset_format
        self.last_result: Optional[DataQuality] = None

    # ----------------------------------------------------------------------
    # Single-dataset path (Java: FantailConnector.computeMetafeatures)
    # ----------------------------------------------------------------------

    def process(self, did: int) -> DataQuality:
        """Compute qualities for one dataset and upload. Returns the
        ``DataQuality`` (Java uploads as a side effect — same here)."""
        info = get_data_and_meta_information_from_did(
            did, dataset_type=self.dataset_format, base_url=self.client.base_url
        )
        data_quality = load_qualities(
            info,
            data_format=self.dataset_format,
            did=did,
            evaluation_engine_id=EVALUATION_ENGINE_ID,
        )

        # Java's "all" set runs landmarkers on top of the base statistical
        # characterizers. We append sklearn-landmarker Quality objects to the
        # same DataQuality before upload.
        if self.characterizer_set == "all" and not data_quality.error:
            self._append_landmarkers(data_quality, info)

        if data_quality.qualities and not data_quality.error:
            try:
                self.client.data_qualities_upload(data_quality)
            except OpenmlApiError:
                # Java's FantailConnector lets upload errors propagate; we
                # match that — the caller (CLI / poll) decides what to do.
                raise

        self.last_result = data_quality
        return data_quality

    def _append_landmarkers(
        self,
        data_quality: DataQuality,
        info: DatasetDownloadInfo,
    ) -> None:
        """Compute all landmarkers and append to ``data_quality.qualities``.

        Silent on failure: if X/y can't be built (no target, unparseable),
        we leave the base qualities intact and skip — Java's GenericLandmarker
        returns null per-landmarker on exception; here the whole landmarker
        pass is all-or-nothing since they share one X/y."""
        from src.helpers import normalize_target_names

        try:
            attributes, rows = DataLoader(self.dataset_format).load(info)
            target_names = normalize_target_names(info.default_target_attribute)
            X, y = build_xy(attributes, rows, target_names)
            landmark_values = compute_all_landmarkers(X, y)
        except Exception:  # noqa: BLE001 — Java parity: landmarkers fail soft
            return

        for name, value in landmark_values.items():
            data_quality.qualities.append(Quality(name=name, value=value))

    # ----------------------------------------------------------------------
    # Polling loop (Java: FantailConnector.start with id == null)
    # ----------------------------------------------------------------------

    def poll(self) -> None:
        """Loop ``data_qualities_unprocessed``, processing each dataset id
        until the server returns "No unprocessed" (Java's
        ``fetchUnprocessedQualities`` catches the same string)."""
        expected = self._expected_quality_ids()
        while True:
            try:
                dids = self.client.data_qualities_unprocessed(
                    EVALUATION_ENGINE_ID,
                    self.mode,
                    expected,
                    priority_tag=self.priority_tag,
                )
            except OpenmlApiError as e:
                if "No unprocessed" in e.message:
                    return
                raise
            if not dids:
                return
            for did in dids:
                self.process(did)

    def _expected_quality_ids(self) -> list[str]:
        ids = list(_SIMPLE_META_FEATURE_IDS)
        if self.characterizer_set == "all":
            ids.extend(expected_landmarker_ids())
        return ids