import itertools
import numpy as np
from utils.count_table import NoisedCountTable
from utils.query import Query, QueryCondition
import matplotlib.pyplot as plt

## PrivTree
def PrivTree(table, prng, epsilon, fanout, theta=0, workloads=[]):
    range_and_budget_list = privtree(table, prng, epsilon, fanout, theta)
    return NoisedCountTable.from_count_table(table, range_and_budget_list), range_and_budget_list

def privtree(table, prng, epsilon, fanout, theta):
    scale = (2*fanout - 1)/(fanout - 1)*(1/(epsilon/2))
    print("scale: ", scale)
    delta = scale*np.log(fanout)
    print("delta: ", delta)
    domain_names = list(table.domain_dict.keys())
    range_and_budget_list = []

    unvisited_domains = []
    unvisited_domains_per_depth = []
    tree_depth = 0
    unvisited_domains.append(list(table.domain_dict.values()))

    while len(unvisited_domains) > 0:
        unvisited_domains_per_depth.extend(unvisited_domains)
        unvisited_domains = []
        for unvisited_domain in unvisited_domains_per_depth:
            count_query = Query([QueryCondition(domain_name, domain_range[0], domain_range[1]) for domain_name, domain_range in zip(domain_names, unvisited_domain)])
            count = table.run_query(count_query)
            b = count - (delta*tree_depth)
            b = max(b, (theta - delta))
            noise = prng.laplace(0.0, scale=scale)
            noisy_b = b + noise
            if (noisy_b > theta) and can_divide(unvisited_domain):
                unvisited_domains.extend(split_domains(unvisited_domain))
            else:
                range_and_budget_list.append({ 
                    "range": domain_dict_from_domain_list(domain_names, unvisited_domain),
                    "noise": prng.laplace(0.0, scale=1/(epsilon/2)),
                    "mean": count / count_domain_dict(unvisited_domain)
                })
        unvisited_domains_per_depth = []
        tree_depth += 1
    
    return range_and_budget_list

def can_divide(domain):
    return count_domain_dict(domain) > 1

def split_domains(unvisited_domain):
    half_divided_domains = []
    for domain_range in unvisited_domain:
        if domain_range[0] == domain_range[1]:
            half_divided_domains.append([(domain_range[0], domain_range[1])])
        else:
            half_divided_domains.append([(domain_range[0], (domain_range[0]+domain_range[1])//2), ((domain_range[0]+domain_range[1])//2+1, domain_range[1])])
    return list(itertools.product(*half_divided_domains))

def domain_dict_from_domain_list(domain_names, domains):
    return { domain_name: domain for domain_name, domain in zip(domain_names, domains)}

def count_domain_dict(domain):
    """All counts from domain dict
    """
    cardinality_list = []
    for val in domain:
        cardinality_list.append(val[1] - val[0] + 1)
    return np.prod(cardinality_list, dtype=np.float)


def plot_data(data, attrs, noised_data):
    assert len(attrs) == 2, "attrs must be 2 length"
    
    x = data.df[list(attrs)].iloc[:, 0].values
    y = data.df[list(attrs)].iloc[:, 1].values
    fig, ax = plot_locations_xy(x, y)
    print("number of rectangle", len(noised_data.blocks))
    for block in noised_data.blocks:
        domain = list(block.ranges.values())
        plot_rect(ax, domain[0][0]-0.5, domain[0][1]+0.5, domain[1][0]-0.5, domain[1][1]+0.5)

    v0 = [noised_data.domain_dict[attrs[0]][0], noised_data.domain_dict[attrs[0]][1], noised_data.domain_dict[attrs[1]][0], noised_data.domain_dict[attrs[1]][1]]
    print('v0', v0)
    v0_w = v0[1] - v0[0]
    v0_h = v0[3] - v0[2]

    ax.set_xlim(v0[0]-0.05*v0_w, v0[1]+0.05*v0_w)
    ax.set_ylim(v0[2]-0.05*v0_h, v0[3]+0.05*v0_h)
    plt.show()
    
    
def plot_locations_xy(x, y, alpha=0.3, marker='o', marker_size=5, figsize=(16,8), save_plot_dir=None, dpi=250):

    fig, ax = plt.subplots(figsize=figsize)
    ax.set_aspect('equal')
    ax.scatter(x, y,alpha=alpha, marker=marker, s=marker_size)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.grid(alpha=0.8,linewidth=0.5)
        
    return fig, ax

def plot_rect(ax, left, right, bottom, top, fill=False, color='r', edgecolor='r', alpha=1, hatch=None):
    p = plt.Rectangle((left, bottom), right-left, top-bottom, fill=fill, color=color, edgecolor=edgecolor, alpha=alpha, hatch=hatch)
    ax.add_patch(p)