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

import static com.knecon.fforesight.service.layoutparser.processor.services.classification.ClassificationPatterns.AMOUNT_PATTERN;
import static com.knecon.fforesight.service.layoutparser.processor.services.classification.ClassificationPatterns.NUMERIC;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

import org.springframework.stereotype.Service;

import com.knecon.fforesight.service.layoutparser.processor.docstrum.model.UnionFind;
import com.knecon.fforesight.service.layoutparser.processor.model.AbstractPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationDocument;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationPage;
import com.knecon.fforesight.service.layoutparser.processor.model.PageBlockType;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextBlockOnPage;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.text.Word;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPositionSequenceComparator;
import com.knecon.fforesight.service.layoutparser.processor.utils.TextNormalizationUtilities;
import com.knecon.fforesight.service.layoutparser.processor.visualization.LayoutDebugLayer;

import lombok.extern.slf4j.Slf4j;

@Slf4j
@Service
public class TableOfContentsClassificationService {

    private static final int MAX_PAGE_COUNT = 10; // maximum length of a toc to avoid runaway classification
    private static final int SURROUNDING_BLOCKS_RADIUS = 10; // number of blocks to look ahead
    private static final int MINIMUM_MATCHES = 2; // minimum cluster size
    public static final int INTERSECTION_TOLERANCE = 2; // equality threshold for x intersection
    public static final int DENSITY_THRESHOLD_COUNT = 10; // describes the minimum density, at least this many entries per page height are required


    @SuppressWarnings("checkstyle:ModifiedControlVariable")
    public void classifyTableOfContents(ClassificationDocument document) {

        List<TextBlockOnPage> textBlocks = buildBlocksPerPage(document);

        for (int i = 0; i < textBlocks.size(); i++) {
            TextBlockOnPage textBlock = textBlocks.get(i);

            if (!isTOCHeadline(textBlock)) {
                continue;
            }

            int end = identifyTOCItems(i + 1, textBlocks, document);

            if (end > i + 1) {
                if (textBlock.textBlock().getClassification() == null) {
                    textBlock.textBlock().setClassification(PageBlockType.H1);
                }
                i = end;
            }
        }
    }


    private int identifyTOCItems(int start, List<TextBlockOnPage> textBlocks, ClassificationDocument document) {

        ClassificationPage startPage = textBlocks.get(start).page();
        List<TextBlockOnPage> initialLookAhead = textBlocks.subList(start, Math.min(start + SURROUNDING_BLOCKS_RADIUS, textBlocks.size()));
        HashMap<Word, TextBlockOnPage> lookup = new HashMap<>();
        List<Word> numbers = extractNumbers(initialLookAhead, lookup, document.getPages().size());
        TocNumberFinder tocNumberFinder = new TocNumberFinder(numbers, lookup);

        int lastCandidate = start;
        for (int i = start; i < Math.min(lastCandidate + SURROUNDING_BLOCKS_RADIUS, textBlocks.size()); i++) {

            TextBlockOnPage textBlockOnPage = textBlocks.get(i);
            if (textBlockOnPage.page().getPageNumber() - MAX_PAGE_COUNT > startPage.getPageNumber()) {
                break;
            }

            if (textBlockOnPage.textBlock().getClassification() != null //
                && textBlockOnPage.textBlock().getClassification().isHeadline() //
                && !(textBlockOnPage.textBlock().getText().startsWith("TABLES") //
                     || textBlockOnPage.textBlock().getText().startsWith("APPENDICES") //
                     || textBlockOnPage.textBlock().getText().startsWith("FIGURES"))) {
                log.debug("hit an outline headline, stop immediately.");
                lastCandidate = i - 1;
                break;
            }

            List<Word> numbersFromBlock = extractNumbers(textBlockOnPage, lookup, document.getPages().size());

            List<Word> currentRightmostCluster = tocNumberFinder.getCurrentRightmostCluster();

            if (currentRightmostCluster.size() < MINIMUM_MATCHES) {
                log.debug("No numbers indicating a table of contents here.");
                return start;
            }

            if (anyIntersection(currentRightmostCluster, numbersFromBlock, lookup)) {
                lastCandidate = i;
                numbersFromBlock.forEach(tocNumberFinder::add);
            }
        }

        addVisualization(document.getLayoutDebugLayer(), tocNumberFinder, lookup);

        Set<TextBlockOnPage> blocksWithNumberInCluster = tocNumberFinder.getCurrentRightmostCluster()
                .stream()
                .map(lookup::get)
                .collect(Collectors.toSet());

        int lastConfirmed = start;
        for (int i = start; i < lastCandidate + 1; i++) {
            TextBlockOnPage textBlockOnPage = textBlocks.get(i);
            if (blocksWithNumberInCluster.contains(textBlockOnPage)) {
                lastConfirmed = i;
            }
        }

        textBlocks.subList(start, lastConfirmed + 1)
                .stream()
                .filter(block -> (block.textBlock().getClassification() == null || !block.textBlock().getClassification().isHeadline()))
                .forEach(textBlockOnPage -> textBlockOnPage.textBlock().setClassification(PageBlockType.TABLE_OF_CONTENTS_ITEM));

        return lastCandidate;
    }


