package uk.ac.soton.ecs.zmk1g19;


import org.openimaj.data.dataset.VFSGroupDataset;
import org.openimaj.data.dataset.VFSListDataset;
import org.openimaj.experiment.dataset.split.GroupedRandomSplitter;
import org.openimaj.feature.DoubleFV;
import org.openimaj.feature.FeatureExtractor;
import org.openimaj.image.FImage;
import org.openimaj.image.feature.dense.gradient.dsift.DenseSIFT;
import org.openimaj.image.feature.dense.gradient.dsift.PyramidDenseSIFT;
import org.openimaj.ml.annotation.Annotator;
import org.openimaj.ml.annotation.bayes.NaiveBayesAnnotator;
import org.openimaj.ml.annotation.linear.LiblinearAnnotator;
import org.openimaj.ml.clustering.assignment.HardAssigner;
import org.openimaj.ml.kernel.HomogeneousKernelMap;
import org.openimaj.util.pair.IntFloatPair;
import uk.ac.soton.ecs.zmk1g19.classifiers.LiblinearFactory;
import uk.ac.soton.ecs.zmk1g19.classifiers.NaiveBayesFactory;
import uk.ac.soton.ecs.zmk1g19.classifiers.VoteClassifier;
import uk.ac.soton.ecs.zmk1g19.featureExtractors.GISTFeatureExtractor;
import uk.ac.soton.ecs.zmk1g19.featureExtractors.HOGFeatureExtractor;
import uk.ac.soton.ecs.zmk1g19.featureExtractors.PHOWExtractor;
import uk.ac.soton.ecs.zmk1g19.utils.AccuracyTest;
import uk.ac.soton.ecs.zmk1g19.utils.HardAssignerFactory;

import java.io.IOException;
import java.util.HashMap;

/**
 * Run 3: Voting Ensemble using Naive Bayes and LiblinearAnnotator
 * with a combination of feature extractors
 *
 * @author Sofia Kisiala (zmk1g19)
 * @author Harry Nelson (hjn2g19)
 * @author Max Burgess (mwmb1g19)
 * @author Anan Venkatesh (av1u19)
 * @author Fergus Adams (fhwa1g19)
 */

public class Run3 {

    public static void run(VFSGroupDataset<FImage> training, VFSListDataset<FImage> testing) throws IOException {
        GroupedRandomSplitter<String, FImage> split = new GroupedRandomSplitter<>(training, 80, 0, 20);

        // Setup HOG extractor
        HOGFeatureExtractor hogFeatureExtractor = HOGFeatureExtractor.getDefault();


        // Setup GIST extractor
        GISTFeatureExtractor gistFeatureExtractor = GISTFeatureExtractor.getDefault();

        // Setup PHOW extractor
        DenseSIFT dsift = new DenseSIFT(5, 7);
        PyramidDenseSIFT<FImage> pdsift = new PyramidDenseSIFT<FImage>(dsift, 6f, 7);
        HardAssigner<byte[], float[], IntFloatPair> hardAssigner = HardAssignerFactory.trainQuantiser(split.getTrainingDataset(), pdsift);
        PHOWExtractor phowExtractor = new PHOWExtractor(pdsift, hardAssigner);

        // Setup Kernel Map and wrapped extractors
        HomogeneousKernelMap homogeneousKernelMap = new HomogeneousKernelMap(HomogeneousKernelMap.KernelType.Intersection, 0.2, HomogeneousKernelMap.WindowType.Uniform);
        FeatureExtractor<DoubleFV, FImage> wrappedGistExtractor = homogeneousKernelMap.createWrappedExtractor(gistFeatureExtractor);
        FeatureExtractor<DoubleFV, FImage> wrappedHogExtractor = homogeneousKernelMap.createWrappedExtractor(hogFeatureExtractor);
        FeatureExtractor<DoubleFV, FImage> wrappedPHOWExtractor = homogeneousKernelMap.createWrappedExtractor(phowExtractor);

        // Make Liblinear Annotators

        LiblinearAnnotator<FImage, String> liblinearGIST = LiblinearFactory.createLiblinear(wrappedGistExtractor, training);
        LiblinearAnnotator<FImage, String> liblinearHOG = LiblinearFactory.createLiblinear(wrappedHogExtractor, training);
        LiblinearAnnotator<FImage, String> liblinearPHOW = LiblinearFactory.createLiblinear(wrappedPHOWExtractor, training);

        // Make Naive Bayes Annotators

        NaiveBayesAnnotator<FImage, String> bayesUnwrappedGIST = NaiveBayesFactory.createNaiveBayes(gistFeatureExtractor, training);
        NaiveBayesAnnotator<FImage, String> bayesUnrwappedPHOW = NaiveBayesFactory.createNaiveBayes(phowExtractor, training);
        NaiveBayesAnnotator<FImage, String> bayesPHOW = NaiveBayesFactory.createNaiveBayes(wrappedPHOWExtractor, training);

        // Add to HashMap with a weighting
        HashMap<Annotator<FImage, String>, Float> annotators = new HashMap<>();
        annotators.put(liblinearGIST, 0.6f);
        annotators.put(liblinearHOG, 0.5f);
        annotators.put(liblinearPHOW, 0.7f);
        annotators.put(bayesUnwrappedGIST, 0.6f);
        annotators.put(bayesUnrwappedPHOW, 0.6f);
        annotators.put(bayesPHOW, 0.6f);

        // Create vote classifier
        VoteClassifier voteClassifier = VoteClassifier.createVoteClassifier(annotators);
        App.writeAnnotatorOutput(voteClassifier, "run3", testing);
    }
}