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

import static com.knecon.fforesight.service.layoutparser.processor.utils.GeometricComparators.CELL_SIZE_COMPARATOR;
import static com.knecon.fforesight.service.layoutparser.processor.utils.GeometricComparators.RECTANGLE_SIZE_COMPARATOR;

import java.awt.geom.AffineTransform;
import java.awt.geom.Rectangle2D;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.springframework.stereotype.Service;

import com.knecon.fforesight.service.layoutparser.processor.docstrum.model.BoundingBox;
import com.knecon.fforesight.service.layoutparser.processor.model.AbstractPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.ClassificationPage;
import com.knecon.fforesight.service.layoutparser.processor.model.table.Cell;
import com.knecon.fforesight.service.layoutparser.processor.model.table.Ruling;
import com.knecon.fforesight.service.layoutparser.processor.model.table.TablePageBlock;
import com.knecon.fforesight.service.layoutparser.processor.model.text.RedTextPosition;
import com.knecon.fforesight.service.layoutparser.processor.model.text.TextPageBlock;
import com.knecon.fforesight.service.layoutparser.processor.utils.CoordinateTransforms;
import com.knecon.fforesight.service.layoutparser.processor.utils.DoubleComparisons;
import com.knecon.fforesight.service.layoutparser.processor.utils.PageInformation;
import com.knecon.fforesight.service.layoutparser.processor.utils.RectangularIntersectionFinder;
import com.knecon.fforesight.service.layoutparser.processor.utils.SpreadsheetFinder;

import lombok.SneakyThrows;

@Service
public class TableExtractionService {

    private static final int MAX_TABLE_CONTAINED_CELLS_WITH_TEXT = 1;
    private static final double TABLE_UNIFORMITY_THRESHOLD = 0.7;


    /**
     * Finds tables on a page and moves textblocks into cells of the found tables.
     * Note: This algorithm uses Pdf Coordinate System where {0,0} rotated with the page rotation.
     * 0 -> LowerLeft
     * 90 -> UpperLeft
     * 180 -> UpperRight
     * 270 -> LowerRight
     * <p>
     * DirAdj (Text direction adjusted) values can not be used here.
     *
     * @param emptyCells The cells used to build the table.
     * @param page       Page object that contains textblocks and statistics.
     */

    public void extractTables(List<Cell> emptyCells, ClassificationPage page) {

        // sort cells by size (height * width) ascending so that textBlocks are always assigned to the smallest cells that contain them
        emptyCells.sort(CELL_SIZE_COMPARATOR);

        for (AbstractPageBlock abstractPageBlock : page.getTextBlocks()) {
            TextPageBlock textBlock = (TextPageBlock) abstractPageBlock;
            for (Cell cell : emptyCells) {
                if (cell.hasMinimumSize() && doesCellContainTextBlock(cell, textBlock)) {
                    cell.addTextBlock(textBlock);
                    break;
                }
            }
        }

        List<Cell> cells = new ArrayList<>(new HashSet<>(emptyCells));
        DoubleComparisons.sort(cells, BoundingBox.ILL_DEFINED_ORDER);

        List<Rectangle2D> spreadsheetAreas = SpreadsheetFinder.findSpreadsheetsFromCells(cells);
        // sort spreadsheetAreas by size (height * width) ascending so that cells are placed in the smallest tables first
        // this way no cell duplication occurs when tables are contained in other tables and only the most inner table contains the cells
        spreadsheetAreas.sort(RECTANGLE_SIZE_COMPARATOR);

        List<TablePageBlock> tables = new ArrayList<>();
        for (Rectangle2D area : spreadsheetAreas) {

            List<Cell> containedCells = new ArrayList<>();
            for (Cell c : cells) {
                if (c.hasMinimumSize() && area.contains(c.getBBoxInitialUserSpace())) {
                    containedCells.add(c);
                }
            }

            var containedCellsWithText = containedCells.stream()
                    .filter(cell -> !cell.getTextBlocks().isEmpty())
                    .toList();

            // verify if table would contain fewer cells with text than the threshold allows
            if (containedCellsWithText.size() >= MAX_TABLE_CONTAINED_CELLS_WITH_TEXT && checkIfTableCellsAreUniform(containedCells)) {
                tables.add(new TablePageBlock(containedCells, page.getRotation()));
                cells.removeAll(containedCells);
            }
        }

        for (TablePageBlock table : tables) {
            int position = -1;

            for (AbstractPageBlock pageBlock : page.getTextBlocks()) {
                if (pageBlock instanceof TextPageBlock ? table.contains(pageBlock) : table.contains(pageBlock) && position == -1) {
                    position = page.getTextBlocks().indexOf(pageBlock);
                }
            }
            if (position != -1) {
                page.getTextBlocks().add(position, table);

                var toBeRemoved = table.getCells()
                        .stream()
                        .map(Cell::getTextBlocks)
                        .flatMap(List::stream)
                        .toList();
                // remove text blocks from the page that were also added with the table (from its contained cells)
                page.getTextBlocks().removeAll(toBeRemoved);
            }
        }
    }


    private boolean checkIfTableCellsAreUniform(List<Cell> containedCells) {

        if (containedCells.size() <= 2) {
            return true;
        }

        Map<Long, List<Long>> cellsGroupedByRoundedWidth = containedCells.stream()
                .map(BoundingBox::getWidth)
                .map(size -> Math.round(size / 10.0) * 10)
                .collect(Collectors.groupingBy(Long::longValue));

        return (double) cellsGroupedByRoundedWidth.size() / containedCells.size() <= TABLE_UNIFORMITY_THRESHOLD;
    }


    private boolean doesCellContainTextBlock(Cell cell, TextPageBlock textBlock) {

        return cell.contains(textBlock, RedTextPosition.HEIGHT_PADDING);
    }


    @SneakyThrows
    public static List<Cell> findCells(List<Ruling> horizontalRulingLines, List<Ruling> verticalRulingLines, PageInformation pageInformation) {

        AffineTransform affineTransform = CoordinateTransforms.calculateInitialUserSpaceCoordsToImageCoords(pageInformation, 1);
        /*
        switch (pageInformation.rotationDegrees()) {
            case 90 -> affineTransform.translate(RedTextPosition.HEIGHT_PADDING, 0); //although this is wrong, our text coordinates are wrong as well
            case 180 -> affineTransform.translate(0, RedTextPosition.HEIGHT_PADDING);
            case 270 -> affineTransform.translate(-RedTextPosition.HEIGHT_PADDING, 0);
            default -> affineTransform.translate(0, -RedTextPosition.HEIGHT_PADDING);
        }
        */
        return RectangularIntersectionFinder.find(horizontalRulingLines, verticalRulingLines)
                .stream()
                .map(rect -> new Cell(rect, affineTransform))
                .collect(Collectors.toList());
    }

}
