from __future__ import annotations

from sarus_differential_privacy.query import PrivateQuery
from sarus_statistics.ops.median.op import MedianOp

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 MedianBuilder(QueryBuilder):
    """Generate size hyperparameters"""

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

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

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


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

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


def median_builder(dataset: Dataset, query: Query) -> OptimizableMedianBuilder:
    return OptimizableMedianBuilder(dataset, query)