from typing import Callable, List

from sarus_differential_privacy.query import (
    ComposedQuery,
    LaplaceQuery,
    PrivateQuery,
    SampledGaussianMechanismQuery,
)

from sarus_data_spec.constants import PUBLIC
from sarus_data_spec.path import path
from sarus_data_spec.typing import Dataset


def links_queries(
    dataset: Dataset,
    hist_noise_ratio: float,
    bounds_noise_ratio: float,
    quantiles_noise_ratio: float,
    n_quantiles: int,
    quantiles_sampling_ratio: float,
) -> Callable[[float], PrivateQuery]:
    """Returns a Callable that given a noise scale returns the corresponding
    private query"""

    private_queries: List[Callable[[float], PrivateQuery]] = []
    f_ks = dataset.foreign_keys()

    data_type = dataset.schema().data_type()
    for pointing, pointed in f_ks.items():
        path_pointing_struct = path(
            paths=dataset.schema().data_type().get(pointing).structs()
        )
        if data_type.sub_types(path_pointing_struct)[0].properties()[
            PUBLIC
        ] != str(True):
            # histograms
            private_queries.append(
                lambda s: LaplaceQuery(hist_noise_ratio * s)
            )

        path_pointed_struct = path(
            paths=(dataset.schema().data_type().get(pointed).structs())
        )
        if data_type.sub_types(path_pointed_struct)[0].properties()[
            PUBLIC
        ] != str(True):
            # bounds
            private_queries.append(
                lambda s: LaplaceQuery(bounds_noise_ratio * s)
            )
            # quantiles
            private_queries.append(
                lambda s: SampledGaussianMechanismQuery(
                    quantiles_sampling_ratio,
                    quantiles_noise_ratio * s,
                    n_quantiles,
                )
            )

    return lambda scale: ComposedQuery([q(scale) for q in private_queries])