package edu.neu.ccs.pyramid.multilabel_classification.imlgb;

import edu.neu.ccs.pyramid.dataset.MultiLabel;
import edu.neu.ccs.pyramid.dataset.MultiLabelClfDataSet;
import edu.neu.ccs.pyramid.regression.IsotonicRegression;
import org.apache.mahout.math.Vector;
import scala.Serializable;

import java.util.stream.IntStream;

public class IMLGBIsotonicScaling implements Serializable{
    private static final long serialVersionUID = 1L;
    private IsotonicRegression isotonicRegression;
    private IMLGradientBoosting boosting;

    public IMLGBIsotonicScaling(IMLGradientBoosting boosting, MultiLabelClfDataSet multiLabelClfDataSet) {
//        System.out.println("calibrating with isotonic regression");
        this.boosting = boosting;
        double[] locations = new double[multiLabelClfDataSet.getNumDataPoints()];
        double[] binaryLabels = new double[multiLabelClfDataSet.getNumDataPoints()];

        SubsetAccPredictor predictor = new SubsetAccPredictor(boosting);
        IntStream.range(0, multiLabelClfDataSet.getNumDataPoints()).parallel()
                .forEach(i->{
                    MultiLabel pre = predictor.predict(multiLabelClfDataSet.getRow(i));
                    double score = boosting.predictAssignmentProbWithConstraint(multiLabelClfDataSet.getRow(i),pre);
                    locations[i] = score;
                    if (pre.equals(multiLabelClfDataSet.getMultiLabels()[i])){
                        binaryLabels[i] = 1;
                    } else {
                       binaryLabels[i] = 0;
                    }
                });
        isotonicRegression = new IsotonicRegression(locations, binaryLabels);
//        System.out.println("calibration done");
    }

    public double calibratedProb(Vector vector, MultiLabel multiLabel){
        double uncalibrated = boosting.predictAssignmentProbWithConstraint(vector, multiLabel);
        return isotonicRegression.predict(uncalibrated);
    }

    public double calibratedProb(double uncalibratedProb){
        return isotonicRegression.predict(uncalibratedProb);
    }
}

@@ -283,16 +283,39 @@
 zable;%0A%0A
+import java.util.List;%0A
 import j
@@ -759,16 +759,85 @@
 osting;%0A
+        List%3CMultiLabel%3E allAssignments = boosting.getAssignments();%0A
         
@@ -907,16 +907,38 @@
 Points()
+*allAssignments.size()
 %5D;%0A     
@@ -1014,16 +1014,38 @@
 Points()
+*allAssignments.size()
 %5D;%0A%0A    
@@ -1246,108 +1246,22 @@
     
-MultiLabel pre = predictor.predict(multiLabelClfDataSet.getRow(i));%0A                    double score
+double%5B%5D probs
  = b
@@ -1268,32 +1268,35 @@
 oosting.predictA
+llA
 ssignmentProbWit
@@ -1284,32 +1284,33 @@
 llAssignmentProb
+s
 WithConstraint(m
@@ -1342,57 +1342,142 @@
 w(i)
-,pre);%0A                    locations%5Bi%5D = score;%0A
+);%0A                    for (int a=0;a%3Cprobs.length;a++)%7B%0A                        locations%5Bi*allAssignments.size()+a%5D = probs%5Ba%5D;%0A    
     
@@ -1496,19 +1496,37 @@
     if (
-pre
+allAssignments.get(a)
 .equals(
@@ -1585,32 +1585,36 @@
                 
+    
 binaryLabels%5Bi%5D 
@@ -1615,15 +1615,43 @@
 ls%5Bi
+*allAssignments.size()+a
 %5D = 1;%0A
+    
     
@@ -1690,32 +1690,37 @@
                 
+   
+  
 binaryLabels%5Bi%5D 
@@ -1721,14 +1721,64 @@
 ls%5Bi
-%5D = 0;
+*allAssignments.size()+a%5D = 0;%0A                        %7D
 %0A   

set all set probabilities as input to isotonic regression