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

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.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.translate.Batchifier;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;
import com.xindao.ocr.smartjavaai.entity.PlateResult;

import java.util.ArrayList;
import java.util.List;

/**
 * @author dwj
 */
public class CRNNPlateRecTranslator implements Translator<Image, PlateResult> {

    private static final String plateName = "#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品";
    private static final String[] plateColors = {"黑色", "蓝色", "绿色", "白色", "黄色"};
    private static final float MEAN = 0.588f;
    private static final float STD = 0.193f;

    @Override
    public NDList processInput(TranslatorContext ctx, Image input) {
        NDManager manager = ctx.getNDManager();

        // Resize to (168, 48)
        NDArray array = input.toNDArray(manager, Image.Flag.COLOR);
        array = NDImageUtils.resize(array, 168, 48);

        // Normalize
        array = array.toType(DataType.FLOAT32, false)
                .div(255f)
                .sub(MEAN)
                .div(STD);

        // HWC to CHW
        array = array.transpose(2, 0, 1);
        array = array.expandDims(0); // batch dimension

        return new NDList(array);
    }

    @Override
    public PlateResult processOutput(TranslatorContext ctx, NDList list) {
        NDArray plateOutput = list.get(0);  // shape: [1, T, num_classes]
        NDArray colorOutput = list.get(1);  // shape: [1, num_colors]

        int[] plateIdx = plateOutput.argMax(-1)
                .toType(DataType.INT32, false)
                .toIntArray();
        int colorIdx = colorOutput.argMax(1).toType(DataType.INT32, false).toIntArray()[0];

        String plateNo = decodePlate(plateIdx);
        String plateColor = plateColors[colorIdx];

        return new PlateResult(plateNo, plateColor);
    }

    private String decodePlate(int[] preds) {
        int pre = 0;
        List<Integer> newPreds = new ArrayList<>();
        for (int idx : preds) {
            if (idx != 0 && idx != pre) {
                newPreds.add(idx);
            }
            pre = idx;
        }

        StringBuilder sb = new StringBuilder();
        for (int i : newPreds) {
            if (i >= 0 && i < plateName.length()) {
                sb.append(plateName.charAt(i));
            }
        }
        return sb.toString();
    }

    @Override
    public Batchifier getBatchifier() {
        return null; // 非批量任务
    }
}