from abc import ABC, abstractmethod
from logging import getLogger
from typing import Callable, Any

import jax
import numpy as np
import pixeltable as pxt
from pixeltable import Error, Query
from pixeltable.catalog import Table

import torch
import jax.numpy as jnp

from coreax import SupervisedData, ScalarValuedKernel
from coreax.kernels import LinearKernel
from coreax.solvers import GreedyKernelPoints
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

from goldener.clusterize import get_random_chunk_assignment
from goldener.pxt_utils import (
    set_value_to_idx_rows,
    GoldPxtTorchDataset,
    get_expr_from_column_name,
    get_valid_table,
    make_batch_ready_for_table,
    check_pxt_table_has_primary_key,
    get_sample_row_from_idx,
)
from goldener.reduce import GoldReductionTool, GoldReductionToolWithFit
from goldener.torch_utils import get_dataset_sample_dict
from goldener.utils import (
    filter_batch_from_indices,
    get_indices_with_labels,
    split_sampling_among_chunks,
    get_sampling_count_from_ratios,
    get_ratios_for_counts,
)

logger = getLogger(__name__)


class GoldSelectionTool(ABC):
    """Run selection of a subset of data points from vectorized samples using a coresubset selection algorithm."""

    def _validate_input(self, x: torch.Tensor) -> None:
        """Validate that input is already vectorized (2D torch.Tensor).

        Args:
            x: Input tensor to validate.

        Raises:
            ValueError: If input is not a 2D tensor.
        """
        if x.ndim != 2:
            raise ValueError(
                f"GoldSelectionTool only accepts 2D tensors (num_vectors, feature_dim). "
                f"Got shape {x.shape}. Please ensure your input is already vectorized."
            )

    @abstractmethod
    def select(
        self,
        x: torch.Tensor,
        k: int,
    ) -> list[int]:
        """Select a subset of data points from the input data.

        Args:
            x: A 2D torch.Tensor where each row represents a data point.
            k: The number of data points to select.

        Returns:
            A list of indices corresponding to the selected data points.
        """


class GoldGreedyKernelPointsSelectionTool(GoldSelectionTool):
    """Coresubset selection using Coreax's GreedyKernelPoints solver.

    The solver is created at each request depending on the specified k points to select.

    Attributes:
        feature_kernel: The kernel to use for the GreedyKernelPoints solver.
        random_key: Random key for reproducibility in the GreedyKernelPoints solver.
    """

    def __init__(
        self,
        feature_kernel: ScalarValuedKernel,
        random_key: int = 42,
    ) -> None:
        self.random_key = random_key
        self.feature_kernel = feature_kernel

    def select(
        self,
        x: torch.Tensor,
        k: int,
    ) -> list[int]:
        """Select a subset of data points from the input data using greedy kernel point coresubset sampling.

        Args:
            x: A 2D torch.Tensor where each row represents a data point.
            k: The number of data points to select.

        Returns:
            A list of indices corresponding to the selected data points.
        """
        self._validate_input(x)
        in_data = SupervisedData(
            data=jnp.array(x.numpy(force=True)), supervision=jnp.array(np.ones(len(x)))
        )
        solver: GreedyKernelPoints = GreedyKernelPoints(
            coreset_size=k,
            random_key=jax.random.key(self.random_key),
            feature_kernel=self.feature_kernel,
        )
        coreset, _ = solver.reduce(in_data)

        return coreset.unweighted_indices.tolist()


class GoldGreedyKCenterSelectionTool(GoldSelectionTool):
    """Create a coresubset from selecting with the K Center greedy algorithm.

    This is a greedy algorithm that selects iteratively the farthest to the already selected points
    among the not selected points. The first center is selected as the point with
    the largest aggregated distance to all the other points.

    Attributes:
        device: The torch device to use for computations.
    """

    def __init__(self, device: torch.device | str) -> None:
        self.device = device

    def select(
        self,
        x: torch.Tensor,
        k: int,
    ) -> list[int]:
        """Select a subset of data points from the input data by choosing iteratively the point
        with the farthest nearest neighbors among the not selected points.

        Args:
            x: A 2D torch.Tensor where each row represents a data point.
            k: The number of data points to select.

        Returns:
            A list of indices corresponding to the selected data points.
        """
        self._validate_input(x)
        x_len = len(x)
        if k > x_len:
            raise ValueError("k cannot be greater than the number of data points in x.")

        if k == x_len:
            return list(range(x_len))

        x = x.to(self.device)

        # the first point is the one with the largest aggregated distance to all the other points
        sample_to_sample_distance = torch.cdist(x, x)
        first_idx = int(sample_to_sample_distance.norm(dim=1).argmax().item())
        selected_indices = [first_idx]
        if k == 1:
            return selected_indices

        # the next points are selected iteratively as the ones
        # with the largest distance to the already selected points
        distances = sample_to_sample_distance[first_idx]
        for _ in range(k - 1):
            next_idx = int(distances.argmax().item())
            selected_indices.append(next_idx)

            distances_to_next = sample_to_sample_distance[next_idx]
            distances = torch.stack([distances, distances_to_next]).min(dim=0).values

        return selected_indices


