from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import line_search
from matplotlib import rcParams

rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16


def fun(x, y):
    return 5*x**2 - 6*x*y + 5*y**2


def f(x):
    x1 = x[0]
    x2 = x[1]
    return fun(x1, x2)


def g(x):
    x1 = x[0]
    x2 = x[1]
    g1 = 10*x1 - 6*x2
    g2 = -6*x1 + 10*x2
    return [g1, g2]


X, Y = np.mgrid[-1.5:1.5:101j, -1.5:1.5:101j]
Z = fun(X, Y)
levels = np.linspace(0.5, 5, 6)
plt.figure(figsize=(10, 10))
CS = plt.contour(X, Y, Z, levels=levels, colors='k')
plt.clabel(CS)
xk = np.array([-0.5, -1])
pk = np.array([0., 1.])
for k in range(10):
    res = line_search(f, g, xk, pk)
    alpha = res[0]
    xnew = xk + 0.6*alpha*pk
    pk = pk[::-1]
    plt.arrow(xk[0], xk[1], xnew[0] - xk[0], xnew[1] - xk[1],
              fc='r', ec='r',
              length_includes_head=True)
    xk = xnew

plt.axis('image')
plt.xlabel(r'$x$', size=18)
plt.ylabel(r'$y$', size=18)
plt.title(r'$f(x,y) = 5x^2 - 6xy + 5y^2$')
plt.savefig('coordinate descent.svg')
plt.show()