from __future__ import annotations

import numpy as np
from sarus_data_spec.typing import Dataset
from sarus_differential_privacy.query import PrivateQuery
from sarus_statistics.tasks.size.base import SizeParameters
from sarus_statistics.tasks.size.visitor import default_size

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 Query


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

    def __init__(self, dataset: Dataset):
        self._dataset = dataset
        self._schema = dataset.schema()
        self.is_big_data = self.dataset.manager().is_big_data(dataset)

    def build_query(self, input_parameter: Query.Size) -> Task:
        size_tree = SizeParameters(
            default_size(self._schema.data_type()),
            is_big_data=self.is_big_data,
        )
        noise = input_parameter.noise
        size_tree.set_noise(noise)
        proto = size_tree.protobuf()
        return proto

    def private_query(self, out: Task) -> PrivateQuery:
        return SizeParameters(
            out, is_big_data=self.is_big_data
        ).private_query()


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

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


def size_builder(dataset: Dataset, query: Query) -> OptimizableSizeBuilder:
    return OptimizableSizeBuilder(dataset, query)