package biz.k11i.xgboost.learner;

import java.util.HashMap;
import java.util.Map;

/**
 * Objective function implementations.
 */
public class ObjFunction {
    private static final Map<String, ObjFunction> FUNCTIIONS;

    static {
        ObjFunction nop = new ObjFunction();
        ObjFunction regLogistic = new RegLossObjLogistic();

        FUNCTIIONS = new HashMap<>();
        FUNCTIIONS.put("binary:logistic", regLogistic);
        FUNCTIIONS.put("binary:logitraw", nop);
        FUNCTIIONS.put("multi:softmax", new SoftmaxMultiClassObjClassify());
        FUNCTIIONS.put("multi:softprob", new SoftmaxMultiClassObjProb());
    }

    /**
     * Gets {@link ObjFunction} from given name.
     *
     * @param name name of objective function
     * @return objective function
     */
    public static ObjFunction fromName(String name) {
        ObjFunction result = FUNCTIIONS.get(name);
        if (result == null) {
            throw new IllegalArgumentException(name + " is not supported objective function.");
        }
        return result;
    }

    /**
     * Transforms prediction values.
     *
     * @param preds prediction
     * @return transformed values
     */
    public double[] predTransform(double[] preds) {
        // do nothing
        return preds;
    }

    /**
     * Logistic regression.
     */
    static class RegLossObjLogistic extends ObjFunction {
        @Override
        public double[] predTransform(double[] preds) {
            for (int i = 0; i < preds.length; i++) {
                preds[i] = sigmoid(preds[i]);
            }
            return preds;
        }

        double sigmoid(double x) {
            return (1 / (1 + Math.exp(-x)));
        }
    }

    /**
     * Multiclass classification.
     */
    static class SoftmaxMultiClassObjClassify extends ObjFunction {
        @Override
        public double[] predTransform(double[] preds) {
            int maxIndex = 0;
            double max = preds[0];
            for (int i = 1; i < preds.length; i++) {
                if (max < preds[i]) {
                    maxIndex = i;
                    max = preds[i];
                }
            }

            return new double[]{maxIndex};
        }
    }

    /**
     * Multiclass classification (predicted probability).
     */
    static class SoftmaxMultiClassObjProb extends ObjFunction {
        @Override
        public double[] predTransform(double[] preds) {
            double max = preds[0];
            for (int i = 1; i < preds.length; i++) {
                max = Math.max(preds[i], max);
            }

            double sum = 0;
            for (int i = 0; i < preds.length; i++) {
                preds[i] = Math.exp(preds[i] - max);
                sum += preds[i];
            }

            for (int i = 0; i < preds.length; i++) {
                preds[i] /= (float) sum;
            }

            return preds;
        }
    }
}

@@ -207,20 +207,37 @@
 %3E FUNCTI
-I
 ONS
+ = new HashMap%3C%3E()
 ;%0A%0A    s
@@ -256,78 +256,35 @@
     
-ObjFunction nop = new ObjFunction();%0A        ObjFunction regL
+register(%22binary:l
 ogistic
- =
+%22,
  new
@@ -308,175 +308,83 @@
 ic()
-;%0A%0A        FUNCTIIONS = new HashMap%3C%3E();%0A        FUNCTIIONS.put(%22binary:logistic%22, regLogistic);%0A        FUNCTIIONS.put(%22binary:logitraw%22, nop);%0A        FUNCTIIONS.put
+);%0A        register(%22binary:logitraw%22, new ObjFunction());%0A        register
 (%22mu
@@ -446,22 +446,16 @@
     
-FUNCTIIONS.put
+register
 (%22mu
@@ -750,17 +750,16 @@
 = FUNCTI
-I
 ONS.get(
@@ -923,32 +923,324 @@
  result;%0A    %7D%0A%0A
+    /**%0A     * Register an %7B@link ObjFunction%7D for a given name.%0A     *%0A     * @param name name of objective function%0A     * @param objFunction objective function%0A     */%0A    public static void register(String name, ObjFunction objFunction) %7B%0A        FUNCTIONS.put(name, objFunction);%0A    %7D%0A%0A
     /**%0A     * T

Allow users to register their ObjFunction