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 edu.neu.ccs.pyramid.util.Pair;
import org.apache.mahout.math.Vector;

import java.io.Serializable;
import java.util.*;
import java.util.stream.IntStream;
import java.util.stream.Stream;

public class CardinalityCalibrator implements Serializable{
    private static final long serialVersionUID = 1L;
    Map<Integer,IsotonicRegression> calibrations;
    private IMLGradientBoosting boosting;

    public CardinalityCalibrator(IMLGradientBoosting boosting, MultiLabelClfDataSet multiLabelClfDataSet) {
        this.calibrations = new HashMap<>();
        List<MultiLabel> support = boosting.getAssignments();
        Set<Integer> cardinalities = new HashSet<>();
        for (MultiLabel multiLabel: support){
            cardinalities.add(multiLabel.getNumMatchedLabels());
        }

        for (int cardinality: cardinalities){
            List<MultiLabel> allAssignments = boosting.getAssignments();
            Stream<Pair<Double,Integer>> stream =  IntStream.range(0, multiLabelClfDataSet.getNumDataPoints()).parallel()
                    .boxed().flatMap(i-> {
                        double[] probs = boosting.predictAllAssignmentProbsWithConstraint(multiLabelClfDataSet.getRow(i));
                        Stream<Pair<Double,Integer>> pairs = IntStream.range(0, probs.length)
                                .filter(a->allAssignments.get(a).getNumMatchedLabels()==cardinality)
                                .mapToObj(a -> {
                            Pair<Double, Integer> pair = new Pair<>();
                            pair.setFirst(probs[a]);
                            pair.setSecond(0);
                            if (allAssignments.get(a).equals(multiLabelClfDataSet.getMultiLabels()[i])) {
                                pair.setSecond(1);
                            }
                            return pair;
                        });
                        return pairs;
                    });
            calibrations.put(cardinality, new IsotonicRegression(stream));
        }



//        System.out.println("calibrating with isotonic regression");
        this.boosting = boosting;

    }



    public double calibrate(double uncalibrated, int cardinality){
        return calibrations.get(cardinality).predict(uncalibrated);
    }


    public double calibratedProb(Vector vector, MultiLabel multiLabel){
        double uncalibrated = boosting.predictAssignmentProbWithConstraint(vector, multiLabel);
        int cardinality = multiLabel.getNumMatchedLabels();
        return calibrations.get(cardinality).predict(uncalibrated);
    }


}

@@ -217,16 +217,56 @@
 ession;%0A
+import edu.neu.ccs.pyramid.util.ArgMin;%0A
 import e
@@ -2469,43 +2469,395 @@
     
-return calibrations.get(cardinality
+//deal with unseen cardinality%0A        List%3CInteger%3E cards = new ArrayList%3C%3E(calibrations.keySet());%0A        Collections.sort(cards);%0A        double%5B%5D diff = new double%5Bcards.size()%5D;%0A        for (int i=0;i%3Ccards.size();i++)%7B%0A            diff%5Bi%5D = Math.abs(cardinality-cards.get(i));%0A        %7D%0A        int closest = cards.get(ArgMin.argMin(diff));%0A        return calibrations.get(closest
 ).pr
@@ -3140,50 +3140,35 @@
 brat
-ions.get(cardinality).predict(uncalibrated
+e(uncalibrated, cardinality
 );%0A 

deal with new cardinality