package org.jeecg.modules.opencv.utils;

import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;
import org.opencv.videoio.VideoCapture;
import org.opencv.videoio.Videoio;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

public class ExtractionVideoKeyFrames {
    Mat computeColorHist3D(Mat im, Size size)
    {
        Mat dst = new Mat(), hist = new Mat();
        Imgproc.resize(im, dst, size, 0, 0, Imgproc.INTER_NEAREST);
        Mat hsv_src = new Mat();
        Imgproc.cvtColor(im, hsv_src, Imgproc.COLOR_BGR2HSV);
        List<Mat> images=new ArrayList<Mat>(3);
        images.add(hsv_src);

        MatOfInt channels= new MatOfInt(0);
        MatOfInt hSize = new MatOfInt(256);
        MatOfFloat ranges= new MatOfFloat(0,255);
        //获取图像直方图
        Imgproc.calcHist(images, channels, new Mat(), hist, hSize, ranges);
        return hist;
    }

    Mat computeColorHist(Mat im, Size size)
    {
        Mat dst = new Mat(), hist = new Mat();
        Imgproc.resize(im, dst, size, 0, 0, Imgproc.INTER_NEAREST);
        Mat gray = new Mat(im.size(),CvType.CV_8UC1);
        Imgproc.cvtColor(im,gray,Imgproc.COLOR_RGB2GRAY);

        List<Mat> images = new ArrayList<Mat>();
        images.add(gray);

        MatOfInt channels= new MatOfInt(0);
        MatOfInt hSize = new MatOfInt(128);
        MatOfFloat ranges= new MatOfFloat(0,128);
        Imgproc.calcHist(images, channels, new Mat(), hist, hSize, ranges);
        return hist;
    }

    double computeASE(Mat a, Mat b) {
        return  1 - Imgproc.compareHist(a, b, 0);
    }

    int ExtractKeyframes(String videoPath, double rate, String outImgPath){
        VideoCapture cap = new VideoCapture();
        if (!cap.open(videoPath)) {
            return -1;
        }

        File file = new File(outImgPath);
        if(!file.exists()){
            file.mkdirs();
        }
        //  Videoio.CAP_PROP_FRAME_COUNT
        int width = (int)cap.get(Videoio.CAP_PROP_FRAME_WIDTH);
        int height = (int)cap.get(Videoio.CAP_PROP_FRAME_HEIGHT);
        float fps = (float)cap.get(Videoio.CAP_PROP_FPS);
        int nframes = (int)cap.get(Videoio.CAP_PROP_FRAME_COUNT);

        int nscale = 1;
        if (width%4==0 && height%4==0){
            nscale = 4;
        }else if (width%2==0 && height%2==0){
            nscale = 2;
        }
        Size size = new Size(width/nscale, height/nscale);

        Mat frame = new Mat();
        Mat lastFrame = new Mat();
        Mat histA = new Mat();
        Mat histB = new Mat();
        long nCount = 0;
        int frameID = 1;
        int lastFrameID = 1;
        double tic = (double)Core.getTickCount();
        while (cap.read(frame)) {
            if (frame.empty()){
                break;
            }

            if (lastFrameID==1) {
                frameID = (int)cap.get(Videoio.CAP_PROP_POS_FRAMES);
                lastFrameID = frameID;
                frame.copyTo(lastFrame);
                continue;
            }
            if (++nCount % 15 != 0){
                continue;	// Processing one frame in each second
            }

            frameID = (int)cap.get(Videoio.CAP_PROP_POS_FRAMES);
            histA = computeColorHist3D(frame, size);

            double aseV = rate;
            if (histB.empty()){
            } else{
                aseV = computeASE(histA, histB);
            }

            if (aseV >= rate) {
                String imagePath = outImgPath+"/"+lastFrameID+".jpg";
                if (!Imgcodecs.imwrite(imagePath, lastFrame)){
                    return -2;
                }
                lastFrameID = frameID;
                frame.copyTo(lastFrame);
            }

            histA.copyTo(histB);
        }

        if (!lastFrame.empty()) {
            String imagePath=outImgPath+"/"+ lastFrameID+".jpg";
            if (!Imgcodecs.imwrite(imagePath, lastFrame)){
                return -2;
            }
        }

        double toc = (double)Core.getTickCount();
        return 0;
    }
}
