#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ---------------------------------------------------------------------------------- #
#                                  "Tracking Minuit"
# ---------------------------------------------------------------------------------- #
# 
#  Python macro for showing how Minuit jumps around in the function space.
#  Shows the Rosenbrock function (https://en.wikipedia.org/wiki/Rosenbrock_function) 
#  which is known to be diffucult to minimize. 
#  Also shows how to track Minuits progress in space. 
#
#  Author: Christian Michelsen (NBI)
#  Email:  mnv794@alumni.ku.dk
#  Date:   5st of December 2017
# 
# ---------------------------------------------------------------------------------- #


from __future__ import division, print_function
import numpy as np
import matplotlib.pyplot as plt
from iminuit import Minuit
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter

plt.close('all')


a = 1


# The function we want to miminize, in this case the Rosenbrock function
def f(x, y):   # The rosenbrock function, has minimum at (x, y) = (a, a^2)
    return 0.5 * (a - x)**2 + 100*(y - x**2)**2

# Create a helper object which keeps track of Minuits progress
class TrackMinuit:
    steps_taken = []
    def __init__(self, function):
#        self.data_x = data_x
        self.function = function
        TrackMinuit.steps_taken = []
    def __call__(self, x, y):
        fval = self.function(x, y)
        TrackMinuit.steps_taken.append([x, y])
        return fval
    def GetSteps(self):
        return np.array(TrackMinuit.steps_taken)

# Run the actual fit
TM = TrackMinuit(f)
m = Minuit(TM, pedantic=False, x=-0.5, y=2.5)
m.migrad()
# see the fitted values
print(m.values)  
#print(m.errors)


# get the steps
steps = TM.GetSteps()
x_steps, y_step = steps.T



# Plot Minuit's progress in the function space
fig = plt.figure(figsize=(14, 8))
ax = fig.gca(projection='3d')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.set_title('Function space of the Rosenbrock function')

X = np.linspace(-2, 2, 50)  
Y = np.linspace(-0.5, 3, 50)  
    
X, Y = np.meshgrid(X, Y)
Z = f(X, Y)

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
         linewidth=0, antialiased=False, alpha=0.8)  #Try coolwarm vs jet

ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

fig.colorbar(surf, shrink=0.5, aspect=5)
ax.plot(x_steps, y_step, f(x_steps, y_step), 'wo-',  label='Minimizer Steps')
ax.scatter(a, a**2, Z.min(), c='red', s=100, label='Theoretical Minimimum')
ax.legend()

fig.tight_layout()
plt.show(block=False)


# Finally, ensure that the program does not termine (and the plot disappears), before you press enter:
try:
    __IPYTHON__
except:
    raw_input('Press Enter to exit')