from __future__ import annotations

from sarus_differential_privacy.query import PrivateQuery
from sarus_statistics.ops.std.op import StdOp

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


def std_task(query: Query.Std) -> GenericTask:
    return GenericTask(
        parameters={
            'noise_mean': query.noise_mean,
            'noise_square': query.noise_square,
            'noise_count': query.noise_count,
        }
    )


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

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

    def build_query(self, input_parameter: Query.Std) -> Task:
        return std_task(input_parameter)

    def private_query(self, out: Task) -> PrivateQuery:
        return StdOp(
            self.dataset,
            out.parameters['noise_mean'],
            out.parameters['noise_square'],
            out.parameters['noise_count'],
        ).private_query()


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

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


def std_builder(dataset: Dataset, query: Query) -> OptimizableStdBuilder:
    return OptimizableStdBuilder(dataset, query)