package org.nd4j.linalg.learning;

import lombok.Data;
import lombok.NoArgsConstructor;
import org.nd4j.linalg.api.buffer.DoubleBuffer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;

import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;

/**
 * Nesterov's momentum.
 * Keep track of the previous layer's gradient
 * and use it as a way of updating the gradient.
 *
 * @author Adam Gibson
 */
@Data
@NoArgsConstructor
public class Nesterovs implements Serializable,GradientUpdater {
    private double momentum = 0.5;
    protected Map<Integer,Double> momentumAfter = new HashMap<>();
    private INDArray v;
    private double lr = 0.1;


    public Nesterovs(double momentum, Map<Integer,Double> momentumAfter, double lr) {
        this.momentum = momentum;
        this.momentumAfter = momentumAfter;
        this.lr = lr;
    }

    public Nesterovs(double momentum, double lr) {
        this.momentum = momentum;
        this.lr = lr;
    }

    public Nesterovs(double momentum, Map<Integer,Double> momentumAfter) {
        this.momentum = momentum;
        this.momentumAfter = momentumAfter;
    }

    public Nesterovs(double momentum) {
        this.momentum = momentum;

    }

    /**
     * Get the nesterov update
     * @param gradient the gradient to get the update for
     * @param iteration
     * @return
     */
    @Override
    public INDArray getGradient(INDArray gradient, int iteration) {
        if(v == null)
            v = Nd4j.zeros(gradient.shape());
        if(momentumAfter !=null)
            momentum = (momentumAfter.containsKey(iteration)) ? momentumAfter.get(iteration) : momentum;
        INDArray vPrev = v;
        v = vPrev.mul(momentum).subi(gradient.mul(lr));
        //reference https://cs231n.github.io/neural-networks-3/#sgd 2nd equation
        //DL4J default is negative step function thus we flipped the signs:
        // x += mu * v_prev + (-1 - mu) * v
        //i.e., we do params -= updatedGradient, not params += updatedGradient
        
        INDArray ret = vPrev.muli(momentum).addi(v.mul(-momentum - 1));
        return ret;
    }


}

@@ -573,75 +573,8 @@
 .5;%0A
-    protected Map%3CInteger,Double%3E momentumAfter = new HashMap%3C%3E();%0A
     
@@ -626,467 +626,117 @@
 .1;%0A
-%0A%0A
     
-public Nesterovs(double momentum, Map%3CInteger,Double%3E momentumAfter, double lr) %7B%0A        this.momentum = momentum;%0A        this.momentumAfter = momentumAfter;%0A        this.lr = lr;%0A    %7D%0A%0A    public Nesterovs(double momentum, double lr) %7B%0A        this.momentum = momentum;%0A        this.lr = lr;%0A    %7D%0A%0A    public Nesterovs(double momentum, Map%3CInteger,Double%3E momentumAfter) %7B%0A        this.momentum = momentum;%0A        this.momentumAfter = momentumAfte
+%0A    public Nesterovs(double momentum, double lr) %7B%0A        this.momentum = momentum;%0A        this.lr = l
 r;%0A 
@@ -1121,146 +1121,8 @@
 ));%0A
-        if(momentumAfter !=null)%0A            momentum = (momentumAfter.containsKey(iteration)) ? momentumAfter.get(iteration) : momentum;%0A
     

Move momentumAfter back into dl4j