###############################################################################
### WORST-CASE FUNCTION FOR RK METHOD
###############################################################################
# Objective : plot one of the function, and initial points satisfying worst-case scenario

# For a given Runge-Kutta method :
# x_2 = x_1 + gamma*sum(beta_i*g(x_i1))
# x_i1 = x_1 + gamma*sum(alpha_j*g(x_j1))
#
# 1. Compute the worst-case contraction rate and the iterates (Gram matrix G and function values F)
# 2. Given the lyapunov at the optimum, compute the iterates and the function f while minimizing the dimension of f.
# 3. Plot the fuction evaluated at a given point
# We use Performance Estimation Problems, that is semi-definite programing on a worst-case formulation of the problem.

import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt

from worst_case_functions_RK import worst_case_RK, B_stability_RK_rate_G_F


## Class of functions
L = 1. # smoothness
mu = 0. # strong-convexity

## Parameter for binary-search
epsilon = 10**(-3)

## EXAMPLES
# name of the scheme
scheme = 'Two stage 1'
# step size
gamma = 1/L

if scheme == 'Two stage 1':
    alpha = np.zeros((2, 2))
    alpha[1][0] = .5
    beta = np.zeros(2)
    beta[1] = 1.
if scheme == 'Euler explicit':
    alpha = np.zeros((1, 1))
    beta = np.zeros(1)
    beta[0] = 1.
if scheme == 'Euler implicit':
    alpha = np.zeros((1, 1))
    alpha[0][0] = 1.
    beta = np.zeros(1)
    beta[0] = 1.
if scheme == 'Two stage explicit':
    alpha = np.zeros((2, 2))
    alpha[1][0] = .5
    beta = .5 * np.ones(2)
if scheme == 'Trapezoidal':
    alpha = np.zeros((2, 2))
    alpha[1][0] = 1.
    beta = .5 * np.ones(2)

### COMPUTE THE WORST-CASE FUNCTION IN A POINT x        
## Compute the contraction rate, and the Gram matrix/function values associated with this contraction rate
tau, F, G = B_stability_RK_rate_G_F(L, mu, gamma, alpha, beta)

## Suppose we have two principal dimensions : 
N = 2 

## Evaluate the function in a new point
x = np.array([0, 0])
eps = 10**-2
f_, X_pts, G_pts, g_ = worst_case_RK(L, mu, gamma, alpha, beta, G, F, x, eps = eps) # get the points generated by the algorithm


### PLOT THE WHOLE FUNCTION
## Plot the shape of the function, and the iterated
# parameters
nn = 21 # number of points per dimension (Nb = nn*nn)
X,Y = np.linspace(-2, 2, nn),np.linspace(-1.5, 1.5, nn) # generate a grid
eps = 0.05

# Evaluate the function on a grid
sol = np.zeros((nn,nn))
for i in range(nn):
    for j in range(nn):
        xx = np.array([X[i], Y[j]])
        sol[i][j], _, _, _ = worst_case_RK(L, mu, gamma, alpha, beta, G, F, xx, eps = eps)

# Plot the worst-case iterates
X_pts=np.array(X_pts)
X_pts1 = np.array([X_pts[2*i] for i in range(len(X_pts)//2)])
X_pts2 = np.array([X_pts[2*i+1] for i in range(len(X_pts)//2)])
plt.subplots(1,1,figsize=(11,7))
plt.plot(X_pts1[:,0], X_pts1[:,1], '--x', color = 'red')
plt.plot(X_pts2[:,0], X_pts2[:,1], '--x', color = 'red')

## Points from the grid
plt.pcolor(X, Y, sol.T)
plt.colorbar()
plt.xlabel('x')
plt.ylabel('y')
plt.show()



