"""
Solution to the exercise 3b in the set 6 (dogleg path)
"""

from sympy import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches


init_printing()

Delta = symbols('Delta',positive=True)
x1,x2,tau = symbols('x_1 x_2 tau')

f = Rational(1,2)*x1**2 + x2**2
x0= Matrix([1,1])

# compute the gradient and evaluate it at x0
g  = Matrix([diff(f,x1),diff(f,x2)])
g0 = g.subs([(x1,x0[0]),(x2,x0[1])])

# compute the Hessian and evaluate it at x0
B  = Matrix([diff(g,x1).T,diff(g,x2).T])
B0 = B.subs([(x1,x0[0]),(x2,x0[1])])

# find the unconstrained minimum
p_B = B0.LUsolve(-g0)
Delta_B2 = p_B.dot(p_B)

# find the steepest-descent unit step
p_U = -g0.dot(g0)/g0.dot(B0*g0) * g0
Delta_U2 = p_U.dot(p_U)

p_leg0 = Delta/sqrt(Delta_U2)*p_U
p_leg2 = p_B
print('If Delta >= %s then solution is:' % pretty(sqrt(Delta_B2)))
pprint(p_leg2.T)
print('If Delta <= %s then solution is:' % (sqrt(Delta_U2)))
pprint(p_leg0.T)

# in between the two, we need to solve the dogleg problem
p_tau = p_U + tau*(p_B-p_U)
roots = solve(p_tau.dot(p_tau)-Delta**2,tau)
# select the positive root
Delta_mid = (sqrt(Delta_B2)+sqrt(Delta_U2))/2
if roots[0].subs(Delta,Delta_mid).evalf()>=0:
    tau_dogleg = roots[0]
else:
    assert(roots[1].subs(Delta,Delta_mid).evalf()>=0)
    tau_dogleg = roots[1]

p_leg1 = p_tau.subs(tau,tau_dogleg)
print('For Delta in between, the solution is:')
pprint(p_leg1)


# Plot the dogleg trajectory
Delta_plot = np.linspace(0,2,1000)
p_star_plot= np.zeros((2,len(Delta_plot)))

Delta_B = sqrt(Delta_B2).evalf()
Delta_U = sqrt(Delta_U2).evalf()

leg0 = lambdify(Delta,p_leg0)
leg1 = lambdify(Delta,p_leg1)
leg2 = lambdify(Delta,p_leg2)
for i in range(len(Delta_plot)):
    if Delta_plot[i] <= Delta_U:
        p_star_plot[:,i] = leg0(Delta_plot[i]).flatten()
    elif Delta_plot[i] >= Delta_B:
        p_star_plot[:,i] = leg2(Delta_plot[i]).flatten()
    else:
        p_star_plot[:,i] = leg1(Delta_plot[i]).flatten()

plt.plot(p_star_plot[0,:],p_star_plot[1,:])
plt.title('Dogleg trajectory for $0 \leq \Delta \leq 2$')
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')


# For comparison, we also plot the exact solution to the trust region problem
# it is given by p = -[1/(1+lambda),2/(2+lambda)] where lambda is found
# from the equation |p|^2 = Delta^2
from scipy.optimize import brentq

lam = symbols('lambda',positive=True)
p_lam = (B0+lam*eye(len(g0))).LUsolve(-g0)
p_lam2= lambdify(lam,p_lam.dot(p_lam))
p_ex  = lambdify(lam,p_lam)

p_exact_plot= np.zeros((2,len(Delta_plot)))
for i in range(len(Delta_plot)):
    if Delta_plot[i] == 0.0:
        # this special case is only needed so that Brent's method
        # does not complain for this trivial case
        p_exact_plot[:,i] = 0.0
    elif Delta_plot[i] >= Delta_B:
        # if unconstrained minimizer feasible - return it
        p_exact_plot[:,i] = leg2(Delta_plot[i]).flatten()
    else:
        # find the lambda such that the point is on the boundary of the
        # trust region
        # we solve the problem numerically, as analytical solution
        # is too slow.  here one can use e.g. bisection, but we use
        # Brent's method, which is faster
        lam_exact = brentq(lambda lam: p_lam2(lam)-Delta_plot[i]**2,0,1.0E+10)
        p_exact_plot[:,i] = p_ex(lam_exact).flatten()

plt.plot(p_exact_plot[0,:],p_exact_plot[1,:])


# plot the TR with delta=5/6
plt.gca().add_patch( patches.Circle([0,0],5/6., linestyle='dashed', \
                                    fill=False, edgecolor='magenta') )


# plot the level sets of the model
g00 = np.array(g0).flatten()
B00 = np.array(B0)
model = lambda xy: float(g00.dot(xy) + 0.5*xy.dot(B00.dot(xy)))
dxy = 0.05
x = np.arange(-1.2, 0.1, dxy)
y = np.arange(-1.2, 0.1, dxy)
X, Y = np.meshgrid(x, y)
it = np.nditer([X, Y, None])
for x,y,z in it:
    z[...] = model(np.array([x,y]))
Z = it.operands[2]
# compute the levels: model values corresponding to different delta
levels = []
for delta in np.arange(1,9)/6.:
    lam_exact = brentq(lambda lam: p_lam2(lam)-delta**2,0,1.0E+10)
    p_exact = p_ex(lam_exact).flatten()
    levels.append(model(p_exact))
levels.extend([0.0,model(np.array([-1.,-1]))])
#
#CS = plt.contour(X, Y, Z,levels = sorted(levels), linestyles='dotted')
CS = plt.contourf(X, Y, Z,levels = sorted(levels),cmap='RdYlBu')
#plt.clabel(CS, inline=1, fontsize=10)

plt.gca().set_xlim([-1.2,0.1])
plt.gca().set_ylim([-1.2,0.1])


plt.legend(['Dogleg','Exact','TR=5/6'])
plt.show()