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_none, linesearch_one, linesearch_simple
from tklib.tksci.tkoptimizeobject import linesearch_exact
from tklib.tksci.tkoptimizeobject import linesearch_golden,linesearch_armijo
from tklib.tksci.tkoptimizeobject import tkOptimizeData
from tklib.tkutils import lvlprint, merge_attributes, joinf


def gradient(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_alphaeps = None, ls_feps = None,
            ls_nmaxiter = None):

    """
    Non-linear LSQ by gradient descent methods
    Supported: sd, cg, newton dfp, dfpB, bfgs, bfgsB, broyden, broydenB
       'B' denotes the approximation for Hessian matrix, o/w that for inverse of Hessian matrix
    """

    n = len(x0)

    if searchdir_func is None:
        searchdir_func = searchdir_sd
        if algorism == 'cg':
            searchdir_func = searchdir_cg
        elif algorism == 'sd':
            searchdir_func = searchdir_sd
        elif algorism == 'newton':
            searchdir_func = searchdir_newton
        elif algorism == 'dfp':
            searchdir_func = searchdir_dfp
        elif algorism == 'dfpH':
            searchdir_func = searchdir_dfpH
        elif algorism == 'bfgs':
            searchdir_func = searchdir_bfgs
        elif algorism == 'bfgsH':
            searchdir_func = searchdir_bfgsH
        elif algorism == 'broyden':
            searchdir_func = searchdir_broyden
        elif algorism == 'broydenH':
            searchdir_func = searchdir_broydenH

    if ls_func is None:
        if algorism == 'newton':
            ls_func = linesearch_newton
        else:
            ls_func = linesearch_armijo

        if lsmode == 'simple':
            ls_func = linesearch_simple
        if lsmode == 'one':
            ls_func = linesearch_one
        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
    optdata2 = tkOptimizeData(x0, itmax = nmaxiter, tolx = tolx, tolf = tolf, 
                func = func, 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 = merge_attributes(optdata, optdata2)

    f = func(x0, optdata)
#    x0str = joinf(x0, "%12.6g", ", ")
#    lvlprint(print_level, 1, "x0 = ({}): f = {}".format(x0, f))


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

        dk, icg, gradfkm, dkm = searchdir_func(x0, diff1func, icg, gradfkm, dkm, optdata)
        x0str = joinf(x0, "%12.6g", ", ")
        dxstr = joinf(dk, "%10.4g", ", ")
        lvlprint(print_level, 1, "iter: {:4d}  x0=({})  dx=({})  fprev={:12.6g}"
                    .format(iter, x0str, dxstr, fprev))

# 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, alphaeps = ls_alphaeps, feps = ls_feps,
                    optdata = optdata)
        f = func(x0, optdata)
        dxmax = max(abs(dx))
        if callback is not None:
            optdata.insiter = insiter
            optdata.x0     = x0
            optdata.x      = 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

        x0str = joinf(x0, "%12.6g", ", ")
        dxstr = joinf(dx, "%10.4g", ", ")
        lvlprint(print_level, 1, "   ls({:3d}): x0=({})  dx=({})".format(insiter, x0str, dxstr))
        lvlprint(print_level, 1, 
                "            dxmax={:12.6g}  alpha2={:12.6g}  f={:14.10g}".format(dxmax,alpha2, f))

        if dxmax < tolx or (fprev is not None and abs(f - fprev) < tolf):
            if dxmax < tolx:
                lvlprint(print_level, 0, "Converged due to  dx={:10.4g} < tolx={:12.6g}".format(dxmax, tolx))
            else:
                lvlprint(print_level, 0, "Converged due to  df={:10.4g} < tolf={:12.6g}".format(abs(f - fprev), tolf))
            x0str = joinf(x0, "%12.6g", ", ")
            lvlprint(print_level, 0, "   at x = ({})   f = {:12.6g}".format(x0str, f))
            return x0, f, optdata

        fprev = f

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



if __name__ == "__main__":
    main()

