package nl.byelex.logodetectie.node.debug.recognizer;

import com.google.common.base.Throwables;
import nl.byelex.logodetectie.lib.commondomain.ClusteringAlgorithm;
import nl.byelex.logodetectie.lib.commondomain.LogoType;
import nl.byelex.logodetectie.lib.commondomain.PointInt;
import nl.byelex.logodetectie.lib.keypoints.LogoKeypointsGenerator;
import nl.byelex.logodetectie.lib.keypoints.LogoTrainData;
import nl.byelex.logodetectie.lib.mask.Mask;
import nl.byelex.logodetectie.lib.opencvutils.ImgUtils;
import nl.byelex.logodetectie.lib.xml.xmodel.ImageKeypointData;
import nl.byelex.logodetectie.node.processing.ExtractKeypointsAndFeatures;
import nl.byelex.logodetectie.node.processing.logodetector.ISingleLogoDetector;
import nl.byelex.logodetectie.node.processing.logodetector.LogoDetectorFactory;
import nl.byelex.logodetectie.node.processing.logodetector.LogoMatchCandidate;
import nl.byelex.logodetectie.opencvutils.InitOpenCv;
import org.opencv.core.Mat;
import org.opencv.core.MatOfKeyPoint;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

public class KeypointDetectTest {

    private static final Logger logger = LoggerFactory.getLogger(KeypointDetectTest.class);
    public final static String queryDir = "/mnt/drive/temp/main/040_task2777_nn_training_data/frames/";
    public final static String trainDir = "/mnt/drive/temp/kpn_logo_train";
    private long pointSize = 0;
    private long time = 0;
    private long timeMatch = 0;
    private long logoSize = 0;

    private ThreadLocal<ExtractKeypointsAndFeatures.ExtractionContext> context = ThreadLocal.withInitial(ExtractKeypointsAndFeatures.ExtractionContext::new);

    public void run() throws Exception {
        List<LogoDetectorFactory> builders = new ArrayList<>();
        List<Path> trainFilesList = Files.list(Paths.get(trainDir)).collect(Collectors.toList());
        for (int i = 0; i < trainFilesList.size(); i++) {
            LogoDetectorFactory factory = new LogoDetectorFactory();
            addTrainLogo(i, trainFilesList.get(i), factory, ClusteringAlgorithm.DBSCAN);
            builders.add(factory);
        }

        List<Path> queryFilesList = Files.list(Paths.get(queryDir)).collect(Collectors.toList());
        for (int i = 0; i < queryFilesList.size(); i++) {
            processFrameImpl(queryFilesList.get(i), builders);
        }
        System.out.print("time:  " + time + " keypoints: " + pointSize + " logos: " + logoSize + " files: " + queryFilesList.size() + " match time: " + timeMatch + " ");
        System.out.println("mean time:  " + (((double) time) / queryFilesList.size()) + " keypoints: " + ((double) pointSize) / queryFilesList.size() + " logos " + ((double) logoSize) / queryFilesList.size() + " " + "mean time match:  " + (((double) timeMatch) / queryFilesList.size()));

    }


    private void addTrainLogo(long id, Path filePath, LogoDetectorFactory factory, @Nonnull ClusteringAlgorithm clusteringAlgorithm) {

        LogoTrainData keypoints = LogoKeypointsGenerator.generateKeypoints(filePath);

        Mat features = null;
        try {
            BufferedImage img = ImageIO.read(filePath.toFile());


            Mask mask = new Mask(new PointInt(0, 0), new PointInt(0, img.getHeight()), new PointInt(img.getWidth(), img.getHeight()), new PointInt(img.getWidth(), img.getHeight()));


            features = keypoints.buildFeaturesMat();
            factory.identifiers("oo", id + 1)
                    .clusteringAlgorithm(clusteringAlgorithm)
                    .keypoints(keypoints.getKpa(), features)
                    .mask(mask.getPolygon())
                    .trainingSizes(keypoints.getImageWidth(), keypoints.getImageHeight())
                    .logoType(LogoType.IMAGE)
                    .keepRejected();
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (features != null)
                features.release();
        }
    }

    public static void main(String[] args) {
        InitOpenCv.init();
        try {
            new KeypointDetectTest().run();
        } catch (Exception e) {
            logger.warn(Throwables.getRootCause(e).getMessage(), e);
        }
    }

    @SuppressWarnings("Duplicates")
    private void processFrameImpl(Path queryFile, List<LogoDetectorFactory> builders) throws IOException {
        final ExtractKeypointsAndFeatures.ExtractionContext ctx = context.get();
        MatOfKeyPoint kpm = null;
        Mat features = new Mat();
        Mat m = ImgUtils.imread(queryFile);
        try {


            long startTime = System.currentTimeMillis();
            kpm = new MatOfKeyPoint();
//                Pair<MatOfKeyPoint,Mat> pair = ctx.getKeypointsDetector().detectKeypointsAndExtractFeatures(m, null, null);
//            kpm = pair.getKey();
//            features = pair.getValue();

            ctx.getKeypointsDetector().detectKeypointsAndExtractFeatures(m, kpm, features);


            time += (System.currentTimeMillis() - startTime);

            pointSize += kpm.toList().size();


            if (features.type() != 5) {
                float[] matrix = matToMatrixB(features);
                int rows = features.rows();
                features = ImageKeypointData.arrToFeatureMatF(matrix, rows);
            }


            startTime = System.currentTimeMillis();
            List<LogoMatchCandidate> candidates = new ArrayList<>();
            for (LogoDetectorFactory builder : builders) {
                ISingleLogoDetector detector = builder.createDetector(kpm.toList(), features, m.width(), m.height());
                for (LogoMatchCandidate c : detector.detect()) {
                    if (c.isRejected()) {
                        // final LogoVariant logoVariant = brandLogoVariants.getLogoVariant(c.getVariantId());
                        // frameRejected0.addMatch(new LogoMatch(frameRejected0, logoVariant, MatchStage.RECOGNITION, c.getMaskRegion(), c.getTransformation()));
                    } else
                        candidates.add(c);
                }
            }
            timeMatch += (System.currentTimeMillis() - startTime);
            logoSize += candidates.size();
            //
//            new ResolveLogoDuplicates(brandLogoVariants, candidates).resolve();
//
//            for (LogoMatchCandidate rr : candidates) {
//                final LogoVariant logoVariant = brandLogoVariants.getLogoVariant(rr.getVariantId());
//                final Frame frame = rr.isRejected() ? frameRejected1 : frameResult;
//                frame.addMatch(new LogoMatch(frame, logoVariant, MatchStage.RECOGNITION, rr.getMaskRegion(), rr.getTransformation()));
//            }

        } finally {
            kpm.release();
            features.release();
            m.release();
        }
    }

    private static float[] matToMatrixB(Mat mat) {
        int c = mat.cols();
        int r = mat.rows();
        byte d[] = new byte[c * r];
        mat.get(0, 0, d);
        int s = 0;
        float[] matrix = new float[r * c];
        for (int i = 0; i < r; i++) {
            for (int j = 0; j < c; j++) {
                matrix[s] = d[s];
//                System.out.println("f "+matrix[i][j]+ " "+d[s]+arr[2]);
                s++;
            }
        }
        return matrix;
    }


}