from pure_ldp.core import FreqOracleClient

import numpy as np
from scipy.linalg import hadamard
import math
import random

class CMSClient(FreqOracleClient):
    def __init__(self, epsilon, hash_funcs, m, is_hadamard=False):
        """
        Apple's Count Mean Sketch (CMS) Algorithm
        Args:
            epsilon (float): Privacy Budget Epsilon
            hash_funcs (list of funcs): A list of hash function mapping data to {0...m-1} (can be generated by CMSServer)
            m (int): The length of the hash domain
            is_hadamard (optional bool): If true, uses Hadamard Count Mean Sketch (HCMS)
        """
        super().__init__(epsilon, None)
        self.update_params(hash_funcs, m, epsilon)
        self.is_hadamard = is_hadamard

        if self.is_hadamard:
            self.had = hadamard(self.m)

    def update_params(self, hash_funcs=None, m=None, epsilon=None, index_mapper=None):
        """
        Updates parameters
        Args:
            hash_funcs (optional list): List of k hash functions mapping data to {0...m-1}
            m (optional int): Length of hash domain
            epsilon (optional int): Privacy Budget
            index_mapper (optional function): Index mapper function
        """
        if hash_funcs is not None:
            self.hash_funcs = hash_funcs
            self.k = len(self.hash_funcs)

        self.epsilon = epsilon if epsilon is not None else self.epsilon
        self.m = m if m is not None else self.m

        if epsilon is not None:
            self.prob = 1 / (1 + math.pow(math.e, self.epsilon / 2))

    def _one_hot(self, data):
        """
        Used internally to perturb data
        Args:
            data: arbitrary data

        Returns: perturbed vector v and hash index j

        """
        j = random.randint(0, self.k-1)
        h_j = self.hash_funcs[j]
        v = [0] * self.m if self.is_hadamard else np.full(self.m, -1)
        v[h_j(data)] = 1
        return v, j

    def _cms_perturb(self, data):
        """
        Used internally for peturbing data using the CMS algorithm
        Args:
            data: data to be perturbed

        Returns: peturbed cms data

        """
        v, j = self._one_hot(data)
        v[np.random.rand(*v.shape) < self.prob] *= -1 # "flip" bits with prob
        # return np.multiply(v, b), j # Used to generate a random vector b using np.random.choice but it was 3x slower than the above line
        return v, j

    def _hcms_perturb(self, data):
        """
        Used internally for perturbing data using HCMS
        Args:
            data: data to be perturbed

        Returns: peturbed hcms data

        """
        if not (self.m & (self.m - 1)) == 0:
            raise ValueError("m must be a positive integer, and m must be a power of 2 to use hcms")

        v, j = self._one_hot(data)
        b = random.choices([-1, 1], k=1, weights=[self.prob, 1 - self.prob])
        h_j = self.hash_funcs[j]
        w = self.had[:, h_j(data)]
        l = random.randint(0, self.m-1)
        return b[0] * w[l], j, l  # Return (b*w_l, index j, index l)

    def privatise(self, data):
        """
        Privatises data item using CMS/HCMS

        Args:
            data: item to be privatised

        Returns: Privatised data

        """
        data = str(data)
        if self.is_hadamard:
            return self._hcms_perturb(data)
        else:
            return self._cms_perturb(data)