import pylab
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from scipy.optimize import minimize
import sympy as sp
from scipy.optimize import rosen, rosen_der, rosen_hess, rosen_hess_prod
from scipy.optimize import Bounds

def makeData():
    x = np.arange(-1, 1, 0.1)
    y = np.arange(-1, 1, 0.1)
    xgrid, ygrid = np.meshgrid(x, y)
    zgrid = -xgrid**2 + xgrid*ygrid - 4*ygrid**2
    return xgrid, ygrid, zgrid

def makeDataG():
    x = np.arange(-1, 1, 0.1)
    xgrid= np.meshgrid(x)
    y1grid = 0.5 * xgrid - 1
    y2grid = 0.5* xgrid + 0.6
    return xgrid, y1grid, y2grid

def f(x):
    return -1 * x[0]*x[0] + x[0]*x[1] - 4*x[1]*x[1]



eq_cons = {'type': 'eq',
           'fun': lambda x: np.array ([x[1] - 0.5 *x[0] + 1, 0.5*x[0] - x[1] + 0.6]),
           'jac': lambda x: np.array ([[-0.5, 1.0],[0.5, -1.0]])
          }

x1, y1, z = makeData()
x2, y2, y3 = makeDataG()

fig = pylab.figure()
axes = Axes3D(fig)

fig1 = pylab.figure()
axes1 = Axes3D(fig1)


bnds = Bounds ([-np.inf, -np.inf], [np.inf, np.inf])

axes.plot_surface(x1, y1, z)
axes1.plot_surface(x2, y2)
axes1.plot_surface(x2, y3)
x0 = np.array([-1.0, -1.0])
res = minimize(rosen, x0, method='SLSQP', jac=rosen_der,
               constraints=[eq_cons], options={'ftol': 1e-9, 'disp': True},
               bounds=bnds)


# axes.scatter(t.x[0], t.x[1], f(t.x), s = 480, c = 'red')

print('x= ' + str(t.x[0]) + ', y= ' + str(t.x[1]) + ', z(x,y)=' + str(f(t.x)))

pylab.show()