package ai.panoptic.vital_sign_camera;

import static com.google.mlkit.vision.face.FaceContour.FACE;

import android.graphics.PointF;
import android.media.Image;
import android.util.ArraySet;

import com.google.android.gms.tasks.Task;
import com.google.mlkit.vision.common.InputImage;
import com.google.mlkit.vision.face.Face;
import com.google.mlkit.vision.face.FaceContour;
import com.google.mlkit.vision.face.FaceDetection;
import com.google.mlkit.vision.face.FaceDetector;
import com.google.mlkit.vision.face.FaceDetectorOptions;

import java.util.List;
import java.util.Set;

import ai.panoptic.wvs.engine.EyesOpenProbability;
import ai.panoptic.wvs.engine.FaceDetectorFeature;
import ai.panoptic.wvs.engine.HeadEulerAngle;
import ai.panoptic.wvs.engine.NormalizedFacebox;
import ai.panoptic.wvs.engine.VitalSignFaceDetector;

public class MLKitFaceDetector implements VitalSignFaceDetector {

    ai.panoptic.wvs.engine.Face face = null;

    FaceDetectorOptions options =
        new FaceDetectorOptions.Builder()
            .setPerformanceMode(FaceDetectorOptions.PERFORMANCE_MODE_FAST)
            .setContourMode(FaceDetectorOptions.CONTOUR_MODE_NONE)
            .setClassificationMode(FaceDetectorOptions.CLASSIFICATION_MODE_ALL)
            .setMinFaceSize(0.15f)
            .build();

    FaceDetector faceDetector = FaceDetection.getClient(options);

    public ai.panoptic.wvs.engine.Face detectFace(Image mediaImage, int rotationDegrees) {

        if (mediaImage == null) {
            return null;
        }

        float width = mediaImage.getWidth();
        float height = mediaImage.getHeight();
        if (rotationDegrees == 270 || rotationDegrees == 90) {
            width = mediaImage.getHeight();
            height = mediaImage.getWidth();
        }
        final float _width = width;
        final float _height = height;
        InputImage image = InputImage.fromMediaImage(mediaImage, rotationDegrees);

        Task<List<Face>> task = faceDetector.process(image);

        task.addOnCompleteListener( result -> {

            if (result.isCanceled()) {
                return;
            }
            if (!result.isSuccessful()) {
                return;
            }
            if (!result.isComplete()) {
                return;
            }
            try {
                List<Face> faces = result.getResult();
                if (faces.isEmpty()) {
                    this.face = null;
                    return;
                }
                Face face = faces.get(0);
                this.face = createFace(face, _width, _height);

            } catch (Exception e) {
                e.printStackTrace();
                this.face = null;
            }

        });

        return face;
    }

    static private ai.panoptic.wvs.engine.Face createFace(Face face, float _width, float _height) {

        FaceContour contour = face.getContour(FACE);

        NormalizedFacebox facebox = new NormalizedFacebox();
        float minX = _width;
        float minY = _height;
        float maxX = 0;
        float maxY = 0;

        if (contour != null) {

            for (PointF p : contour.getPoints()) {
                if (p.x < minX) {
                    minX = p.x;
                }
                if (p.y < minY) {
                    minY = p.y;
                }
                if (p.x > maxX) {
                    maxX = p.x;
                }
                if (p.y > maxY) {
                    maxY = p.y;
                }
            }

        } else {
            minX = face.getBoundingBox().left;
            minY = face.getBoundingBox().top;
            maxX = face.getBoundingBox().right;
            maxY = face.getBoundingBox().bottom;
        }

        facebox.xCenter = (minX + (maxX - minX) / 2.0f) / _width;
        facebox.yCenter = (minY + (maxY - minY) / 2.0f) / _height;
        facebox.width = (maxX - minX) / _width;
        facebox.height = (maxY - minY) / _height;

        ai.panoptic.wvs.engine.Face result = new ai.panoptic.wvs.engine.Face(facebox);

        result.setEulerAngle(
                new HeadEulerAngle(
                        face.getHeadEulerAngleX(),
                        face.getHeadEulerAngleY(),
                        face.getHeadEulerAngleY()
                )
        );

        if (face.getLeftEyeOpenProbability() != null && face.getRightEyeOpenProbability() != null) {
            result.setEyesOpenProbability(
                    new EyesOpenProbability(
                            face.getLeftEyeOpenProbability(),
                            face.getRightEyeOpenProbability()
                    )
            );
        }

        return result;
    }

    @Override
    public Set<FaceDetectorFeature> supportedFeatures() {
        ArraySet<FaceDetectorFeature> features = new ArraySet<FaceDetectorFeature>();
        features.add(FaceDetectorFeature.eyeOpenProbability);
        features.add(FaceDetectorFeature.eulerAngle);
        return features;
    }

    @Override
    public NormalizedFacebox detect(Image mediaImage, int rotationDegrees) {
        detectFace(mediaImage, rotationDegrees);
        if (face == null) {
            return null;
        }
        return face.getFacebox();
    }
}
