package model.supervised.boosting.gradiantboost;

import data.DataSet;
import data.core.Label;
import model.Predictable;
import model.Trainable;
import model.supervised.boosting.Boost;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import performance.ClassificationEvaluator;
import performance.Evaluator;
import utils.random.RandomUtils;
import utils.sort.SortIntDoubleUtils;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;

/**
 * Created by hanxuan on 11/3/15 for machine_learning.
 */
public class GradientBoostClassification implements Predictable, Trainable, Boost {

    public static boolean NEED_REPORT = false;

    public static int MAX_THREADS = 4;

    private static Logger log = LogManager.getLogger(GradientBoostClassification.class);

    private DataSet trainData = null;

    private DataSet testData = null;

    private GradientBoostRegression[] regressions = null;

    private int classCount = -1;

    private ExecutorService service = null;

    private CountDownLatch countDownLatch = null;

    private double kl = -1;

    private double testError = -1;

    private double trainError = -1;

    @Override
    public double predict(double[] feature) {

        double[] probs = new double[classCount];
        for (int i = 0; i < classCount; i++) {
            probs[i] = regressions[i].predict(feature);
        }
        int[] indexes = RandomUtils.getIndexes(classCount);
        SortIntDoubleUtils.sort(indexes, probs);
        return indexes[classCount - 1];
    }

    @Override
    public void train() {

        service = Executors.newFixedThreadPool(MAX_THREADS);
        countDownLatch = new CountDownLatch(classCount);
        IntStream.range(0, classCount).forEach(i -> service.submit(() -> {
            try {
                doTask(i);
            }catch (Throwable t) {
                log.error(t.getMessage(), t);
            }
            log.info("task: {}/{} finished...", i, classCount);
            countDownLatch.countDown();
        }));
        try {
            TimeUnit.SECONDS.sleep(10);
            countDownLatch.await();
        }catch (Throwable t) {
            log.error(t.getMessage(), t);
        }
        service.shutdown();

        log.info("GradientBoostClassification training finished ...");

        if(NEED_REPORT) {
            statisticReport();
            printRoundReport();
        }
    }

    @Override
    public void initialize(DataSet d) {
        this.trainData = d;
        classCount = d.getLabels().getIndexClassMap().size();
    }


    @Override
    public void boostConfig(int iteration, String boosterClassName, Evaluator evaluator, DataSet testData) throws Exception {
        regressions = new GradientBoostRegression[classCount];
        for (int i = 0; i < classCount; i++) {
            regressions[i] = new GradientBoostRegression();
            regressions[i].boostConfig(iteration, boosterClassName, evaluator, testData);
        }
        this.testData = testData;
        log.info("GradientBoostClassification config finished, classCount {}, regressor iteration {}", classCount, iteration);
    }

    @Override
    public void printRoundReport() {
        log.info("============ report ============");
        log.info("KL-divergence {}", kl);
        log.info("test error {}", testError);
        log.info("train error {}", trainError);
        log.info("============ ====== ============");
    }

    private void doTask(int classId) {

        Label classLabel = makeClassLabel(classId);
        DataSet classData = new DataSet(trainData.getFeatureMatrix(), classLabel);
        regressions[classId].initialize(classData);
        regressions[classId].train();
    }

    private Label makeClassLabel(int classId) {

        float[] labels = new float[trainData.getInstanceLength()];
        for (int i = 0; i < labels.length; i++) {
            if (trainData.getLabel(i) == classId) {
                labels[i] = 1.0F;
            }
        }
        return new Label(labels, trainData.getLabels().getClassIndexMap());
    }

    private void statisticReport() {

        ClassificationEvaluator evaluator = new ClassificationEvaluator();
        evaluator.initialize(trainData, this);
        evaluator.getPredictLabel();
        trainError = 1 - evaluator.evaluate();

        evaluator.initialize(testData, this);
        evaluator.getPredictLabel();
        testError = 1 - evaluator.evaluate();
    }

//    private void klCalculation() {
//
//    }
}

@@ -332,24 +332,57 @@
 .Evaluator;%0A
+import utils.array.ArraySumUtil;%0A
 import utils
@@ -1393,24 +1393,267 @@
  feature) %7B%0A
+        double%5B%5D probs = probs(feature);%0A        int%5B%5D indexes = RandomUtils.getIndexes(classCount);%0A        SortIntDoubleUtils.sort(indexes, probs);%0A        return indexes%5BclassCount - 1%5D;%0A    %7D%0A%0A    public double%5B%5D probs (double%5B%5D feature) %7B
 %0A        dou
@@ -1815,147 +1815,44 @@
     
-int%5B%5D indexes = RandomUtils.getIndexes(classCount);%0A        SortIntDoubleUtils.sort(indexes, probs);%0A        return indexes%5BclassCount - 1%5D
+return ArraySumUtil.normalize(probs)
 ;%0A  

add probs predict