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.awt.Rectangle;
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;
    private int patch = 24;   // receptive-field / training patch (24 or 32); set via setPatch to match the model. By Claude on 06/16/2026
    private int outCh = -1;   // model output channels: 6 = continuous-velocity (reg) head {det,Vx,Vy,logvar,dx,dy}; else grid. By Claude on 06/17/2026

    public CuasDnnInfer(String modelSpec, int nframes) throws Exception {
        String local = resolveModel(modelSpec);
        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"
        // Derive temporal depth (N frames) from the model's input shape [B, N, H, W] so an
        // 8- vs 9-frame model swaps at runtime; the constructor arg is only a fallback hint.
        // By Claude on 06/15/2026
        int modelN = nframes;
        try {
            ai.onnxruntime.NodeInfo ni = session.getInputInfo().get(inputName);
            long[] shape = ((ai.onnxruntime.TensorInfo) ni.getInfo()).getShape();
            if ((shape != null) && (shape.length >= 2) && (shape[1] > 0)) modelN = (int) shape[1];
        } catch (Exception e) { /* keep the hint value */ }
        this.nframes = modelN;
        // Output channel count [B, out_ch, H, W]: 6 = continuous-velocity (reg) head, else grid (1+vdim^2+2). By Claude on 06/17/2026
        try {
            String outName = session.getOutputNames().iterator().next();
            long[] osh = ((ai.onnxruntime.TensorInfo) session.getOutputInfo().get(outName).getInfo()).getShape();
            if ((osh != null) && (osh.length >= 2) && (osh[1] > 0)) this.outCh = (int) osh[1];
        } catch (Exception e) { /* leave -1 (grid assumed) */ }
        System.out.println("CuasDnnInfer: loaded " + local + "  input=" + inputName
                + " outputs=" + session.getOutputNames() + " nframes=" + this.nframes + " outCh=" + this.outCh
                + (isReg() ? " (REG head: S,Vx,Vy,sigma)" : "")
                + ((modelN != nframes) ? " (hint was " + nframes + ")" : ""));
    }

    /** Temporal depth the loaded model expects (N frames), read from its input shape. By Claude on 06/15/2026 */
    public int getNFrames() { return nframes; }

    /** Model output channel count (6 = reg head; else grid). By Claude on 06/17/2026 */
    public int getOutCh() { return outCh; }
    /** True if the loaded model is the continuous-velocity (reg) head: out = {det,Vx,Vy,logvar,dx,dy}. */
    public boolean isReg() { return outCh == 6; }

    /** Set the receptive-field / training patch size (24 or 32) the model was trained at;
     *  inferROI extracts this P×P patch per ROI pixel (half=P/2 = the output-pixel index). By Claude on 06/16/2026 */
    public void setPatch(int p) { if (p > 0) this.patch = p; }

    /** 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;
        }
    }

    /** Run the FCN per-pixel (stride-1, matching the ±0.5px-center training) over an ROI of the # By Claude on 06/13/2026
     *  conditioned N-frame stack. frames=[N][width*height], newest first. Extracts a 24x24 patch
     *  centered on each ROI pixel (zero-padded at image edges), batches them, runs ORT once, and
     *  returns the C5P-shaped field [roi_npix][1][vdim*vdim] = softmax(vel)*s (confidence-gated
     *  velocity posterior per pixel) - plugs straight into showConvKernel5d / the recurrent layer. */
    public double[][][] inferROI(float[][] frames, int width, int height, Rectangle roi, int velRadius, double sThresh) throws Exception {
        return inferROI(frames, width, height, roi, velRadius, sThresh, null);
    }

    /** As inferROI, but if offOut != null also fills offOut[roi_pix] = {dx, dy, s} from the two // By Claude on 06/14/2026
     *  offset channels (op[1+nvel], op[1+nvel+1]) - the DNN's per-pixel sub-pixel position estimate. */
    public double[][][] inferROI(float[][] frames, int width, int height, Rectangle roi, int velRadius, double sThresh, double[][] offOut) throws Exception {
        final int N = frames.length;
        final int P = patch, half = P / 2;              // receptive field == training patch (24 or 32, via setPatch) // By Claude on 06/16/2026
        final int vdim = 2 * velRadius + 1, nvel = vdim * vdim;
        final int rw = roi.width, rh = roi.height, npix = rw * rh;
        final int CHUNK = 2048;   // tile the batch so ORT activations stay small (B*32*22*22)
        double[][][] field = new double[npix][1][nvel];
        for (int base = 0; base < npix; base += CHUNK) {
            int bs = Math.min(CHUNK, npix - base);
            FloatBuffer fb = FloatBuffer.allocate(bs * N * P * P);
            for (int q = 0; q < bs; q++) {
                int p = base + q;
                int cy = roi.y + (p / rw), cx = roi.x + (p % rw);
                for (int n = 0; n < N; n++) {
                    float[] fr = frames[n];
                    for (int j = 0; j < P; j++) {
                        int iy = cy + j - half;
                        boolean yok = (iy >= 0 && iy < height);
                        for (int i = 0; i < P; i++) {
                            int ix = cx + i - half;
                            fb.put((yok && ix >= 0 && ix < width) ? fr[iy * width + ix] : 0f);
                        }
                    }
                }
            }
            fb.rewind();
            try (OnnxTensor in = OnnxTensor.createTensor(env, fb, new long[] {bs, N, P, P});
                 OrtSession.Result res = session.run(Collections.singletonMap(inputName, in))) {
                float[][][][] out = (float[][][][]) res.get(0).getValue();   // [bs, outCh, 1, 1]
                for (int q = 0; q < bs; q++) {
                    float[][][] op = out[q];
                    float s = sigmoid(op[0][0][0]);
                    if (offOut != null) offOut[base + q] = new double[] {op[1 + nvel][0][0], op[1 + nvel + 1][0][0], s}; // {dx,dy,s} for all pixels // By Claude on 06/14/2026
                    if (s < sThresh) continue;                       // below confidence -> leave field 0 // By Claude on 06/13/2026
                    double mx = Double.NEGATIVE_INFINITY;
                    for (int v = 0; v < nvel; v++) mx = Math.max(mx, op[1 + v][0][0]);
                    double sum = 0; double[] e = new double[nvel];
                    for (int v = 0; v < nvel; v++) { e[v] = Math.exp(op[1 + v][0][0] - mx); sum += e[v]; }
                    double[] fv = field[base + q][0];
                    for (int v = 0; v < nvel; v++) fv[v] = (e[v] / sum) * s;
                }
            }
            if ((((base / CHUNK) & 7) == 7) || (base + CHUNK >= npix))    // periodic progress // By Claude on 06/13/2026
                System.out.println("      DNN inferROI " + Math.min(base + CHUNK, npix) + "/" + npix + " px");
        }
        return field;
    }

    /** Continuous-velocity (reg) inference over an ROI. For each pixel returns
     *  {S, Vx, Vy, sigma, dx, dy} from the 6-ch reg head (det, Vx, Vy, logvar, dx, dy):
     *  S=sigmoid(det), Vx/Vy already bounded by the model, sigma=exp(logvar/2). By Claude on 06/17/2026 */
    public double[][] inferROIReg(float[][] frames, int width, int height, Rectangle roi) throws Exception {
        final int N = frames.length;
        final int P = patch, half = P / 2;
        final int rw = roi.width, npix = rw * roi.height;
        final int CHUNK = 2048;
        double[][] out6 = new double[npix][6];
        for (int base = 0; base < npix; base += CHUNK) {
            int bs = Math.min(CHUNK, npix - base);
            FloatBuffer fb = FloatBuffer.allocate(bs * N * P * P);
            for (int q = 0; q < bs; q++) {
                int p = base + q;
                int cy = roi.y + (p / rw), cx = roi.x + (p % rw);
                for (int n = 0; n < N; n++) {
                    float[] fr = frames[n];
                    for (int j = 0; j < P; j++) {
                        int iy = cy + j - half;
                        boolean yok = (iy >= 0 && iy < height);
                        for (int i = 0; i < P; i++) {
                            int ix = cx + i - half;
                            fb.put((yok && ix >= 0 && ix < width) ? fr[iy * width + ix] : 0f);
                        }
                    }
                }
            }
            fb.rewind();
            try (OnnxTensor in = OnnxTensor.createTensor(env, fb, new long[] {bs, N, P, P});
                 OrtSession.Result res = session.run(Collections.singletonMap(inputName, in))) {
                float[][][][] out = (float[][][][]) res.get(0).getValue();   // [bs, 6, 1, 1]
                for (int q = 0; q < bs; q++) {
                    float[][][] op = out[q];
                    double[] o = out6[base + q];
                    o[0] = sigmoid(op[0][0][0]);                 // S
                    o[1] = op[1][0][0]; o[2] = op[2][0][0];      // Vx, Vy (model-bounded)
                    o[3] = Math.exp(0.5 * op[3][0][0]);          // sigma = exp(logvar/2)
                    o[4] = op[4][0][0]; o[5] = op[5][0][0];      // dx, dy
                }
            }
            if ((((base / CHUNK) & 7) == 7) || (base + CHUNK >= npix))
                System.out.println("      DNN inferROIReg " + Math.min(base + CHUNK, npix) + "/" + npix + " px");
        }
        return out6;
    }

    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 {
        if (args[0].equals("roitest")) {                 // roitest model testvec_in.bin testvec_out.bin
            int N = 8, P = 24, vr = 5, nvel = (2 * vr + 1) * (2 * vr + 1);
            float[] flat = readFloatsLE(args[2]);
            float[][] frames = new float[N][P * P];      // one P*P "image" per frame
            for (int n = 0; n < N; n++) System.arraycopy(flat, n * P * P, frames[n], 0, P * P);
            try (CuasDnnInfer infer = new CuasDnnInfer(args[1], N)) {
                // 1x1 ROI at the image center -> extracted patch == the full P*P frame
                double[][][] field = infer.inferROI(frames, P, P, new Rectangle(P / 2, P / 2, 1, 1), vr, 0.0);
                float[] exp = readFloatsLE(args[3]);     // raw [124]; build softmax(vel)*s
                float s = sigmoid(exp[0]); double mx = -1e30, sum = 0; double[] e = new double[nvel];
                for (int v = 0; v < nvel; v++) mx = Math.max(mx, exp[1 + v]);
                for (int v = 0; v < nvel; v++) { e[v] = Math.exp(exp[1 + v] - mx); sum += e[v]; }
                double maxd = 0; for (int v = 0; v < nvel; v++) maxd = Math.max(maxd, Math.abs(field[0][0][v] - e[v] / sum * s));
                System.out.printf("inferROI roitest: max abs diff vs single-patch softmax*s = %.6g  %s%n",
                        maxd, maxd < 1e-4 ? "PASS" : "MISMATCH");
            }
            return;
        }
        // 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");
            }
        }
    }
}
