from __future__ import annotations

from sarus_differential_privacy.query import PrivateQuery
from sarus_statistics.ops.mean.op import MeanOp

from sarus_query_builder.core.core import OptimizableQueryBuilder, QueryBuilder
from sarus_query_builder.core.typing import Task
from sarus_query_builder.protobuf.query_pb2 import GenericTask, Query


class MeanBuilder(QueryBuilder):
    """Generate size hyperparameters"""

    def __init__(self, dataset: Dataset):
        self._dataset = dataset

    def build_query(self, input_parameter: Query.Mean) -> Task:
        return GenericTask(parameters={'noise': input_parameter.noise})

    def private_query(self, out: Task) -> PrivateQuery:
        return MeanOp(self.dataset, out.parameters['noise']).private_query()


class OptimizableMeanBuilder(OptimizableQueryBuilder):
    def __init__(self, dataset: Dataset, query: Query):
        self._dataset = dataset
        self.query = query
        self._builders = [MeanBuilder(dataset)]

    def build_query(self, input_parameter: float) -> Task:
        query = self.query
        if input_parameter:
            query.mean.noise = 1 / input_parameter
        else:
            query.mean.noise = np.inf
        return self.builders[0].build_query(query.mean)


def mean_builder(dataset: Dataset, query: Query) -> OptimizableMeanBuilder:
    return OptimizableMeanBuilder(dataset, query)