import numpy as np

# Utilities
def gen_iid_truncated_normal_utils(n, mean, var):
    pass

def gen_iid_scaled_beta_utils(n=1000, a=1, b=1, low=-1, high=1):
    return np.random.beta(a, b, size=n) * (high-low) + low

def gen_iid_uniform_utils(n, low=-1, high=1):
    return np.random.random(size=n) * (high-low) + low

# Gammas
def gen_uniform_gamma(n):
	# Generate a single uniform distribution
    x = np.random.random(size=n)
    return x / x.sum()

def gen_normed_beta_gamma(n, a=1, b=1):
	# Generate a single beta normed distribution
    x = np.random.beta(a, b, size=n)
    return x / x.sum()

def gen_uniform_gamma_cb(n, ut, s):
	# Generate a single cost bounded uniform distribution
	# Calculate the E[C(A)]
	# while E[C(A)] > S:
	#	D = E[C(A)] - S
	# 	Increase utility of each element of x by D
	#	Renormalize x
	# return x
	x = np.random.random(size=n)
	x = x / x.sum()
	E_cost = np.dot(x, ut)
	print "E_cost: %s, s: %s" % (E_cost, s)
	while E_cost > s:
		diff = E_cost - s
		for i, g in enumerate(x):
			delta = diff / ut[i]
			x[i] += delta
		x = x / x.sum()
	return x

# def gen_normed_beta_gamma_cb(n, ut, a=1, b=1):
#	# Generate a single beta normed distribution
#    x = np.random.beta(a, b, size=n)
#    return x / x.sum()

# Lists
def generate_utils_list(n):
    uts = []
    uts += [gen_iid_scaled_beta_utils(n).tolist()]
    uts += [gen_iid_scaled_beta_utils(n, a=10, b=10).tolist()]
    uts += [gen_iid_scaled_beta_utils(n, a=10).tolist()]
    return uts

def generate_gamma_list(n):
	# Generate list of gamma distributions
    gs = []
    gs += [gen_uniform_gamma(n).tolist()]
    gs += [gen_normed_beta_gamma(n, a=3).tolist()]
    return gs

def generate_gamma_list_ut(n, ut, s):
	# Generate list of gamma distributions
	# number of worlds, utility, cost bound
    gs = []
    gs += [gen_uniform_gamma_cb(n, ut, s).tolist()]
#    gs += [gen_normed_beta_gamma_cb(n, ut, a=3, s).tolist()]
    gs += [gen_normed_beta_gamma(n, a=3).tolist()]
    return gs