package org.nd4j.linalg.lossfunctions.impl;

import lombok.EqualsAndHashCode;
import org.apache.commons.math3.util.Pair;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.lossfunctions.ILossFunction;
import org.nd4j.linalg.ops.transforms.Transforms;

/**
 * Created by susaneraly on 8/15/16.
 */
@EqualsAndHashCode
public class LossMSLE implements ILossFunction {
    public INDArray scoreArray(INDArray labels, INDArray preOutput, String activationFn, INDArray mask) {
        INDArray scoreArr;
        INDArray postOutput = Nd4j.getExecutioner().execAndReturn(Nd4j.getOpFactory().createTransform(activationFn, preOutput.dup()));
        scoreArr = Transforms.log(postOutput);
        scoreArr.subi(Transforms.log(labels));
        scoreArr = scoreArr.mul(scoreArr);
        if (mask != null) scoreArr.muliColumnVector(mask);
        return scoreArr;
    }

    @Override
    public double computeScore(INDArray labels, INDArray preOutput, String activationFn, INDArray mask, boolean average) {
        INDArray scoreArr = scoreArray(labels, preOutput, activationFn, mask);

        double score = scoreArr.sumNumber().doubleValue();

        if(average) score /= scoreArr.size(0);

        return score;
    }

    @Override
    public INDArray computeScoreArray(INDArray labels, INDArray preOutput, String activationFn, INDArray mask) {
        INDArray scoreArr = scoreArray(labels, preOutput, activationFn, mask);
        return scoreArr.sum(1);
    }

    @Override
    public INDArray computeGradient(INDArray labels, INDArray preOutput, String activationFn, INDArray mask) {
        INDArray postOutput = Nd4j.getExecutioner().execAndReturn(Nd4j.getOpFactory().createTransform(activationFn, preOutput.dup()));
        INDArray postOutDer = Nd4j.getExecutioner().execAndReturn(Nd4j.getOpFactory().createTransform(activationFn,preOutput.dup()).derivative());

        INDArray logyHat = Transforms.log(postOutput);
        INDArray gradients = logyHat.sub(Transforms.log(labels));
        gradients.muli(2);
        gradients.muli(postOutDer);
        gradients.divi(postOutput.addi(Nd4j.EPS_THRESHOLD));

        if(mask != null){
            gradients.muliColumnVector(mask);
        }

        return gradients;
    }

    @Override
    public org.apache.commons.math3.util.Pair<Double, INDArray> computeGradientAndScore(INDArray labels, INDArray preOutput, String activationFn, INDArray mask, boolean average) {
        //TODO: probably a more efficient way to do this...
        //Yes - will implement in round two. Just want to get done now.

        return new Pair<>(
                computeScore(labels, preOutput, activationFn, mask, average),
                computeGradient(labels, preOutput, activationFn, mask));
    }

    @Override
    public String toString(){
        return "LossMSLE()";
    }

}

@@ -554,37 +554,33 @@
        INDArray 
-postO
+o
 utput = Nd4j.get
@@ -714,65 +714,52 @@
 log(
-postOutput);%0A        scoreArr.subi(Transforms.log(labels)
+output.addi(1.0).divi(labels.add(1.0)),false
 );%0A 
@@ -788,16 +788,17 @@
 eArr.mul
+i
 (scoreAr
@@ -799,16 +799,37 @@
 coreArr)
+.divi(labels.size(1))
 ;%0A      
@@ -1654,21 +1654,17 @@
 NDArray 
-postO
+o
 utput = 
@@ -1785,26 +1785,27 @@
 NDArray 
-postOutDer
+sigmaPrimeZ
  = Nd4j.
@@ -1938,233 +1938,213 @@
 ray 
-logyHat = Transforms.log(postOutput);%0A        INDArray gradients = logyHat.sub(Transforms.log(labels));%0A        gradients.muli(2);%0A        gradients.muli(postOutDer);%0A        gradients.divi(postOutput.addi(Nd4j.EPS_THRESHOLD)
+p1 = output.add(1.0);%0A        INDArray gradients = sigmaPrimeZ.divi(p1).muli(2.0/labels.size(1));%0A        INDArray logRatio = Transforms.log(p1.divi(labels.add(1.0)),false);%0A        gradients.muli(logRatio
 );%0A%0A

Fix LossMSLE score and gradient