package com.amlan.javacvfacerecognization.utils;

import android.util.Log;

import org.bytedeco.javacpp.opencv_core.Mat;
import org.bytedeco.javacpp.opencv_core.MatVector;
import org.bytedeco.javacpp.opencv_face.FaceRecognizer;

import java.io.File;
import java.io.FilenameFilter;
import java.nio.IntBuffer;

import static org.bytedeco.javacpp.opencv_core.CV_32SC1;
import static org.bytedeco.javacpp.opencv_face.createEigenFaceRecognizer;
import static org.bytedeco.javacpp.opencv_face.createFisherFaceRecognizer;
import static org.bytedeco.javacpp.opencv_imgcodecs.CV_LOAD_IMAGE_GRAYSCALE;
import static org.bytedeco.javacpp.opencv_imgcodecs.imread;

public class OpenCVFaceRecognizer {

    private FaceRecognizer faceRecognizer;

    public void trainModel(String trainingDir) {

        File root = new File(trainingDir);

        FilenameFilter imgFilter = new FilenameFilter() {
            public boolean accept(File dir, String name) {
                name = name.toLowerCase();
                return name.endsWith(".jpg") || name.endsWith(".pgm") || name.endsWith(".png");
            }
        };

        File[] imageFiles = root.listFiles(imgFilter);

        MatVector images = new MatVector(imageFiles.length);

        Mat labels = new Mat(imageFiles.length, 1, CV_32SC1);
        IntBuffer labelsBuf = labels.createBuffer();

        int counter = 0;

        for (File image : imageFiles) {
            Mat img = imread(image.getAbsolutePath(), CV_LOAD_IMAGE_GRAYSCALE);

            String parentFolderName = image.getParent();
            String folderName = parentFolderName.substring(parentFolderName.lastIndexOf("s") + 1);

            // lable is a int which denotes the type. Basically associate a set of image to a lable
            int label = Integer.parseInt(folderName);

            images.put(counter, img);

            labelsBuf.put(counter, label);

            counter++;
        }


        // We are using Eigenfaces Algorithm
        faceRecognizer = createEigenFaceRecognizer();
        faceRecognizer.train(images, labels);


    }

    public int preditImage(String fileName) {

        Mat testImage = imread(fileName, CV_LOAD_IMAGE_GRAYSCALE);
        int predictedLabel = faceRecognizer.predict(testImage);
        System.out.println("Predicted label: " + predictedLabel);
        Log.e("amlan","Prediction :::: "+predictedLabel);
        return predictedLabel;

    }
}