from __future__ import annotations

from sarus_differential_privacy.query import PrivateQuery
from sarus_statistics.ops.sum.op import SumOp

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

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

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

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


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

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


def sum_builder(dataset: Dataset, query: Query) -> OptimizableSumBuilder:
    return OptimizableSumBuilder(dataset, query)