package org.nd4j.linalg.learning;

import lombok.Data;
import lombok.NoArgsConstructor;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.api.shape.Shape;

import java.io.Serializable;

/**
 * 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 {
    public static final double DEFAULT_NESTEROV_MOMENTUM = 0.9;

    private double momentum = DEFAULT_NESTEROV_MOMENTUM;
    private volatile INDArray v;
    private double learningRate = 0.1;

    @Override
    public int stateSizeForInputSize(int inputSize) {
        return inputSize;
    }

    @Override
    public void setStateViewArray(INDArray viewArray, int[] gradientShape, char gradientOrder, boolean initialize) {
        if (!viewArray.isRowVector())
            throw new IllegalArgumentException("Invalid input: expect row vector input");
        if (initialize)
            viewArray.assign(0);
        this.v = viewArray;

        //Reshape to match the expected shape of the input gradient arrays
        this.v = Shape.newShapeNoCopy(this.v, gradientShape, gradientOrder == 'f');
        if (v == null)
            throw new IllegalStateException("Could not correctly reshape gradient view array");
    }

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

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

    }

    @Override
    public void update(Object... args) {
        if (args.length > 0) {
            learningRate = (Double) args[0];
            momentum = (Double) args[1];
        }

    }


    /**
     * 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)
            throw new IllegalStateException("Updater has not been initialized with view state");

        INDArray vPrev = v;
        v = vPrev.mul(momentum).subi(gradient.mul(learningRate));
        //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));
        gradient.assign(ret);

        return gradient;
    }

    @Override
    public GradientUpdaterAggregator getAggregator(boolean addThis) {
        NesterovsAggregator ag = new NesterovsAggregator();
        if (addThis)
            ag.aggregate(this);
        return ag;
    }

    public static class NesterovsAggregator implements GradientUpdaterAggregator {
        private INDArray vSum;
        private double lrSum;
        private double momentumSum;
        private int count = 0;

        @Override
        public GradientUpdater getUpdater() {
            Nesterovs nesterovs = new Nesterovs(momentumSum / count, lrSum / count);
            nesterovs.setV(vSum.div(count));
            return nesterovs;
        }

        @Override
        public void aggregate(GradientUpdater updater) {
            if (!(updater instanceof Nesterovs))
                throw new UnsupportedOperationException("Cannot aggregate Nesterovs with updater: " + updater);
            Nesterovs nesterovs = (Nesterovs) updater;
            if (vSum == null) {
                vSum = nesterovs.v.dup();
                lrSum = nesterovs.learningRate;
                momentumSum = nesterovs.momentum;
            } else {
                vSum.addi(nesterovs.v);
                lrSum += nesterovs.learningRate;
                momentumSum += nesterovs.momentum;
            }
            count++;
        }

        @Override
        public GradientUpdaterAggregator combine(GradientUpdaterAggregator other) {
            if (!(other instanceof NesterovsAggregator))
                throw new IllegalArgumentException("Cannot combine NesterovsAggregator with aggregator: " + other);
            NesterovsAggregator aggregator = (NesterovsAggregator) other;
            vSum.addi(aggregator.vSum);
            lrSum += aggregator.lrSum;
            momentumSum += aggregator.momentumSum;
            count += aggregator.count;
            return this;
        }
    }
}

@@ -157,19 +157,121 @@
 api.
-shape.Shape
+ops.impl.transforms.arithmetic.AddOp;%0Aimport org.nd4j.linalg.api.shape.Shape;%0Aimport org.nd4j.linalg.factory.Nd4j
 ;%0A%0Ai
@@ -2246,102 +2246,8 @@
 );%0A%0A
-        INDArray vPrev = v;%0A        v = vPrev.mul(momentum).subi(gradient.mul(learningRate));%0A
     
@@ -2535,99 +2535,432 @@
     
-INDArray ret = vPrev.muli(momentum).addi(v.mul(-momentum - 1));%0A        gradient.assign(ret
+//v = mu * v - lr * gradient%0A        INDArray vPrev = v.dup();%0A        v.muli(momentum).subi(gradient.mul(learningRate));              //Modify state array in-place%0A%0A        /*%0A        Next line is equivalent to:%0A        INDArray ret = vPrev.muli(momentum).addi(v.mul(-momentum - 1));%0A        gradient.assign(ret);%0A        */%0A        Nd4j.getExecutioner().exec(new AddOp(vPrev.muli(momentum), v.mul(-momentum - 1), gradient)
 );%0A%0A

Fix nesterov state array