from __future__ import division
from scitbx.array_family import flex
from libtbx.test_utils import approx_equal

def weighted_correlation(w, x, y, derivatives_wrt_y_depth=0):
  "http://en.wikipedia.org/wiki/Pearson_product-moment_correlation_coefficient#Calculating_a_weighted_correlation"
  assert derivatives_wrt_y_depth in [0,1,2]
  sum_w = flex.sum(w)
  assert sum_w != 0
  wxm = flex.sum(w * x) / sum_w
  wym = flex.sum(w * y) / sum_w
  xc = x - wxm
  yc = y - wym
  sum_wxy = flex.sum(w * xc * yc)
  sum_wxx = flex.sum(w * xc * xc)
  sum_wyy = flex.sum(w * yc * yc)
  cc_den_sq = sum_wxx * sum_wyy
  cc_den = cc_den_sq**0.5
  cc = sum_wxy / cc_den
  if (derivatives_wrt_y_depth == 0):
    return cc
  #
  # 1st derivatives w.r.t. y
  d_sum_wxy = w * xc
  d_sum_wyy = 2 * w * yc
  d_cc_den_sq = sum_wxx * d_sum_wyy
  d_cc_den = 1/2 * d_cc_den_sq / cc_den
  d_cc_num = d_sum_wxy * cc_den - sum_wxy * d_cc_den
  d_cc = d_cc_num / cc_den_sq
  if (derivatives_wrt_y_depth == 1):
    return cc, d_cc
  #
  # 2nd derivatives w.r.t. y
  d_yc = 1 - w / sum_w
  d2_sum_wyy = 2 * w * d_yc
  d2_cc_den_sq = sum_wxx * d2_sum_wyy
  d2_cc_den = 1/2 * (d2_cc_den_sq / cc_den - d_cc_den_sq * d_cc_den / cc_den_sq)
  d2_cc_num = -sum_wxy * d2_cc_den
  d2_cc = d2_cc_num / cc_den_sq - d_cc_num * d_cc_den_sq / cc_den_sq**2
  return cc, d_cc, d2_cc

def finite_difference_derivatives(w, x, y, depth, eps=1e-6):
  assert depth in [1,2]
  result = flex.double()
  for i in xrange(len(y)):
    fs = []
    y_orig = y[i]
    for signed_eps in [eps, -eps]:
      y[i] = y_orig + signed_eps
      if (depth == 1):
        fs.append(weighted_correlation(w, x, y))
      else:
        _, d_cc = weighted_correlation(w, x, y, derivatives_wrt_y_depth=1)
        fs.append(d_cc[i])
    y[i] = y_orig
    result.append((fs[0]-fs[1])/(2*eps))
  return result

def exercise():
  mt = flex.mersenne_twister(seed=0)
  sz = 12
  for i_trial in xrange(10):
    x = mt.random_double(size=sz)*5-1
    y = mt.random_double(size=sz)*3-1
    for i_w,w in enumerate([flex.double(sz, 1), mt.random_double(size=sz)*7]):
      cc, d_cc, d2_cc = weighted_correlation(
        w, x, y, derivatives_wrt_y_depth=2)
      if (i_w == 0):
        cc_w1 = flex.linear_correlation(x, y).coefficient()
        assert approx_equal(cc, cc_w1)
      d_cc_fd = finite_difference_derivatives(w, x, y, depth=1)
      assert approx_equal(d_cc, d_cc_fd)
      d2_cc_fd = finite_difference_derivatives(w, x, y, depth=2)
      assert approx_equal(d2_cc, d2_cc_fd)

def run(args):
  assert len(args) == 0
  exercise()
  print "OK"

if (__name__ == "__main__"):
  import sys
  run(args=sys.argv[1:])