import numpy as np


###########################################################################
##################    Gradient Evaluations     ############################

def x_d_eps(x1, pm, eps=1e-8):
    return x1 + pm * eps


def eval_grad(f, x1, f1=False, eps=1e-8, evaltype="forward", direction=False, alpha=1., component=False):
    x0 = tuple(x1)
    if  f1 is False:
        if direction is not False:
            f1 = f(x1 + alpha * direction)
        else:
            f1 = f(x1)
    
    if   evaltype == "forward" :
        ##  Differentiation with forward finite difference method
        if component is not False:
            x1[component] = x_d_eps( x1[component], 1., eps=eps )
            x2 = x1
        else:
            x2 = x_d_eps( x1, 1., eps=eps )
        if direction is not False:
            alpha2 = x_d_eps(alpha, 1., eps=eps)
            f2 = f( x1 + alpha2 * direction )
            print("alpha1=%s alpha2=%s dir=%s f1=%s, f2=%s "%(alpha, alpha2, direction, f1, f2))
        else:
            f2 = f(x2)
        #print("f2=%s, f1=%s, f2-f1=%s" %(f2, f1, f2-f1))
        g  = (f2 - f1) / eps

    elif evaltype == "backward":
        ##  Do differentiation with backward finite difference method
        if component is not False:
            x1[component] = x_d_eps( x1[component], -1., eps=eps )
        else:
            x2 = x_d_eps( x1, -1., eps=eps )
        if direction is not False:
            alpha2 = x_d_eps(alpha, -1., eps=eps)
            f2 = f( x1 + alpha2 * direction )
        else:
            f2 = f(x2)
        g  = (f1 - f2) / eps

    elif evaltype == "central" :
        ##  Do differentiation with central finite difference method
        if component is not False:
            x1 = np.asarray(x0)
            x1[component] = x_d_eps( x1[component], 1., eps=eps )
            x2p = x1
            x1 = np.asarray(x0)
            x1[component] = x_d_eps( x1[component], -1., eps=eps )
            x2m = x1
        else:
            x2p = x_d_eps( x1,  1., eps=eps )
            x2m = x_d_eps( x1, -1., eps=eps )
        if direction is not False:
            alphap = x_d_eps(alpha, 1., eps=eps)
            alpham = x_d_eps(alpha, -1., eps=eps)
            f2p = f( x1 + alphap * direction )
            f2m = f( x1 + alpham * direction )
        else:
            f2p = f(x2p)
            f2m = f(x2m)
        g   = (f2p -  f2m) / (2 * eps)
    #print("evalgrad", g)
    return g



def grad_f(f1, x1, f2, x2):
    """
    This calculates the gradient of the function. 
    f is a vector of function evaluations fi at the matrix of vector points x
    """
    if type(x1) == float or type(x2) == float:
        ##  If x1 or x2 is just a scalar
        return (f2 - f1) / (x2 - x1)
    else:
        if shape.x1 != shape.x2:
            print("The shapes of the input vectors are not equal!\n    Exiting...")
            raise ValueError

        print("\n GRAD_F \n Evaluating the gradient at point \n x0 = %.8f \n  "%( (x2 + x1)/2. ) )

        g_f     = np.zeros(x1.shape)
        deltax  =  x2 - x1
        df      =  f2 - f1

        for ind, dx in enumerate(deltax):
            g_f[ind] = df / dx

        return g_f

def eval_f_df(f, f0, df0, a):
    if  f0    is False:
        f0    =  f(a)
    if  df0   is False:
        df0   =  eval_grad( f,  a, f1=f0  )
    return f0, df0
   

def df(f, x, f1=False, ev_type="forward", eps=1e-8):
    df = np.zeros(x.shape)
    for i, xi in enumerate(x):
        df[i] = eval_grad(f, x, f1=f1, evaltype=ev_type, eps=eps, component=i)
    return df

#########################################################################################
##################     Curvature and Hessian Evaluations     ############################

