%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt

"""
Explanation: Optimization Exercise 1
Imports
End of explanation
"""


def hat(x,a,b):
    return -a*x**2+b*x**4

assert hat(0.0, 1.0, 1.0)==0.0
assert hat(0.0, 1.0, 1.0)==0.0
assert hat(1.0, 10.0, 1.0)==-9.0

"""
Explanation: Hat potential
The following potential is often used in Physics and other fields to describe symmetry breaking and is often known as the "hat potential":
$$ V(x) = -a x^2 + b x^4 $$
Write a function hat(x,a,b) that returns the value of this function:
End of explanation
"""


a = 5.0
b = 1.0

x=np.linspace(-3,3,100)
plt.figure(figsize=(9,6))
plt.xlabel('Range'), plt.ylabel('V(x)'), plt.title('Hat Potential')
plt.plot(x, hat(x,a,b))
plt.box(False)
plt.grid(True)
plt.tick_params(axis='x', top='off', direction='out')
plt.tick_params(axis='y', right='off', direction='out');

assert True # leave this to grade the plot

"""
Explanation: Plot this function over the range $x\in\left[-3,3\right]$ with $b=1.0$ and $a=5.0$:
End of explanation
"""


res1 = opt.minimize_scalar(hat, bounds=(-3,0), args=(a,b), method='bounded')
res2 = opt.minimize_scalar(hat, bounds=(0,3), args=(a,b), method='bounded')
print('Local minima: %f, %f' % (res1.x, res2.x))
plt.figure(figsize=(9,6))
plt.xlabel('Range'), plt.ylabel('V(x)')
plt.plot(x, hat(x,a,b), label="Potential")
plt.scatter(res1.x, res1.fun, marker="o", color="r")
plt.scatter(res2.x, res2.fun, marker="o", color="r")
plt.title('Finding Local Minima of Hat Potential')
plt.box(False), plt.grid(True), plt.xlim(-2.5,2.5), plt.ylim(-8,4)
plt.tick_params(axis='x', top='off', direction='out')
plt.tick_params(axis='y', right='off', direction='out')
plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.);

assert True # leave this for grading the plot

"""
Explanation: Write code that finds the two local minima of this function for $b=1.0$ and $a=5.0$.

Use scipy.optimize.minimize to find the minima. You will have to think carefully about how to get this function to find both minima.
Print the x values of the minima.
Plot the function as a blue line.
On the same axes, show the minima as red circles.
Customize your visualization to make it beatiful and effective.
End of explanation
"""


