from __future__ import division
from scitbx.array_family import flex
from scitbx import sparse
from scitbx.lstbx import normal_eqns, normal_eqns_solving
from libtbx.test_utils import approx_equal, Exception_expected

from scitbx.lstbx.tests import test_problems

def exercise_linear_normal_equations():
  py_eqs = [ ( 1, (-1,  0,  0),  1),
             ( 2, ( 2, -1,  0),  3),
             (-1, ( 0,  2,  1),  2),
             (-2, ( 0,  1,  0), -2),
             ]

  eqs_0 = normal_eqns.linear_ls(3)
  for b, a, w in py_eqs:
    eqs_0.add_equation(right_hand_side=b,
                       design_matrix_row=flex.double(a),
                       weight=w)

  eqs_1 = normal_eqns.linear_ls(3)
  b = flex.double()
  w = flex.double()
  a = sparse.matrix(len(py_eqs), 3)
  for i, (b_, a_, w_) in enumerate(py_eqs):
    b.append(b_)
    w.append(w_)
    for j in xrange(3):
      if a_[j]: a[i, j] = a_[j]
  eqs_1.add_equations(right_hand_side=b, design_matrix=a, weights=w)

  assert approx_equal(
    eqs_0.normal_matrix_packed_u(), eqs_1.normal_matrix_packed_u(), eps=1e-15)
  assert approx_equal(
    eqs_0.right_hand_side(), eqs_1.right_hand_side(), eps=1e-15)
  assert approx_equal(
    list(eqs_0.normal_matrix_packed_u()), [ 13, -6, 0, 9, 4, 2 ], eps=1e-15)
  assert approx_equal(
    list(eqs_0.right_hand_side()), [ 11, -6, -2 ], eps=1e-15)



def exercise_non_linear_ls_with_separable_scale_factor():
  test = test_problems.polynomial_fit(normalised=False)
  test.build_up()

  assert test.n_equations == test.n_data;

  # Reference values computed in tst_normal_equations.nb
  eps = 5e-14
  assert approx_equal(test.optimal_scale_factor(), 0.6148971786833856, eps)
  assert approx_equal(test.objective(), 0.039642707534326034, eps)
  assert approx_equal(test.chi_sq(), 0.011326487866950296, eps)


  assert not test.step_equations().solved
  try:
    test.step_equations().cholesky_factor_packed_u()
    raise Exception_expected
  except RuntimeError:
    pass
  try:
    test.step_equations().solution()
    raise Exception_expected
  except RuntimeError:
    pass

  assert approx_equal(
    list(test.step_equations().normal_matrix_packed_u()),
    [ 0.371944193675858, 0.39066546997866547  , 0.10797294655500618,
                           0.41859250354804045, 0.08077629438075473,
                                                0.19767268057900367 ],
    eps)
  assert approx_equal(
    list(test.step_equations().right_hand_side()),
    [ 0.12149917297914861, 0.13803759252793774, -0.025190641142579157 ],
    eps)

  test.step_equations().solve()

  assert test.step_equations().solved
  try:
    test.step_equations().normal_matrix_packed_u()
    raise Exception_expected
  except RuntimeError:
    pass
  try:
    test.step_equations().right_hand_side()
    raise Exception_expected
  except RuntimeError:
    pass

  assert approx_equal(
    list(test.step_equations().cholesky_factor_packed_u()),
    [ 0.6098722765266986, 0.6405693208478925   ,  0.1770418999366983 ,
                            0.09090351333425013, -0.3589664912436558 ,
                                                  0.19357661121640218 ],
    eps)
  assert approx_equal(
    list(test.step_equations().solution()),
    [ 1.2878697604109028, -0.7727798877778043, -0.5151113342942297 ],
    eps=1e-12)

  test_bis = test_problems.polynomial_fit(normalised=True)
  test_bis.build_up()
  assert approx_equal(test_bis.objective(),
                      test.objective()/test.sum_w_yo_sq(),
                      eps=1e-15)
  assert approx_equal(test_bis.chi_sq(), test.chi_sq(), eps=1e-15)


def exercise_non_linear_ls_with_separable_scale_factor_plus_penalty():
  test = test_problems.polynomial_fit_with_penalty(normalised=False)
  test.build_up()
  assert test.n_equations == test.n_data + 1

  eps = 5e-14
  # reference values from tst_normal_equations.nb again

  assert approx_equal(test.optimal_scale_factor(), 0.6148971786833856, eps)
  redu = test.reduced_problem()
  assert test.objective() == redu.objective()
  assert test.step_equations().right_hand_side()\
         .all_eq(redu.step_equations().right_hand_side())
  assert test.step_equations().normal_matrix_packed_u()\
         .all_eq(redu.step_equations().normal_matrix_packed_u())

  assert approx_equal(test.objective(), 1.3196427075343262, eps)
  assert approx_equal(test.chi_sq(), 0.32991067688358156, eps)
  assert approx_equal(
    test.step_equations().right_hand_side(),
    (1.7214991729791487, -1.4619624074720623, 1.5748093588574208),
    eps)
  assert approx_equal(
    test.step_equations().normal_matrix_packed_u(),
    (1.371944193675858, -0.6093345300213344,  1.107972946555006,
                         1.4185925035480405, -0.9192237056192452,
                                              1.1976726805790037),
    eps)

  test_bis = test_problems.polynomial_fit_with_penalty(normalised=True)
  test_bis.build_up()
  assert approx_equal(test_bis.chi_sq(), test.chi_sq(), eps=1e-15)

  n_equations = test.n_equations
  test.build_up()
  assert test.n_equations == n_equations

def exercise_levenberg_marquardt(non_linear_ls, plot=False):
  non_linear_ls.restart()
  iterations = normal_eqns_solving.levenberg_marquardt_iterations(
    non_linear_ls,
    track_all=True,
    gradient_threshold=1e-8,
    step_threshold=1e-8,
    tau=1e-4,
    n_max_iterations=200)
  assert non_linear_ls.n_equations == non_linear_ls.n_data
  assert approx_equal(non_linear_ls.x, non_linear_ls.arg_min, eps=5e-4)
  print "L-M: %i iterations" % iterations.n_iterations
  if plot:
    f = open('plot.nb', 'w')
    print >>f, "g=%s;" % iterations.gradient_norm_history.mathematica_form()
    print >>f, "\[Mu]=%s;" % iterations.mu_history.mathematica_form()
    print >>f, "ListLogPlot[{g,\[Mu]},Joined->True]"
    f.close()

def run():
  import sys
  plot = '--plot' in sys.argv[1:]
  t = test_problems.exponential_fit()
  exercise_levenberg_marquardt(t, plot)
  exercise_linear_normal_equations()
  exercise_non_linear_ls_with_separable_scale_factor()
  exercise_non_linear_ls_with_separable_scale_factor_plus_penalty()
  print 'OK'

if __name__ == '__main__':
  run()