import numpy as np
# function to return the gradient of the quadratic function
def g_f(x):

  return np.array([3*x[0]+x[2]-3,4*x[1]+2*x[2], x[0]+2*x[1]+3*x[2]-1])

# The function conjugate gradient : for finding the minimum value for quadratic function.
def conjugate_gradient(x,Q,eps):

  i = 0
  lst_x = []
  lst_g = []
  sum_g = [] 
  g = g_f(x) # getting the gradient
  lst_x.append(x) 
  lst_g.append(g)
  sum_g.append(abs(np.sum(g)))
  
  d = -g # specifying the direction of descent
  
  while abs(np.sum(g)) > eps: # specifing exit instance from the executing of the loop
    
    i +=1 # specifying a counter to keep a count on the number of executions.

    alpha = - np.divide(np.dot(g,d),np.dot(d,np.dot(Q,d))) # calculating the step size.
    x = x + alpha*d #calculating the next value of x
    lst_x.append(x) 

    g = g_f(x) 
    lst_g.append(g) 
    beta = np.divide(np.dot(g,np.dot(Q,d)),np.dot(d,np.dot(Q,d))) # calculating beta, to calculate the Q-conjugate descent direction
    d = -g + beta*d 
    sum_g.append(abs(np.sum(g)))


  return x,lst_x,lst_g,i,sum_g

# Initializing values of x, Q and function call 

x = np.array([0,0,0])
Q = np.array([[3,0,1],[0,4,2],[1,2,3]])
x,lst,lst_g,i,sum_g = conjugate_gradient(x,Q,eps=1e-4)

# Defining the exact line search method , to calculate alpha for every stage.

def f(alpha,x,d):
  
  return (1-(x[0]-alpha*d[0]))**2 + 100*((x[1]-alpha*d[1])-(x[0]-alpha*d[0])**2)**2
 
# derivative of the objective function 

def grad_alpha(f,alpha,x,d):
  h = 0.00001
  return (f(alpha+h,x,d)-f(alpha,x,d))/h

def bracket_minimum(alpha,x,d,s,k):

  a, fa = alpha,  f(alpha,x,d)
  b, fb = a+s , f(a+s,x,d)

  if fb > fa:
    a,b = b,a 
    fa, fb = fb,fb
    s = -s

  while True:
    c, fc = b+s, f(b+s,x,d)
    if fc > fb:
      if a < c:
        return [a,c]
      else:
        return [c,a]
    a,fa,b,fb = b,fb,c,fc
    s*=k

def bisection_root_finding(a,b,x,d,eps):

  if a > b : 
    a,b = b,a 
  fa, fb = grad_alpha(f,a,x,d), grad_alpha(f,b,x,d)
  if fa == 0:
    b = a
    return (a)

  if fb == 0:
    a = b
    return (b)
  iter = 0
 
  while abs(b - a) > eps:
    iter +=1
    c = (a+b)/2
    y = grad_alpha(f,c,x,d)

    if y == 0:
      a,b = c, c
      break

    if np.sign(y)  < 0  :
      a = c
    else:
      b = c

  return (a+b)/2

def line_search(alpha,x,d):
  
  a,b = bracket_minimum(alpha,x,d,s=0.01,k=2.0)
  alpha = bisection_root_finding(a,b,x,d,eps=1e-5)
  return alpha

# We will use the conjugate gradient method to arrive at the minimum of the Rosenbrock's Banana Function
def fun(x,a):
  pow = a
  if pow == 0:
    return (1 - x[0])**2 + 100*(x[1] - x[0]**2)**2
   
  else:
    h = 0.00001
    x1 = np.array([x[0]+h,x[1]])
    x2 = np.array([x[0],x[1]+h])
    # The gradient Vector
    g = np.array([(fun(x1,0)-fun(x,0))/h,(fun(x2,0)-fun(x,0))/h])
    return g

# The function conjugate gradient : for finding the minimum value for quadratic function.
def conjugate_gradient_nonquad(x,alpha,eps):

  # Initialize counter and empty lists to hold x and function values
  i = 0
  lst_x = []
  lst_f = []

  # assigning the function and gradient value.
  f, g = fun(x,0), fun(x,1) 
  
  lst_x.append(x) 
  lst_f.append(f)

  # specifying the direction of descent
  d = np.array(g/np.linalg.norm(g))
  
  # specifing exit instance from the executing of the loop
  while abs(f) > eps: 
    
    #specifying a counter to keep a count on the number of executions.
    i +=1 

    # replace the calculating on alpha with exact line search, similar to our other first - order descent methods
    alpha = line_search(alpha,x,d)
   
    #calculating the next value of x
    x = x - alpha*d 
    f = fun(x,0)
    g1= fun(x,1) 
    
    # This is the step, which differentiates from the above conjugate gradient method.
    beta = np.divide(np.dot(g1.T,g1-g),np.dot(g.T,g) )# calculating beta, using the Polak Ribiere formula 
    beta = max(0,beta)  # to assure convergence, we add this step

    # updating the direction vector
    d1 = np.array(g1/np.linalg.norm(g1))
    d = d1 + beta*d  
    g = g1

    # appending values to our lists.
    lst_x.append(x)
    lst_f.append(f) 

  print("Number of iterations",i)
  print("The value of the Rosenbrock banana function at its minimum:",fun(x,0))
  return x,lst_x,lst_f,i