package com.knecon.fforesight.service.layoutparser.processor.utils;

import java.awt.geom.Rectangle2D;
import java.awt.geom.RectangularShape;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import com.knecon.fforesight.service.layoutparser.processor.docstrum.model.AngleFilter;
import com.knecon.fforesight.service.layoutparser.processor.docstrum.model.TextBoundingBox;
import com.knecon.fforesight.service.layoutparser.processor.docstrum.model.UnionFind;
import com.knecon.fforesight.service.layoutparser.processor.docstrum.utils.DoubleUtils;
import com.knecon.fforesight.service.layoutparser.processor.docstrum.utils.FastAtan2;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPositionSequence;

import lombok.experimental.UtilityClass;

@UtilityClass
public class TextPositionOperations {

    public static final double ANGLE_TOLERANCE = Math.PI / 35;
    public static final AngleFilter ANGLE_FILTER = new AngleFilter(-ANGLE_TOLERANCE, ANGLE_TOLERANCE);
    public static final double MAX_LINE_HEIGHT_FACTOR = 0.66; // multiplied with max word height
    public static final double MAX_WORD_DISTANCE_FACTOR = 3.5; // multiplied with max word width

    private static final double Y_THRESHOLD = 5;
    private static final double X_THRESHOLD = 3;
    private static final Comparator<TextBoundingBox> COMPARATOR_DIR_ADJ = //
            Comparator.comparing(TextBoundingBox::getDir)
                    .thenComparing(TextBoundingBox::getMaxYDirAdj, (o1, o2) -> DoubleUtils.compareDouble(o1, o2, Y_THRESHOLD))
                    .thenComparing(TextBoundingBox::getXDirAdj, (o1, o2) -> DoubleUtils.compareDouble(o1, o2, X_THRESHOLD));


    public List<TextPositionSequence> mergeAndSort(List<TextPageBlock> textBlocks) {

        var sequences = textBlocks.stream()
                .flatMap(tb -> tb.getSequences()
                        .stream())
                .collect(Collectors.toSet());
        return sortUsingLineDetection(sequences);
    }


    public List<TextPositionSequence> sort(List<TextPositionSequence> sequences) {

        return sortUsingLineDetection(new HashSet<>(sequences));
    }


    private List<TextPositionSequence> sortUsingLineDetection(Set<TextPositionSequence> sequences) {

        return sortLines(groupByLine(sequences));

    }


    public List<TextPositionSequence> sortLines(Collection<Set<TextPositionSequence>> lines) {

        return lines.stream()
                .map(TextPositionOperations::sortByXDirAdj)
                .filter(line -> !line.isEmpty())
                .sorted(Comparator.comparing(line -> line.get(0), COMPARATOR_DIR_ADJ))
                .flatMap(Collection::stream)
                .toList();
    }


    private List<TextPositionSequence> sortByXDirAdj(Set<TextPositionSequence> line) {

        return line.stream()
                .sorted(Comparator.comparing(TextPositionSequence::getXDirAdj))
                .toList();
    }


    public Collection<Set<TextPositionSequence>> groupByLine(Set<TextPositionSequence> sequences) {

        double maxLineDistance = sequences.stream()
                                         .map(TextPositionSequence::getBBoxDirAdj)
                                         .mapToDouble(RectangularShape::getHeight).average()
                                         .orElse(10) * MAX_LINE_HEIGHT_FACTOR;
        double maxXGap = sequences.stream()
                                 .map(TextPositionSequence::getBBoxDirAdj)
                                 .mapToDouble(RectangularShape::getWidth).average()
                                 .orElse(75) * MAX_WORD_DISTANCE_FACTOR;

        UnionFind<TextPositionSequence> unionFind = new UnionFind<>(sequences);

        for (TextPositionSequence sequence : sequences) {
            for (TextPositionSequence sequence2 : sequences) {

                if (sequence.equals(sequence2) || unionFind.inSameSet(sequence, sequence2)) {
                    continue;
                }

                double angle = computeAngle(sequence.getBBoxDirAdj(), sequence2.getBBoxDirAdj());

                double normalizedVerticalDistance = Math.abs(sequence.getBBoxDirAdj().getCenterY() - sequence2.getBBoxDirAdj().getCenterY()) / maxLineDistance;
                double normalizedHorizontalDistance = Math.abs(sequence.getBBoxDirAdj().getCenterX() - sequence2.getBBoxDirAdj().getCenterX()) / maxXGap;

                if (sequence.getDir() != sequence2.getDir()
                    || Math.abs(sequence.getFontSize() - sequence2.getFontSize()) > 0.5 * Math.min(sequence.getFontSize(),
                                                                                                   sequence2.getFontSize())
                    || Math.pow(normalizedVerticalDistance, 2) + Math.pow(normalizedHorizontalDistance, 2) > 1
                    || !ANGLE_FILTER.matches(angle)) {
                    continue;
                }

                unionFind.union(sequence, sequence2);
            }
        }

        return unionFind.getGroups();
    }


    public double computeAngle(Rectangle2D rect1, Rectangle2D rect2) {

        double rect1CentroidX = rect1.getCenterX();
        double rect1CentroidY = rect1.getCenterY();
        double rect2CentroidX = rect2.getCenterX();
        double rect2CentroidY = rect2.getCenterY();

        double deltaX = rect2CentroidX - rect1CentroidX;
        double deltaY = rect2CentroidY - rect1CentroidY;

        return FastAtan2.fastAtan2(deltaY, deltaX);
    }


    public List<TextPositionSequence> merge(List<TextPageBlock> textBlocks) {

        return textBlocks.stream()
                .map(TextPageBlock::getSequences)
                .flatMap(Collection::stream)
                .collect(Collectors.toList());
    }

}
