import numpy as np
import warnings
from pylearn2.utils import sharedX
from theano import config
from theano import function
from theano.printing import Print
from theano.printing import min_informative_str

import theano.tensor as T

def norm_sq(s):
    return np.square(s.get_value()).sum()

def scale(s, a):
    s.set_value(s.get_value() * np.cast[config.floatX](a))

class BatchGradientDescent:
    """ A class for minimizing a function via the method of steepest
    descent.
    """

    def __init__(self, objective, params, inputs = None,
            param_constrainers = None, max_iter = -1,
            lr_scalers = None):
        """ objective: a theano expression to be minimized
                       should be a function of params and,
                       if provided, inputs
            params: A list of theano shared variables.
                    These are the optimization variables
            inputs: (Optional) A list of theano variables
                    to serve as inputs to the graph.
            param_constrainers: (Optional) A list of callables
                    to be called on all updates dictionaries to
                    be applied to params. This is how you implement
                    constrained optimization.

            Calling the ``minimize'' method with values for
            for ``inputs'' will update ``params'' to minimize
            ``objective''.
        """

        self.max_iter = max_iter

        if inputs is None:
            inputs = []

        if param_constrainers is None:
            param_constrainers = []

        obj = objective

        self.verbose = False

        param_to_grad_sym = {}
        param_to_grad_shared = {}
        updates = {}

        self.params = [ param for param in params ]

        for param in params:
            grad = T.grad(objective, param)
            param_to_grad_sym[param] = grad
            grad_shared = sharedX( param.get_value() )
            param_to_grad_shared[param] = grad_shared
            updates[grad_shared] = grad

        self.param_to_grad_shared = param_to_grad_shared

        if self.verbose:
            print 'batch gradient class compiling gradient function'
        self._compute_grad = function(inputs, updates = updates )
        if self.verbose:
            print 'done'

        if self.verbose:
            print 'batch gradient class compiling objective function'
        self.obj = function(inputs, obj)
        if self.verbose:
            print 'done'

        self.param_to_cache = {}
        alpha = T.scalar(name = 'alpha')
        cache_updates = {}
        goto_updates = {}
        for param in params:
            self.param_to_cache[param] = sharedX(param.get_value(borrow=False))
            cache_updates[self.param_to_cache[param]] = param
            cached = self.param_to_cache[param]
            grad = self.param_to_grad_shared[param]
            if lr_scalers is not None and param in lr_scalers:
                scaled_alpha = alpha * lr_scalers[param]
            else:
                scaled_alpha = alpha
            mul = scaled_alpha * grad
            diff = cached - mul
            goto_updates[param] = diff
        self._cache_values = function([],updates = cache_updates)
        for param_constrainer in param_constrainers:
            param_constrainer(goto_updates)
        self._goto_alpha = function([alpha], updates = goto_updates)

    def _normalize_grad(self):

        n = sum([norm_sq(elem) for elem in self.param_to_grad_shared.values()])
        n = np.sqrt(n)

        for grad_shared in self.param_to_grad_shared.values():
            scale(grad_shared, 1./n)

        return n

    def minimize(self, * inputs ):

        if self.verbose:
            print 'minimizing'
        alpha_list = [ .001, .005, .01, .05, .1 ]

        orig_obj = self.obj(*inputs)

        if self.verbose:
            print orig_obj


        iters = 0
        while iters != self.max_iter:
            iters += 1
            best_obj, best_alpha, best_alpha_ind = self.obj( * inputs), 0., -1
            self._cache_values()
            self._compute_grad(*inputs)
            norm = self._normalize_grad()

            prev_best_obj = best_obj

            for ind, alpha in enumerate(alpha_list):
                self._goto_alpha(alpha)
                obj = self.obj(*inputs)
                if self.verbose:
                    print '\t',alpha,obj

                #Use <= rather than = so if there are ties
                #the bigger step size wins
                if obj <= best_obj:
                    best_obj = obj
                    best_alpha = alpha
                    best_alpha_ind = ind
                #end if obj
            #end for ind, alpha


            if self.verbose:
                print best_obj

            assert not np.isnan(best_obj)
            assert best_obj <= prev_best_obj
            self._goto_alpha(best_alpha)

            #if best_obj == prev_best_obj and alpha_list[0] < 1e-5:
            #    break

            if best_alpha_ind < 1 and alpha_list[0] > 3e-7:
                alpha_list = [ alpha / 3. for alpha in alpha_list ]
                if self.verbose:
                    print 'shrinking the step size'
            elif best_alpha_ind > len(alpha_list) -2:
                alpha_list = [ alpha * 2. for alpha in alpha_list ]
                if self.verbose:
                    print 'growing the step size'
            elif best_alpha_ind == -1 and alpha_list[0] <= 3e-7:
                if alpha_list[-1] > 1:
                    if self.verbose:
                        print 'converged'
                    break
                if self.verbose:
                    print 'expanding the range of step sizes'
                for i in xrange(len(alpha_list)):
                    for j in xrange(i,len(alpha_list)):
                        alpha_list[j] *= 1.5
                    #end for j
                #end for i
            #end check on alpha_ind
        #end while

        if norm > 1e-2:
            warnings.warn(str(norm)+" seems pretty big for a gradient at convergence...")

        return best_obj

@@ -622,19 +622,36 @@
 s = None
+, verbose = False
 ):%0A
-
         
@@ -1640,19 +1640,21 @@
 rbose = 
-Fal
+verbo
 se%0A%0A    
@@ -3423,24 +3423,136 @@
 o_updates)%0A%0A
+        if objective.dtype == %22float32%22:%0A            self.tol = 1e-6%0A        else:%0A            self.tol = 3e-7%0A%0A
     def _nor
@@ -5145,17 +5145,16 @@
   break%0A
-%0A
         
@@ -5199,20 +5199,24 @@
 st%5B0%5D %3E 
-3e-7
+self.tol
 :%0A      
@@ -5626,20 +5626,24 @@
 t%5B0%5D %3C= 
-3e-7
+self.tol
 :%0A      

Make BatchGradientDescent use lower tolerance for float32 objective