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

import ai.djl.Model;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.index.NDIndex;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import ai.djl.util.Utils;
import cn.smartjavaai.common.entity.Point;
import com.xindao.ocr.smartjavaai.entity.OcrBox;
import com.xindao.ocr.smartjavaai.entity.OcrItem;
import com.xindao.ocr.smartjavaai.entity.TableStructureResult;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

/**
 * 表格识别的前后处理
 */
public class TableStructTranslator implements Translator<Image, TableStructureResult> {

    private final int maxLength = 488;
    private int height;
    private int width;
    private float scale = 1.0f;
    private float xScale;
    private float yScale;
    private List<String> dict;
    private String beg_str = "sos";
    private String end_str = "eos";
    private List<String> td_token = new ArrayList<>();

    @Override
    public void prepare(TranslatorContext ctx) throws IOException {
        Model model = ctx.getModel();
        try (InputStream is = model.getArtifact("table_structure_dict_ch.txt").openStream()) {
            dict = Utils.readLines(is, false);
            dict.add(0,beg_str);
            if(dict.contains("<td>"))
                dict.remove("<td>");
            if(!dict.contains("<td></td>"))
                dict.add("<td></td>");
            dict.add(end_str);
        }

        td_token.add("<td>");
        td_token.add("<td");
        td_token.add("<td></td>");
    }

    @Override
    public NDList processInput(TranslatorContext ctx, Image input) {
        NDArray img = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
        height = input.getHeight();
        width = input.getWidth();

        img = ResizeTableImage(img, height, width, maxLength);
        img = PaddingTableImage(ctx, img, maxLength);

        img = img.transpose(2, 0, 1).div(255).flip(0);
        img = NDImageUtils.normalize(
                        img, new float[]{0.485f, 0.456f, 0.406f}, new float[]{0.229f, 0.224f, 0.225f});
        img = img.expandDims(0);
        return new NDList(img);
    }

    @Override
    public TableStructureResult processOutput(TranslatorContext ctx, NDList list) {
        NDArray bbox_preds = list.get(0);
        NDArray structure_probs = list.get(1);

        NDArray structure_idx = structure_probs.argMax(2);
        structure_probs = structure_probs.max(new int[]{2});

        List<List<String>> structure_batch_list = new ArrayList<>();
        List<List<NDArray>> bbox_batch_list = new ArrayList<>();
        List<List<NDArray>> result_score_list = new ArrayList<>();

        // get ignored tokens
        int beg_idx = dict.indexOf(beg_str);
        int end_idx = dict.indexOf(end_str);

        long batch_size = structure_idx.size(0);
        for (int batch_idx = 0; batch_idx < batch_size; batch_idx++) {
            List<String> structure_list = new ArrayList<>();
            List<NDArray> bbox_list = new ArrayList<>();
            List<NDArray> score_list = new ArrayList<>();

            long len = structure_idx.get(batch_idx).size();
            for (int idx = 0; idx < len; idx++) {
                int char_idx = (int) structure_idx.get(batch_idx).get(idx).toLongArray()[0];
                if (idx > 0 && char_idx == end_idx) {
                    break;
                }
//                if (char_idx == beg_idx || char_idx == end_idx) {
//                    continue;
//                }
                String text = dict.get(char_idx);
                if(td_token.indexOf(text)>-1){
                    NDArray bbox = bbox_preds.get(batch_idx, idx);
//                    bbox.set(new NDIndex("0::2"), bbox.get(new NDIndex("0::2")));
//                    bbox.set(new NDIndex("1::2"), bbox.get(new NDIndex("1::2")));
                    bbox_list.add(bbox);
                }
                structure_list.add(text);
                score_list.add(structure_probs.get(batch_idx, idx));
            }

            structure_batch_list.add(structure_list); // structure_str
            bbox_batch_list.add(bbox_list);
            result_score_list.add(score_list);
        }
        List<String> structure_str_list =structure_batch_list.get(0);
        List<NDArray> bbox_list = bbox_batch_list.get(0);
        List<NDArray> score_list = result_score_list.get(0);

        structure_str_list.add(0,"<html>");
        structure_str_list.add(1,"<body>");
        structure_str_list.add(2,"<table>");
        structure_str_list.add("</table>");
        structure_str_list.add("</body>");
        structure_str_list.add("</html>");

        List<OcrItem> ocrItemList = new ArrayList<>();

        for (int i = 0; i < bbox_list.size(); i++) {
            NDArray box = bbox_list.get(i);
            float[] arr = new float[4];
            arr[0] = box.get(new NDIndex("0::2")).min().toFloatArray()[0];
            arr[1] = box.get(new NDIndex("1::2")).min().toFloatArray()[0];
            arr[2] = box.get(new NDIndex("0::2")).max().toFloatArray()[0];
            arr[3] = box.get(new NDIndex("1::2")).max().toFloatArray()[0];

            Point topLeft = new Point(arr[0] * xScale * width, arr[1] * yScale * height);
            Point topRight = new Point(arr[2] * xScale * width, arr[1] * yScale * height);
            Point bottomRight = new Point(arr[2] * xScale * width, arr[3] * yScale * height);
            Point bottomLeft = new Point(arr[0] * xScale * width, arr[3] * yScale * height);


            OcrBox ocrBox = new OcrBox(topLeft, topRight, bottomRight, bottomLeft);
            //String tag = structure_str_list.get(i + 3); // 前面加了<html><body><table> 所以偏移+3
            float score = score_list.get(i).toFloatArray()[0]; // 获取每个结构token的得分
            OcrItem item = new OcrItem();
            item.setOcrBox(ocrBox);
            item.setScore(score);
            //item.setTableTag(tag);
            ocrItemList.add(item);
        }
        return new TableStructureResult(ocrItemList, structure_str_list);
    }

    @Override
    public Batchifier getBatchifier() {
        return null;
    }

    private NDArray ResizeTableImage(NDArray img, int height, int width, int maxLen) {
        int localMax = Math.max(height, width);
        float ratio = maxLen * 1.0f / localMax;
        int resize_h = (int) (height * ratio);
        int resize_w = (int) (width * ratio);
        scale = ratio;

        if(width > height){
            xScale = 1f;
            yScale = (float)width /(float)height;
        } else{
            xScale = (float)height /(float)width;
            yScale = 1f;
        }

        img = NDImageUtils.resize(img, resize_w, resize_h);
        return img;
    }

    private NDArray PaddingTableImage(TranslatorContext ctx, NDArray img, int maxLen) {
        NDArray paddingImg = ctx.getNDManager().zeros(new Shape(maxLen, maxLen, 3), DataType.UINT8);
        paddingImg.set(
                new NDIndex("0:" + img.getShape().get(0) + ",0:" + img.getShape().get(1) + ",:"), img);
        return paddingImg;
    }

}