class GoldGreedyClosestPointSelectionTool(GoldSelectionTool):
    """Create a coresubset from selecting the closest points iteratively.

    This is a greedy algorithm that selects iteratively the point with the closest nearest neighbors
    among the not selected points.

    Attributes:
        device: The torch device to use for computations.
    """

    def __init__(self, device: torch.device | str) -> None:
        self.device = device

    def select(
        self,
        x: torch.Tensor,
        k: int,
    ) -> list[int]:
        """Select a subset of data points from the input data by choosing iteratively the point
        with the closest nearest neighbors among the not selected points.

        Args:
            x: A 2D torch.Tensor where each row represents a data point.
            k: The number of data points to select.

        Returns:
            A list of indices corresponding to the selected data points.
        """
        self._validate_input(x)
        x_len = len(x)
        if k > x_len:
            raise ValueError("k cannot be greater than the number of data points in x.")

        if k == x_len:
            return list(range(x_len))

        x = x.to(self.device)

        selected_indices: list[int] = []
        remaining_indices = set(range(len(x)))

        for _ in range(k):
            remaining_indices_as_list = list(remaining_indices)
            remaining_vectors = x[remaining_indices_as_list]

            distance = torch.cdist(remaining_vectors, remaining_vectors)
            distance = (
                distance
                + torch.eye(len(remaining_vectors), device=x.device, dtype=x.dtype)
                * distance.max()
            )  # set self-distance to max

            point_with_closest = int(distance.min(dim=1).values.argmin().item())
            index_in_original = remaining_indices_as_list[point_with_closest]

            selected_indices.append(index_in_original)
            remaining_indices.remove(index_in_original)

        return selected_indices


class GoldGreedyFarthestPointSelectionTool(GoldSelectionTool):
    """Create a coresubset from selecting the farthest points iteratively.

    This is a greedy algorithm that selects iteratively the point with the farthest nearest neighbors
    among the not selected points.

    Attributes:
        device: The torch device to use for computations.
    """

    def __init__(self, device: torch.device | str) -> None:
        self.device = device

    def select(
        self,
        x: torch.Tensor,
        k: int,
    ) -> list[int]:
        """Select a subset of data points from the input data by choosing iteratively the point
        with the farthest nearest neighbors among the not selected points.

        Args:
            x: A 2D torch.Tensor where each row represents a data point.
            k: The number of data points to select.

        Returns:
            A list of indices corresponding to the selected data points.
        """
        self._validate_input(x)
        x_len = len(x)
        if k > x_len:
            raise ValueError("k cannot be greater than the number of data points in x.")

        if k == x_len:
            return list(range(x_len))

        x = x.to(self.device)

        selected_indices: list[int] = []
        remaining_indices = set(range(len(x)))

        for _ in range(k):
            remaining_indices_as_list = list(remaining_indices)
            remaining_vectors = x[remaining_indices_as_list]

            distance = torch.cdist(remaining_vectors, remaining_vectors)
            distance = (
                distance
                + torch.eye(len(remaining_vectors), device=x.device, dtype=x.dtype)
                * distance.max()
            )  # set self-distance to max
            point_with_farthest = int(distance.min(dim=1).values.argmax().item())
            index_in_original = remaining_indices_as_list[point_with_farthest]

            selected_indices.append(index_in_original)
            remaining_indices.remove(index_in_original)

        return selected_indices


