import typing as t

from sarus_differential_privacy.query import PrivateQuery
import numpy as np

import sarus_data_spec.typing as st


class StatisticsOp:
    """Object describing a statistics operation
    Takes dataset & parameters as input, and can compute:

    - the DP-result of the operation with `.value()`,
    - the associated private_query with `.private_query()`.

    Note that the returned private queries are an upper bound, and the actual
    consumption might be lower."""

    def dataset(self) -> st.Dataset:
        raise NotImplementedError

    def value(
        self,
        random_generator: t.Optional[np.random.Generator] = None,
        *args: t.Any,
        **kwargs: t.Any,
    ) -> t.Any:
        raise NotImplementedError

    def private_query(self) -> PrivateQuery:
        raise NotImplementedError

    def _pandas(  # or _local() if we want to do it in arrow ?
        self,
        random_generator: t.Optional[np.random.Generator] = None,
        *args: t.Any,
        **kwargs: t.Any,
    ) -> t.Any: ...

    def _sql(
        self,
        random_generator: t.Optional[np.random.Generator] = None,
        *args: t.Any,
        **kwargs: t.Any,
    ) -> t.Any: ...