#!/usr/bin/env python
#
# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
# Copyright (c) 1997-2015 California Institute of Technology.
# License: 3-clause BSD.  The full license text is available at:
#  - http://trac.mystic.cacr.caltech.edu/project/mystic/browser/mystic/LICENSE
"""
Methods to support discrete measures
"""
# what about sample_unsafe ?
from mystic.math.stats import *
from mystic.math.samples import *
from mystic.math.integrate import *
from mystic.symbolic import generate_solvers, generate_constraint, solve
from mystic.symbolic import generate_conditions, generate_penalty
from mystic.math import almostEqual

def weighted_select(samples, weights, mass=1.0):
  """randomly select a sample from weighted set of samples

Inputs:
    samples -- a list of sample points
    weights -- a list of sample weights
    mass -- sum of normalized weights
"""
  from numpy import sum, array
  from mystic.tools import random_state
  rand = random_state().random
  # generate a list representing the weighted distribution
  wts = normalize(weights, mass)
  wts = array([sum(wts[:i+1]) for i in range(len(wts))])
  # correct for any rounding error
  wts[-1] = mass
  # generate a random weight
  w = mass * rand()
  # select samples that corresponds to randomly selected weight
  selected = len(wts[ wts <= w ])
  return samples[selected]

##### calculate methods #####
def spread(samples):
  """calculate the range of a list of points   [range(x) = min(x) - min(x)]

Inputs:
    samples -- a list of sample points
"""
  return max(samples) - min(samples)

def norm(weights):
  """calculate the norm of a list of points   [norm(x) = mean(x)]

Inputs:
    weights -- a list of sample weights
"""
  return mean(weights)

def maximum(f, samples):
  """calculate the max of function for the given list of points

Inputs: 
    f -- a function that takes a list and returns a number
    samples -- a list of sample points
"""
  y = [f(x) for x in samples]
  return max(y)

def ess_maximum(f, samples, weights=None, tol=0.):
  """calculate the max of function for support on the given list of points

Inputs: 
    f -- a function that takes a list and returns a number
    samples -- a list of sample points
    weights -- a list of sample weights
    tol -- weight tolerance, where any weight <= tol is considered zero
"""
  if weights is None:
    return maximum(f, samples)
  return maximum(f, support(samples, weights, tol))

def minimum(f, samples):
  """calculate the min of function for the given list of points

Inputs: 
    f -- a function that takes a list and returns a number
    samples -- a list of sample points
"""
  y = [f(x) for x in samples]
  return min(y)

def ess_minimum(f, samples, weights=None, tol=0.):
  """calculate the min of function for support on the given list of points

Inputs: 
    f -- a function that takes a list and returns a number
    samples -- a list of sample points
    weights -- a list of sample weights
    tol -- weight tolerance, where any weight <= tol is considered zero
"""
  if weights is None:
    return minimum(f, samples)
  return minimum(f, support(samples, weights, tol))

def expectation(f, samples, weights=None, tol=0.0):
  """calculate the (weighted) expectation of a function for a list of points

Inputs:
    f -- a function that takes a list and returns a number
    samples -- a list of sample points
    weights -- a list of sample weights
    tol -- weight tolerance, where any weight <= tol is ignored
"""
  if weights is None:
    y = [f(x) for x in samples]
    return mean(y, weights)
  # contributed by TJS #
  # to prevent function evaluation if weight is "too small":
  # skip evaluation of f(x) if the corresponding weight <= tol
  yw = [(f(x),w) for (x,w) in zip(samples, weights) if abs(w) > tol]
  return mean(*zip(*yw))
  ##XXX: at around len(samples) == 150, the following is faster
  #aw = asarray(weights)
  #ax = asarray(samples)
  #w = aw > tol
  #return mean(ax[w], aw[w])

def mean(samples, weights=None):
  """calculate the (weighted) mean for a list of points

Inputs:
    samples -- a list of sample points
    weights -- a list of sample weights
"""
  if weights is None:
    weights = [1.0/float(len(samples))] * len(samples)
  # get weighted sum
  ssum = sum(i*j for i,j in zip(samples, weights))
  # normalize by sum of the weights
  wts = float(sum(weights))
  if wts: return ssum / wts
  from numpy import inf
  return ssum * inf  # protect against ZeroDivision

def support_index(weights, tol=0):
  """get the indicies of the positions which have non-zero weight

Inputs:
    weights -- a list of sample weights
    tol -- weight tolerance, where any weight <= tol is considered zero
"""
  return [i for (i,w) in enumerate(weights) if w > tol]

def support(samples, weights, tol=0):
  """get the positions which have non-zero weight

Inputs:
    samples -- a list of sample points
    weights -- a list of sample weights
    tol -- weight tolerance, where any weight <= tol is considered zero
"""
  return [samples[i] for (i,w) in enumerate(weights) if w > tol]

def variance(samples, weights=None): #, _mean=None):
  """calculate the (weighted) variance for a list of points

Inputs:
    samples -- a list of sample points
    weights -- a list of sample weights
"""
  if weights is None:
    weights = [1.0/float(len(samples))] * len(samples)
 #if _mean is None:
  _mean = mean(samples, weights)
  svar = [abs(s - _mean)**2 for s in samples]
  return mean(svar, weights)

def std(samples, weights=None): #, _mean=None):
  """calculate the (weighted) standard deviation for a list of points

Inputs:
    samples -- a list of sample points
    weights -- a list of sample weights
"""
  from numpy import sqrt
  return sqrt(variance(samples, weights)) # _mean)


##### coordinate shift methods #####
from numpy import asarray
def impose_mean(m, samples, weights=None):
  """impose a mean on a list of (weighted) points
  (this function is 'range-preserving' and 'variance-preserving')

Inputs:
    m -- the target mean
    samples -- a list of sample points
    weights -- a list of sample weights
"""
 #XXX: this is as expected... mean(impose_mean(2.0, samples, weights), weights)
 #XXX: this is unexpected?... mean(impose_mean(2.0, samples, weights))
  samples = asarray(list(samples)) #XXX: faster to use x = array(x, copy=True) ?
  shift = m - mean(samples, weights)
  samples = samples + shift  #NOTE: is "range-preserving"
  return list(samples)


def impose_variance(v, samples, weights=None):
  """impose a variance on a list of (weighted) points
  (this function is 'mean-preserving')

Inputs:
    v -- the target variance
    samples -- a list of sample points
    weights -- a list of sample weights
"""
  m = mean(samples, weights)
  samples = asarray(list(samples)) #XXX: faster to use x = array(x, copy=True) ?
  sv = variance(samples,weights) #,m)
  if not sv:  # protect against ZeroDivision when variance = 0
    from numpy import nan
    return [nan]*len(samples) #XXX: better to space pts evenly across range?
  from numpy import sqrt
  scale = sqrt(float(v) / sv)
  samples = samples * scale  #NOTE: not "mean-preserving", until the next line
  return impose_mean(m, samples, weights) #NOTE: not range preserving

#FIXME: for range and variance to be 'mutually preserving'...
#       must reconcile scaling by sqrt(v2/v1) & (r2/r1)
#       ...so likely, must scale the weights... or scale each point differently

def impose_std(s, samples, weights=None):
  """impose a standard deviation on a list of (weighted) points
  (this function is 'mean-preserving')

Inputs:
    s -- the target standard deviation
    samples -- a list of sample points
    weights -- a list of sample weights
"""
  return impose_variance(s**2, samples, weights)

def impose_spread(r, samples, weights=None): #FIXME: fails if len(samples) = 1
  """impose a range on a list of (weighted) points