from __future__ import annotations

import numpy as np
from sarus_data_spec.typing import Dataset
from sarus_differential_privacy.query import (
    EpsilonQuery,
    GaussianQuery,
    LaplaceQuery,
    PrivateQuery,
)

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


# Laplace
class LaplaceBuilder(QueryBuilder):
    """Generate Bounds hyperparameters"""

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

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

    def private_query(self, out: Task) -> PrivateQuery:
        # if not isinstance(out, GenericTask):
        #    raise TypeError("Expected GenericTask task")
        return LaplaceQuery(out.parameters['noise'])


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

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


# Gaussian
class GaussianBuilder(QueryBuilder):
    """Generate Bounds hyperparameters"""

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

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

    def private_query(self, out: Task) -> PrivateQuery:
        # if not isinstance(out, GenericTask):
        #    raise TypeError("Expected GenericTask task")
        return GaussianQuery(out.parameters['noise'])


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

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


# Epsilon
class EpsilonBuilder(QueryBuilder):
    """Generate Bounds hyperparameters"""

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

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

    def private_query(self, out: Task) -> PrivateQuery:
        # if not isinstance(out, GenericTask):
        #    raise TypeError("Expected GenericTask task")
        return EpsilonQuery(out.parameters['noise'])


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

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


def laplace_builder(
    dataset: Dataset, query: Query
) -> OptimizableLaplaceBuilder:
    return OptimizableLaplaceBuilder(dataset, query)


def gaussian_builder(
    dataset: Dataset, query: Query
) -> OptimizableGaussianBuilder:
    return OptimizableGaussianBuilder(dataset, query)


def epsilon_builder(
    dataset: Dataset, query: Query
) -> OptimizableEpsilonBuilder:
    return OptimizableEpsilonBuilder(dataset, query)