package com.knecon.fforesight.service.layoutparser.processor.services.graphics;

import java.awt.geom.AffineTransform;
import java.awt.geom.Rectangle2D;
import java.awt.image.BufferedImage;
import java.awt.image.DataBufferByte;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.springframework.stereotype.Service;

import lombok.SneakyThrows;

@Service
public class FindGraphicsRaster {

    // Pixels that are lighter then this threshold are ignored
    private final static int THRESHOLD = 240;

    // DPI to render the image at, in practice sub-72 seems to risk pixels being lost
    private final static int DPI = 72;


    @SneakyThrows
    public List<Box> findCCBoundingBoxes(PDDocument doc, List<Rectangle2D> remove, PageInformation pageInformation) {

        var renderer = new PDFRenderer(doc);
        var img = renderer.renderImageWithDPI(pageInformation.number() - 1, DPI, ImageType.GRAY);
        var imageCtm = getImageCTM(pageInformation, img.getWidth());
        return findCCBoundingBoxes(img, remove, THRESHOLD, DPI / 72, imageCtm);
    }


    @SneakyThrows
    private List<Box> findCCBoundingBoxes(BufferedImage image, List<Rectangle2D> remove, int grayScaleTresh, int rescale, AffineTransform imageCTM) {

        var inverseCTM = imageCTM.createInverse();

        var h = image.getHeight();
        var w = image.getWidth();
        var pixels = new int[w * h];
        image.getRaster().getPixels(0, 0, w, h, pixels);
        remove.stream().map(rect -> inverseCTM.createTransformedShape(rect).getBounds2D()).forEach(box -> {
            for (int y = (int) Math.floor(box.getMinY() / rescale); y <= (int) Math.min(Math.ceil(box.getMaxY() / rescale), h); y++) {
                for (int x = (int) Math.floor(box.getMinX() / rescale); x <= (int) Math.min(Math.ceil(box.getMaxX() / rescale), w); x++) {
                    pixels[w * y + x] = grayScaleTresh;
                }
            }
        });

//        var image2 = createImageFromMatrix(pixels, w, h);

        return findCCBoundingBoxes(pixels, w, h, grayScaleTresh, rescale, imageCTM);
    }


    public static BufferedImage createImageFromMatrix(int[] matrix, int width, int height) {

        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_BYTE_GRAY);

        byte[] pixelData = ((DataBufferByte) image.getRaster().getDataBuffer()).getData();

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                int index = y * width + x;
                int pixel = matrix[index]; // Assuming each element in the matrix represents a pixel color
                pixelData[index] = (byte) pixel;
            }
        }

        return image;
    }


    private List<Box> findCCBoundingBoxes(int[] pixels, int w, int h, int pixThreshold, int rescale, AffineTransform imageCTM) {

        DistinctQueue<Integer> pixelsToExplore = new DistinctQueue<>();
        var boundingBoxes = new ArrayList<Box>();

        for (int y = 0; y < h; y++) {
            for (int x = 0; x < w; x++) {
                var pixelIndex = x + y * w;
                if (pixels[pixelIndex] < pixThreshold) {
                    var minX = x;
                    var maxX = x;
                    var minY = y;
                    var maxY = y;
                    pixelsToExplore.enqueue(pixelIndex);
                    while (!pixelsToExplore.isEmpty()) {
                        var currentPixel = pixelsToExplore.dequeue();
                        if (currentPixel > w) {
                            var lowerPixel = currentPixel - w;
                            if (pixels[lowerPixel] < pixThreshold) {
                                pixelsToExplore.enqueue(currentPixel - w);
                                minY = Math.min(minY, lowerPixel / w);
                            }
                        }
                        if (currentPixel < pixels.length - w) {
                            var upperPixel = currentPixel + w;
                            if (pixels[upperPixel] < pixThreshold) {
                                pixelsToExplore.enqueue(upperPixel);
                                maxY = Math.max(maxY, upperPixel / w);
                            }
                        }
                        if (currentPixel % w != 0) {
                            var leftPixel = currentPixel - 1;
                            if (pixels[leftPixel] < pixThreshold) {
                                pixelsToExplore.enqueue(leftPixel);
                                minX = Math.min(minX, leftPixel % w);
                            }
                        }
                        if ((currentPixel + 1) % w != 0) {
                            var rightPixel = currentPixel + 1;
                            if (pixels[rightPixel] < pixThreshold) {
                                pixelsToExplore.enqueue(rightPixel + 1);
                                maxX = Math.max(maxX, rightPixel % w);
                            }
                        }
                        // Set the current pixel to white so we don't visit it again.
                        pixels[currentPixel] = pixThreshold;
                    }
                    boundingBoxes.add(new Box(minX * rescale, minY * rescale, maxX * rescale, maxY * rescale));
                }
            }
        }
        return boundingBoxes.stream().filter(box -> box.area() > 0).map(box -> box.transform(imageCTM)).collect(Collectors.toList());
    }


    public AffineTransform getImageCTM(PageInformation pageInformation, int imageWidth) {

        double scalingFactor = calculateScalingFactor(pageInformation, imageWidth);
        AffineTransform imageToCropBoxScaling = new AffineTransform(scalingFactor, 0, 0, scalingFactor, -pageInformation.minX(), -pageInformation.minY());

        AffineTransform mirrorMatrix = new AffineTransform(1, 0, 0, -1, 0, pageInformation.height());

        AffineTransform rotationMatrix = switch (pageInformation.rotationDegrees()) {
            case 90 -> new AffineTransform(0, 1, -1, 0, pageInformation.height(), 0);
            case 180 -> new AffineTransform(-1, 0, 0, -1, pageInformation.width(), pageInformation.height());
            case 270 -> new AffineTransform(0, -1, 1, 0, pageInformation.width() - pageInformation.height(), pageInformation.height()); // results from 90 + 180 rotations
            default -> new AffineTransform();
        };

        // matrix multiplication is performed from right to left, so the order is reversed.
        // scaling -> mirror -> rotation
        AffineTransform resultMatrix = new AffineTransform();

        resultMatrix.concatenate(rotationMatrix);
        resultMatrix.concatenate(mirrorMatrix);
        resultMatrix.concatenate(imageToCropBoxScaling);
        return resultMatrix;
    }


    private double calculateScalingFactor(PageInformation pageInformation, int imageWidth) {

        // PDFBox always returns page height and width based on rotation
        double pageWidth;
        if (pageInformation.rotationDegrees() == 90 || pageInformation.rotationDegrees() == 270) {
            pageWidth = pageInformation.height();
        } else {
            pageWidth = pageInformation.width();
        }

        return pageWidth / imageWidth;
    }

}
