# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: CC-BY-NC-4.0
import jax.numpy as np
from jax import jit
from math import sqrt
import string


def get_sensitivity(numqueries):
    def get_sensitity_k(D):
        n, d = D.shape
        return sqrt(numqueries) / n

    return get_sensitity_k


# k way correlations --- i.e. for all k-tuple q in queries we compute the value \sum_t (\prod_j D[t,j], j in q)
def preserve_statistic(queries):
    # hardcoding threshold to 1 here

    letters_idx = string.ascii_lowercase
    row_idx, col_idxs = letters_idx[0], letters_idx[1:]

    assert len(queries[0]) <= len(col_idxs)

    # construct string for k-th order statistic with einsum, e.g., for 3-th marginal 'ab,ac,ad->bcd'
    eins_string = "ij,ik,il->jkl"

    @jit
    def compute_statistic(D):
        return (
            np.concatenate(
                [
                    D.shape[0]
                    - np.einsum(
                        eins_string, *[(1 - D[:, idx_q]) for idx_q in q]
                    ).flatten()
                    for q in queries
                ]
            )
            / D.shape[0]
        )

    return compute_statistic


def preserve_subset_statistic(queries):
    @jit
    def compute_statistic(D):
        np.concatenate(
            [(1 - np.prod(1 - D[:, q], 2)).sum(0) for q in np.array_split(queries, 10)]
        )
        return (
            np.concatenate(
                [
                    (1 - np.prod(1 - D[:, q], 2)).sum(0)
                    for q in np.array_split(queries, 10)
                ]
            )
            / D.shape[0]
        )

    return compute_statistic
