package edu.neu.ccs.pyramid.calibration;

import edu.neu.ccs.pyramid.dataset.MultiLabel;
import edu.neu.ccs.pyramid.multilabel_classification.DynamicProgramming;
import edu.neu.ccs.pyramid.multilabel_classification.MultiLabelClassifier;
import edu.neu.ccs.pyramid.multilabel_classification.cbm.BMDistribution;
import edu.neu.ccs.pyramid.multilabel_classification.cbm.CBM;
import edu.neu.ccs.pyramid.util.Pair;
import org.apache.mahout.math.Vector;

import java.util.*;
import java.util.stream.Collectors;


/**
 * find top K sets by calibrated probability
 */
public class TopKFinder {

    public static List<Pair<MultiLabel,Double>> topK(Vector x, MultiLabelClassifier.ClassProbEstimator classProbEstimator, LabelCalibrator labelCalibrator,
                                                     VectorCalibrator vectorCalibrator, PredictionFeatureExtractor predictionFeatureExtractor,
                                                     int top){
        List<Pair<MultiLabel,Double>> list = new ArrayList<>();
        double[] marginals = labelCalibrator.calibratedClassProbs(classProbEstimator.predictClassProbs(x));
        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);

        for (int i=0;i<top;i++) {
            MultiLabel candidate = dynamicProgramming.nextHighestVector();
            PredictionCandidate predictionCandidate = new PredictionCandidate();
            predictionCandidate.x = x;
            predictionCandidate.labelProbs = marginals;
            predictionCandidate.multiLabel = candidate;

            Vector feature = predictionFeatureExtractor.extractFeatures(predictionCandidate);
            double pro = vectorCalibrator.calibrate(feature);
            list.add(new Pair<>(candidate,pro));
        }

        Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
        return list.stream().sorted(comparator.reversed()).collect(Collectors.toList());
    }

    public static List<Pair<MultiLabel,Double>> topKinSupport(Vector x, MultiLabelClassifier.ClassProbEstimator classProbEstimator, LabelCalibrator labelCalibrator,
                                                     VectorCalibrator vectorCalibrator, PredictionFeatureExtractor predictionFeatureExtractor,
                                                     List<MultiLabel> support,
                                                     int top){
        double[] marginals = labelCalibrator.calibratedClassProbs(classProbEstimator.predictClassProbs(x));
        Comparator<Pair<MultiLabel,Double>> comparator = Comparator.comparing(pair->pair.getSecond());
        return support.stream().map(multiLabel -> {
            PredictionCandidate predictionCandidate = new PredictionCandidate();
            predictionCandidate.x = x;
            predictionCandidate.labelProbs = marginals;
            predictionCandidate.multiLabel = multiLabel;
            Vector feature = predictionFeatureExtractor.extractFeatures(predictionCandidate);
            double pro = vectorCalibrator.calibrate(feature);
            return new Pair<>(multiLabel,pro);
        }).sorted(comparator.reversed()).limit(top).collect(Collectors.toList());
    }


}

@@ -365,16 +365,97 @@
 bm.CBM;%0A
+import edu.neu.ccs.pyramid.util.ArgMax;%0Aimport edu.neu.ccs.pyramid.util.ArgSort;%0A
 import e
@@ -2032,24 +2032,2034 @@
 ());%0A    %7D%0A%0A
+%0A    public static List%3CPair%3CMultiLabel,Double%3E%3E topK(Vector x, MultiLabelClassifier.ClassProbEstimator classProbEstimator, LabelCalibrator labelCalibrator,%0A                                                     VectorCalibrator vectorCalibrator, PredictionFeatureExtractor predictionFeatureExtractor,%0A                                                     int minSetSize, int maxSetSize,%0A                                                     int top)%7B%0A        List%3CPair%3CMultiLabel,Double%3E%3E list = new ArrayList%3C%3E();%0A        double%5B%5D marginals = labelCalibrator.calibratedClassProbs(classProbEstimator.predictClassProbs(x));%0A        DynamicProgramming dynamicProgramming = new DynamicProgramming(marginals);%0A%0A        List%3CMultiLabel%3E candidates = new ArrayList%3C%3E();%0A        for (int i=0;i%3Ctop;i++) %7B%0A            MultiLabel candidate = dynamicProgramming.nextHighestVector();%0A            if (candidate.getNumMatchedLabels()%3E=minSetSize&&candidate.getNumMatchedLabels()%3C=maxSetSize)%7B%0A                candidates.add(candidate);%0A            %7D%0A        %7D%0A%0A        if (candidates.isEmpty())%7B%0A            int%5B%5D sorted = ArgSort.argSortDescending(marginals);%0A            MultiLabel candidate = new MultiLabel();%0A            for (int i=0;i%3CminSetSize;i++)%7B%0A                candidate.addLabel(sorted%5Bi%5D);%0A            %7D%0A        %7D%0A%0A        for (MultiLabel candidate: candidates)%7B%0A            PredictionCandidate predictionCandidate = new PredictionCandidate();%0A            predictionCandidate.x = x;%0A            predictionCandidate.labelProbs = marginals;%0A            predictionCandidate.multiLabel = candidate;%0A%0A            Vector feature = predictionFeatureExtractor.extractFeatures(predictionCandidate);%0A            double pro = vectorCalibrator.calibrate(feature);%0A            list.add(new Pair%3C%3E(candidate,pro));%0A        %7D%0A%0A        Comparator%3CPair%3CMultiLabel,Double%3E%3E comparator = Comparator.comparing(pair-%3Epair.getSecond());%0A        return list.stream().sorted(comparator.reversed()).collect(Collectors.toList());%0A    %7D%0A%0A
     public s

enforce set size constraints in top K sets