def get_hess(f, x, epsi=1e-8, epsj=1e-8):
    n  = x.shape
    x0 = x
    f0 = f(x0)
    x0 = tuple(x)
    H = np.zeros((n,n))
    for i in range(n):
        x    = np.asarray(x0)
        xi   = x_d_eps( x0[i],  1. , eps=epsi)
        x[i] = xi
        f_ei = f(x)
        for j in range(n):
            x    = np.asarray(x0)
            x[j] = x_d_eps( x[j],  1., eps=epsj )
            f_ej = f(x)
            x[i] = xi
            f_eij= f(x) 
            H[i][j] = ( f_eij - f_ei - f_ej +  f0 ) / ( epsi * epsj )
    return H
    

def eval_curv(f, x1, f1=False, evaltype="central"):

    if  f1 is False:
        f1  = f(x1)
    
    if   evaltype == "forward" :
        ##  Differentiation with forward finite difference met
        x2 = x_d_eps( x1, 1. )
        x3 = x_d_eps( x2, 1. )
        f2 = f(x2)
        f3 = f(x3)
        g  = (f3 - 2 * f2 + f1) / eps**2

    elif evaltype == "backward":
        ##  Do differentiation with backward finite difference method
        x2 = x_d_eps( x1, -1. )
        x3 = x_d_eps( x2, -1. )
        f2 = f(x2)
        f3 = f(x3)
        g  = (f1 - 2 * f2 + f3) / eps**2 

    elif evaltype == "central" :
        ##  Do differentiation with central finite difference method
        x2p = x_d_eps( x1,  1. )
        x2m = x_d_eps( x1, -1. )        
        f2p = f(x2p)
        f2m = f(x2m)
        g   = (f2p  -  2 * f1  +  f2m) /  eps**2
        
    return g




##################################################################################
############################    Line Search     ##################################

def line_search_wolfe(f, xk, pk, a_max, f0   = False, df0   = False,
                                        f1   = False, f_alo = False,
                                        c1   = 1e-4,  c2    = 0.9,  eps=1e-8, maxit=100):
    ##  This is a line search algorithm to find the best step length a_k_best
    a0 = 0.
    aj = (a_max + a0) /2.
    
    f0    = f(xk)
    ajm   = a0
    f_ajm = f0
    f_aj  = f(xk + aj*pk)
    a_st  = False

    df0   = eval_grad(f, xk, f1=f0, direction=pk, alpha=a0) 
    print(f0, df0, f_aj, pk)
    for i in range(maxit):
        cond = (f_aj  > f0  +  c1 * aj * df0) or ( f_aj > f_ajm )

        if cond:
            a_st = ( ajm + aj ) / 2.
            a_st = zoom(f, xk, a_st, pk, alo=ajm, ahi=aj)

        else:
            df_aj = eval_grad(f, xk, f1=f_aj, direction=pk, alpha=aj) 
            if abs(df_aj) <= - c2 * df0:
                a_st = aj
            if df_aj > 0.:
                a_st = ( ajm + aj ) / 2.
                a_st = zoom(f, xk, a_st, pk, alo=a0, ahi=aj)
            df0  = df_aj        
        aj   = (aj + a_max) / 2.
        fajm = f_aj
        f_aj = f( xk +  aj * pk )
        if a_st is not False:
            return a_st
        if i == maxit -1:
            print("Hit the maximum number of iterations %s for the line search" %(maxit))
            print("Setting a_st = %s" %(aj))
            return a_st



