def test_solve_constraint():

  constraints = """
  spread([x0,x1]) - 1.0 = mean([x0,x1])   
  mean([x0,x1,x2]) = x2"""

  from mystic.math.measures import mean, spread
  _constraints = solve(constraints)
  solv = generate_solvers(_constraints)
  constraint = generate_constraint(solv)
  x = constraint([1.0, 2.0, 3.0])
  assert all(x) == all([1.0, 5.0, 3.0])
  assert mean(x) == x[2]
  assert spread(x[:-1]) - 1.0 == mean(x[:-1])

def test_solve_constraint():

  # sympy can no longer do "spread([x0,x1])"... so use "x1 - x0"
  constraints = """
  (x1 - x0) - 1.0 = mean([x0,x1])   
  mean([x0,x1,x2]) = x2"""

  from mystic.math.measures import mean
  _constraints = solve(constraints)
  solv = generate_solvers(_constraints)
  constraint = generate_constraint(solv)
  x = constraint([1.0, 2.0, 3.0])
  assert all(x) == all([1.0, 5.0, 3.0])
  assert mean(x) == x[2]
  assert (x[1] - x[0]) - 1.0 == mean(x[:-1])

def test_penalize():

  from mystic.math.measures import mean, spread
  def mean_constraint(x, target):
    return mean(x) - target

  def range_constraint(x, target):
    return spread(x) - target

  @quadratic_equality(condition=range_constraint, kwds={'target':5.0})
  @quadratic_equality(condition=mean_constraint, kwds={'target':5.0})
  def penalty(x):
    return 0.0

  def cost(x):
    return abs(sum(x) - 5.0)

  from mystic.solvers import fmin
  from numpy import array
  x = array([1,2,3,4,5])
  y = fmin(cost, x, penalty=penalty, disp=False)

  assert round(mean(y)) == 5.0
  assert round(spread(y)) == 5.0
  assert round(cost(y)) == 4*(5.0)

def test_constrain():

  from mystic.math.measures import mean, spread
  from mystic.math.measures import impose_mean, impose_spread
  def mean_constraint(x, mean=0.0):
    return impose_mean(mean, x)

  def range_constraint(x, spread=1.0):
    return impose_spread(spread, x)

  @inner(inner=range_constraint, kwds={'spread':5.0})
  @inner(inner=mean_constraint, kwds={'mean':5.0})
  def constraints(x):
    return x

  def cost(x):
    return abs(sum(x) - 5.0)

  from mystic.solvers import fmin_powell
  from numpy import array
  x = array([1,2,3,4,5])
  y = fmin_powell(cost, x, constraints=constraints, disp=False)

  assert mean(y) == 5.0
  assert spread(y) == 5.0
  assert almostEqual(cost(y), 4*(5.0))

 def constraints(x):
     # constrain the last x_i to be the same value as the first x_i
     x[-1] = x[0]
     # constrain x such that mean(x) == target
     if not almostEqual(mean(x), target):
         x = impose_mean(target, x)
     return x

def test_solve_constraint():

  from mystic.math.measures import mean
  @with_mean(1.0)
  def constraint(x):
    x[-1] = x[0]
    return x

  x = solve(constraint, guess=[2,3,1])

  assert almostEqual(mean(x), 1.0, tol=1e-15)
  assert x[-1] == x[0]
  assert issolution(constraint, x)

def test_generate_constraint():

  constraints = """
  spread([x0, x1, x2]) = 10.0
  mean([x0, x1, x2]) = 5.0"""

  from mystic.math.measures import mean, spread
  solv = generate_solvers(constraints)
  assert almostEqual(mean(solv[0]([1,2,3])), 5.0)
  assert almostEqual(spread(solv[1]([1,2,3])), 10.0)

  constraint = generate_constraint(solv)
  assert almostEqual(constraint([1,2,3]), [0.0,5.0,10.0], 1e-10)

def test_with_constraint():

  from mystic.math.measures import mean, impose_mean

  @with_constraint(inner, kwds={'target':5.0})
  def mean_of_squared(x, target):
    return impose_mean(target, [i**2 for i in x])

  from numpy import array
  x = array([1,2,3,4,5])
  y = impose_mean(5, [i**2 for i in x])
  assert mean(y) == 5.0
  assert mean_of_squared(x) == y

def test_simplify():
  constraints = """
  mean([x0, x1, x2]) <= 5.0
  x0 <= x1 + x2"""

  from mystic.math.measures import mean
  _constraints = simplify(constraints)
  solv = generate_solvers(_constraints)
  constraint = generate_constraint(solv)
  x = constraint([1.0, -2.0, -3.0])
  assert all(x) == all([-5.0, -2.0, -3.0])
  assert mean(x) <= 5.0
  assert x[0] <= x[1] + x[2]

def test_with_mean():

  from mystic.math.measures import mean, impose_mean

  @with_mean(5.0)
  def mean_of_squared(x):
    return [i**2 for i in x]

  from numpy import array
  x = array([1,2,3,4,5])
  y = impose_mean(5, [i**2 for i in x])
  assert mean(y) == 5.0
  assert mean_of_squared(x) == y

