from __future__ import annotations

from sarus_data_spec.typing import Dataset
from sarus_differential_privacy.query import PrivateQuery
from sarus_statistics.ops.covariance.op import CovarianceOp

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

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

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

    def private_query(self, out: Task) -> PrivateQuery:
        return CovarianceOp(
            self.dataset,
            out.parameters['epsilon'],
            out.parameters['dims'],
        ).private_query()


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

    def build_query(self, input_parameter: float) -> Task:
        query = self.query
        query.covariance.epsilon = input_parameter
        return self.builders[0].build_query(query.covariance)


def simple_covariance_builder(
    dataset: Dataset, query: Query
) -> OptimizableCovarianceBuilder:
    return OptimizableCovarianceBuilder(dataset, query)