package com.elphel.imagej.cuas.rt;

import java.nio.FloatBuffer;
import java.util.Collections;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtSession;

/**
 * C5P DNN v2 Stage-2 (Hough-vote refine) ONNX runner. By Claude on 06/18/2026
 * Loads the learned refine net (3-channel vote accumulator [accS, accVx, accVy] -> det logit + Vx,Vy consensus)
 * and runs it over a whole ROI-sized accumulator in one pass (full-field conv, NOT patch-based like CuasDnnInfer).
 * The vote scatter itself is done in Java (CuasDetectRT.voteScatter); this only runs the conv refine.
 */
public class CuasStage2Infer {
    private final OrtEnvironment env;
    private final OrtSession session;
    private final String inputName;

    public CuasStage2Infer(String model) throws Exception {
        env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions opts = new OrtSession.SessionOptions();
        try { opts.setSessionLogLevel(ai.onnxruntime.OrtLoggingLevel.ORT_LOGGING_LEVEL_ERROR); } catch (Exception e) {}
        session = env.createSession(model, opts);
        inputName = session.getInputNames().iterator().next();
        System.out.println("CuasStage2Infer: loaded " + model + "  input=" + inputName);
    }

    /** acc = [3][w*h] (accS, accVx, accVy), row-major. Returns [3][w*h]: {sigmoid(det), Vx, Vy} refined. */
    public double[][] refine(double[][] acc, int w, int h) throws Exception {
        final int npix = w * h;
        FloatBuffer fb = FloatBuffer.allocate(3 * npix);
        for (int c = 0; c < 3; c++) for (int i = 0; i < npix; i++) fb.put((float) acc[c][i]);
        fb.rewind();
        double[][] out = new double[3][npix];
        try (OnnxTensor in = OnnxTensor.createTensor(env, fb, new long[] {1, 3, h, w});
             OrtSession.Result res = session.run(Collections.singletonMap(inputName, in))) {
            float[][][][] o = (float[][][][]) res.get(0).getValue();   // [1,3,h,w]
            for (int y = 0; y < h; y++) for (int x = 0; x < w; x++) {
                int i = y * w + x;
                out[0][i] = 1.0 / (1.0 + Math.exp(-o[0][0][y][x]));     // det -> sigmoid
                out[1][i] = o[0][1][y][x];                              // Vx consensus
                out[2][i] = o[0][2][y][x];                              // Vy consensus
            }
        }
        return out;
    }

    public void close() { try { if (session != null) session.close(); } catch (Exception e) {} }
}