    private static void addVisualization(LayoutDebugLayer layoutDebugLayer, TocNumberFinder tocNumberFinder, Map<Word, TextBlockOnPage> lookup) {

        tocNumberFinder.getCurrentRightmostCluster()
                .stream()
                .collect(Collectors.groupingBy(key -> lookup.get(key).page().getPageNumber()))
                .forEach((pageNumber, number) -> layoutDebugLayer.addTocPages(number, pageNumber));
    }


    private static boolean anyIntersection(Collection<Word> numbers1,
                                           Collection<Word> numbers2,
                                           Map<Word, TextBlockOnPage> lookup) {

        return numbers1.stream()
                .anyMatch(numberFromCluster -> numbers2.stream()
                        .anyMatch(numberFromBlock -> matches(numberFromBlock, numberFromCluster, lookup)));
    }


    private static List<Word> extractNumbers(List<TextBlockOnPage> textBlocks, Map<Word, TextBlockOnPage> lookup, int numberOfPages) {

        List<Word> blocks = new LinkedList<>();
        for (TextBlockOnPage textBlock : textBlocks) {
            blocks.addAll(extractNumbers(textBlock, lookup, numberOfPages));
        }
        return blocks;
    }


    private static List<Word> extractNumbers(TextBlockOnPage textBlock, Map<Word, TextBlockOnPage> lookup, int numberOfPages) {

        List<Word> blocks = new LinkedList<>();
        TextPageBlock block = textBlock.textBlock();
        List<Word> sequences = block.getWords();
        for (int i = 0; i < sequences.size(); i++) {

            Word word = sequences.get(i);

            if (!NUMERIC.matcher(word).matches()) {
                continue;
            }

            if (AMOUNT_PATTERN.matcher(getSurroundingString(i, sequences)).matches()) {
                continue;
            }

            try {
                int pageNumber = Integer.parseInt(word.toString());
                if (0 >= pageNumber || pageNumber > numberOfPages) {
                    continue;
                }
                lookup.put(word, textBlock);
                blocks.add(word);
            } catch (NumberFormatException e) {
                log.debug("That wasn't a number! Should not happen, due to numeric check beforehand.");
            }
        }
        return blocks;
    }


    private static CharSequence getSurroundingString(int i, List<Word> sequences) {

        int end = Math.min(i + 5, sequences.size());
        return sequences.subList(i, end)
                .stream()
                .map(Word::toString)
                .collect(Collectors.joining(" "));
    }


    private static boolean matches(Word number1, Word number2, Map<Word, TextBlockOnPage> lookup) {

        if (number1.getDir() != number2.getDir()) {
            return false;
        }

        return number1.intersectsXDirAdj(number2, INTERSECTION_TOLERANCE);
    }


    private boolean isTOCHeadline(TextBlockOnPage textBlock) {

        if (textBlock.textBlock().getText().length() > 50) {
            return false;
        }
        String text = TextNormalizationUtilities.removeAllWhitespaces(textBlock.textBlock().getText().toLowerCase(Locale.ENGLISH));
        return (text.contains("content") && text.length() < "content".length() + 6) //
               || (text.contains("tableofcontent") && text.length() < "tableofcontent".length() + DENSITY_THRESHOLD_COUNT)//
               || text.equals("tables")//
               || text.equals("appendices")//
               || text.equals("figures");
    }


    private List<TextBlockOnPage> buildBlocksPerPage(ClassificationDocument document) {

        List<TextBlockOnPage> blocks = new ArrayList<>();
        for (ClassificationPage page : document.getPages()) {
            for (AbstractPageBlock abstractPageBlock : page.getTextBlocks()) {
                if (abstractPageBlock instanceof TextPageBlock textBlock) {
                    if (textBlock.getClassification() != null && (textBlock.getClassification().equals(PageBlockType.HEADER) //
                                                                  || textBlock.getClassification().equals(PageBlockType.FOOTER))) {
                        continue;
                    }
                    blocks.add(new TextBlockOnPage(textBlock, page));
                }
            }
        }
        return blocks;
    }


