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.ops.links.local import links_query
from sarus_statistics.protobuf.links_pb2 import (
    LinksParameters,
    RandomLinksParameters,
)

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

NB_QUANTILES = 32


class LinksBuilder(QueryBuilder):
    """Generate Links hyperparameters"""

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

    def build_query(self, input_parameter: Query.Links) -> Task:
        nb_quantiles = input_parameter.nb_quantiles or NB_QUANTILES
        isotonic_regression = True

        random_links = RandomLinksParameters(
            noise=input_parameter.noise,
            nb_quantiles=nb_quantiles,
            isotonic_regression=isotonic_regression,
        )

        links_params = LinksParameters(random_links=random_links)
        return links_params

    def private_query(self, out: Task) -> PrivateQuery:
        # if not isinstance(out, LinksParameters):
        #    raise TypeError("Expected LinksParameters task")
        return links_query(self._dataset, out.random_links.noise)


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

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


def links_builder(dataset: Dataset, query: Query) -> OptimizableLinksBuilder:
    return OptimizableLinksBuilder(dataset, query)