from __future__ import absolute_import

import numpy as nm
import numpy.linalg as nla

from sfepy.base.base import output, get_default, pause, Struct
from sfepy.base.log import Log, get_logging_conf
from sfepy.base.timing import Timer
from sfepy.solvers.solvers import OptimizationSolver

import scipy.optimize as sopt
import scipy.optimize.linesearch as linesearch
import six
from six.moves import range

def conv_test(conf, it, of, of0, ofg_norm=None):
    """
    Returns
    -------
    flag : int
        * -1 ... continue
        *  0 ... small OF -> stop
        *  1 ... i_max reached -> stop
        *  2 ... small OFG -> stop
        *  3 ... small relative decrase of OF
     """

    status = -1
    output('opt: iter: %d, of: %e (||ofg||: %e)' % (it, of, ofg_norm))

    if (abs(of) < conf.eps_of):
        status = 0
    elif ofg_norm and (ofg_norm < conf.eps_ofg):
        status = 2
    elif (it > 0) and (abs(of0 - of) < (conf.eps_rd * abs(of0))):
        status = 3

    if (status == -1) and (it >= conf.i_max):
        status = 1

    return status

def wrap_function(function, args):
    ncalls = [0]
    times = []
    timer = Timer()
    def function_wrapper(x):
        ncalls[0] += 1
        timer.start()
        out = function(x, *args)
        times.append(timer.stop())
        return out
    return ncalls, times, function_wrapper

def check_gradient(xit, aofg, fn_of, delta, check):

    dofg = nm.zeros_like(aofg)
    xd = xit.copy()
    for ii in range(xit.shape[0]):
        xd[ii] = xit[ii] + delta
        ofp = fn_of(xd)

        xd[ii] = xit[ii] - delta
        ofm = fn_of(xd)

        xd[ii] = xit[ii]

        dofg[ii] = 0.5 * (ofp - ofm) / delta

        output('**********', ii, aofg[ii], dofg[ii])

    diff = abs(aofg - dofg)
    aux = nm.concatenate((aofg[:,nm.newaxis], dofg[:,nm.newaxis],
                          diff[:,nm.newaxis]), 1)
    output(aux)
    output(nla.norm(diff, nm.Inf))
    aofg.tofile('aofg.txt', ' ')
    dofg.tofile('dofg.txt', ' ')
    diff.tofile('diff.txt', ' ')
    if check == 2:
        import pylab
        pylab.plot(aofg)
        pylab.plot(dofg)
        pylab.legend(('analytical', 'finite difference'))
        pylab.show()
    pause('gradient checking done')

class FMinSteepestDescent(OptimizationSolver):
    """
    Steepest descent optimization solver.
    """
    name = 'opt.fmin_sd'

    _parameters = [
        ('i_max', 'int', 10, False,
         'The maximum number of iterations.'),
        ('eps_rd', 'float', 1e-5, False,
         'The relative delta of the objective function.'),
        ('eps_of', 'float', 1e-4, False,
         'The tolerance for the objective function.'),
        ('eps_ofg', 'float', 1e-8, False,
         'The tolerance for the objective function gradient.'),
        ('norm', 'numpy norm', nm.Inf, False,
         'The norm to be used.'),
        ('ls', 'bool', True, False,
         'If True, use a line-search.'),
        ('ls_method', "{'backtracking', 'full'}", 'backtracking', False,
         'The line-search method.'),
        ('ls_on', 'float', 0.99999, False,
         """Start the backtracking line-search by reducing the step, if
            :math:`||f(x^i)|| / ||f(x^{i-1})||` is larger than `ls_on`."""),
        ('ls0', '0.0 < float < 1.0', 1.0, False,
         'The initial step.'),
        ('ls_red', '0.0 < float < 1.0', 0.5, False,
         'The step reduction factor in case of correct residual assembling.'),
        ('ls_red_warp', '0.0 < float < 1.0', 0.1, False,
         """The step reduction factor in case of failed residual assembling
            (e.g. the "warp violation" error caused by a negative volume
            element resulting from too large deformations)."""),
        ('ls_min', '0.0 < float < 1.0', 1e-5, False,
         'The minimum step reduction factor.'),
        ('check', '0, 1 or 2', 0, False,
         """If >= 1, check the tangent matrix using finite differences.  If 2,
            plot the resulting sparsity patterns."""),
        ('delta', 'float', 1e-6, False,
         r"""If `check >= 1`, the finite difference matrix is taken as
            :math:`A_{ij} = \frac{f_i(x_j + \delta) - f_i(x_j - \delta)}{2
            \delta}`."""),
        ('output', 'function', None, False,
         """If given, use it instead of :func:`output()
            <sfepy.base.base.output()>` function."""),
        ('yscales', 'list of str', ['linear', 'log', 'log', 'linear'], False,
         'The list of four convergence log subplot scales.'),
        ('log', 'dict or None', None, False,
         """If not None, log the convergence according to the configuration in
            the following form: ``{'text' : 'log.txt', 'plot' : 'log.pdf'}``.
            Each of the dict items can be None."""),
    ]

    def __init__(self, conf, **kwargs):
        OptimizationSolver.__init__(self, conf, **kwargs)

        conf = self.conf

        log = get_logging_conf(conf)
        conf.log = log = 
Struct(name='log_conf', **log)