def test_with_mean_spread():

  from mystic.math.measures import mean, spread, impose_mean, impose_spread

  @with_spread(50.0)
  @with_mean(5.0)
  def constrained_squared(x):
    return [i**2 for i in x]

  from numpy import array
  x = array([1,2,3,4,5])
  y = impose_spread(50.0, impose_mean(5.0,[i**2 for i in x]))
  assert almostEqual(mean(y), 5.0, tol=1e-15)
  assert almostEqual(spread(y), 50.0, tol=1e-15)
  assert constrained_squared(x) == y

def test_as_constraint():

  from mystic.math.measures import mean, spread
  def mean_constraint(x, target):
    return mean(x) - target

  def range_constraint(x, target):
    return spread(x) - target

  @quadratic_equality(condition=range_constraint, kwds={'target':5.0})
  @quadratic_equality(condition=mean_constraint, kwds={'target':5.0})
  def penalty(x):
    return 0.0

  ndim = 3
  constraints = as_constraint(penalty, solver='fmin')
  #XXX: this is expensive to evaluate, as there are nested optimizations

  from numpy import arange
  x = arange(ndim)
  _x = constraints(x)
  
  assert round(mean(_x)) == 5.0
  assert round(spread(_x)) == 5.0
  assert round(penalty(_x)) == 0.0

  def cost(x):
    return abs(sum(x) - 5.0)

  npop = ndim*3
  from mystic.solvers import diffev
  y = diffev(cost, x, npop, constraints=constraints, disp=False, gtol=10)

  assert round(mean(y)) == 5.0
  assert round(spread(y)) == 5.0
  assert round(cost(y)) == 5.0*(ndim-1)

def test_outer_constraint():

  from mystic.math.measures import impose_mean, mean

  def impose_constraints(x, mean, weights=None):
    return impose_mean(mean, x, weights)

  @outer(outer=impose_constraints, kwds={'mean':5.0})
  def mean_of_squared(x):
    return [i**2 for i in x]

  from numpy import array
  x = array([1,2,3,4,5])
  y = impose_mean(5, [i**2 for i in x])
  assert mean(y) == 5.0
  assert mean_of_squared(x) == y

def test_with_penalty():

  from mystic.math.measures import mean, spread
  @with_penalty(quadratic_equality, kwds={'target':5.0})
  def penalty(x, target):
    return mean(x) - target

  def cost(x):
    return abs(sum(x) - 5.0)

  from mystic.solvers import fmin
  from numpy import array
  x = array([1,2,3,4,5])
  y = fmin(cost, x, penalty=penalty, disp=False)

  assert round(mean(y)) == 5.0
  assert round(cost(y)) == 4*(5.0)

 def constrain(rv):
   "constrain:  y >= m  and  sum(wi)_{k} = 1 for each k in K"
   pm = scenario()
   pm.load(rv, pts)      # here rv is param: w,x,y
   #impose: sum(wi)_{k} = 1 for each k in K
   norm = 1.0
   for i in range(len(pm)):
     w = pm[i].weights
     w[-1] = norm - sum(w[:-1])
     pm[i].weights = w
   #impose: y >= m 
   values, weights = pm.values, pm.weights
   y = float(mean(values, weights))
   if not (y >= float(target[0])):
     pm.values = impose_mean(target[0]+target[1], values, weights)
   rv = pm.flatten(all=True) 
   return rv

def test_as_penalty():

  from mystic.math.measures import mean, spread
  @with_spread(5.0)
  @with_mean(5.0)
  def constraint(x):
    return x

  penalty = as_penalty(constraint)

  from numpy import array
  x = array([1,2,3,4,5])
  
  def cost(x):
    return abs(sum(x) - 5.0)

  from mystic.solvers import fmin
  y = fmin(cost, x, penalty=penalty, disp=False)

  assert round(mean(y)) == 5.0
  assert round(spread(y)) == 5.0
  assert round(cost(y)) == 4*(5.0)

 def mean_value(self):  # get mean of y's
   """calculate the mean of the associated values for a scenario"""
   from mystic.math.measures import mean
   return mean(self.values, self.weights)

  from numpy import sum
  ans = sum(lipschitz_distance(L, pm, _data))
  print "original: %s @ %s\n" % (ans, a)
 #print "pm: %s" % pm
 #print "data: %s" % data
  #---
  lb = [0,.5,-100,-100,  0,.5,-100,-100,  0,.5,-100,-100,   0,0,0,0,0,0,0,0]
  ub = [.5,1, 100, 100,  .5,1, 100, 100,  .5,1, 100, 100,   9,9,9,9,9,9,9,9]
  bounds = (lb,ub)

  _constrain = mean_y_norm_wts_constraintsFactory((y_mean,y_buffer), pts)
  results = impose_feasible(feasability, data, guess=pts, tol=deviation, \
                            bounds=bounds, constraints=_constrain)
  from mystic.math.measures import mean
  print "solved: %s" % results.flatten(all=True)
  print "mean(y): %s >= %s" % (mean(results.values, results.weights), y_mean)
  print "sum(wi): %s == 1.0" % [sum(w) for w in results.wts]

  print "\n---------------------------------------------------\n"

  bc = bc[:-2]
  ids = ['1','2','3']
  t = dataset()
  t.load(bc, map(model, bc), ids)
  t.update(t.coords, map(model, t.coords))
