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.regex.Matcher;
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.SectionIdentifier;
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.TocNumberComparator;
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.TABLE_OF_CONTENTS_HEADLINE);
                }
                i = end;
            }
        }
    }


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

        if (start >= textBlocks.size()) {
            return start;
        }
        ClassificationPage startPage = textBlocks.get(start).page();
        List<TextBlockOnPage> initialLookAhead = textBlocks.subList(start, Math.min(start + SURROUNDING_BLOCKS_RADIUS, textBlocks.size()));
        HashMap<NumberWord, TextBlockOnPage> numberToBlockLookup = new HashMap<>();
        List<NumberWord> numbers = extractNumbers(initialLookAhead, numberToBlockLookup, document.getPages().size());
        TocNumberFinder tocNumberFinder = new TocNumberFinder(numbers, numberToBlockLookup);

        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<NumberWord> numbersFromBlock = extractNumbers(textBlockOnPage, numberToBlockLookup, document.getPages().size());

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

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

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

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

        addVisualization(document.getLayoutDebugLayer(), tocNumberFinder, numberToBlockLookup, blocksWithNumberInCluster, textBlocks.get(start - 1));

        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<NumberWord, TextBlockOnPage> lookup,
                                         Set<TextBlockOnPage> blocksWithNumberInCluster,
                                         TextBlockOnPage startingHeadline) {

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


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

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


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

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


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

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

            Word word = words.get(i);
            if (!wordIsEndOfLine(i, words)) {
                continue;
            }

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

            Matcher matcher = SectionIdentifier.numericalIdentifierPattern.matcher(word.toString());
            if (matcher.find() && matcher.group(2) != null) {
                continue;
            }

            Matcher numberFinder = NUMERIC.matcher(word);
            if (!numberFinder.find() || word.length() > 5) {
                continue;
            }

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


    private static boolean wordIsEndOfLine(int i, List<Word> words) {

        if (i == words.size() - 1) {
            return true;
        }
        Word word = words.get(i);
        Word nextWord = words.get(i + 1);
        return !nextWord.rightOf(word);
    }


    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(NumberWord number1, NumberWord number2, Map<NumberWord, TextBlockOnPage> lookup) {

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

        return number1.word().intersectsXDirAdj(number2.word(), 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<NumberWord> numberClusters;
        final HashMap<NumberWord, TextBlockOnPage> lookup;


        TocNumberFinder(List<NumberWord> blocks, HashMap<NumberWord, 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(NumberWord number) {

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

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


        public List<NumberWord> getCurrentRightmostCluster() {

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


        private List<NumberWord> removeOnNonConsecutivePages(List<NumberWord> numbers) {

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

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

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

                if (Math.abs(prev - curr) > 1) {
                    break;
                } else {
                    result.add(numbers.get(i));
                }
            }
            return result;
        }


        private int getPageNumber(List<NumberWord> numbers, int i) {

            return lookup.get(numbers.get(i)).page().getPageNumber();
        }


        private List<NumberWord> filterByWordNearTopOfPage(List<NumberWord> numbers) {

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

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

            for (int i = 1; i < numbers.size(); i++) {
                NumberWord prev = numbers.get(i - 1);
                NumberWord curr = numbers.get(i);
                ClassificationPage prevPage = lookup.get(prev).page();
                ClassificationPage currPage = lookup.get(curr).page();
                if (prevPage.equals(currPage)) {
                    result.add(curr);
                } else if (curr.word().getBBox().getMinY() < currPage.getPageHeight() * 0.33) {
                    result.add(curr);
                }
            }
            return result;
        }


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

            List<NumberWord> confirmedClusterNumbers = new ArrayList<>();

            confirmedClusterNumbers.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)) {
                    confirmedClusterNumbers.add(numbers.get(i));
                }
            }
            if (getNumberAsInt(numbers, numbers.size() - 1) >= getLatestNumber(confirmedClusterNumbers)) {
                confirmedClusterNumbers.add(numbers.get(numbers.size() - 1));
            }

            return confirmedClusterNumbers;
        }


        private static int getLatestNumber(List<NumberWord> confirmedClusterNumbers) {

            return confirmedClusterNumbers.get(confirmedClusterNumbers.size() - 1).number();
        }


        // Helper method to check if removing the current number results in a better order
        public static boolean isBetterWithout(List<NumberWord> 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<NumberWord> numbers, int i) {

        return numbers.get(i).number();
    }

}
