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

import java.awt.geom.Rectangle2D;
import java.awt.geom.RectangularShape;
import java.util.ArrayList;
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.Word;

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<Word> mergeAndSort(List<TextPageBlock> textBlocks) {

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


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

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


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

        return sortLines(groupByLine(sequences));
    }


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

        List<List<Word>> lineBlocks = new ArrayList<>();
        for (Set<Word> line : lines) {
            List<Word> sortedLine = sortByXDirAdj(line);
            if (!sortedLine.isEmpty()) {
                lineBlocks.add(sortedLine);
            }
        }
        // need to use old sorting, since COMPARATOR_DIR_ADJ is not transitive
        QuickSort.sort(lineBlocks, Comparator.comparing(line -> line.get(0), COMPARATOR_DIR_ADJ));

        List<Word> list = new ArrayList<>();
        for (List<Word> words : lineBlocks) {
            list.addAll(words);
        }
        return list;
    }


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

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


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

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

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

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

                if (sequence.equals(sequence2)) { // || unionFind.inSameSet(sequence, sequence2)) doing this is actually slower than not doing it
                    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()) {
                    continue;
                }
                if (Math.abs(sequence.getFontSize() - sequence2.getFontSize()) > 0.5 * Math.max(sequence.getFontSize(), sequence2.getFontSize())) {
                    continue;
                }
                if (Math.pow(normalizedVerticalDistance, 2) + Math.pow(normalizedHorizontalDistance, 2) > 1) {
                    continue;
                }
                if (!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<Word> merge(List<TextPageBlock> textBlocks) {

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

}
