import numpy as np
from numpy import sin, cos, tan, pi, exp, sqrt
import sys

from tklib.tksci.tkoptimizeobject import searchdir_sd, searchdir_cg, searchdir_newton
from tklib.tksci.tkoptimizeobject import linesearch_one, linesearch_simple, linesearch_exact, linesearch_newton
from tklib.tksci.tkoptimizeobject import linesearch_golden,linesearch_armijo
from tklib.tksci.tkoptimizeobject import tkOptimizeData
from tklib.tkutils import lvlprint, mergeattributes


def cg(func = None, x0 = None, diff1func = None, diff2func = None, diff2arrayfunc = None,
            nmaxiter = 100, tolx = 1.0e-5, tolf = 1.0e-5,
            optdata = None, callback = None, print_level = 1, iprintinterval = 10, 
            algorism = None, searchdir_func = None, 
            lsmode = None, ls_func = None,
            ls_xrange = None, ls_h = None, ls_alpha = None, ls_dump = None, 
            ls_nmaxiter = None, ls_alphaeps = None):

    n = len(x0)

    if searchdir_func is None:
        searchdir_func = searchdir_sd
        if algorism == 'cg':
            searchdir_func = searchdir_cg
        elif algorism == 'newton':
            searchdir_func = searchdir_newton

    if ls_func is None:
        ls_func = linesearch_one
        if lsmode == 'simple':
            ls_func = linesearch_simple
        elif lsmode == 'exact':
            ls_func = linesearch_exact
        elif lsmode == 'newton':
            ls_func = linesearch_newton
        elif lsmode == 'golden':
            ls_func = linesearch_golden
        elif lsmode == 'armijo':
            ls_func = linesearch_armijo

    gradfkm = None
    dkm     = None
    f = func(x0)
    lvlprint(print_level, 1, "x0 = ({}, {}): f = {}".format(x0[0], x0[1], f))

    optdata2 = tkOptimizeData(x0, itmax = nmaxiter, tolx = tolx, tolf = tolf, 
                diff1func = diff1func, diff2func = diff2func, diff2arrayfunc = diff2arrayfunc,
                algorism = algorism, searchdir_func = searchdir_func, 
                lsmode = lsmode, ls_func = ls_func,
                callback = callback)
    if optdata is None:
        optdata = optdata2
    else:
        optdata = mergeattributes(optdata, optdata2)


# counter of cg steps
    icg = 0
# optimization start
    fprev = None
    for iter in range(nmaxiter):
        optdata.iter = iter

        dk, icg, gradfkm, dkm = searchdir_func(x0, diff1func, icg, gradfkm, dkm, optdata)

# pass the first derivative of the target function, i.e., the negative value of the search direction dk, -dk
        insiter, x0, dx, alpha2 = ls_func(func, x0, dk, 
                    xrange = ls_xrange, h = ls_h, alpha = ls_alpha, dump = ls_dump, 
                    nmaxiter = ls_nmaxiter, eps = ls_alphaeps)
        f = func(x0)
        dxmax = max(abs(dx))
        if callback is not None:
            optdata.insiter = insiter
            optdata.x0     = x0
            optdata.f      = f
            optdata.dx     = dx
            optdata.dxmax  = dxmax
            optdata.alpha2 = alpha2
            ret = callback(optdata)
            if ret > 0:
                var.status = ret
                return None, None, optdata

        lvlprint(print_level, 1, "{:4d}({:3d}): x = ({:14.6g}, {:14.6g}) dx = {:14.6g}: f = {:14.6g}".format(
                        iter, insiter, x0[0], x0[1], dxmax, f))

        if dxmax < tolx or (fprev is not None and abs(f - fprev) < tolf):
            lvlprint(print_level, 0, "Converged at x = ", x0, " dx = {}   f = {}".format(dxmax, f))
            return f, x0, optdata

        fprev = f

    lvlprint(print_lvel, 0, "Not converged")
    return f, x0, optdata



if __name__ == "__main__":
    main()