# r = dataset()
# r.load(t.coords, t.values)
# L = [0.1, 0.0, 0.0]
  print "%s" % t
  print "L: %s" % L
  print "shortness:"

def test_expect(constrain=False):
  G = marc_surr  #XXX: uses the above-provided test function
  function_name = G.__name__

  _mean = 06.0   #NOTE: SET THE mean HERE!
  _range = 00.5  #NOTE: SET THE range HERE!
  nx = 3  #NOTE: SET THE NUMBER OF 'h' POINTS HERE!
  ny = 3  #NOTE: SET THE NUMBER OF 'a' POINTS HERE!
  nz = 3  #NOTE: SET THE NUMBER OF 'v' POINTS HERE!

  h_lower = [60.0];  a_lower = [0.0];  v_lower = [2.1]
  h_upper = [105.0]; a_upper = [30.0]; v_upper = [2.8]

  lower_bounds = (nx * h_lower) + (ny * a_lower) + (nz * v_lower)
  upper_bounds = (nx * h_upper) + (ny * a_upper) + (nz * v_upper)
  bounds = (lower_bounds,upper_bounds)

  if debug:
    print(" model: f(x) = %s(x)" % function_name)
    print(" mean: %s" % _mean)
    print(" range: %s" % _range)
    print("..............\n")

  if debug:
    param_string = "["
    for i in range(nx):
      param_string += "'x%s', " % str(i+1)
    for i in range(ny):
      param_string += "'y%s', " % str(i+1)
    for i in range(nz):
      param_string += "'z%s', " % str(i+1)
    param_string = param_string[:-2] + "]"

    print(" parameters: %s" % param_string)
    print(" lower bounds: %s" % lower_bounds)
    print(" upper bounds: %s" % upper_bounds)
  # print(" ...")

  wx = [1.0 / float(nx)] * nx
  wy = [1.0 / float(ny)] * ny
  wz = [1.0 / float(nz)] * nz

  from mystic.math.measures import _pack, _unpack
  wts = _pack([wx,wy,wz])
  weights = [i[0]*i[1]*i[2] for i in wts]

  if not constrain:
    constraints = None
  else:  # impose a mean constraint on 'thickness'
    h_mean = (h_upper[0] + h_lower[0]) / 2.0
    h_error = 1.0
    v_mean = (v_upper[0] + v_lower[0]) / 2.0
    v_error = 0.05
    if debug:
      print("impose: mean[x] = %s +/- %s" % (str(h_mean),str(h_error)))
      print("impose: mean[z] = %s +/- %s" % (str(v_mean),str(v_error)))
    def constraints(x, w):
      from mystic.math.discrete import compose, decompose
      c = compose(x,w)
      E = float(c[0].mean)
      if not (E <= float(h_mean+h_error)) or not (float(h_mean-h_error) <= E):
        c[0].mean = h_mean
      E = float(c[2].mean)
      if not (E <= float(v_mean+v_error)) or not (float(v_mean-v_error) <= E):
        c[2].mean = v_mean
      return decompose(c)[0]

  from mystic.math.measures import mean, expectation, impose_expectation
  samples = impose_expectation(_mean, G, (nx,ny,nz), bounds, weights, \
                               tol=_range, constraints=constraints)

  smp = _unpack(samples,(nx,ny,nz))
  if debug:
    from numpy import array
    # rv = [xi]*nx + [yi]*ny + [zi]*nz
    print("\nsolved [x]: %s" % array( smp[0] ))
    print("solved [y]: %s" % array( smp[1] ))
    print("solved [z]: %s" % array( smp[2] ))
    #print("solved: %s" % smp)
  mx = mean(smp[0])
  my = mean(smp[1])
  mz = mean(smp[2])
  if debug:
    print("\nmean[x]: %s" % mx)  # weights are all equal
    print("mean[y]: %s" % my)  # weights are all equal
    print("mean[z]: %s\n" % mz)  # weights are all equal
  if constrain:
    assert almostEqual(mx, h_mean, tol=h_error)
    assert almostEqual(mz, v_mean, tol=v_error)

  Ex = expectation(G, samples, weights)
  cost = (Ex - _mean)**2
  if debug:
    print("expect: %s" % Ex)
    print("cost = (E[G] - m)^2: %s" % cost)
  assert almostEqual(cost, 0.0, 0.01)

 def __mean(self):
   from mystic.math.measures import mean
   return mean(self.positions, self.weights)