def zoom(f, xk, aj, pk, alo=0., ahi=10.,
            f0   = False, df0   = False,
            f_aj = False, f_alo = False,
            c1   = 1e-4,  c2    = 0.9,  eps=1e-8, maxit=100):
    """
    Interpolate (using quadratic, cubic, or bisection) to find
    a trial step length a_j between a_lo and a_hi
    """
    if  f0    is False:
        f0    =  f(xk)
    if  df0   is False:
        df0   =  eval_grad( f, xk, f1=f0 )
    if  f_aj  is False:
        f_aj  =  f(xk + aj*pk )
    if  f_alo is False:
        f_alo =  f(xk + alo*pk)

    df0 = eval_grad(f, xk, f1=f0, direction=pk, alpha=0., eps=eps)
        
    ast = False
    for i in range(maxit):
        cond = (f_aj  > f0  +  c1 * aj * df0) or ( f_aj > f_alo )

        if cond:
            ahi = aj
        else:
            print("In the else statement: Evaluating Gradient")
            print("f_aj =%s, aj = %s, alo=%s, ahi=%s" %(f_aj, aj, alo, ahi))
            df_aj = eval_grad(f, xk, f1=f_aj, direction=pk, alpha=aj, eps=eps)
            if abs(df_aj) <= -c2 * df0:
                ast = aj
            if df_aj * (ahi - alo) >= 0:
                ahi = alo
            alo = aj
        f_aj = f(xk + aj * pk)
        if ast is not False:
            break
        if i == maxit - 1:
            print("Warning: Maximum number of iterations for finding step length has been exceeded.")
            print("Setting a_s to current a_j:\n    a_s = %s" %(aj))
            ast = aj
    return ast



        

###########################################################################
##################     Conjugate Gradient      ############################

def fr_cg(f, xk, a_max=10., maxiter=100, tol=1e-5, fr=True, c1 =1e-4, c2=0.4, eps=1e-8):
    fk  = f(xk)
    print("First fk", fk)
    dfk = df(f, xk, f1=fk, eps=eps)
    print("df0 = %s"%(dfk))
    pk  = -dfk
    for i in range(maxiter):
        ak   = line_search_wolfe(f, xk, pk, a_max, f0=fk, df0=dfk, maxit=maxiter, c1=c1, c2=c2, eps=eps)
        print("CG: Wolfe step length    ak = %s" %(ak))
        xk1  = xk  +  ak * pk
        dfk1 = df(f, xk1, eps=eps)        
        if fr:
            print("\n\nFletcher-Reeves Conjugate gradient: Iteration %s" %(i) )
            bk1  = ( dfk1.dot(dfk1) ) / ( dfk.dot(dfk)  )
        else:
            print("Polak-Ribere Conjugate gradient:    Iteration %s" %(i) )
            bk1  = ( dfk1.dot(dfk1-dfk) ) / ( dfk.dot(dfk)  )
        print("    bk1 = %s" %(bk1))            
        print("    xk = %s" %(xk))
        print("    xk1 = %s\n" %(xk1))
        pk  = -dfk1 + bk1 * pk
        xk  = xk1
        fk  = f(xk)        
        dfk = dfk1
        if np.all(np.abs(dfk1) < tol):
            print("END conditon dfk=",dfk)
            break
    print("Return    xk = %s" %(xk))
    return xk
        
        

def linear_cg( A, x_k, b, tol=1e-5):
    #This solves the equation Ax = b
    r_k = A.dot(x_k) - b
    p_k = -r_k
    while np.all(r_k > tol):
        a_k = ( r_k.dot(p_k) )/( p_k.dot( A.dot( p_k ) ) )
        x_k += a_k * p_k
        r_kp1 = r_k + a_k * A.dot( p_k )
        beta_k = r_kp1.dot(r_kp1)/r_k.dot(r_k)
        p_k = -r_kp1 + beta_k*p_k 
        r_k = r_kp1
    return x_k


def testf(x):
    tf = np.sum((x-6)*(x-6))
    return tf

#xk = np.array([1,2,3])
#eps = 1e-5
#fr_cg(testf, xk, eps=eps)

        
"""
############     From BOP   ##########
            
            gg = sum(ftot(:,:nd)*ftot(:,:nd))

            if ((abs(ggold) < 1.0e-6_dp) .or. (mod(iter,5) == 0)) then
               gamma = 0.0_dp
            else
               gamma = gg/ggold
            endif
            
            if (.not. quiet) write(9,'("gamma = ",g12.5)') gamma
            ggold = gg
            
            do
               cg(:,:nd) = gamma * cg(:,:nd) + ftot(:,:nd) 
               newad(:,:nd) = ad(:,:nd)
               
               call safemin(newad,cg,ein,eout,status)

               if ((status /= 1) .or. abs(ggold) <= 1.0e-6_dp) exit
               
               ggold = 0.0_dp               
               gamma = 0.0_dp
               
            end do

"""
