from __future__ import annotations

import numpy as np
from sarus_data_spec.typing import Dataset
from sarus_differential_privacy.query import LaplaceQuery, PrivateQuery
from sarus_statistics.ops.bounds.op import BoundOp
from sarus_statistics.tasks.bounds.base import BoundsParameters
from sarus_statistics.tasks.bounds.visitor import default_bounds

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 BoundsBuilder(QueryBuilder):
    """Generate Bounds hyperparameters"""

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

    def build_query(self, input_parameter: Query.Bounds) -> Task:
        bounds_tree = BoundsParameters(
            default_bounds(self._schema.data_type())
        )
        bounds_tree.set_noise(input_parameter.noise)
        return bounds_tree.protobuf()

    def private_query(self, out: Task) -> PrivateQuery:
        return BoundsParameters(out).private_query()


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

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


def bounds_builder(dataset: Dataset, query: Query) -> OptimizableBoundsBuilder:
    return OptimizableBoundsBuilder(dataset, query)


# Simple bound builder for one column
class SimpleBoundsBuilder(QueryBuilder):
    """Generate Bounds hyperparameters"""

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

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

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


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

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


def simple_bounds_builder(
    dataset: Dataset, query: Query
) -> OptimizableBoundsBuilder:
    return OptimizableSimpleBoundsBuilder(dataset, query)