    private static class TocNumberFinder {

        final UnionFind<Word> numberClusters;
        final HashMap<Word, TextBlockOnPage> lookup;


        TocNumberFinder(List<Word> blocks, HashMap<Word, TextBlockOnPage> lookup) {

            this.numberClusters = new UnionFind<>(new HashSet<>(blocks));
            for (int i = 0; i < blocks.size(); i++) {
                for (int j = i + 1; j < blocks.size(); j++) {
                    if (matches(blocks.get(i), blocks.get(j), lookup)) {
                        numberClusters.union(blocks.get(i), blocks.get(j));
                    }
                }
            }
            this.lookup = lookup;
        }


        public void add(Word number) {

            if (numberClusters.getElements().contains(number)) {
                return;
            }

            numberClusters.addElement(number);
            for (Word element : numberClusters.getElements()) {
                if (matches(number, element, lookup)) {
                    numberClusters.union(element, number);
                }
            }
        }


        public List<Word> getCurrentRightmostCluster() {

            return numberClusters.getGroups()
                    .stream()
                    .filter(cluster -> cluster.size() > MINIMUM_MATCHES)
                    .map(cluster -> cluster.stream()
                            .sorted(new TextPositionSequenceComparator(lookup))
                            .toList())
                    .map(this::removeOutliers)
//                    .map(this::filterByMinimumDensity)
                    .filter(cluster -> cluster.size() > MINIMUM_MATCHES)
                    .max(Comparator.comparingDouble(cluster -> cluster.get(0).getBBox().getMaxX())).orElse(Collections.emptyList());
        }

// does not seem to be doing much, ideally instead of using the height of the blocks i would like to use the height, beginning from the MainBody top,
// but as the MainBody is often wrong, this results in some numbers being discarded even though they are correct.
//        private List<TextPositionSequence> filterByMinimumDensity(List<TextPositionSequence> numbers) {
//
//            Map<ClassificationPage, List<TextPositionSequence>> clustersPerPage = numbers.stream()
//                    .collect(Collectors.groupingBy(number -> lookup.get(number).page()));
//
//            List<TextPositionSequence> result = new ArrayList<>(numbers.size());
//            clustersPerPage.keySet()
//                    .stream()
//                    .sorted(Comparator.comparingInt(ClassificationPage::getPageNumber))
//                    .forEach(page -> {
//                        var numbersOnPage = clustersPerPage.get(page);
//
//                        double height = numbersOnPage.stream()
//                                .map(BoundingBox::getBBox)
//                                .collect(RectangleTransformations.collectBBox()).getHeight();
//
//                        double count = numbersOnPage.size();
//
//                        if ((count / height) >= (DENSITY_THRESHOLD_COUNT / page.getPageHeight())) {
//                            result.addAll(numbers);
//                        }
//                    });
//            return result;
//        }


        public List<Word> removeOutliers(List<Word> numbers) {

            List<Word> result = new ArrayList<>();

            result.add(numbers.get(0));

            for (int i = 1; i < numbers.size() - 1; i++) {
                int prev = getNumberAsInt(numbers, i - 1);
                int curr = getNumberAsInt(numbers, i);
                int next = getNumberAsInt(numbers, i + 1);

                if (!(curr <= prev || curr >= next) || !isBetterWithout(numbers, i)) {
                    result.add(numbers.get(i));
                }
            }
            if (getNumberAsInt(numbers, numbers.size() - 1) >= getNumberAsInt(numbers, Math.max(0, numbers.size() - 2))) {
                result.add(numbers.get(numbers.size() - 1));
            }

            return result;
        }


        // Helper method to check if removing the current number results in a better order
        public static boolean isBetterWithout(List<Word> numbers, int i) {

            if (i == 0 || i == numbers.size() - 1) {
                return false;
            }

            int prev = getNumberAsInt(numbers, i - 1);
            int curr = getNumberAsInt(numbers, i);
            int next = getNumberAsInt(numbers, i + 1);

            return (prev <= next) && (Math.abs(prev - next) < Math.abs(prev - curr) + Math.abs(curr - next));
        }

    }


    private static int getNumberAsInt(List<Word> numbers, int i) {

        return Integer.parseInt(numbers.get(i).toString());
    }

}
