from itertools import izip

class HFModel(object):
    """
    This is a wrapper class around models we want to optimize with
    HF. The model instance we wrap only needs to be able to handle a
    single batch at a time, whereas this object will handle lists of
    batches. Generally, a batch contains both inputs and targets,
    although this class doesn't assume anything about that. The model
    that gets passed to HFModel needs to be able to deal with the
    contents of a batch. In the case of an RNN, a batch might be a
    triple of inputs, targets, and a mask showing which predictions
    contribute to the error. Or something even more exotic.
    """
    def __init__(self, model):
        self._model = model

    def getParamVect(self):
        return self._model.getParamVect()

    def setParamVect(self, vect):
        return self._model.setParamVect(vect)
    
    def gradEtc(self, batches, computeGradSquared = False):
        assert(len(batches) > 0)
        grad, gradSquared, error = self._model.gradEtc(batches[0], computeGradSquared)
        for b in batches[1:]:
            curGrad, curGradSquared, curError = self._model.gradEtc(b, computeGradSquared)
            grad += curGrad
            if computeGradSquared:
                gradSquared += curGradSquared
            error += curError
        return grad, gradSquared, error
    
    def loss(self, batches):
        return sum(self._model.loss(b) for b in batches)

    def getFpropState(self, batches):
        """
        The targets in each batch in batches are not used, but must be
        present (and can be anything).
        """
        return [self._model.getFpropState(b) for b in batches]
    
    def Gv(self, batches, vect, fpropStates = None):
        assert(len(batches) > 0)
        if fpropStates == None:
            fpropStates = [None for i in range(len(batches))]
        batchesAndStates = izip(batches, fpropStates)
        b, s = batchesAndStates.next()
        Gv = self._model.Gv(b, vect, s)
        for batch, fpropState in batchesAndStates:
            Gv += self._model.Gv(batch, vect, fpropState)
        return Gv
                                      
