package com.xindao.ocr.smartjavaai.model.table;

import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import cn.smartjavaai.common.entity.DetectionRectangle;
import cn.smartjavaai.common.entity.R;
import cn.smartjavaai.common.utils.FileUtils;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import com.xindao.ocr.smartjavaai.config.OcrRecOptions;
import com.xindao.ocr.smartjavaai.entity.OcrBox;
import com.xindao.ocr.smartjavaai.entity.OcrInfo;
import com.xindao.ocr.smartjavaai.entity.OcrItem;
import com.xindao.ocr.smartjavaai.entity.TableStructureResult;
import com.xindao.ocr.smartjavaai.exception.OcrException;
import com.xindao.ocr.smartjavaai.model.common.detect.OcrCommonDetModel;
import com.xindao.ocr.smartjavaai.model.common.direction.OcrDirectionModel;
import com.xindao.ocr.smartjavaai.model.common.recognize.OcrCommonRecModel;
import com.xindao.ocr.smartjavaai.utils.ConvertHtml2Excel;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.opencv.core.Mat;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Paths;
import java.util.*;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 表格内容识别器
 * @author dwj
 */
@Slf4j
public class TableRecognizer {

    private OcrCommonDetModel textDetector;
    private TableStructureModel tableStructureModel;
    private OcrCommonRecModel textRecModel;
    private OcrDirectionModel directionModel;

    private TableRecognizer(Builder builder) {
        this.tableStructureModel = builder.tableStructureModel;
        this.textRecModel = builder.textRecModel;
        this.directionModel = builder.directionModel;
        this.textDetector = builder.textDetector;
        textRecModel.setTextDetModel(textDetector);
        textRecModel.setDirectionModel(directionModel);
    }

    public static Builder builder() {
        return new Builder();
    }

    // 链式设置文本识别模型
    public TableRecognizer withTextRecModel(OcrCommonRecModel textRecModel) {
        this.textRecModel = textRecModel;
        return this;
    }

    // 链式设置表格结构模型
    public TableRecognizer withStructureModel(TableStructureModel tableStructureModel) {
        this.tableStructureModel = tableStructureModel;
        return this;
    }

    /**
     * 表格识别
     * @param image
     * @return
     */
    public R<TableStructureResult> recognize(Image image) {
        //表格结构识别
        R<TableStructureResult> result = tableStructureModel.detect(image);
        if(!result.isSuccess()){
            return R.fail(result.getCode(), result.getMessage());
        }
        //文本检测+文字识别
        boolean enableDirectionCorrect = directionModel == null ? false : true;
        OcrRecOptions options = new OcrRecOptions(enableDirectionCorrect, false);
        OcrInfo ocrInfo = textRecModel.recognize(image, options);
        List<String> tableContentList = buildTable(result.getData(), ocrInfo);
        String html = convertHtml(result.getData().getTableTagList(), tableContentList);
        result.getData().setHtml(html);
        return result;
    }


    /**
     * 表格识别
     * @param image
     * @return
     */
    public R<TableStructureResult> recognize(BufferedImage image) {
        if(!ImageUtils.isImageValid(image)){
            return R.fail(R.Status.INVALID_IMAGE);
        }
        Image img = null;
        try {
            img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
            return recognize(img);
        } catch (Exception e) {
            throw new OcrException(e);
        } finally {
            if(Objects.nonNull(img)){
                ((Mat)img.getWrappedImage()).release();
            }
        }
    }

    /**
     * 表格识别
     * @param imagePath
     * @return
     */
    public R<TableStructureResult> recognize(String imagePath) {
        if(!FileUtils.isFileExists(imagePath)){
            return R.fail(R.Status.FILE_NOT_FOUND);
        }
        Image img = null;
        try {
            img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
            return recognize(img);
        } catch (IOException e) {
            throw new OcrException("无效的图片", e);
        } finally {
            if (Objects.nonNull(img)){
                ((Mat)img.getWrappedImage()).release();
            }
        }
    }