class GoldSelector:
    """Select a subset of data points from vectorized samples.

    The GoldSelector processes a dataset or PixelTable table to perform coresubset selection using a
    selection algorithm on already vectorized representations. The selection results are stored
    in a local PixelTable table (specified by `table_path`) so that the selection process is idempotent:
    calling the same operation multiple times will not duplicate or recompute selections that are already
    present in the table.

    If the dataset is too big to fit into memory or the coresubset selection algorithm is too
    computationally expensive, the coresubset selection can be performed in chunks.
    All torch tensors will be converted to numpy arrays before saving.

    The selection can operate in a sequential (single-process) mode or a
    distributed mode (not implemented). The table schema is created/validated automatically and will include
    minimal indexing columns (`idx`, `idx_vector`) required to link selected samples back to
    their originating data points.

    Attributes:
        table_path: Path to the PixelTable table where selection results will be stored locally.
        selection_tool: The GoldSelectionTool implementing the selection algorithm.
        reducer: Optional GoldReductionTool instance for dimensionality reduction before selection.
        chunk: Optional chunk size for processing data in chunks to reduce memory consumption.
        collate_fn: Optional function to collate dataset samples into batches composed of
            dictionaries with at least the key specified by `vectorized_key` returning a PyTorch Tensor.
            If None, the dataset is expected to directly provide such batches.
        vectorized_key: Key in the batch dictionary that contains the vectorized data for selection. Default is "vectorized".
        include_vectorized_in_table: Whether to include the vectorized data in the selection table. Defaults to False.
        It is only applied if the cluster table is created from a Table (it is forced anyway for Dataset).
        selection_key: Column name to store the selection value in the PixelTable table. Default is "selected".
        label_key: Optional key for label-based stratified selection.
        exclude_labels: Optional set of label strings to exclude from selection. Default is None.
        to_keep_schema: Optional dictionary defining additional columns to keep from the original dataset/table
            into the selection table. The keys are the column names and the values are the PixelTable types.
        min_pxt_insert_size: Minimum number of rows to accumulate before inserting into PixelTable. Default is 100.
        batch_size: Batch size used when iterating over the data.
        num_workers: Number of workers for the PyTorch DataLoader during iteration on data.
        allow_existing: If False, an error will be raised when the table already exists. Default is True.
        distribute: Whether to use distributed processing for selection and table population. Not implemented yet. Default is False.
        drop_table: Whether to drop the selection table after creating the dataset with selection results. It is only applied
            when using `select_in_dataset`. Default is False.
        max_batches: Optional maximum number of batches to process. Useful for testing on a small subset of the dataset.
        random_state: Optional random state for reproducibility during chunk assignment. Default is None.
    """

    _MINIMAL_SCHEMA: dict[str, type] = {
        "idx": pxt.Required[pxt.Int],
        "idx_vector": pxt.Required[pxt.Int],
    }

    def __init__(
        self,
        table_path: str,
        selection_tool: GoldSelectionTool = GoldGreedyKernelPointsSelectionTool(
            feature_kernel=LinearKernel(output_scale=1, constant=0)
        ),
        reducer: GoldReductionTool | None = None,
        chunk: int | None = None,
        collate_fn: Callable | None = None,
        vectorized_key: str = "vectorized",
        include_vectorized_in_table: bool = False,
        selection_key: str = "selected",
        label_key: str | None = None,
        exclude_labels: set[str] | None = None,
        to_keep_schema: dict[str, type] | None = None,
        min_pxt_insert_size: int = 100,
        batch_size: int = 1,
        num_workers: int = 0,
        allow_existing: bool = True,
        distribute: bool = False,
        drop_table: bool = False,
        max_batches: int | None = None,
        random_state: int | None = None,
    ) -> None:
        """Initialize the GoldSelector.

        Args:
            table_path: Path to store the PixelTable table.
            selection_tool: The GoldSelectionTool implementing the selection algorithm.
            reducer: Optional GoldReductionTool instance for dimensionality reduction before selection.
            chunk: Optional chunk size for processing data in chunks.
            collate_fn: Optional collate function for the DataLoader.
            vectorized_key: Key pointing to the vector for selection. Defaults to "vectorized".
            include_vectorized_in_table: Whether to include the vectorized data in the selection table. Defaults to False.
                It is only applied if the selection table is created from a Table (it is forced anyway for Dataset).
            selection_key: Key for storing selection values. Defaults to "selected".
            label_key: Optional key for label stratification.
            exclude_labels: Optional set of label strings to exclude from selection. Default is None.
            to_keep_schema: Optional schema for additional columns to preserve.
            min_pxt_insert_size: Minimum number of rows to accumulate before inserting into PixelTable. Default is 100.
            batch_size: Batch size used when iterating over the data.
            num_workers: Number of workers for the PyTorch DataLoader during iteration on data.
            allow_existing: Whether to allow using an existing table. Defaults to True.
            distribute: Whether to use distributed selection. Defaults to False.
            drop_table: Whether to drop the table after dataset creation. Defaults to False.
            max_batches: Optional maximum number of batches to process.
            random_state: Optional random state for reproducibility during chunk assignment. Default is None.
        """
        self.table_path = table_path
        self.selection_tool = selection_tool
        self.reducer = reducer
        self.chunk = chunk
        self.collate_fn = collate_fn
        self.vectorized_key = vectorized_key
        self.include_vectorized_in_table = include_vectorized_in_table
        self.selection_key = selection_key
        if exclude_labels is not None and label_key is None:
            raise ValueError(
                "If exclude_labels is provided, label_key must also be provided."
            )
        self.label_key = label_key
        self.exclude_labels = exclude_labels
        self.to_keep_schema = to_keep_schema
        self.min_pxt_insert_size = min_pxt_insert_size
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.allow_existing = allow_existing
        self.distribute = distribute
        self.drop_table = drop_table
        self.max_batches = max_batches
        self.random_state = random_state

    def setup_selection_table(self, select_from: Dataset | Table) -> Table:
        """Set up the selection table based on the source dataset or table.

        Args:
            select_from: Dataset or Table to select from. If a Dataset is provided, each item should be a
                dictionary with at least the `vectorized_key` and `idx` keys after applying the collate_fn.
                If the collate_fn is None, the dataset is expected to directly provide such batches.
                If a Table is provided, it should contain at least the `vectorized_key`,
                `idx` and `idx_vector` columns.

        Returns: A PixelTable Table with the appropriate schema for selection,
            populated with rows corresponding to the source dataset or table.
        """
        logger.info(f"Loading the existing selection table from {self.table_path}")
        try:
            old_selection_table = pxt.get_table(
                self.table_path,
                if_not_exists="ignore",
            )
        except Error:
            logger.info(f"No existing selection table from {self.table_path}")
            old_selection_table = None

        # selection table are expected to have a primary key allowing to idempotent updates
        if old_selection_table is not None:
            check_pxt_table_has_primary_key(old_selection_table, set(["idx_vector"]))

        if not self.allow_existing and old_selection_table is not None:
            raise ValueError(
                f"Table at path {self.table_path} already exists and "
                "allow_existing is set to False."
            )

        if isinstance(select_from, Table):
            selection_table = self._selection_table_from_table(
                select_from=select_from,
                old_selection_table=old_selection_table,
            )
        else:
            selection_table = self._selection_table_from_dataset(
                select_from, old_selection_table
            )

        return selection_table

    def select_in_dataset(
        self, select_from: Dataset | Table, select_count: int, value: str
    ) -> GoldPxtTorchDataset:
        """Select a subset of samples using coresubset selection and return results as a GoldPxtTorchDataset.

        The selection process applies a coresubset selection algorithm on already vectorized
        representations of the data points and stores results in a PixelTable table specified by `table_path`.
        When the chunk attribute is set, the selection is performed in chunks to reduce memory consumption.
        If a reducer is provided, the vectors are reduced in dimension before applying the coresubset selection.

        This is a convenience wrapper that runs `select_in_table` to populate
        (or resume populating) the PixelTable table, then wraps the table into a
        `GoldPxtTorchDataset` for downstream consumption. If `drop_table` is True,
        the table will be removed after the dataset is created.

        Args:
            select_from: Dataset or Table to select from. If a Dataset is provided, each item should be a
                dictionary with at least the `vectorized_key` and `idx` keys after applying the collate_fn.
                If the collate_fn is None, the dataset is expected to directly provide such batches.
                If a Table is provided, it should contain at least the `vectorized_key`, `idx` and `idx_vector` columns.
            select_count: Number of data points to select.
            value: Value to set in the `selection_key` column for selected samples.

        Returns:
            A GoldPxtTorchDataset containing at least the selection information in the `selection_key` key
                and `idx` (index of the sample) and `idx_vector` (index of the vector) keys as well.
        """

        selected_table = self.select_in_table(select_from, select_count, value)

        selected_dataset = GoldPxtTorchDataset(selected_table, keep_cache=True)

        if self.drop_table:
            pxt.drop_table(selected_table)

        return selected_dataset

    def select_in_table(
        self, select_from: Dataset | Table, select_size: int | float, value: str | None
    ) -> Table:
        """Select a subset of samples using coresubset selection and store results in a PixelTable table.

        The selection process applies a coresubset selection algorithm on already vectorized
        representations of the data points and stores results in a PixelTable table specified by `table_path`.
        When the chunk attribute is set, the selection is performed in chunks to reduce memory consumption.
        If a reducer is provided, the vectors are reduced in dimension before applying the coresubset selection.

        This method is idempotent (i.e., failure proof), meaning that if it is called
        multiple times on the same dataset or table, it will restart the selection process
        based on the vectors already present in the PixelTable table.

        Args:
            select_from: Dataset or Table to select from. If a Dataset is provided, each item should be a
                dictionary with at least the `vectorized_key` and `idx` keys after applying the collate_fn.
                If the collate_fn is None, the dataset is expected to directly provide such batches.
                If a Table is provided, it should contain at least the `vectorized_key`, `idx` and `idx_vector` columns.
            select_size: Number or ratio (between 0 and 1) of data points to select.
            value: Value to set in the `selection_key` column for selected samples.

        Returns:
            A PixelTable Table containing at least the selection information in the `selection_key` column
                and `idx` (index of the sample) and `idx_vector` (index of the vector) columns as well.

        Raises:
            ValueError: If select_size is a float not in the range (0.0, 1.0).
        """
        if isinstance(select_size, float):
            if not (0.0 < select_size < 1.0):
                raise ValueError(
                    "When select_size is a float, it must be in the range (0.0, 1.0)."
                )
        else:
            if select_size <= 0:
                raise ValueError("select_size must be a positive integer.")

        selection_table = self.setup_selection_table(select_from)
        if isinstance(select_from, Dataset):
            select_from = selection_table

        assert isinstance(select_from, Table)

        # define the number of element to sample
        total_size = select_from.select(select_from.idx).distinct().count()
        select_count = (
            select_size
            if isinstance(select_size, int)
            else int(select_size * total_size)
        )
        if select_count == 0 and isinstance(select_size, float):
            select_count = 1  # at least one sample

        selection_indices = self.get_selection_indices(
            selection_table, value, self.selection_key
        )
        already_selected_count = len(selection_indices)
        if already_selected_count > 0:
            # if new vectors have been added to the selection table
            # make sure already selected samples are already flagged out
            set_value_to_idx_rows(
                table=selection_table,
                col_expr=get_expr_from_column_name(selection_table, self.selection_key),
                idx_expr=selection_table.idx,
                indices=selection_indices,
                value=value,
            )

        still_to_select_count = select_count - already_selected_count
        if still_to_select_count == 0:
            logger.info(
                f"Selection table already fully filled out for {value} from {self.table_path}"
            )
            return selection_table
        elif still_to_select_count < 0:
            raise ValueError(
                f"Selection table already contains {already_selected_count} selected samples for {value}, "
                f"which is more than the requested {select_count} samples."
            )
        elif self.distribute:
            self._distributed_select(select_from, selection_table, select_count, value)
        else:
            self._sequential_select(select_from, selection_table, select_count, value)

        logger.info(
            f"Selection table populated {
                len(
                    self.get_selection_indices(
                        selection_table, value, self.selection_key
                    )
                )
            } rows with value {value} at {self.table_path}"
        )

        return selection_table

    def _selection_table_from_table(
        self, select_from: Table, old_selection_table: Table | None
    ) -> Table:
        """Create or validate the selection table schema from a PixelTable table.

        This private method sets up the table structure with necessary columns for tracking
        selection status and ensures all rows from the source table are represented.

        Args:
            select_from: The source PixelTable table to select from.
            old_selection_table: Existing selection table if resuming, or None.

        Returns:
            The selection table with proper schema and initial rows.
        """
        minimal_schema = self._MINIMAL_SCHEMA.copy()

        if self.to_keep_schema is not None:
            minimal_schema |= self.to_keep_schema

        if self.label_key is not None:
            minimal_schema[self.label_key] = pxt.String

        selection_table = get_valid_table(
            table=old_selection_table
            if old_selection_table is not None
            else self.table_path,
            minimal_schema=minimal_schema,
            primary_key="idx_vector",
        )

        if self.vectorized_key not in select_from.columns():
            raise ValueError(
                f"Table at path {self.table_path} does not contain "
                f"the required column {self.vectorized_key}."
            )

        selection_table_cols = selection_table.columns()
        if self.selection_key not in selection_table_cols:
            selection_table.add_column(
                if_exists="error", **{self.selection_key: pxt.String}
            )

        if (
            self.include_vectorized_in_table
            and self.vectorized_key not in selection_table_cols
        ):
            idx_vectors = [row["idx_vector"] for row in select_from.sample(1).collect()]
            if not idx_vectors:
                raise ValueError(
                    f"Source table at {self.table_path} is empty, cannot sample vectorized data for schema inference."
                )
            sample = get_sample_row_from_idx(
                select_from,
                idx_key="idx_vector",
                # make sure to take an existing vector
                idx=idx_vectors[0],
                collate_fn=self.collate_fn,
                expected_keys=[self.vectorized_key],
            )

            vectorized_value = sample[self.vectorized_key]
            selection_table.add_column(
                **{
                    self.vectorized_key: pxt.Array[  # type: ignore[misc]
                        vectorized_value.shape, pxt.Float
                    ]
                }
            )

        if selection_table.count() > 0:
            to_select_indices = set(
                [
                    row["idx_vector"]
                    for row in select_from.select(select_from.idx_vector)
                    .distinct()
                    .collect()
                ]
            )
            already_in_selection = set(
                [
                    row["idx_vector"]
                    for row in selection_table.select(selection_table.idx_vector)
                    .distinct()
                    .collect()
                ]
            )
            still_to_select = to_select_indices.difference(already_in_selection)
            if not still_to_select:
                logger.info(
                    f"The selection table is already initialized in {self.table_path}"
                )
                return selection_table

        col_list = ["idx", self.vectorized_key]
        if self.to_keep_schema is not None:
            col_list.extend(list(self.to_keep_schema.keys()))

        if self.selection_key in select_from.columns():
            col_list.append(self.selection_key)

        if self.label_key is not None and self.label_key in select_from.columns():
            col_list.append(self.label_key)

        self._add_rows_to_selection_table_from_dataset(
            select_from=GoldPxtTorchDataset(
                select_from.select(
                    *[
                        get_expr_from_column_name(select_from, col)
                        for col in col_list + ["idx_vector"]
                    ]
                ),
                keep_cache=False,
            ),
            selection_table=selection_table,
            include_vectorized=self.include_vectorized_in_table,
        )

        return selection_table

    def _selection_table_from_dataset(
        self, select_from: Dataset, old_selection_table: Table | None
    ) -> Table:
        """Create or validate the selection table schema from a PyTorch Dataset.

        This private method sets up the table structure with necessary columns
        and validate that the vectorized column is present in the dataset.

        Args:
            select_from: The source PyTorch Dataset to select from.
            old_selection_table: Existing selection table if resuming, or None.

        Returns:
            The selection table with proper schema.
        """
        minimal_schema = self._MINIMAL_SCHEMA.copy()
        if self.to_keep_schema is not None:
            minimal_schema |= self.to_keep_schema

        if self.label_key is not None:
            minimal_schema[self.label_key] = pxt.String

        selection_table = get_valid_table(
            table=old_selection_table
            if old_selection_table is not None
            else self.table_path,
            minimal_schema=minimal_schema,
            primary_key="idx_vector",
        )
        selection_table_cols = selection_table.columns()
        if self.vectorized_key not in selection_table_cols:
            sample = get_dataset_sample_dict(
                select_from,
                collate_fn=self.collate_fn,
                expected=[self.vectorized_key],
            )

            vectorized_value = (
                sample[self.vectorized_key].squeeze(0).detach().cpu().numpy()
            )
            selection_table.add_column(
                **{
                    self.vectorized_key: pxt.Array[  # type: ignore[misc]
                        vectorized_value.shape, pxt.Float
                    ]
                }
            )

        if self.selection_key not in selection_table_cols:
            selection_table.add_column(
                if_exists="error", **{self.selection_key: pxt.String}
            )

        self._add_rows_to_selection_table_from_dataset(
            select_from, selection_table, include_vectorized=True
        )

        return selection_table

    def _add_rows_to_selection_table_from_dataset(
        self,
        select_from: Dataset,
        selection_table: Table,
        include_vectorized: bool = False,
    ) -> None:
        """Add rows from the source dataset to the selection table.

        This private method iterates through the dataset in batches and populates the
        selection table with vectorized data and metadata, skipping already processed samples.

        Args:
            select_from: The source PyTorch Dataset.
            selection_table: The selection table to populate.
            include_vectorized: Whether to include the vectorized data in the selection table.
        """
        dataloader = DataLoader(
            select_from,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            collate_fn=self.collate_fn,
        )

        # the vectors already present in the selection table will be excluded
        # from the batches to process to avoid duplicates and speed up processing
        # as well as the ones with excluded labels if specified
        to_exclude_from_batch = set(
            [
                row["idx_vector"]
                for row in selection_table.select(selection_table.idx_vector).collect()
            ]
        )

        ready_to_insert: list[dict[str, Any]] = []

        for batch_idx, batch in tqdm(
            enumerate(dataloader),
            desc="Initializing rows for the selection table",
            total=(
                None if hasattr(select_from, "__len__") is False else len(dataloader)
            ),
        ):
            # Stop if we've processed enough batches
            if self.max_batches is not None and batch_idx >= self.max_batches:
                break

            if "idx_vector" not in batch:
                starts = batch_idx * self.batch_size
                batch["idx_vector"] = [
                    starts + idx for idx in range(len(batch[self.vectorized_key]))
                ]

            if "idx" not in batch:
                batch["idx"] = batch["idx_vector"]

            if self.exclude_labels is not None and self.label_key in batch:
                assert self.label_key is not None
                to_exclude_from_batch.update(
                    get_indices_with_labels(
                        batch,
                        self.label_key,
                        self.exclude_labels,
                        index_key="idx_vector",
                    )
                )

            # Remove the vectors already present in the selection table to avoid duplicates and speed up processing
            # and as well the ones with excluded labels
            if to_exclude_from_batch:
                batch = filter_batch_from_indices(
                    batch,
                    to_exclude_from_batch,
                    index_key="idx_vector",
                )

                if len(batch) == 0:
                    continue  # all samples already described

            if self.selection_key not in batch:
                batch[self.selection_key] = [
                    None for _ in range(len(batch[self.vectorized_key]))
                ]

            # add newly added vectors to the exclusion set to avoid duplicates in the next batches
            to_exclude_from_batch.update(
                [
                    idx.item() if isinstance(idx, torch.Tensor) else idx
                    for idx in batch["idx_vector"]
                ]
            )
            to_insert_keys = ["idx", self.selection_key]
            if self.to_keep_schema is not None:
                to_insert_keys.extend(list(self.to_keep_schema.keys()))
            if self.label_key is not None and self.label_key in batch:
                to_insert_keys.append(self.label_key)
            if include_vectorized:
                to_insert_keys.append(self.vectorized_key)

            batch_as_list = make_batch_ready_for_table(
                batch,
                to_insert_keys,
                "idx_vector",
            )

            ready_to_insert.extend(batch_as_list)
            if len(ready_to_insert) >= self.min_pxt_insert_size:
                selection_table.insert(ready_to_insert)
                ready_to_insert = []

        if ready_to_insert:
            selection_table.insert(ready_to_insert)

    @staticmethod
    def get_selection_pxt_query(
        table: Table,
        value: str | None,
        selection_key: str,
        label_key: str | None = None,
        label_value: str | None = None,
        idx_key: str = "idx",
    ) -> Query:
        """Get the Pixeltable query to access samples selected with a given value.

        Args:
            table: PixelTable table to query.
            value: Value in the selection_key column to filter selected samples.
            selection_key: Column name used to store selection values.
            label_key: Optional column name used to filter samples by label.
            label_value: Optional label value to filter samples by label.
            idx_key: Column name used to get sample indices.
        """
        idx_col = get_expr_from_column_name(table, idx_key)
        selection_col = get_expr_from_column_name(table, selection_key)
        if label_value is not None and label_key is not None:
            label_col = get_expr_from_column_name(table, label_key)
            query = (selection_col == value) & (label_col == label_value)  # noqa: E712
        else:
            if label_key is not None or label_value is not None:
                raise ValueError("label_key and label_value must be set together.")
            query = selection_col == value  # noqa: E712

        return table.where(query).select(idx_col).distinct()

    @staticmethod
    def get_selection_indices(
        table: Table,
        value: str | None,
        selection_key: str,
        label_key: str | None = None,
        label_value: str | None = None,
        idx_key: str = "idx",
    ) -> set[int]:
        """Get the indices of samples selected with a given value.

        Args:
            table: PixelTable table to query.
            value: Value in the selection_key column to filter selected samples.
            selection_key: Column name used to store selection values.
            label_key: Optional column name used to filter samples by label.
            label_value: Optional label value to filter samples by label.
            idx_key: Column name used to get sample indices.
        """
        return set(
            [
                row[idx_key]
                for row in GoldSelector.get_selection_pxt_query(
                    table,
                    value,
                    selection_key,
                    label_key,
                    label_value,
                    idx_key,
                ).collect()
            ]
        )

    @staticmethod
    def get_selection_count(
        table: Table,
        value: str | None,
        selection_key: str,
        label_key: str | None = None,
        label_value: str | None = None,
        idx_key: str = "idx",
    ) -> int:
        """Get the number of samples selected with a given value.

        Args:
            table: PixelTable table to query.
            value: Value in the selection_key column to filter selected samples.
            selection_key: Column name used to store selection values.
            label_key: Optional column name used to filter samples by label.
            label_value: Optional label value to filter samples by label.
            idx_key: Column name used to get sample indices.
        """
        return GoldSelector.get_selection_pxt_query(
            table,
            value,
            selection_key,
            label_key,
            label_value,
            idx_key,
        ).count()

    def _sequential_select(
        self,
        select_from: Table,
        selection_table: Table,
        select_count: int,
        value: str | None,
    ) -> None:
        """Run sequential (single-process) selection process.

        This private method handles label-stratified selection if a label_key is configured,
        otherwise performs selection on the full dataset. It delegates the actual coresubset
        selection to _label_select.

        Args:
            select_from: The source table with vectorized data.
            selection_table: The table to store selection results.
            select_count: Number of samples to select.
            value: Value to assign to selected samples in the selection_key column.
        """
        if self.label_key is not None:
            label_col = get_expr_from_column_name(selection_table, self.label_key)

            already_selected_count = self.get_selection_count(
                selection_table,
                value,
                self.selection_key,
            )

            # The selection is done iteratively per label
            # The first labels might as well include the next labels in case of multi labels per sample
            # For the first selection loop, the initial target count is limiting the selection
            # (allowing to maximize the idempotence), but if not enough samples have been selected,
            # a new round of selection is done allowing some labels to not be sampled
            # (the selection count ends up to 0 for the label) and the selection
            # is going on ignoring the count selected in the previous round.
            force_all_labels = True
            labels = sorted(
                set(
                    row[self.label_key]
                    for row in selection_table.select(label_col).distinct().collect()
                )
            )
            while already_selected_count < select_count:
                # For the first loop, the initial count is the target and all labels are forced to be selected.
                # For the next loops, only the missing count is targeted
                loop_select_count = (
                    select_count  # force_all_labels means 1st initial selection loop
                    if force_all_labels
                    else select_count
                    - already_selected_count  # otherwise, select only the remaining samples to select after counting the already selected ones for the label
                )
                # During the loop over the labels, the status of each label is computed and this values
                # allows to define how many samples are still to select for the label
                # For the first loop, the initial population is considered as null in order to maximize the idempotence.
                # For the next loops, the already selected samples are the basis for the selection
                label_counts_init = (
                    {label: 0 for label in labels}
                    if force_all_labels
                    else {
                        label: self.get_selection_count(
                            selection_table,
                            value,
                            self.selection_key,
                            self.label_key,
                            label,
                        )
                        for label in labels
                    }
                )

                selection_label_counts = self._compute_label_counts_for_selection(
                    selection_table=selection_table,
                    labels=labels,
                    select_count=loop_select_count,
                    value=value,
                    force_all_labels=force_all_labels,
                )
                for label_idx, (label_value, label_count) in enumerate(
                    selection_label_counts.items()
                ):
                    if label_count == 0:
                        if force_all_labels:
                            raise ValueError(
                                f"Not enough sample for Label '{label_value}' to run the selection"
                            )

                        continue

                    # if the label is already selected for at least the target count defined for the label,
                    # skip the label (over selection is not an issue)
                    label_count_status = (
                        self.get_selection_count(
                            selection_table,
                            value,
                            self.selection_key,
                            self.label_key,
                            label_value,
                        )
                        - label_counts_init[label_value]
                    )
                    if label_count_status >= label_count:
                        continue

                    # select the missing ones based on the current status of the label
                    # in order to reach the target count for the label
                    logger.info(
                        f"Selecting {label_count} samples for label '{label_value}' for value '{value}'"
                    )
                    self._select_label(
                        select_from,
                        selection_table,
                        (
                            label_count
                            if force_all_labels
                            else label_count - label_count_status
                        ),
                        value,
                        label_value=label_value,
                        from_already=force_all_labels,
                    )

                # after the first loop, labels with lower population might are allowed to be sampled again.
                force_all_labels = False
                already_selected_count = self.get_selection_count(
                    selection_table, value=value, selection_key=self.selection_key
                )

        else:
            logger.info(f"Selecting {select_count} samples for value '{value}'")
            self._select_label(
                select_from,
                selection_table,
                select_count,
                value,
            )

    def _select_label(
        self,
        select_from: Table,
        selection_table: Table,
        select_count: int,
        value: str | None,
        label_value: str | None = None,
        from_already: bool = True,
    ) -> None:
        """Perform coresubset selection for a specific label or all data.

        This private method implements chunked coresubset selection.
        It processes data in chunks to manage memory, applies optional dimensionality reduction,
        and updates the selection table with selected samples.

        Args:
            select_from: The source table with vectorized data.
            selection_table: The table to store selection results.
            select_count: Number of samples to select.
            value: Value to assign to selected samples in the selection_key column.
            label_value: Optional label value to filter samples by label.
            from_already: Whether to restart selection from the already selected ones for the label.
                If False, the selection is done to select the full select_count even if some samples have already
                been selected for the label. If True, the selection is done to select only the remaining samples
                 to reach select_count when counting the already selected ones for the label.
        """
        selection_col = get_expr_from_column_name(selection_table, self.selection_key)
        vectorized_col = get_expr_from_column_name(select_from, self.vectorized_key)

        # define the number of samples to select from the already sampled ones for the label
        # when more_than_already is False, the selection is done to select only the remaining samples
        # to reach select_count. When it is True, the selection is done to select the full select_count
        # even if some samples have already been selected
        already_selected_count = self.get_selection_count(
            table=selection_table,
            value=value,
            selection_key=self.selection_key,
            label_key=self.label_key,
            label_value=label_value,
        )

        # validate that there is still enough samples to select from
        current_selected_count = already_selected_count if from_already else 0
        available_samples_for_selection_count = self.get_selection_count(
            table=selection_table,
            value=None,
            selection_key=self.selection_key,
            label_key=self.label_key,
            label_value=label_value,
        )
        if available_samples_for_selection_count < (
            select_count - current_selected_count
        ):
            raise ValueError(
                "Cannot select more unique data points than available in the dataset."
            )

        # The coresubset selection is done from all the vectors (after filtering) of all data points
        # (depending on data, a data point can have multiple vectors).
        # Then, the same data point can be selected multiple times if it has multiple vectors selected.
        # To achieve select_count of unique data points, we loop until we have enough unique data points selected.
        if label_value is not None:
            assert self.label_key is not None
            label_col = get_expr_from_column_name(selection_table, self.label_key)
            available_query = (selection_col == None) & (label_col == label_value)  # noqa: E712 E711
        else:
            available_query = selection_col == None  # noqa: E711
        while current_selected_count < select_count:
            loop_start = (
                current_selected_count  # validate some samples have been selected
            )
            selection_pool = selection_table.where(available_query)

            # when the number of samples to select is the same as the number of samples available for selection,
            # the coresubset selection can be skipped and all the available samples are selected
            still_selectable = (
                selection_pool.select(selection_table.idx).distinct().count()
            )
            if still_selectable == 0:
                if from_already:
                    raise ValueError(
                        "No more sample available for selection, but the selection is not complete. "
                        "This might be due to an issue in the selection process or the data."
                    )

                break

            loop_select_count = select_count - current_selected_count
            if loop_select_count == still_selectable:
                set_value_to_idx_rows(
                    table=selection_table,
                    col_expr=selection_col,
                    idx_expr=selection_table.idx,
                    indices=set(
                        row["idx"]
                        for row in selection_pool.select(selection_table.idx)
                        .distinct()
                        .collect()
                    ),
                    value=value,
                )
                break

            # assign the vectors to chunks for chunked selection
            to_select_vector_indices = [
                row["idx_vector"]
                for row in selection_pool.select(selection_table.idx_vector).collect()
            ]
            chunk_assignment = get_random_chunk_assignment(
                to_chunk=to_select_vector_indices,
                max_chunk_size=self.chunk,
                random_state=self.random_state,
            )
            select_count_per_chunk = split_sampling_among_chunks(
                to_split=loop_select_count,
                chunk_sizes=[len(chunk) for chunk in chunk_assignment],
            )

            # make coresubset selection per chunk
            for chunk_indices, chunk_select_count in zip(
                chunk_assignment, select_count_per_chunk
            ):
                if chunk_select_count == 0:
                    continue

                # update the vector available for selection because the selection from previous chunks
                # might have selected some of the samples in the current chunk, making them no longer
                # available for selection in the current chunk
                chunk_vector_indices_not_selected = [
                    row["idx_vector"]
                    for row in selection_table.where(
                        (selection_table.idx_vector.isin(chunk_indices))
                        & (selection_col == None)  # noqa: E711
                    )
                    .select(selection_table.idx_vector)
                    .collect()
                ]
                if len(chunk_vector_indices_not_selected) == 0:
                    continue

                # load the vectors and the corresponding indices for the chunk
                to_select_from = select_from.where(
                    select_from.idx_vector.isin(chunk_vector_indices_not_selected)
                ).select(select_from.idx, vectorized_col)
                to_select_for_chunk = [
                    (
                        torch.from_numpy(sample[self.vectorized_key]),
                        torch.tensor(sample["idx"]).unsqueeze(0),
                    )
                    for sample in to_select_from.collect()
                ]
                vectors_list, indices_list = map(list, zip(*to_select_for_chunk))
                vectors = torch.stack(vectors_list, dim=0)
                indices = torch.cat(indices_list, dim=0)

                if len(to_select_for_chunk) == chunk_select_count:
                    # take all the indices
                    # It can happen when the selection size is close to the total size
                    coresubset_indices = set(indices.tolist())
                else:
                    # make coresubset selection for the chunk
                    if self.reducer is not None:
                        vectors = (
                            self.reducer.fit_transform(vectors)
                            if isinstance(self.reducer, GoldReductionToolWithFit)
                            else self.reducer.transform(vectors)
                        )

                    coresubset_indices = self._coresubset_selection(
                        vectors, chunk_select_count, indices
                    )

                # update table with selected indices
                set_value_to_idx_rows(
                    table=selection_table,
                    col_expr=selection_col,
                    idx_expr=selection_table.idx,
                    indices=coresubset_indices,
                    value=value,
                )

                current_selected_count = self.get_selection_count(
                    table=selection_table,
                    value=value,
                    selection_key=self.selection_key,
                    label_key=self.label_key,
                    label_value=label_value,
                )
                if not from_already:
                    current_selected_count -= already_selected_count  # remove the initial selection count to follow the current one

            if current_selected_count == loop_start:
                raise ValueError(
                    "No new samples were selected in the iteration, but the selection is not complete. "
                    "This might be due to an issue in the selection process or the data."
                )

    def _distributed_select(
        self,
        select_from: Table,
        selection_table: Table,
        select_count: int,
        value: str | None,
    ) -> None:
        """Run distributed selection process (not implemented).

        Args:
            select_from: The source table with vectorized data.
            selection_table: The table to store selection results.
            select_count: Number of samples to select.
            value: Value to assign to selected samples in the selection_key column.

        Raises:
            NotImplementedError: Always raised as distributed mode is not yet implemented.
        """
        raise NotImplementedError("Distributed selection is not implemented yet.")

    def _coresubset_selection(
        self, x: torch.Tensor, select_count: int, indices: torch.Tensor
    ) -> set[int]:
        """Apply the selection algorithm.

        Args:
            x: Input vectors to select from.
            select_count: Number of vectors to select.
            indices: Original indices corresponding to each vector.

        Returns:
            Set of selected indices from the original index space.
        """
        selected_indices = self.selection_tool.select(x, select_count)

        return set(indices[torch.tensor(selected_indices)].tolist())

    def _compute_label_counts_for_selection(
        self,
        selection_table: Table,
        labels: list[str],
        select_count: int,
        value: str | None = None,
        force_all_labels: bool = True,
    ) -> dict[str, int]:
        assert self.label_key is not None
        # When all labels are required, all the not selected samples and the selected samples for the current
        # value are included in the count of the label.
        # This allows to maximize the idempotence of the selection process
        label_counts = {}
        for label in labels:
            count = self.get_selection_count(
                table=selection_table,
                value=None,
                selection_key=self.selection_key,
                label_key=self.label_key,
                label_value=label,
            )
            if force_all_labels:
                count += self.get_selection_count(
                    table=selection_table,
                    value=value,
                    selection_key=self.selection_key,
                    label_key=self.label_key,
                    label_value=label,
                )

            if count > 0:
                label_counts[label] = count
            elif force_all_labels:
                raise ValueError(
                    f"Label '{label}' has no more selectable sample for value '{value}', but force_all_labels is True. "
                    f"This might be due to an issue in the selection process or the data."
                )

        label_ratios = {
            value: ratio
            for value, ratio in zip(
                label_counts.keys(),
                get_ratios_for_counts(list(label_counts.values())),
            )
        }

        return get_sampling_count_from_ratios(
            label_ratios, select_count, force_non_zero=force_all_labels
        )
