package com.elphel.imagej.cuas.rt;
// C5P DNN front-end inference via ONNX Runtime (CPU). // By Claude on 06/13/2026
// Loads a PyTorch-trained, ONNX-exported all-conv FCN that maps an N-frame patch stack to a
// per-pixel (Vx,Vy,s) output. Phase 1: CPU, fed from Java float arrays (the convenient testing
// path, like the old TensorFlow integration). Phase 2 (later): onnxruntime_gpu CUDA/TensorRT EP,
// device chosen via the existing GPU-detection logic; eventual zero-copy from JCuda CUdeviceptr.
//
// Model location (curt_dnn_model): resolved as a local path, scp "user@host:/path", or http(s)
// URL; remote specs are fetched to ~/.cache/c5p_dnn/ together with the external-data sibling
// (model.onnx + model.onnx.data, which torch 2.9's exporter splits).

import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.net.URL;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.file.Files;
import java.util.Collections;

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

public class CuasDnnInfer implements AutoCloseable {
    private final OrtEnvironment env;
    private final OrtSession session;
    private final String inputName;
    private final int nframes;

    public CuasDnnInfer(String modelSpec, int nframes) throws Exception {
        String local = resolveModel(modelSpec);
        this.nframes = nframes;
        this.env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions opts = new OrtSession.SessionOptions();
        // CPU EP by default. GPU later: opts.addCUDA(deviceId) once cuDNN 9 is present.
        this.session = env.createSession(local, opts);
        this.inputName = session.getInputNames().iterator().next();  // "frames"
        System.out.println("CuasDnnInfer: loaded " + local + "  input=" + inputName
                + " outputs=" + session.getOutputNames());
    }

    /** Run one patch. frames = [N][H][W] (newest first). Returns the raw output channels
     *  [outCh] for the center output pixel: [0]=det logit, [1..V*V]=vel logits, last 2 = offset. */
    public float[] inferRaw(float[][][] frames) throws Exception {
        int N = frames.length, H = frames[0].length, W = frames[0][0].length;
        FloatBuffer fb = FloatBuffer.allocate(N * H * W);
        for (int n = 0; n < N; n++)
            for (int y = 0; y < H; y++)
                fb.put(frames[n][y]);
        fb.rewind();
        try (OnnxTensor in = OnnxTensor.createTensor(env, fb, new long[] {1, N, H, W});
             OrtSession.Result res = session.run(Collections.singletonMap(inputName, in))) {
            float[][][][] out = (float[][][][]) res.get(0).getValue();   // [1, outCh, 1, 1]
            int outCh = out[0].length;
            float[] flat = new float[outCh];
            for (int c = 0; c < outCh; c++) flat[c] = out[0][c][0][0];
            return flat;
        }
    }

    public static float sigmoid(float x) { return (float) (1.0 / (1.0 + Math.exp(-x))); }

    @Override public void close() throws Exception { session.close(); }

    // ---- model location resolver (mirrors the tile_processor_gpu default/override scheme) ----
    static String resolveModel(String spec) throws Exception {
        if (spec == null || spec.trim().isEmpty())
            throw new IllegalArgumentException("CuasDnnInfer: empty model spec (curt_dnn_model)");
        spec = spec.trim();
        if (new File(spec).isFile()) return spec;                    // local path
        String cacheDir = System.getProperty("user.home") + "/.cache/c5p_dnn/";
        new File(cacheDir).mkdirs();
        String name = spec.substring(spec.lastIndexOf('/') + 1);
        String localOnnx = cacheDir + name;
        if (spec.startsWith("http://") || spec.startsWith("https://")) {
            httpFetch(spec, localOnnx);
            try { httpFetch(spec + ".data", localOnnx + ".data"); } catch (Exception e) { /* no external data */ }
        } else if (spec.matches("[^@/]+@[^:]+:.*")) {                 // user@host:/path (scp/rsync)
            rsyncFetch(spec, localOnnx);
            try { rsyncFetch(spec + ".data", localOnnx + ".data"); } catch (Exception e) { /* none */ }
        } else {
            throw new java.io.FileNotFoundException("CuasDnnInfer: model not found / unrecognized spec: " + spec);
        }
        return localOnnx;
    }

    static void rsyncFetch(String remote, String dest) throws Exception {
        Process p = new ProcessBuilder("rsync", "-az", remote, dest).inheritIO().start();
        if (p.waitFor() != 0) throw new RuntimeException("rsync failed for " + remote);
    }

    static void httpFetch(String url, String dest) throws Exception {
        try (InputStream is = new URL(url).openStream()) {
            Files.copy(is, new File(dest).toPath(), java.nio.file.StandardCopyOption.REPLACE_EXISTING);
        }
    }

    // ---- standalone verification: Java ORT output vs the saved PyTorch raw output ----
    static float[] readFloatsLE(String path) throws Exception {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        try (DataInputStream in = new DataInputStream(new FileInputStream(path))) {
            byte[] buf = new byte[8192]; int n;
            while ((n = in.read(buf)) > 0) bos.write(buf, 0, n);
        }
        ByteBuffer bb = ByteBuffer.wrap(bos.toByteArray()).order(ByteOrder.LITTLE_ENDIAN);
        float[] f = new float[bb.remaining() / 4];
        for (int i = 0; i < f.length; i++) f[i] = bb.getFloat();
        return f;
    }

    public static void main(String[] args) throws Exception {
        // args: model(path/spec) testvec_in.bin [testvec_out.bin] [N H W]
        String model = args[0], invec = args[1];
        String outvec = args.length > 2 ? args[2] : null;
        int N = args.length > 3 ? Integer.parseInt(args[3]) : 8;
        int H = args.length > 4 ? Integer.parseInt(args[4]) : 24;
        int W = args.length > 5 ? Integer.parseInt(args[5]) : 24;
        float[] flat = readFloatsLE(invec);
        if (flat.length != N * H * W) throw new RuntimeException("in size " + flat.length + " != " + (N * H * W));
        float[][][] frames = new float[N][H][W];
        int k = 0;
        for (int n = 0; n < N; n++) for (int y = 0; y < H; y++) for (int x = 0; x < W; x++) frames[n][y][x] = flat[k++];
        try (CuasDnnInfer infer = new CuasDnnInfer(model, N)) {
            float[] out = infer.inferRaw(frames);
            System.out.println("output len=" + out.length + "  det_logit=" + out[0] + "  s=" + sigmoid(out[0]));
            if (outvec != null) {
                float[] exp = readFloatsLE(outvec);
                double maxabs = 0; int amax = 0;
                for (int i = 0; i < Math.min(out.length, exp.length); i++) {
                    double d = Math.abs(out[i] - exp[i]); if (d > maxabs) { maxabs = d; amax = i; }
                }
                System.out.printf("max abs diff vs PyTorch = %.6g (at ch %d: java %.5f vs torch %.5f)%n",
                        maxabs, amax, out[amax], exp[amax]);
                System.out.println(maxabs < 1e-3 ? "PASS - Java ORT matches PyTorch" : "MISMATCH - investigate");
            }
        }
    }
}