    /**
     * 表格识别
     * @param imageData
     * @return
     */
    public R<TableStructureResult> recognize(byte[] imageData) {
        if(Objects.isNull(imageData)){
            return R.fail(R.Status.INVALID_IMAGE);
        }
        try {
            BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData));
            return recognize(image);
        } catch (IOException e) {
            throw new OcrException("错误的图像", e);
        }
    }

    /**
     * 绘制表格
     * @param tableStructureResult
     * @param image
     * @param savePath
     */
    public void drawTable(TableStructureResult tableStructureResult, BufferedImage image, String savePath){
        if(Objects.isNull(tableStructureResult) || CollectionUtils.isEmpty(tableStructureResult.getTableTagList())){
            throw new OcrException("表格结构为空");
        }
        for (int i = 0; i < tableStructureResult.getOcrItemList().size(); i++){
            OcrItem item = tableStructureResult.getOcrItemList().get(i);
            DetectionRectangle detectionRectangle = item.getOcrBox().toDetectionRectangle();
            ImageUtils.drawImageRectWithText(image, detectionRectangle, i + "", Color.RED);
        }
        ImageUtils.saveImage(image, savePath);
    }


    /**
     * 绘制表格
     * @param tableStructureResult
     * @param image
     * @return
     */
    public BufferedImage drawTable(TableStructureResult tableStructureResult, BufferedImage image){
        if(Objects.isNull(tableStructureResult) || CollectionUtils.isEmpty(tableStructureResult.getTableTagList())){
            throw new OcrException("表格结构为空");
        }
        for (int i = 0; i < tableStructureResult.getOcrItemList().size(); i++){
            OcrItem item = tableStructureResult.getOcrItemList().get(i);
            DetectionRectangle detectionRectangle = item.getOcrBox().toDetectionRectangle();
            ImageUtils.drawImageRectWithText(image, detectionRectangle, i + "", Color.RED);
        }
        return image;
    }

    /**
     * 删除 HTML 中第一个 <style> ... </style> 段落
     * @param html 原始 HTML
     * @return 去掉 <style> 的 HTML
     */
    public static String removeStyleBlock(String html) {
        String lowerHtml = html.toLowerCase();
        int styleStart = lowerHtml.indexOf("<style");
        if (styleStart == -1) {
            return html; // 没有 style，返回原文
        }
        int styleEnd = lowerHtml.indexOf("</style>", styleStart);
        if (styleEnd == -1) {
            return html; // 没闭合标签，不处理
        }
        styleEnd += "</style>".length();
        // 去掉 style 块
        return html.substring(0, styleStart) + html.substring(styleEnd);
    }


    /**
     * 导出 Excel
     * @param html
     * @param out
     */
    public void exportExcel(String html, OutputStream out){
        String content = removeStyleBlock(html);
        content = content.replace("<html><body>", "");
        content = content.replace("</body></html>", "");
        try (HSSFWorkbook workbook = ConvertHtml2Excel.table2Excel(content)){
            workbook.write(out);
            out.flush();
        } catch (Exception e) {
            throw new OcrException("导出excel失败，请检查表结构是否识别正确");
        }
    }

    /**
     * 导出 Excel
     * @param html
     * @param savePath
     */
    public void exportExcel(String html, String savePath){
        String content = removeStyleBlock(html);
        content = content.replace("<html><body>", "");
        content = content.replace("</body></html>", "");
        try (HSSFWorkbook workbook = ConvertHtml2Excel.table2Excel(content)){
            workbook.write(new File(savePath));
        } catch (Exception e) {
            throw new OcrException("导出excel失败，请检查表结构是否识别正确");
        }
    }


    /**
     * 构建表格
     * @param tableStructureResult
     * @param ocrInfo
     * @return
     */
    public List<String> buildTable(TableStructureResult tableStructureResult, OcrInfo ocrInfo) {
        // 获取 Cell 与 文本检测框 的对应关系(1:N)。
        Map<Integer, List<Integer>> matched = new ConcurrentHashMap<>();
        List<OcrItem> ocrItems = ocrInfo.getOcrItemList();

        for (int i = 0; i < ocrItems.size(); i++) {
            OcrBox ocrBox = ocrItems.get(i).getOcrBox();
            int[] box_1 = {
                    (int)ocrBox.getTopLeft().getX(),
                    (int)ocrBox.getTopLeft().getY(),
                    (int)ocrBox.getBottomRight().getX(),
                    (int)ocrBox.getBottomRight().getY()
            };
            // 获取两两cell之间的L1距离和 1- IOU
            List<Pair<Float, Float>> distances = new ArrayList<>();
            for (OcrItem cell : tableStructureResult.getOcrItemList()) {
                OcrBox cellBox = cell.getOcrBox();
                int[] box_2 = {
                        (int)cellBox.getTopLeft().getX(),
                        (int)cellBox.getTopLeft().getY(),
                        (int)cellBox.getBottomRight().getX(),
                        (int)cellBox.getBottomRight().getY()
                };
                float distance = distance(box_1, box_2);
                float iou = 1 - computeIou(box_1, box_2);
                distances.add(Pair.of(distance, iou));
            }
            // 根据距离和IOU挑选最"近"的cell
            Pair<Float, Float> nearest = sorted(distances);

            // 获取最小距离对应的下标id，也等价于cell的下标id  （distances列表是根据遍历cells生成的）
            int id = 0;
            for (int idx = 0; idx < distances.size(); idx++) {
                Pair<Float, Float> current = distances.get(idx);
                if (current.getLeft().floatValue() == nearest.getLeft().floatValue()
                        && current.getRight().floatValue() == nearest.getRight().floatValue()) {
                    id = idx;
                    break;
                }
            }
            if (!matched.containsKey(id)) {
                List<Integer> textIds = new ArrayList<>();
                textIds.add(i);
                // cell id, text id list (dt_boxes index list)
                matched.put(id, textIds);
            } else {
                matched.get(id).add(i);
            }
        }

        List<String> cell_contents = new ArrayList<>();
        List<Double> probs = new ArrayList<>();
        for (int i = 0; i < tableStructureResult.getOcrItemList().size(); i++) {
            List<Integer> textIds = matched.get(i);
            List<String> contents = new ArrayList<>();
            String content = "";
            if (textIds != null) {
                for (Integer id : textIds) {
                    contents.add(ocrItems.get(id).getText());
                }
                content = StringUtils.join(contents, " ");
            }
            cell_contents.add(content);
            probs.add(-1.0);
        }
        return cell_contents;
    }

    /**
     * 计算欧式距离
     * Calculate L1 distance
     *
     * @param box_1
     * @param box_2
     * @return
     */
    private int distance(int[] box_1, int[] box_2) {
        int x1 = box_1[0];
        int y1 = box_1[1];
        int x2 = box_1[2];
        int y2 = box_1[3];
        int x3 = box_2[0];
        int y3 = box_2[1];
        int x4 = box_2[2];
        int y4 = box_2[3];
        int dis = Math.abs(x3 - x1) + Math.abs(y3 - y1) + Math.abs(x4 - x2) + Math.abs(y4 - y2);
        int dis_2 = Math.abs(x3 - x1) + Math.abs(y3 - y1);
        int dis_3 = Math.abs(x4 - x2) + Math.abs(y4 - y2);
        return dis + Math.min(dis_2, dis_3);
    }

    /**
     * 计算交并比
     * computing IoU
     *
     * @param rec1: (y0, x0, y1, x1), which reflects (top, left, bottom, right)
     * @param rec2: (y0, x0, y1, x1)
     * @return scala value of IoU
     */
    private float computeIou(int[] rec1, int[] rec2) {
        // computing area of each rectangles
        int S_rec1 = (rec1[2] - rec1[0]) * (rec1[3] - rec1[1]);
        int S_rec2 = (rec2[2] - rec2[0]) * (rec2[3] - rec2[1]);

        // computing the sum_area
        int sum_area = S_rec1 + S_rec2;

        // find the each edge of intersect rectangle
        int left_line = Math.max(rec1[1], rec2[1]);
        int right_line = Math.min(rec1[3], rec2[3]);
        int top_line = Math.max(rec1[0], rec2[0]);
        int bottom_line = Math.min(rec1[2], rec2[2]);

        // judge if there is an intersect
        if (left_line >= right_line || top_line >= bottom_line) {
            return 0.0f;
        } else {
            float intersect = (right_line - left_line) * (bottom_line - top_line);
            return (intersect / (sum_area - intersect)) * 1.0f;
        }
    }

    /**
     * 距离排序
     * Distance sorted
     *
     * @param distances
     * @return
     */
    private Pair<Float, Float> sorted(List<Pair<Float, Float>> distances) {
        Comparator<Pair<Float, Float>> comparator =
                new Comparator<Pair<Float, Float>>() {
                    @Override
                    public int compare(Pair<Float, Float> a1, Pair<Float, Float> a2) {
                        // 首先根据IoU排序
                        if (a1.getRight().floatValue() > a2.getRight().floatValue()) {
                            return 1;
                        } else if (a1.getRight().floatValue() == a2.getRight().floatValue()) {
                            // 然后根据L1距离排序
                            if (a1.getLeft().floatValue() > a2.getLeft().floatValue()) {
                                return 1;
                            }
                            return -1;
                        }
                        return -1;
                    }
                };

        // 距离排序
        List<Pair<Float, Float>> newDistances = new ArrayList<>();
        CollectionUtils.addAll(newDistances, new Object[distances.size()]);
        Collections.copy(newDistances, distances);
        Collections.sort(newDistances, comparator);
        return newDistances.get(0);
    }

    /**
     * 生成表格html
     * Generate table html
     *
     * @param pred_structures
     * @param cell_contents
     * @return
     */
    public String convertHtml(List<String> pred_structures, List<String> cell_contents) {
        StringBuffer html = new StringBuffer();
        // 添加统一的样式（可选放到<head>中）
        html.append("<style>\n");
        html.append("table { border-collapse: collapse; }\n");
        html.append("td, th, table { border: 1px solid black; padding: 5px; }\n");
        html.append("</style>\n");
        int td_index = 0;
        for (String tag : pred_structures) {
            if (tag.contains("<td></td>")) {
                String content = cell_contents.get(td_index);
                html.append("<td>");
                html.append(content);
                html.append("</td>");
                td_index++;
                continue;
            }
            html.append(tag);
        }
        return html.toString();
    }


    public static class Builder {
        private TableStructureModel tableStructureModel;
        private OcrCommonRecModel textRecModel;
        private OcrDirectionModel directionModel;
        private OcrCommonDetModel textDetector;

        public Builder withStructureModel(TableStructureModel model) {
            this.tableStructureModel = model;
            return this;
        }

        public Builder withTextRecModel(OcrCommonRecModel model) {
            this.textRecModel = model;
            return this;
        }

        public Builder withDirectionModel(OcrDirectionModel model) {
            this.directionModel = model;
            return this;
        }

        public Builder withTextDetModel(OcrCommonDetModel model) {
            this.textDetector = model;
            return this;
        }

        public TableRecognizer build() {
            if (this.tableStructureModel == null) {
                throw new IllegalStateException("tableStructureModel 未设置");
            }
            if (this.textDetector == null) {
                throw new IllegalStateException("textDetector 未设置");
            }
            if (this.textRecModel == null) {
                throw new IllegalStateException("textRecModel 未设置");
            }
            return new TableRecognizer(this);
        }
    }

}