package uk.ac.soton.ecs.zmk1g19;

import de.bwaldvogel.liblinear.SolverType;
import org.openimaj.data.DataSource;
import org.openimaj.data.dataset.VFSGroupDataset;
import org.openimaj.data.dataset.VFSListDataset;
import org.openimaj.experiment.dataset.split.GroupedRandomSplitter;
import org.openimaj.experiment.evaluation.classification.ClassificationEvaluator;
import org.openimaj.experiment.evaluation.classification.ClassificationResult;
import org.openimaj.experiment.evaluation.classification.analysers.confusionmatrix.CMAnalyser;
import org.openimaj.experiment.evaluation.classification.analysers.confusionmatrix.CMResult;
import org.openimaj.feature.DoubleFV;
import org.openimaj.feature.FeatureExtractor;
import org.openimaj.feature.local.data.LocalFeatureListDataSource;
import org.openimaj.feature.local.list.LocalFeatureList;
import org.openimaj.image.FImage;
import org.openimaj.image.feature.local.keypoints.FloatKeypoint;
import org.openimaj.ml.annotation.linear.LiblinearAnnotator;
import org.openimaj.ml.clustering.FloatCentroidsResult;
import org.openimaj.ml.clustering.assignment.HardAssigner;
import org.openimaj.ml.clustering.kmeans.FloatKMeans;
import org.openimaj.util.pair.IntFloatPair;
import uk.ac.soton.ecs.zmk1g19.featureExtractors.PixelPatchFeatureExtractor;
import uk.ac.soton.ecs.zmk1g19.utils.Run2Engine;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * Run 2: Fixed Size Densely Sampled Pixel Patches
 *
 * @author Sofia Kisiala    (zmk1g19)
 * @author Harry Nelson     (hjn2g19)
 * @author Max Burgess      (mwmb1g19)
 * @author Anan Venkatesh   (av1u19)
 * @author Fergus Adams     (fhwa1g19)
 */
public class Run2 {
    /**
     * Main method used to run our Second classifier.
     *
     * @param training Image set used to train classifier
     * @param testing  Image set used to test classifier
     */
    public static void run(VFSGroupDataset<FImage> training, VFSListDataset<FImage> testing) {
        GroupedRandomSplitter<String, FImage> split = new GroupedRandomSplitter<>(training, 1, 0, 1);

        // The sampling thing to extract features based on fixed size patches, 8x8 size, every 4 pixels x and y
        Run2Engine engine = new Run2Engine();

        // Get the data to train the quantiser, from the training set
        System.out.println("Starting quantiser training");
        HardAssigner<float[], float[], IntFloatPair> assigner = trainQuantiser(training, engine);
        System.out.println("finished quantiser training");

        //
        System.out.println("start feature extractor");
        FeatureExtractor<DoubleFV, FImage> extractor = new PixelPatchFeatureExtractor(assigner, engine);
        //FeatureExtractor<DoubleFV, FImage> extractor = new PixelPatchExtractor(assigner, 8, 4);
        System.out.println("finish feature extractor");

        // The linear classifier specified
        System.out.println("start make annotator");
        LiblinearAnnotator<FImage, String> annotator = new LiblinearAnnotator<FImage, String>(extractor, LiblinearAnnotator.Mode.MULTICLASS, SolverType.L2R_L2LOSS_SVC, 1.0, 0.00001);
        System.out.println("finish make annotator");

        // train the annotator
        System.out.println("start training annotator");
        annotator.train(split.getTrainingDataset());
        System.out.println("finish training annotator");

        // write to file
        App.writeAnnotatorOutput(annotator, "run2", testing);

        System.out.println("start classification eval");
        ClassificationEvaluator<CMResult<String>, String, FImage> eval =
                new ClassificationEvaluator<CMResult<String>, String, FImage>(
                        annotator, split.getTestDataset(), new CMAnalyser<FImage, String>(CMAnalyser.Strategy.SINGLE));
        System.out.println("finish classification eval");

        System.out.println("start eval");
        Map<FImage, ClassificationResult<String>> guesses = eval.evaluate();
        System.out.println("finish eval");

        CMResult<String> result = eval.analyse(guesses);
        System.out.println(result.getDetailReport());
    }

    /**
     * @param sample Set of images the Quantiser will be trained using
     * @param engine Class that implements findFeatures (Run2Engine) which returns a feature list given an image
     * @return Trained HardAssigner
     */
    static HardAssigner<float[], float[], IntFloatPair> trainQuantiser(VFSGroupDataset<FImage> sample, Run2Engine engine) {
        // HEAVILY BASED ON CH12 FROM TUTORIAL
        List<LocalFeatureList<FloatKeypoint>> allKeys = new ArrayList<LocalFeatureList<FloatKeypoint>>();

        for (FImage img : sample) {
            allKeys.add(engine.findFeatures(img));
        }

        //prune the keypoint list if it exceeds length 10000
        if (allKeys.size() > 10000) {
            allKeys = allKeys.subList(0, 10000);
        }

        // k means clustering
        FloatKMeans km = FloatKMeans.createKDTreeEnsemble(500);
        DataSource<float[]> datasource = new LocalFeatureListDataSource<FloatKeypoint, float[]>(allKeys);
        FloatCentroidsResult result = km.cluster(datasource);

        return result.defaultHardAssigner();
    }
}