import conjugateGrad as cg
import gnumpy as gnp
import numpy as num
import sys

from query_memUsage import memUsage

#the interface to the model:
#the model must know how to get gradEtc (gradient, gradient squared, loss)
#the model must also be able to compute Gv and finally
#the model must compute loss by itself
#####@TODO: rethink this part of the interface
#model.getParamVect() must return a copy of the model parameters as a single packed monolithic vector
#model.setParamVect(W) must set the model parameters to be equal to
#  the vector W (not a reference to W) which may involve unpacking W

class ShouldNeverHappen(Exception):
    pass

def dampingHeuristic(damping, rho):
    if rho < 0.25:
        return 3*damping/2.
    if rho > 0.75:
        return 2*damping/3.
    return damping

class HF(object):
    def __init__(self, log, initialDamping, useDiagPrecond = True):
        self.damping = initialDamping
        self.log = log

        self.cgStartVect = None #can't be set until we have a model

        self.cacheFpropState = True
        
        #this needs to be adapted, but how?!
        #we will let people adapt it between HF iterations outside the optimizer, we just give up
        self.maxCGIters = 250 #Ilya says 200-300 is OK on curves if you use the diagonal preconditioner
        
        #preconditioning stuff
        self.precondPower = 0.75
        self.useDiagPrecond = useDiagPrecond
        
        #hopefully none of these settings will need adjustment
        self.cgStopTol = 5e-4
        self.cgGapRatio = 0.1
        self.minCGTestGap = 10
        self.minCGIters = 1 #meaningless if less than self.minCGTestGap
        self.doCGBacktracking = True
        self.cgBacktrackTestInterval = 10 #meaningless if not self.doCGBacktracking
        self.cgResidualRefresh = None #seems to work well to never refresh the residual in CG
        
        self.lineSearchShrinkFactor = 0.8
        self.maxLineSearchShrinks = 50 #any more than this we just reject the update
        
        #factor we shrink the previous final CG vector before using it to initialize the next CG run
        self.cgStartVectDecay = 0.95
        
    def lineSearch(self, model, W, searchDirection, lineSearchBatch, lineSearchBatchSize):
        """
        Perform a backtracking line-search. In other words, compare the
        loss at W+a**j*searchDirection (for 0<a<1 and j=0,1,2...) to
        the loss at W and increase j (shrink the update vector) until
        we beat the starting loss.
        
        Nota Bene: We assume that model.getParamVect() == W when this method gets called.
        """
        origLoss = model.loss(lineSearchBatch)/lineSearchBatchSize
        normSearchDir = num.sqrt((searchDirection*searchDirection).sum())
        print >>self.log, "Starting line search, |dir|=%f" % (normSearchDir)
        for i in range(self.maxLineSearchShrinks):
            newW = W + self.lineSearchShrinkFactor**i*searchDirection
            model.setParamVect(newW)
            newLoss = model.loss(lineSearchBatch)/lineSearchBatchSize
            print >>self.log, "line search: dist=%f, loss=%f" % (self.lineSearchShrinkFactor**i, newLoss)
            if newLoss < origLoss:
                print >>self.log, "Line search complete"
                #at this point, model.getParamVect() == newW
                return newW
        print >>self.log, "Line search failed to reduce loss, restoring old weights and zeroing self.cgStartVect"
        #line below not needed since model.setParamVect(W) gets called on our return value
        #model.setParamVect(W) #restore the original weights, we have rejected this search direction
        self.cgStartVect = gnp.zeros(W.shape)
        return W
    
    def getSearchDirection(self, cgStartVect, gradient, model, curvatureBatch, curvatureBatchSize, \
                           applyPreconditioner, backtrackBatch, backtrackBatchSize = None):
        """
        This method finds a new search direction using CG. It returns
        the search direction, the reduction in the CG cost function,
        and the final CG vector, which may be different than the
        search direction because of CG backtracking.
        
        The way we do CG backtracking is different from Ilya's
        implementation. Ours will be greedier, but his will be more
        likely to take a later CG vector.
        """
        if backtrackBatchSize == None:
            backtrackBatchSize = curvatureBatchSize
        cgStoppingCriterion = cg.CGStoppingCriterion(self.maxCGIters, self.minCGIters, self.cgStopTol, \
                                                      self.cgGapRatio, self.minCGTestGap, 1e-10, log = self.log)
        b = -gradient
        curveFpropState = None
        if self.cacheFpropState:
            print >>self.log, "Caching fprop state on curvature batch."
            curveFpropState = model.getFpropState(curvatureBatch)
        A = lambda v: model.Gv(curvatureBatch, v, curveFpropState)/curvatureBatchSize + self.damping*v
        x0 = cgStartVect
        invM = applyPreconditioner
        cgIters = cg.precondConjGrad(A, b, x0, invM, cgStoppingCriterion, self.cgResidualRefresh, log = self.log)
        bestXCGCost = cg.CGObjectiveFromScratch(A, b, x0) #@TODO: inefficient, this will be computed by CG
        if self.doCGBacktracking:
            btrackBatch_grad, unused_gradSq, unused_loss = model.gradEtc(backtrackBatch, False)
            btrackBatch_grad /= backtrackBatchSize
            bb = -btrackBatch_grad
            btackFpropState = None
            #We no longer allow caching fprop state on the
            #backtracking batch since it isn't useful enough to make a
            #separate option
            #if self.cacheFpropState:
            #    btackFpropState = model.getFpropState(backtrackBatch)
            AA = lambda v: model.Gv(backtrackBatch, v, btackFpropState)/backtrackBatchSize + self.damping*v
            backtrackDesirabiliy = lambda xx : -cg.CGObjectiveFromScratch(AA, bb, xx)
            bestX = x0 #the one we backtrack to
            bestXQuality = backtrackDesirabiliy(x0)
        btrackId = -1
        for j,x in enumerate(cgIters):
            self.totalCGSteps += 1
            if self.doCGBacktracking and \
                   j % self.cgBacktrackTestInterval == (self.maxCGIters-1)%self.cgBacktrackTestInterval:
                newXQuality = backtrackDesirabiliy(x)
                if newXQuality >= bestXQuality:
                    bestX = x.copy()
                    bestXQuality = newXQuality
                    bestXCGCost = cgStoppingCriterion.CGCosts[-1]
                    btrackId = j
        if not self.doCGBacktracking:
            bestX = x.copy()
            bestXCGCost = cgStoppingCriterion.CGCosts[-1]
        finalCGVect = x.copy()
        #the way we compute CGCost doesn't include a constant term so
        #CGCost(0)==0 therefore cgCostReduction is just 0-bestXCGCost
        cgCostReduction = -bestXCGCost #used later to compute rho
        print >>self.log, "CGBackTrackID=", j
        return bestX, cgCostReduction, finalCGVect
    
    def buildDiagPreconditioner(self, g2, otherPrecondDampingTerm):
        #otherPrecondDampingTerm should be something like the L2 cost for the model
        MAsVect = (g2 + self.damping + otherPrecondDampingTerm)**self.precondPower
        invM = lambda v: v/MAsVect
        return invM
    
    def computeRho(self, searchDirection, model, W, gradBatch, gradBatchSize, cgCostReduction):
        """
        We do NOT include damping in the computation of rho so we have
        to subtract it out of cgCostReduction.
        """
        origLoss = model.loss(gradBatch)/gradBatchSize
        model.setParamVect(W+searchDirection)
        newLoss = model.loss(gradBatch)/gradBatchSize
        model.setParamVect(W)
        lossReduction = origLoss - newLoss
        #remove the damping term that is in cgCostReduction
        cgCostReduction += 0.5*self.damping*(searchDirection*searchDirection).sum()
        if cgCostReduction == 0:
            raise ShouldNeverHappen("How did CG stop early enough to make no reduction in the CG cost?")
        if lossReduction < 0:
            return -num.inf
        print >>self.log, "rho_numer=", lossReduction, "rho_denom=", cgCostReduction
        rho = lossReduction/cgCostReduction
        return rho
    
    def optimize(self, model, data, iters, otherPrecondDampingTerm, curIter = 0):
        """
        otherPrecondDampingTerm should be something like the L2 cost for the model.
        """
        self.totalCGSteps = 0
        W = model.getParamVect() #reconsider this interface?
        self.cgStartVect = gnp.zeros(W.shape)
        print >>self.log, "Memory usage before entering main optimizer loop: %f (MB)" % memUsage()
        for i in range(curIter, iters):
            print >>self.log, "*HF* %d: damping=%.9f, totalCGSteps=%d" % (i, self.damping, self.totalCGSteps)
            sepCGBtrackBatch = self.doCGBacktracking #for now we will always use separate data if we do backtracking
            
            print >>self.log, "Memory usage before obtaining gradBatch for HF iter %d: %f (MB)" % (i, memUsage())
            gradBatch = data.obtainBatch(data.gradBatchSize)
            print >>self.log, "Memory usage after obtaining gradBatch for HF iter %d: %f (MB)" % (i, memUsage())
            g, g2, loss = model.gradEtc(gradBatch, self.useDiagPrecond)
            
            self.log.flush()
            if self.useDiagPrecond:
                g2 = g2/data.gradBatchSize
            g, loss = g/data.gradBatchSize, loss/data.gradBatchSize
            self.latestLoss = loss
            print >>self.log, ">HF %d: loss = %f" % (i, loss)
            print >>self.log, ">HF %d: |g|=%f" % (i, num.sqrt((g*g).sum()))
            
            #create the preconditioning function for CG
            #  at the moment, we will assume that we are using the diagonal preconditioner or nothing
            invM = lambda vv : vv
            if self.useDiagPrecond:
                invM = self.buildDiagPreconditioner(g2, otherPrecondDampingTerm) #this is where we need the gradient squared

            print >>self.log, "Memory usage before obtaining curvatureBatch for HF iter %d: %f (MB)" % (i, memUsage())
            curvatureBatch = data.obtainBatch(data.curvatureBatchSize)
            cgBacktrackBatch = data.obtainBatch(data.curvatureBatchSize) if sepCGBtrackBatch else curvatureBatch
            print >>self.log, "Memory usage before getSearchDirection for HF iter %d: %f (MB)" % (i, memUsage())
            searchDirection, cgCostReduction, finalCGVect = self.getSearchDirection(self.cgStartVect, g, model, curvatureBatch,\
                                                                                    data.curvatureBatchSize, invM, cgBacktrackBatch)
            curvatureBatch = None #allow garbage collection for curvatureBatch
            print >>self.log, "Memory usage after getSearchDirection for HF iter %d: %f (MB)" % (i, memUsage())
            
            self.rho = self.computeRho(searchDirection, model, W, gradBatch, data.gradBatchSize, cgCostReduction)
            gradBatch = None #allow garbage collection for gradBatch
            print >>self.log, "rho =", self.rho
            
            prevDamping = self.damping
            self.damping = dampingHeuristic(self.damping, self.rho)
            print >>self.log, ">HF %d: damping=%f (was %f)" % (i, self.damping, prevDamping)
            print >>self.log, ">HF %d: rho=%f " % (i, self.rho)
            
            self.cgStartVect = self.cgStartVectDecay*finalCGVect
            
            #line search may potentially zero-out self.cgStartVect
            lineSearchBatch = data.obtainBatch(data.lineSearchBatchSize)
            W = self.lineSearch(model, W, searchDirection, lineSearchBatch, data.lineSearchBatchSize)
            lineSearchBatch = None #allow garbage collection for gradBatch
            
            model.setParamVect(W)
            yield W
