import numpy as np
from scipy.stats import norm, cauchy,uniform,laplace
import ray
from DPQuantile import DPQuantile
from tqdm import tqdm
import numpy as np

def generate_data(dist_type, tau , n_samples):
    """Generate a data stream of the specified distribution and return (data, true quantile)"""
    if dist_type == 'normal':
        data = np.random.normal(0, 1, n_samples)
        true_q = norm.ppf(tau)
    elif dist_type == 'uniform':
        data = np.random.uniform(-1, 1, n_samples)
        true_q = -1 + 2 * tau 
    elif dist_type == 'cauchy':
        data = np.random.standard_cauchy(n_samples)
        true_q = cauchy.ppf(tau)
    elif dist_type == 'laplace':
        data = np.random.laplace(0, 1, n_samples)
        true_q = laplace.ppf(tau)
    else:
        raise ValueError("Unsupported distribution type")
    return data, true_q

def generate_true_q(dist_type, tau):
    """Generate the true quantile for the specified distribution"""
    if dist_type == 'normal':
        true_q = norm.ppf(tau)
    elif dist_type == 'uniform':
        true_q = -1 + 2 * tau
    elif dist_type == 'cauchy':
        true_q = cauchy.ppf(tau)
    elif dist_type == 'laplace':
        true_q = laplace.ppf(tau)
    else:
        raise ValueError("Unsupported distribution type")
    return true_q

def get_true_var(dist_type,r,tau):
    """Generate the true quantile, response rate r and density for the specified distribution, and return the theoretical variance"""
    if dist_type == 'normal':
        true_q = norm.ppf(tau)
        f_q = norm.pdf(true_q)
    elif dist_type == 'uniform':
        true_q = -1 + 2 * tau
        f_q = uniform.pdf(true_q, loc=-1, scale=2)
    elif dist_type == 'cauchy':
        true_q = cauchy.ppf(tau)
        f_q = cauchy.pdf(true_q)
    elif dist_type == 'laplace':
        from scipy.stats import laplace
        true_q = laplace.ppf(tau)
        f_q = laplace.pdf(true_q)
    else:
        raise ValueError("Unsupported distribution type")
    return (1-r**2*(1-2*(1-tau))**2)/(4*f_q**2*r**2)   

def distribute_data(data, n_clients):
    """Randomly shuffle the data and distribute it to clients in proportion"""
    data = np.random.permutation(data) 
    return np.split(data, n_clients)

def calculate_rae(estimated_vars, true_var):
    """
    Calculate the relative absolute error between the estimated and true variance

    Parameters:
    estimated_vars (numpy.ndarray): Estimated variance
    true_var (float): True variance

    Returns:
    float: Relative absolute error
    """
    return np.abs(estimated_vars - true_var) / true_var