package uk.ac.soton.ecs.zmk1g19;

import org.apache.commons.vfs2.FileObject;
import org.openimaj.data.dataset.VFSGroupDataset;
import org.openimaj.data.dataset.VFSListDataset;
import org.openimaj.experiment.evaluation.classification.ClassificationResult;
import org.openimaj.image.FImage;
import org.openimaj.image.ImageUtilities;
import org.openimaj.ml.annotation.Annotator;

import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.FileSystems;
import java.nio.file.Path;
import java.util.Set;

/**
 * Main class used to preprocess data, run all classifiers, and print output data to text file.
 *
 * @author Sofia Kisiala    (zmk1g19)
 * @author Harry Nelson     (hjn2g19)
 * @author Max Burgess      (mwmb1g19)
 * @author Anan Venkatesh   (av1u19)
 * @author Fergus Adams     (fhwa1g19)
 */
public class App {
    public static void main(String[] args) throws IOException {
        VFSGroupDataset<FImage> training;
        VFSListDataset<FImage> testing;
        Path dir = FileSystems.getDefault().getPath("").toAbsolutePath();
        // DatasetEmbiggener.transformImages();
        //import images from folders
        try {
            training = new VFSGroupDataset<>(dir + "/bigSet", ImageUtilities.FIMAGE_READER);
            testing = new VFSListDataset<>(dir + "/testing", ImageUtilities.FIMAGE_READER);
        } catch (Exception e) {
            e.printStackTrace();
            return;
        }

        //Run1.run(training, testing);
        //Run2.run(training, testing);
        Run3.run(training, testing);
    }

    /**
     * File writing utility function; provide an annotator, a file name and a testing dataset and get a brand new output file today!
     *
     * @param ann     Annotator
     * @param outFile Output file name (without extension)
     * @param testing Testing dataset to use
     */
    public static void writeAnnotatorOutput(Annotator<FImage, String> ann, String outFile, VFSListDataset<FImage> testing) {
        FileObject[] files = testing.getFileObjects();
        StringBuilder sb = new StringBuilder();
        //go through all images in testing folder and match names to files in output
        for (int i = 0; i < testing.size(); i++) {
            String name = files[i].getName().getBaseName();
            FImage image = testing.get(i);
            //use annotator to generate classification
            ClassificationResult<String> classificationResult = ann.classify(image);
            Set<String> predictedClasses = classificationResult.getPredictedClasses();
            //need to use iterator to extract classification from Set format
            String classification = predictedClasses.iterator().next();
            sb.append(name).append(" ").append(classification).append(System.lineSeparator());
        }
        //push stringbuilder to text file
        try (FileOutputStream fos = new FileOutputStream(outFile + ".txt")) {
            fos.write(sb.toString().getBytes());
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}