package com.xindao.ocr.smartjavaai.model.common.recognize.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 java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * 文字识别前后处理
 *
 */
public class PPOCRRecTranslator implements Translator<Image, String> {
    private List<String> table;
    private final boolean use_space_char;

    private String batchifier;

    public PPOCRRecTranslator(Map<String, ?> arguments) {
        use_space_char =
                arguments.containsKey("use_space_char")
                        ? Boolean.parseBoolean(arguments.get("use_space_char").toString())
                        : true;
        batchifier =  arguments.containsKey("batchifier")
                ? arguments.get("batchifier").toString()
                : "padding";
    }

    @Override
    public void prepare(TranslatorContext ctx) throws IOException {
        Model model = ctx.getModel();
        try (InputStream is = model.getArtifact("dict.txt").openStream()) {
            table = Utils.readLines(is, true);
            table.add(0, "blank");
            if(use_space_char){
                table.add(" ");
                table.add(" ");
            }
            else{
                table.add("");
                table.add("");
            }

        }
    }

    @Override
    public String processOutput(TranslatorContext ctx, NDList list) throws IOException {
        StringBuilder sb = new StringBuilder();
        NDArray tokens = list.singletonOrThrow();

//        long[] indices = tokens.get(0).argMax(1).toLongArray();
        long[] indices = tokens.argMax(1).toLongArray();
        boolean[] selection = new boolean[indices.length];
        Arrays.fill(selection, true);
        for (int i = 1; i < indices.length; i++) {
            if (indices[i] == indices[i - 1]) {
                selection[i] = false;
            }
        }

        // 字符置信度
//        float[] probs = new float[indices.length];
//        for (int row = 0; row < indices.length; row++) {
//            NDArray value = tokens.get(0).get(new NDIndex(""+ row +":" + (row + 1) +"," + indices[row] +":" + ( indices[row] + 1)));
//            probs[row] = value.toFloatArray()[0];
//        }

        int lastIdx = 0;
        for (int i = 0; i < indices.length; i++) {
            if (selection[i] == true && indices[i] > 0 && !(i > 0 && indices[i] == lastIdx)) {
                sb.append(table.get((int) indices[i]));
            }
        }
        return sb.toString();
    }

    @Override
    public NDList processInput(TranslatorContext ctx, Image input) {
        NDArray img = input.toNDArray(ctx.getNDManager(), Image.Flag.COLOR);
        int imgC = 3;
        int imgH = 48;
        int imgW = 320;

        float max_wh_ratio = (float) imgW / (float) imgH;

        int h = input.getHeight();
        int w = input.getWidth();
        float wh_ratio = (float) w / (float) h;

        max_wh_ratio = Math.max(max_wh_ratio,wh_ratio);
        imgW = (int)(imgH * max_wh_ratio);

        int resized_w;
        if (Math.ceil(imgH * wh_ratio) > imgW) {
            resized_w = imgW;
        } else {
            resized_w = (int) (Math.ceil(imgH * wh_ratio));
        }
        NDArray resized_image = NDImageUtils.resize(img, resized_w, imgH);
        resized_image = resized_image.transpose(2, 0, 1).toType(DataType.FLOAT32,false);
        resized_image.divi(255f).subi(0.5f).divi(0.5f);
        NDArray padding_im = ctx.getNDManager().zeros(new Shape(imgC, imgH, imgW), DataType.FLOAT32);
        padding_im.set(new NDIndex(":,:,0:" + resized_w), resized_image);

        padding_im = padding_im.flip(0);
//        padding_im = padding_im.expandDims(0);
        return new NDList(padding_im);
    }

    @Override
    public Batchifier getBatchifier() {
        return Batchifier.fromString(batchifier);
    }

}