package assignment07;

import ij.IJ;
import ij.ImagePlus;
import ij.gui.Overlay;
import ij.gui.Roi;
import ij.gui.ShapeRoi;
import ij.plugin.filter.PlugInFilter;
import ij.process.FloatProcessor;
import ij.process.ImageProcessor;
import javafx.util.Pair;

import java.awt.*;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

// Linear and normalized did not produce good results. Euclidean worked reliably on Ramses image and PCB image.
// Running time highly depends on the size of the area of interest.
public class Correlation_Based_Matching implements PlugInFilter {
    private static final int LIMIT_OF_MARKS = 7;
    @Override
    public int setup(String s, ImagePlus imagePlus) {
        return DOES_8G + DOES_8C + NO_CHANGES;
    }

    @Override
    public void run(ImageProcessor ip) {
        ImageProcessor reference = ip.crop();
        if (reference.getPixelCount() == ip.getPixelCount()) {
            IJ.error("Please select area for the reference image");
            return;
        }

        (new ImagePlus("Reference image", reference)).show();

        FloatProcessor img = ip.convertToFloatProcessor();
        FloatProcessor ref = reference.convertToFloatProcessor();

        float[][] imgArr = img.getFloatArray();
        float[][] refArr = ref.getFloatArray();

        float[][] sqrDiffSumScore = testAndDrawCorrelation(SQR_DIFF_SUM, refArr, imgArr, "Sqr Diff Sum");
        List<Pair<Integer, Integer>> sqrDiffExtremas = findLocalExtrema(sqrDiffSumScore, LIMIT_OF_MARKS, true);
        displayLocalExtrema(sqrDiffExtremas, ip, ref.getWidth(), ref.getHeight(), "Sqr Diff Sum");


        float[][] linearCrossCorrScore = testAndDrawCorrelation(LINEAR_CROSS_CORR, refArr, imgArr, "Linear Cross Corr");
        List<Pair<Integer, Integer>> linearCrossCorrExtremas =
            findLocalExtrema(linearCrossCorrScore, LIMIT_OF_MARKS, false);
        displayLocalExtrema(linearCrossCorrExtremas, ip, ref.getWidth(), ref.getHeight(), "Linear Cross Corr");

        float[][] normCrossCorrScore = testAndDrawCorrelation(NORM_CROSS_CORR, refArr, imgArr, "Norm Cross Corr");
        List<Pair<Integer, Integer>> normCrossCorrExtremas =
            findLocalExtrema(normCrossCorrScore, LIMIT_OF_MARKS, false);
        displayLocalExtrema(normCrossCorrExtremas, ip, ref.getWidth(), ref.getHeight(), "Norm Cross Corr");

    }

    private float[][] testAndDrawCorrelation(CorrelationCalculator calculator, float[][] ref, float[][] img, String imgName) {
        int scoreRows = img.length - ref.length;
        int scoreCols = img[0].length - ref[0].length;

        IJ.showProgress(0, scoreRows * scoreCols);
        float[][] correlationScore = new float[scoreRows][scoreCols];
        for (int i = 0; i < scoreRows; i++) {
            for (int j = 0; j < scoreCols; j++) {
                correlationScore[i][j] = calculator.correlation(ref, img, i, j);
                IJ.showProgress(i * scoreCols + j, scoreRows * scoreCols);
            }
        }
        FloatProcessor correlationImageProcessor = new FloatProcessor(correlationScore);
        (new ImagePlus(imgName, correlationImageProcessor)).show();
        return correlationScore;
    }

    private interface CorrelationCalculator {
        float correlation(float[][] ref, float[][] img, int row, int col);
    }

    private static final CorrelationCalculator SQR_DIFF_SUM =
        (ref, img, row, col) -> {
            float ans = 0;
            for (int i = 0; i < ref.length; i++) {
                for (int j = 0; j < ref[i].length; j++) {
                    ans += sqr(img[row + i][col + j] - ref[i][j]);
                }
            }
            return (float) Math.sqrt(ans);
        };

    private static final CorrelationCalculator LINEAR_CROSS_CORR =
        (ref, img, row, col) -> {
            float ans = 0;
            for (int i = 0; i < ref.length; i++) {
                for (int j = 0; j < ref[i].length; j++) {
                    ans += img[row + i][col + j] * ref[i][j];
                }
            }
            return ans;
        };

    private static final CorrelationCalculator NORM_CROSS_CORR =
        (ref, img, row, col) -> {
            float denRef = 0, denImg = 0;
            for (int i = 0; i < ref.length; i++) {
                for (int j = 0; j < ref[i].length; j++) {
                    denRef += sqr(ref[i][j]);
                    denImg += sqr(img[row + i][col + j]);
                }
            }
            return LINEAR_CROSS_CORR.correlation(ref, img, row, col)
                / ((float) Math.sqrt(denRef) * (float) Math.sqrt(denImg));
        };

    private List<Pair<Integer, Integer>> findLocalExtrema(float[][] scores, int limit, boolean isMinima) {
        int rows = scores.length;
        int cols = scores[0].length;

        ArrayList<Pair<Integer, Integer>> localExtremas = new ArrayList<>();
        for (int i = 0; i < rows; i++) {
            for (int j = 0; j < cols; j++) {
                boolean isExtrema = true;
                for (int di = -1; di <= 1; di++) {
                    for (int dj = -1; dj <= 1; dj++) {
                        if (di != 0 || dj != 0) {
                            int neighb_i = i + di;
                            int neighb_j = j + dj;
                            if (0 <= neighb_i && neighb_i < rows && 0 <= neighb_j && neighb_j < cols) {
                                if (!compare(isMinima, scores[i][j], scores[neighb_i][neighb_j])) {
                                    isExtrema = false;
                                }
                            }
                        }
                    }
                }

                if (isExtrema) {
                    localExtremas.add(new Pair<>(i, j));
                }
            }
        }

        localExtremas =
            localExtremas
                .stream()
                .sorted(Comparator.comparingDouble(p -> scores[p.getKey()][p.getValue()]))
                .limit(limit)
                .collect(Collectors.toCollection(ArrayList::new));
        if (!isMinima) {
            Collections.reverse(localExtremas);
        }

        return localExtremas;
    }

    private void displayLocalExtrema(
            List<Pair<Integer, Integer>> extremas, ImageProcessor ip, int rectW, int rectH, String imgName) {
        Overlay oly = new Overlay();
        for (Pair<Integer, Integer> extrema : extremas) {
            Roi rectRoi = new ShapeRoi(
                new Rectangle(extrema.getKey(), extrema.getValue(), rectW, rectH));
            rectRoi.setStrokeWidth(0.8f);
            rectRoi.setStrokeColor(Color.blue);
            oly.add(rectRoi);
        }

        ImageProcessor copy = ip.duplicate();
        ImagePlus asd = new ImagePlus(imgName, copy);
        asd.setOverlay(oly);
        asd.show();
    }

    private static boolean compare(boolean isMinima, float a, float b) {
        if (isMinima) {
            return a <= b;
        }
        return a >= b;
    }

    private static float sqr(float t) {
        return t * t;
    }
}