import numpy as np
import multiprocessing
import itertools
import functools

from utils.count_table import NoisedCountTable, calc_aggregation_error

"""
HDPView.
"""

def HDPView(table, epsilon, prng, ratio=0.9, alpha=1.6, beta=1.2, gamma=0.9, is_classification=False, label_attr=""):
    return run(table, epsilon, prng, ratio, alpha, beta, gamma, is_classification, label_attr)

def run(table, epsilon, prng, ratio, alpha, beta, gamma, is_classification, label_attr):
    """Run 
        Prepare parameters and execute HDPView.
        1st phase, divide blocks.
        2nd phase, perturbation.

    Args:
        table (CountTable): count table
        epsilon (float): privacy budget
        ratio (float): ubdget ratio of block division and perturbation, 0 to 1 value
        mond_ratio (float): budget ratio of convergence and cutwhere, 0 to 1 value
        prng (np.random.RandomState): random state

    Returns:
        [{
            "range": {int: (int,int)}, 
            "depth": int,
            "perturbation_budget": float,
            "mean": float,
            "aggregation_error": float,
            "noise": float
        }]: 
            List of the noised count values for each divided blocks.
    """
    seed = prng.randint(0, 294967295)
    table.set_random(seed)
    epsilon_r = epsilon * ratio
    epsilon_p = epsilon *(1 - ratio)
    
    N = table.domain_size()
    kappa = np.ceil(np.log2(N)*beta)
    epsilon_cut = (1-gamma)*epsilon_r / kappa
    
    theta = 1/epsilon_p

    scale = ((2*alpha - 1)/(alpha - 1) + 1) * (2/(gamma*epsilon_r))
    delta = scale*np.log(alpha)
    
    # prepare shared memories
    manager = multiprocessing.Manager()
    block_queue = manager.Queue()
    queue_buffer = manager.Queue()
    range_and_budget_list = []

    """
        When using HDPView in classification task, 
        HDPView pre-cut count tensors by all 'label' attributes.
    """
    if is_classification:
        # Pre-cutting by all label attributes
        label_num = table.cardinality_dict[label_attr]
        will_splited = table
        for i in range(label_num-1):
            left_block, right_block = will_splited.split(label_attr, i)
            queue_buffer.put(left_block)
            will_splited = right_block
        queue_buffer.put(will_splited)
    else:
        queue_buffer.put(table)
    
    # prepare thread pool
    MAX_PROCESS = multiprocessing.cpu_count()-1
    pool = multiprocessing.Pool(MAX_PROCESS)
    
    while not queue_buffer.empty():
        queue_move(queue_buffer, block_queue)
        async_results = []
        while not block_queue.empty():
            result = pool.apply_async(
                recursive_bisection, (block_queue.get(), queue_buffer, epsilon_cut, kappa, theta, scale, delta)
            )
            async_results.append(result)
        range_and_budget_list.extend(list(itertools.chain.from_iterable([i.get() for i in async_results]))) # flatten and extend
    
    range_and_budget_list.sort(key=functools.cmp_to_key(range__gt__))
    for d in range_and_budget_list:
        budget = epsilon_p
        d["perturbation_budget"] = budget
        d["mean"], d["aggregation_error"] = calculate_mean_and_aggregation_error(table, d["range"].items())
        d["noise"] = prng.laplace(0.0,  1.0 / budget)
        
    return NoisedCountTable.from_count_table(table, range_and_budget_list), range_and_budget_list

def recursive_bisection(block, block_queue, epsilon_cut, kappa, theta, scale, delta):
    """Recursive bisection by random cut and random converge

    Args:
        block (CountTable): sub-tensor of original count tensor
        block_queue (multiprocessing.Queue): Shared queue to store blocks to be executed

    Returns:
        [{"range": {int: (int,int)}, "depth": int}]
    """

    # Random cut
    if block.depth > kappa:
        axis, index = random_calculate(block)
    else:
        axis, index = exp_calculate_sum(block, epsilon_cut)

    left_block, right_block = block.split(axis, index)

    # Random converge
    convergence_list = []
    if left_block.domain_size() == 1:
        convergence_list.append(blockResult(left_block))
    elif decide_convergence(left_block, theta, scale, delta):
        convergence_list.append(blockResult(left_block))
    else:
        block_queue.put(left_block)

    if right_block.domain_size() == 1:
        convergence_list.append(blockResult(right_block))
    elif decide_convergence(right_block, theta, scale, delta):
        convergence_list.append(blockResult(right_block))
    else:
        block_queue.put(right_block)

    return convergence_list


def blockResult(block):
    """
    Returns:
        {"range": {int: (int,int)}, "depth": int}:
            {"range": {0: (0,3), 1: (2,4)}, "depth": 1}
    """
    return {"range": block.domain_dict, "depth": block.depth}


def queue_move(ori_queue, queue):
    """Move data from ori_queue to queue.
        ori_queue becomes empty.

    Args:
        ori_queue, queue (multiprocessing.Queue): created by multiprocessing.Manager()
    """
    for _ in range(ori_queue.qsize()):
        queue.put(ori_queue.get())


def exp_calculate_sum(block, epsilon_cut):
    scores = []
    targets = []
    for d in block.index:
        cardinality = block.cardinality_dict[d]
        for i in np.arange(cardinality-1):
            scores.append(- np.sum(block.split_values(d, i)))
            targets.append((d, i))

    return exponential_mechanism(
        block.prng,
        epsilon_cut, 
        scores,
        targets,
        2*(2 - 2/block.domain_size())
    )

def exponential_mechanism(prng, eps , scores, targets, sensitivity):
    if sensitivity == 0:
        index = prng.choice(len(targets))
        return targets[index]

    np_scores = np.array(scores)
    score_max = np.max(np_scores)
    weights = np.exp((eps*(np_scores - score_max)) / (2*sensitivity))
    total_weight = np.sum(weights)
    cum_weights = np.cumsum(weights) / total_weight
    index = prng.rand()
    return targets[cum_weights.searchsorted(index)]

def random_calculate(block):
    dim_list = list(block.cardinality_dict.keys())
    while True:
        d = block.prng.choice(dim_list)
        cardinality = block.cardinality_dict[d]
        if cardinality > 1:
            break
    i = block.prng.choice(cardinality-1)
    return (d, i)

def decide_convergence(block, theta, scale, delta):
    ae = calc_aggregation_error(block.values(), block.zero_num())
    b = ae - (delta*block.depth)
    b = max(b, (theta + 2 - delta))
    noise = block.prng.laplace(0.0, scale=scale)
    noisy_b = b + noise
    return noisy_b <= theta

def calculate_mean_and_aggregation_error(table, ranges):
    """Calculate mean and aggregation error for given blocks as table and ranges.

    Args:
        table (CountTable): count table
        ranges ([{int, (int, int)}]): [{dimension, (start, end)}]
    """
    values, zero_num = table.get_sub_block_values(ranges)
    assert len(values) + zero_num > 0, "Error: range is invalid"
    mean = values.sum() / (len(values) + zero_num)
    ae = calc_aggregation_error(values, zero_num)
    return mean, ae

def range__gt__(d, other):
    for dim, dom_range in d["range"].items():
        other_range = other["range"][dim]
        if dom_range > other_range:
            return 1
        elif dom_range < other_range:
            return -1
    return 0