package com.elphel.imagej.gpu.jna;

import java.util.Arrays;

import com.elphel.imagej.gpu.GpuQuad;
import com.elphel.imagej.gpu.GPUTileProcessor;
import com.elphel.imagej.gpu.TpTask;
import com.elphel.imagej.tileprocessor.Correlation2d;
import com.elphel.imagej.tileprocessor.CorrVector;
import com.elphel.imagej.tileprocessor.GeometryCorrection;
import com.elphel.imagej.tileprocessor.IntersceneLmaFloat;
import com.elphel.imagej.tileprocessor.QuadCLT;
import com.sun.jna.Native;
import com.sun.jna.Pointer;

/**
 * Native (JNA) GPU backend — architecture B: subclass {@link GpuQuad}, but allocate NO JCuda memory and
 * create NO JCuda context. The constructor uses GpuQuad's protected no-alloc constructor (config fields
 * only) and stands up the native tile-processor library ({@code libtileproc.so}) via {@link TpJna}:
 * a {@code TpModule} (NVRTC-compiled kernels) and a persistent {@code TpProc} instance that owns all GPU
 * memory. Each GPU-touching method of GpuQuad is overridden to delegate to {@code TpProc}; methods not yet
 * ported throw {@link UnsupportedOperationException} (so anything off the validated path fails loudly
 * rather than silently using the null base JCuda buffers).
 *
 * Selected once at startup (default backend stays JCUDA); JNA mode never initializes JCuda.
 * Validated kernel path (Stages 1–5 + StageProc): geometry, convert_direct, imclt, correlations.
 * By Claude on 2026-06-25.
 */
public class GpuQuadJna extends GpuQuad {
    private TpJna   lib;        // assigned once in initNative() (non-final to share across both ctors)
    private Pointer module;     // TpModule* (NVRTC kernels)
    private Pointer proc;       // TpProc*   (persistent device buffers)
    private boolean closed = false;

    /**
     * Per-scene (sensor) backend.
     * @param quadCLT   the per-scene QuadCLT (geometry, kernels, sensor config)
     * @param srcdir    directory with the kernel .cu/.h sources for NVRTC (e.g. tile_processor_gpu/src)
     * @param devrt     path to libcudadevrt.a (for CDP linking)
     * @param debugLevel gpu debug level
     */
    public GpuQuadJna(QuadCLT quadCLT, String srcdir, String devrt, int debugLevel) {
        super(quadCLT, debugLevel, true); // no-alloc: config fields only, no JCuda context
        initNative(srcdir, devrt);
    }

    /**
     * RECTILINEAR single-array backend (synthetic non-sensor images: no distortion, no aberration kernels,
     * single camera). Mirrors GpuQuad's rectilinear ctor config but stands up its own native TpProc.
     * Built via GpuQuad.createRectilinear() in JNA mode (e.g. the CUAS interscene GpuQuad). By Claude on 06/26/2026.
     */
    public GpuQuadJna(int width, int height, int num_colors, String srcdir, String devrt, int debugLevel) {
        super(width, height, num_colors, debugLevel, true); // no-alloc rectilinear config (num_cams=1, no kernels)
        initNative(srcdir, devrt);
    }

    // Shared native bring-up: load the lib, NVRTC-compile the module, create the persistent TpProc and size
    // its buffers. tp_proc_setup uses the (already-set) config fields uniformly — for the rectilinear ctor
    // num_cams=1 and kernels_hor=kern_tiles=0, so no kernel buffers are allocated.
    private void initNative(String srcdir, String devrt) {
        this.lib = Native.load("tileproc", TpJna.class);
        this.module = lib.tp_create_module(srcdir, devrt);
        if (module == null) {
            throw new IllegalStateException("GpuQuadJna: tp_create_module failed: " + lib.tp_last_error());
        }
        this.proc = lib.tp_proc_create(module);
        if (proc == null) {
            throw new IllegalStateException("GpuQuadJna: tp_proc_create failed: " + lib.tp_last_error());
        }
        int rc = lib.tp_proc_setup(proc, num_cams, num_colors, img_width, img_height,
                                   kernels_hor, kern_tiles);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna: tp_proc_setup rc=" + rc + ": " + lib.tp_last_error());
        }
        // tp_proc_setup_rbg_corr (imclt + correlation buffers) is deferred until the correlation config
        // (num_pairs / sel_pairs / color_weights / corr_out_rad) is known at first use.
    }

    // By Claude on 07/02/2026: native-side equivalent of gpuTileProcessor.bayer_set (that field's owner is null
    // here). Tracks whether the native bayer buffers hold the bound quadCLT's current data, so that
    // execConvertDirect()'s internal setBayerImages(false,...) SKIPS the re-pull exactly like the JCuda base
    // does. Without it every convert unconditionally re-uploaded quadCLT.getResetImageData(), silently
    // clobbering any explicitly uploaded data (setBayerImages(data, true) / setBayerImage()) - found by the
    // curt_cond_test raw baseline coming out bit-identical to the conditioned render.
    private boolean jna_bayer_set = false;
    // Native equivalent of GpuQuad.geometry_correction_set. The reverse-distortion
    // table depends only on GeometryCorrection, not on the per-iteration pose/tasks;
    // keep the upload/table cached until resetGeometryCorrection() invalidates it.
    // By Codex on 07/14/2026 (roadmap rung-3 hoist).
    private boolean jna_geometry_set = false;

    // Switch this shared GpuQuad to a different scene's QuadCLT. Base also clears
    // gpuTileProcessor.bayer_set (null here) — mirrored with jna_bayer_set: the next convert re-uploads.
    // resetGeometryCorrection*/() are flag-only (safe to call).
    @Override public void updateQuadCLT(final QuadCLT quadCLT) {
        this.quadCLT = quadCLT;
        resetGeometryCorrection();
        resetGeometryCorrectionVector();
        jna_bayer_set = false; // By Claude on 07/02/2026: match base (scene switch invalidates uploaded bayer)
    }
    @Override public void resetGeometryCorrection() {
        super.resetGeometryCorrection();
        jna_geometry_set = false;
    }
    @Override public void resetBayer() { jna_bayer_set = false; } // By Claude on 07/02/2026: match base semantics

    // CLT sizing (base derefs null gpu_clt_wh) — full-frame dims; matches TpProc's slice.
    @Override public int getCltSize(boolean use_ref) {
        int tilesX = img_width / getDttSize(), tilesY = img_height / getDttSize();
        return tilesY * tilesX * num_colors * 4 * getDttSize() * getDttSize();
    }
    @Override public int getNumTiles(boolean use_ref) {
        return (img_width / getDttSize()) * (img_height / getDttSize());
    }
    // Push a per-cam CLT slice to the device (e.g. CLT restore via setComboToTD).
    @Override public void setCltData(int ncam, float[] fclt, boolean use_ref) {
        int clt_size = getCltSize(use_ref);
        if (fclt.length != clt_size) {
            System.out.println("GpuQuadJna.setCltData(): wrong size got " + fclt.length + " expected " + clt_size);
            return;
        }
        lib.tp_proc_set_clt(proc, ncam, fclt, use_ref ? 1 : 0);
    }
    // Read all per-cam CLT slices from the device (base derefs null gpu_clt_h on JNA shells).
    // D2H half of the CPU sensor-TD consolidation bridge (CuasTD). By Claude on 07/05/2026.
    @Override public float [][] getCltData(boolean use_ref) {
        int clt_size = getCltSize(use_ref);
        float [][] fclt = new float [num_cams][clt_size];
        for (int ncam = 0; ncam < num_cams; ncam++) {
            lib.tp_proc_get_clt(proc, ncam, use_ref ? 1 : 0, fclt[ncam]);
        }
        return fclt;
    }

    // On-device 16-sensor TD consolidation (tp_consolidate.cu v1 chain) - replaces the CPU
    // bridge (getCltData -> CuasTD.consolidateSensorsTD -> setCltData(0)): the per-sensor TD
    // stays on the device, only the small flattened task streams cross the JNA boundary. The
    // average lands in the cam-0 CLT slot (what the single conj-multiply correlates). Oracle:
    // CuasTD.oobSensorMasks + masked consolidateSensorsTD, validated bit-exact (tol 0) on the
    // real-scene avg_td_oob case through this very export. By Claude on 07/12/2026.
    @Override public boolean execConsolidateSensorsTD(TpTask [][] tp_tasks, double soft_margin, double hard_margin) {
        if ((tp_tasks == null) || (tp_tasks.length < 1) || (tp_tasks[0] == null) || (tp_tasks[0].length == 0)) {
            return false; // nothing to consolidate - let the caller decide
        }
        final int task_size = tp_tasks[0][0].getSize();
        final float [][] ftasks = new float [2][];
        for (int nset = 0; nset < 2; nset++) {
            if ((nset >= tp_tasks.length) || (tp_tasks[nset] == null)) continue;
            final TpTask [] tasks = tp_tasks[nset];
            ftasks[nset] = new float [tasks.length * task_size];
            for (int i = 0; i < tasks.length; i++) {
                tasks[i].asFloatArray(ftasks[nset], i, false); // exact GPU struct layout
            }
        }
        final int [] stats = new int [3]; // {MB pairs, surviving tiles, misaligned}
        int rc = lib.tp_proc_exec_consolidate(proc, ftasks[0], ftasks[1],
                tp_tasks[0].length, task_size, (float) soft_margin, (float) hard_margin, stats);
        if (rc != 0) {
            System.out.println("GpuQuadJna.execConsolidateSensorsTD(): native rc=" + rc +
                    " (" + lib.tp_last_error() + ") - falling back to the CPU bridge");
            return false;
        }
        if (stats[2] != 0) {
            System.out.println("GpuQuadJna.execConsolidateSensorsTD(): WARNING - " + stats[2] +
                    " misaligned MB pairs (task sets not index-aligned), " +
                    stats[0] + " pairs -> " + stats[1] + " tiles");
        }
        return true;
    }

    // Lean-chain fast path: geometry already updated task slots 0/1 in place, so
    // consolidation consumes those device buffers without rebuilding TpTask on
    // the host or uploading the flattened streams again. By Codex on 07/15/2026.
    @Override public boolean execConsolidateResidentTasksTD(
            TpTask [][] tp_tasks, double soft_margin, double hard_margin) {
        if ((tp_tasks == null) || (tp_tasks.length < 1) ||
                (tp_tasks[0] == null) || (tp_tasks[0].length == 0)) return false;
        final int task_size = tp_tasks[0][0].getSize();
        final int slot1 = ((tp_tasks.length > 1) && (tp_tasks[1] != null)) ? 1 : -1;
        final int [] stats = new int [3];
        int rc = lib.tp_proc_exec_consolidate_slots(proc, 0, slot1,
                tp_tasks[0].length, task_size,
                (float) soft_margin, (float) hard_margin, stats);
        if (rc != 0) {
            System.out.println("GpuQuadJna.execConsolidateResidentTasksTD(): native rc=" + rc +
                    " (" + lib.tp_last_error() + ") - falling back to task readback");
            return false;
        }
        if (stats[2] != 0) {
            System.out.println("GpuQuadJna.execConsolidateResidentTasksTD(): WARNING - " + stats[2] +
                    " misaligned MB pairs (task slots not index-aligned), " +
                    stats[0] + " pairs -> " + stats[1] + " tiles");
        }
        return true;
    }

    /** Native handles for the override implementations (added incrementally). */
    protected TpJna   lib()    { return lib; }
    protected Pointer module() { return module; }
    protected Pointer proc()   { return proc; }

    /** Release native GPU memory + module (deterministic — does not rely on GC/finalize). */
    public synchronized void close() {
        if (closed) return;
        closed = true;
        if (proc != null)   lib.tp_proc_destroy(proc);
        if (module != null) lib.tp_destroy_module(module);
    }

    // ---- own caching flags (base uses gpuTileProcessor.* which is null here) ----
    private boolean jna_kernels_set = false;
    private boolean jna_rbg_corr_ready = false;

    // Allocate imclt(RBG)+correlation buffers once. RBG sizing needs only img/num_colors; correlation
    // buffers are sized by num_all_pairs + corr_out_rad. sel_pairs/color_weights here are placeholders
    // (the correlation overrides — added next — will pass the real per-call values).
    private void ensureRbgCorr() {
        if (jna_rbg_corr_ready) return;
        int corr_out_rad = 7; // CORR_OUT_RAD -> 15x15
        float[] cw = (num_colors == 1) ? new float[]{1f,1f,1f} : new float[]{0.25f,0.25f,0.5f};
        // Correlation buffers are sized by num_pairs. The rectilinear single-camera config does INTER-scene
        // correlation, for which getNumPairs(num_cams=1)=0 -> zero-size buffers -> correlate2D_inter writes
        // OOB -> cudaErrorIllegalAddress (deferred to the next CUDA call). The JCuda oracle's rectilinear ctor
        // hardcodes num_pairs=3 for exactly this (GpuQuad.java:732); mirror it. By Claude on 06/26/2026.
        int num_pairs = rectilinear ? 3 : Correlation2d.getNumPairs(num_cams);
        lib.tp_proc_setup_rbg_corr(proc, num_pairs, 0,0,0,0, cw[0],cw[1],cw[2], corr_out_rad);
        jna_rbg_corr_ready = true;
    }

    // ---- geometry ----
    @Override public void setGeometryCorrection() {
        if (!jna_geometry_set) setGeometryCorrection(quadCLT.getGeometryCorrection(), false);
    }
    @Override public void setGeometryCorrection(GeometryCorrection gc, boolean use_java_rByRDist) {
        float[] fgc = gc.expandSensors(GPUTileProcessor.MAX_NUM_CAMS).toFloatArray();
        int rc = lib.tp_proc_set_geometry(proc, fgc, fgc.length);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna.setGeometryCorrection rc="+rc+": "+lib.tp_last_error());
        }
        jna_geometry_set = true;
    }
    @Override public void setExtrinsicsVector(CorrVector cv) {
        double[] dcv = cv.toFullRollArray();
        float[] fcv = new float[dcv.length];
        for (int i = 0; i < dcv.length; i++) fcv[i] = (float) dcv[i];
        lib.tp_proc_set_correction_vector(proc, fcv, fcv.length);
    }

    // ---- aberration kernels (same transpose-flatten as GpuQuad, then native upload) ----
    @Override public void setConvolutionKernels(boolean force) {
        if (jna_kernels_set && !force) return;
        setConvolutionKernels(quadCLT.getCLTKernels(), true);
    }
    @Override public void setConvolutionKernels(double[][][][][][] clt_kernels, boolean force) {
        if (jna_kernels_set && !force) return;
        final boolean transpose = true;
        int num_kernels = clt_kernels[0][0].length * clt_kernels[0][0][0].length * clt_kernels[0].length;
        int kernel_length = num_kernels * 4 * GPUTileProcessor.DTT_SIZE * GPUTileProcessor.DTT_SIZE;
        float[] fkernel  = new float[kernel_length];
        float[] foffsets = new float[num_kernels * 8];
        for (int ncam = 0; ncam < clt_kernels.length; ncam++) {
            int indx = 0;
            for (int ty = 0; ty < clt_kernels[ncam][0].length; ty++)
                for (int tx = 0; tx < clt_kernels[ncam][0][ty].length; tx++)
                    for (int col = 0; col < clt_kernels[ncam].length; col++)
                        for (int p = 0; p < 4; p++) {
                            double[] pa = clt_kernels[ncam][col][ty][tx][p];
                            for (int i0 = 0; i0 < 64; i0++) {
                                int i = transpose ? (((i0 & 7) << 3) + ((i0 >> 3) & 7)) : i0;
                                fkernel[indx++] = (float) pa[i];
                            }
                        }
            indx = 0;
            for (int ty = 0; ty < clt_kernels[ncam][0].length; ty++)
                for (int tx = 0; tx < clt_kernels[ncam][0][ty].length; tx++)
                    for (int col = 0; col < clt_kernels[ncam].length; col++) {
                        double[] pa = clt_kernels[ncam][col][ty][tx][4];
                        for (int i = 0; i < pa.length; i++) foffsets[indx++] = (float) pa[i];
                    }
            lib.tp_proc_set_kernels(proc, ncam, fkernel, fkernel.length);
            lib.tp_proc_set_kernel_offsets(proc, ncam, foffsets, foffsets.length);
        }
        jna_kernels_set = true;
    }

    // ---- bayer / source images (combine channels -> one float[] per cam) ----
    @Override public void setBayerImages(boolean force, boolean center) {
        if (center) {
            double[][] bc = quadCLT.getImageCenter();
            quadCLT.getResetImageCenter();
            float[] f = combineChannels(bc);
            if (f != null) lib.tp_proc_set_center_image(proc, f); // broadcast to all sensors
            return;
        }
        // By Claude on 07/02/2026: same guard as the JCuda base - do NOT re-pull quadCLT data over an
        // explicit upload when the native buffers are current and the quadCLT has no new data.
        if (!force && jna_bayer_set && !quadCLT.hasNewImageData()) {
            return;
        }
        setBayerImages(quadCLT.getResetImageData(), true);
    }
    @Override public void setBayerImages(double[][][] bayer_data, boolean force) {
        for (int ncam = 0; ncam < bayer_data.length; ncam++) {
            float[] f = combineChannels(bayer_data[ncam]);
            if (f != null) lib.tp_proc_set_image(proc, ncam, f);
        }
        jna_bayer_set = true; // By Claude on 07/02/2026: match base (gpuTileProcessor.bayer_set = true)
    }
    @Override public void setBayerImages(float[][] bayer_data, boolean force) {
        if (jna_bayer_set && !force) return;
        for (int ncam = 0; ncam < bayer_data.length; ncam++) {
            if (bayer_data[ncam] != null) lib.tp_proc_set_image(proc, ncam, bayer_data[ncam]);
        }
        jna_bayer_set = true;
    }
    // sum the (1 or 3) split-color channels into one image, as GpuQuad.setBayerImages does
    private static float[] combineChannels(double[][] chans) {
        if (chans == null || chans[0] == null) return null;
        float[] f = new float[chans[0].length];
        for (int i = 0; i < f.length; i++) { double s = chans[0][i]; for (int j = 1; j < chans.length; j++) s += chans[j][i]; f[i] = (float) s; }
        return f;
    }
    // Single pre-combined image upload (the rectilinear/synthetic path: ImageDtt.setRectilinearReferenceTD
    // passes fpixels_ref directly). Base does a JCuda cuMemcpy2D into gpu_bayer_h[ncam]. By Claude on 06/26/2026.
    @Override public void setBayerImage(float[] bayer_image, int ncam) {
        lib.tp_proc_set_image(proc, ncam, bayer_image);
    }
    @Override public void getBayerImage(float[] bayer_image, int ncam) {
        int rc = lib.tp_proc_get_image(proc, ncam, bayer_image);
        if (rc != 0) throw new IllegalStateException("tp_proc_get_image rc="+rc+": "+lib.tp_last_error());
    }
    // Resident LWIR two-map conditioning. Calibration upload is deliberately
    // separate from execution so it is paid once per calibration, never once
    // per scene. By Codex on 07/15/2026.
    @Override public boolean supportsConditioningMaps() {
        return !rectilinear;
    }
    @Override public boolean setConditioningMaps(float[] scale_maps, float[] offset_maps) {
        if ((scale_maps == null) || (offset_maps == null) ||
                (scale_maps.length != num_cams * img_width * img_height) ||
                (offset_maps.length != scale_maps.length)) return false;
        int rc = lib.tp_proc_set_conditioning_maps(proc, scale_maps, offset_maps, img_width * img_height);
        if (rc != 0) {
            System.out.println("GpuQuadJna.setConditioningMaps(): native rc="+rc+" ("+lib.tp_last_error()+")");
            return false;
        }
        return true;
    }
    @Override public boolean execConditioning() {
        int rc = lib.tp_proc_exec_conditioning(proc, 0, 0.0); // production maps already contain fixed DC
        if (rc != 0) {
            System.out.println("GpuQuadJna.execConditioning(): native rc="+rc+" ("+lib.tp_last_error()+")");
            return false;
        }
        return true;
    }
    // CLT buffers (main + ref) are pre-allocated full-frame in tp_proc_setup, so there is nothing to (re)allocate
    // per call. Base allocs JCuda gpu_clt[_ref] here; native is fixed-size -> no-op (nothing newly allocated).
    @Override public boolean reAllocateClt(int[] wh, boolean ref_scene) {
        return false;
    }

    // ---- tasks ----
    @Override public void setTasks(TpTask[] tile_tasks, boolean use_aux, boolean verify) {
        setTasks(tile_tasks, use_aux, verify, 0);
    }
    @Override public void setTasks(TpTask[] tile_tasks, boolean use_aux, boolean verify, int slot) {
        num_task_tiles = tile_tasks.length;
        int task_size = getTaskSize();
        float[] ftasks = new float[task_size * num_task_tiles];
        for (int i = 0; i < num_task_tiles; i++) { tile_tasks[i].task |= 511; tile_tasks[i].asFloatArray(ftasks, i, use_aux); }
        int rc = lib.tp_proc_set_tasks_slot(proc, ftasks, num_task_tiles, task_size * num_task_tiles, slot);
        if (rc != 0) throw new IllegalStateException("tp_proc_set_tasks_slot("+slot+") rc="+rc+
                " ("+lib.tp_last_error()+")");
    }

    @Override public boolean supportsResidentTaskSlots() { return true; }

    // Read tasks back after calculate_tiles_offsets (computed centerXY/disp_dist) -> rebuild TpTask[].
    @Override public void updateTasks(TpTask[] tile_tasks, boolean use_aux) {
        updateTasks(tile_tasks, use_aux, 0);
    }
    @Override public void updateTasks(TpTask[] tile_tasks, boolean use_aux, int slot) {
        num_task_tiles = tile_tasks.length;
        int task_size = getTaskSize();
        float[] ftasks = new float[task_size * num_task_tiles];
        int rc = lib.tp_proc_get_tasks_slot(proc, ftasks, task_size * num_task_tiles, slot);
        if (rc != 0) throw new IllegalStateException("tp_proc_get_tasks_slot("+slot+") rc="+rc+
                " ("+lib.tp_last_error()+")");
        for (int i = 0; i < num_task_tiles; i++)
            tile_tasks[i] = new TpTask(getNumSensors(), ftasks, i, use_aux);
    }
    // Full-frame dims (base returns the null gpu_clt_wh).
    @Override public int[] getWH(boolean use_ref) { return new int[]{img_width, img_height}; }

    // ---- geometry execution (calc_reverse_distortions + rot_derivs + calculate_tiles_offsets) ----
    @Override public void execSetTilesOffsets(boolean uniform_grid) {
        setGeometryCorrection();                                   // gc
        setExtrinsicsVector(quadCLT.getGeometryCorrection().expandSensors(GPUTileProcessor.MAX_NUM_CAMS).getCorrVector()); // cv
        lib.tp_proc_exec_geometry(proc, uniform_grid ? 1 : 0);
    }

    // ---- direct CLT conversion (+ the fragile no_kernels / use_center_image / erase_clt / ref_scene paths) ----
    @Override public void execConvertDirect(boolean ref_scene, int[] wh, int erase_clt, boolean no_kernels, boolean use_center_image) {
        boolean skip_kernels = rectilinear || no_kernels;
        if (!skip_kernels) setConvolutionKernels(false);
        if (!rectilinear)  setBayerImages(false, use_center_image);
        lib.tp_proc_exec_convert_direct(proc, ref_scene ? 1 : 0, erase_clt, no_kernels ? 1 : 0);
    }

    // ---- inverse CLT -> RBG ----
    @Override public void execImcltRbgAll(boolean is_mono, boolean ref_scene, int[] wh) {
        ensureRbgCorr();
        lib.tp_proc_exec_imclt(proc, 1, ref_scene ? 1 : 0); // apply_lpf = 1; honor ref_scene (gpu_clt_ref). By Claude 06/26/2026.
    }
    @Override public float[][] getRBG(int ncam) {
        int out_width = getImageWidth(), out_height = getImageHeight();
        int gpu_width = out_width + GPUTileProcessor.DTT_SIZE, gpu_height = out_height + GPUTileProcessor.DTT_SIZE;
        int gpu_img_size = gpu_width * gpu_height, rslt_img_size = out_width * out_height;
        float[] packed = new float[num_colors * gpu_img_size]; // == TpProc rbg_w x rbg_h (de-pitched)
        lib.tp_proc_get_rbg(proc, ncam, packed);
        float[][] fimg = new float[num_colors][rslt_img_size];
        for (int ncol = 0; ncol < num_colors; ncol++) {
            int tl = (GPUTileProcessor.DTT_SIZE / 2) * (gpu_width + 1) + ncol * gpu_img_size;
            for (int nrow = 0; nrow < out_height; nrow++)
                System.arraycopy(packed, tl + gpu_width * nrow, fimg[ncol], out_width * nrow, out_width);
        }
        return fimg;
    }

    // ---- correlations (delegate to the granular TpProc functions) ----
    private float[] scaleTriplet(double[] scales) {
        int ncol = Math.min(num_colors, 3);
        return new float[]{
            (ncol > 1) ? (float) scales[0] : 1f,
            (ncol > 1) ? (float) scales[1] : 0f,
            (ncol > 2) ? (float) scales[2] : 0f };
    }
    @Override public void execCorr2D_TD(double[] scales, int[] sel_pairs_in) {
        ensureRbgCorr();
        setCorrMask(sel_pairs_in); // base: corr_mask_int + indices (host-only) so getNumUsedPairs() works
        float[] s = scaleTriplet(scales);
        lib.tp_proc_exec_corr2d_td(proc, sel_pairs_in[0], sel_pairs_in[1], sel_pairs_in[2], sel_pairs_in[3], s[0], s[1], s[2]);
    }
    @Override public void execCorr2D_inter_TD(double[] scales) {
        ensureRbgCorr();
        float[] s = scaleTriplet(scales);
        lib.tp_proc_exec_corr2d_inter_td(proc, getSensorMaskInter(), s[0], s[1], s[2]);
    }
    @Override public void execCorr2D_combine(boolean init_corr, int num_pairs_in, int pairs_mask, boolean no_transpose_vertical) {
        int num_pairs = getNumUsedPairs();
        int init_corr_combo = (init_corr ? 1 : 0) + (no_transpose_vertical ? 2 : 0);
        lib.tp_proc_exec_corr2d_combine(proc, init_corr_combo, num_pairs, pairs_mask);
    }
    @Override public void execCorr2D_normalize(boolean combo, double fat_zero, float[] fcorr_weights, int corr_radius) {
        // fcorr_weights (per-tile fat-zero divisor) not yet plumbed -> native uses null (matches the common null path)
        lib.tp_proc_exec_corr2d_normalize(proc, combo ? 1 : 0, fat_zero, corr_radius);
    }
    // ---- peak measurement (tp_peak.cu, roadmap 3-C rung 2) ---- // By Claude on 07/14/2026
    @Override public float[] execCorr2DPeaks(
            int pair_select,
            double abs_min, double rel_min, double min_peak, double eig_sub_frac,
            int refine, double eig_sub_frac1, double scale_axes, double inc_axes,
            boolean use_float) {
        int n = lib.tp_proc_exec_corr2d_peak(proc, pair_select,
                (float) abs_min, (float) rel_min, (float) min_peak, (float) eig_sub_frac,
                refine, (float) eig_sub_frac1, (float) scale_axes, (float) inc_axes,
                use_float ? 1 : 0);
        if (n <= 0) return null; // caller falls back to the Java oracle path
        float[] peaks = new float[n * PEAK_ROW_FLOATS];
        return (lib.tp_proc_get_peaks(proc, peaks) == n) ? peaks : null;
    }
    @Override public boolean setPeakDebias(float[] debias) {
        return lib.tp_proc_set_peak_debias(proc, debias, (debias == null) ? 0 : debias.length) == 0;
    }

    // ---- primitive-float pose fx/J (tp_lma.cu, roadmap rung 3-A4b) ---- // By Codex on 07/15/2026
    private static float[] poseCameraMeta(final IntersceneLmaFloat.Camera camera) {
        return new float[] {
                camera.width, camera.height, camera.focalLength, camera.pixelSize,
                camera.distortionRadius, camera.disparityRadius, camera.lineTime, camera.stepR};
    }
    private static void copyPoseVector3(final float[] vector, final float[] packed, final int offset) {
        if ((vector == null) || (vector.length < 3)) {
            throw new IllegalArgumentException("GpuQuadJna.execPoseFxJacobian: expected 3-vector");
        }
        System.arraycopy(vector, 0, packed, offset, 3);
    }
    @Override public IntersceneLmaFloat.Result execPoseFxJacobian(
            final IntersceneLmaFloat.Camera reference,
            final IntersceneLmaFloat.Camera scene,
            final float[] reference_xyz,
            final float[] reference_atr,
            final float[] scene_xyz,
            final float[] scene_atr,
            final float[] centers,
            final boolean[] selection,
            final boolean calculate_jacobian) {
        if ((reference == null) || (scene == null) || (centers == null) || ((centers.length % 3) != 0)) {
            throw new IllegalArgumentException("GpuQuadJna.execPoseFxJacobian: inconsistent inputs");
        }
        final int numTiles = centers.length / 3;
        if ((numTiles == 0) || ((selection != null) && (selection.length != numTiles))) {
            throw new IllegalArgumentException("GpuQuadJna.execPoseFxJacobian: selection length");
        }
        final float[] poseVectors = new float[12];
        copyPoseVector3(reference_xyz, poseVectors, 0);
        copyPoseVector3(reference_atr, poseVectors, 3);
        copyPoseVector3(scene_xyz, poseVectors, 6);
        copyPoseVector3(scene_atr, poseVectors, 9);
        final byte[] packedSelection = new byte[numTiles];
        for (int tile = 0; tile < numTiles; tile++) {
            packedSelection[tile] = (byte) (((selection == null) || selection[tile]) ? 1 : 0);
        }
        final float[] fx = new float[IntersceneLmaFloat.NUM_COMPONENTS * numTiles];
        final float[] nativeJt = calculate_jacobian ?
                new float[IntersceneLmaFloat.NUM_PARAMS * IntersceneLmaFloat.NUM_COMPONENTS * numTiles] :
                new float[1]; // keep the JNA array argument non-null; native ignores it in fx-only mode
        final byte[] valid = new byte[numTiles];
        final int rc = lib.tp_proc_exec_pose_fx_jacobian(
                proc,
                poseCameraMeta(reference), reference.radial,
                reference.rByRDist, reference.rByRDist.length, reference.ers, reference.ers.length,
                poseCameraMeta(scene), scene.radial,
                scene.rByRDist, scene.rByRDist.length, scene.ers, scene.ers.length,
                poseVectors, centers, packedSelection, numTiles, calculate_jacobian ? 1 : 0,
                fx, nativeJt, valid);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna.execPoseFxJacobian rc="+rc+": "+lib.tp_last_error());
        }
        int numValid = 0;
        for (byte value : valid) if (value != 0) numValid++;
        return new IntersceneLmaFloat.Result(
                numTiles, fx, calculate_jacobian ? nativeJt : null, valid, numValid);
    }

    // ---- fixed float pose LMA candidate (tp_lma.cu, roadmap rung 3-A4d) ---- // By Codex on 07/15/2026
    @Override public IntersceneLmaFloat.LmaStepResult execPoseLmaStep(
            final float lambda,
            final float[] weights,
            final float[] jt,
            final float[] ymfxWeighted,
            final float[] vector) {
        if ((weights == null) || (weights.length == 0) ||
                (jt == null) || (jt.length != IntersceneLmaFloat.NUM_PARAMS * weights.length) ||
                (ymfxWeighted == null) || (ymfxWeighted.length != weights.length) ||
                (vector == null) || (vector.length != IntersceneLmaFloat.NUM_PARAMS)) {
            throw new IllegalArgumentException("GpuQuadJna.execPoseLmaStep: inconsistent inputs");
        }
        final float[] packed = new float[IntersceneLmaFloat.LMA_RESULT_FLOATS];
        final int rc = lib.tp_proc_exec_pose_lma_step(
                proc, lambda, weights, jt, ymfxWeighted, vector, weights.length, packed);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna.execPoseLmaStep rc="+rc+": "+lib.tp_last_error());
        }
        return unpackPoseLmaStep(packed);
    }

    // ---- rung C1: resident prepareLMA (design 2026-07-16) ---- // By Claude on 07/17/2026
    @Override public IntersceneLmaFloat.PrepareResidentResult execPoseLmaPrepareResident(
            final IntersceneLmaFloat.Camera reference,
            final IntersceneLmaFloat.Camera scene,
            final float[] reference_xyz,
            final float[] reference_atr,
            final float[] scene_xyz,
            final float[] scene_atr,
            final float[] centers,
            final int corrSlot,
            final int ntileShift,
            final boolean useEigen,
            final float eigMinSqrt,
            final float eigMaxSqrt,
            final float minConfidence,
            final boolean sameWeights,
            final float[] pullTargets,
            final float[] paramRegweights,
            final boolean capture) {
        if ((reference == null) || (scene == null) || (centers == null) ||
                ((centers.length % 3) != 0) ||
                (pullTargets == null) || (pullTargets.length != IntersceneLmaFloat.NUM_PARAMS) ||
                (paramRegweights == null) || (paramRegweights.length != IntersceneLmaFloat.NUM_PARAMS)) {
            throw new IllegalArgumentException("GpuQuadJna.execPoseLmaPrepareResident: inconsistent inputs");
        }
        final int numTiles = centers.length / 3;
        final int numValues = IntersceneLmaFloat.NUM_COMPONENTS * numTiles +
                IntersceneLmaFloat.NUM_PARAMS;
        final float[] poseVectors = new float[12];
        copyPoseVector3(reference_xyz, poseVectors, 0);
        copyPoseVector3(reference_atr, poseVectors, 3);
        copyPoseVector3(scene_xyz, poseVectors, 6);
        copyPoseVector3(scene_atr, poseVectors, 9);
        final float[] prepResult = new float[IntersceneLmaFloat.PREP_RESULT_FLOATS];
        final float[] outY = capture ? new float[numValues] : null;
        final float[] outWeights = capture ? new float[numValues] : null;
        final float[] outEigen = capture ? new float[4 * numTiles] : null;
        final int rc = lib.tp_proc_exec_pose_lma_prepare_resident(
                proc,
                poseCameraMeta(reference), reference.radial,
                reference.rByRDist, reference.rByRDist.length, reference.ers, reference.ers.length,
                poseCameraMeta(scene), scene.radial,
                scene.rByRDist, scene.rByRDist.length, scene.ers, scene.ers.length,
                poseVectors, centers,
                numTiles, IntersceneLmaFloat.NUM_COMPONENTS,
                corrSlot, ntileShift,
                useEigen ? 1 : 0, eigMinSqrt, eigMaxSqrt,
                minConfidence, sameWeights ? 1 : 0,
                pullTargets, paramRegweights,
                prepResult, capture ? 1 : 0, outY, outWeights, outEigen);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna.execPoseLmaPrepareResident rc="+rc+": "+
                    lib.tp_last_error());
        }
        return new IntersceneLmaFloat.PrepareResidentResult(
                numTiles,
                prepResult[0] != 0.0f,
                (int) prepResult[1],
                prepResult[2], prepResult[3], prepResult[4],
                new float[] {prepResult[5], prepResult[6], prepResult[7]},
                outY, outWeights, outEigen);
    }

    // ---- resident float pose preparation/reduction continuation (roadmap rung 3-A4e) ----
    @Override public IntersceneLmaFloat.ResidentLmaStepResult execPoseLmaResidentStep(
            final float lambda,
            final float[] weights,
            final float[] y,
            final float[] eigen,
            final float[] vector,
            final float pureWeight,
            final float rmsDiff,
            final int numTiles,
            final boolean capturePrepared) {
        final int numValues = IntersceneLmaFloat.NUM_COMPONENTS * numTiles +
                IntersceneLmaFloat.NUM_PARAMS;
        // Rung C1: all-null weights/y/eigen = use the resident buffers assembled by
        // execPoseLmaPrepareResident (no per-step H2D). By Claude on 07/17/2026.
        final boolean usePrepared = (weights == null) && (y == null) && (eigen == null);
        if ((numTiles <= 0) ||
                (!usePrepared && ((weights == null) || (weights.length != numValues) ||
                (y == null) || (y.length != numValues) || (eigen == null) ||
                (eigen.length != 4 * numTiles))) || (vector == null) ||
                (vector.length != IntersceneLmaFloat.NUM_PARAMS) ||
                !Float.isFinite(pureWeight) || !(pureWeight > 0.0f) ||
                !Float.isFinite(rmsDiff)) {
            throw new IllegalArgumentException("GpuQuadJna.execPoseLmaResidentStep: inconsistent inputs");
        }
        final float[] packed = new float[IntersceneLmaFloat.LMA_RESIDENT_RESULT_FLOATS];
        final float[] preparedJt = capturePrepared ?
                new float[IntersceneLmaFloat.NUM_PARAMS * numValues] : null;
        final float[] preparedYmfx = capturePrepared ? new float[numValues] : null;
        final int rc = lib.tp_proc_exec_pose_lma_resident_step(
                proc, lambda, weights, y, eigen, vector, pureWeight, rmsDiff, numTiles,
                capturePrepared ? 1 : 0, packed, preparedJt, preparedYmfx);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna.execPoseLmaResidentStep rc="+rc+": "+
                    lib.tp_last_error());
        }
        final IntersceneLmaFloat.PreparedLma prepared = capturePrepared ?
                new IntersceneLmaFloat.PreparedLma(
                        preparedJt, preparedYmfx,
                        packed[IntersceneLmaFloat.LMA_RESULT_FLOATS],
                        packed[IntersceneLmaFloat.LMA_RESULT_FLOATS + 1]) : null;
        return new IntersceneLmaFloat.ResidentLmaStepResult(
                unpackPoseLmaStep(packed), prepared,
                packed[IntersceneLmaFloat.LMA_RESULT_FLOATS],
                packed[IntersceneLmaFloat.LMA_RESULT_FLOATS + 1],
                packed[IntersceneLmaFloat.LMA_RESULT_FLOATS + 2],
                packed[IntersceneLmaFloat.LMA_RESULT_FLOATS + 3],
                packed[IntersceneLmaFloat.LMA_RESULT_FLOATS + 4] != 0.0f,
                packed[IntersceneLmaFloat.LMA_RESULT_FLOATS + 5] != 0.0f);
    }

    private static IntersceneLmaFloat.LmaStepResult unpackPoseLmaStep(final float[] packed) {
        final int deltaOffset = 1 + IntersceneLmaFloat.LMA_PRODUCTS;
        final int candidateOffset = deltaOffset + IntersceneLmaFloat.NUM_PARAMS;
        return new IntersceneLmaFloat.LmaStepResult(
                packed[0] != 0.0f,
                Arrays.copyOfRange(packed, 1, deltaOffset),
                Arrays.copyOfRange(packed, deltaOffset, candidateOffset),
                Arrays.copyOfRange(packed, candidateOffset, IntersceneLmaFloat.LMA_RESULT_FLOATS));
    }

    // ---- LMA normal-equation products (tp_lma.cu, roadmap rung 3) ---- // By Codex on 07/14/2026
    @Override public double[] execLmaNormalProducts(
            double[] weights, double[][] jt, double[] ymfxWeighted) {
        if ((weights == null) || (jt == null) || (jt.length == 0) ||
                (ymfxWeighted == null) || (ymfxWeighted.length != weights.length)) {
            throw new IllegalArgumentException("GpuQuadJna.execLmaNormalProducts: inconsistent inputs");
        }
        final int numParams = jt.length;
        final int numValues = weights.length;
        final double[] flatJt = new double[numParams * numValues];
        for (int i = 0; i < numParams; i++) {
            if ((jt[i] == null) || (jt[i].length != numValues)) {
                throw new IllegalArgumentException("GpuQuadJna.execLmaNormalProducts: jt["+i+"] length");
            }
            System.arraycopy(jt[i], 0, flatJt, i * numValues, numValues);
        }
        final double[] out = new double[numParams * numParams + numParams];
        final int rc = lib.tp_proc_exec_lma_products(
                proc, weights, flatJt, ymfxWeighted, numParams, numValues, out);
        if (rc != 0) {
            throw new IllegalStateException("GpuQuadJna.execLmaNormalProducts rc="+rc+": "+lib.tp_last_error());
        }
        return out;
    }

    @Override public float[][] getCorr2D(int corr_rad) {
        int corr_size = (2 * corr_rad + 1) * (2 * corr_rad + 1);
        int n = lib.tp_proc_num_corr_tiles(proc);
        float[] flat = new float[n * corr_size];
        lib.tp_proc_get_corr2d(proc, flat, corr_rad);
        float[][] corrs = new float[n][corr_size];
        for (int i = 0; i < n; i++) System.arraycopy(flat, i * corr_size, corrs[i], 0, corr_size);
        return corrs;
    }
    @Override public float[][] getCorr2DCombo(int corr_rad) {
        int corr_size = (2 * corr_rad + 1) * (2 * corr_rad + 1);
        int n = lib.tp_proc_num_corr_combo(proc);
        num_corr_combo_tiles = n; // NOT YET TESTED (latent: no live consumer on validated path) — side-effect parity w/ base, cf. f6dcc90f. By Claude 06/27/2026.
        float[] flat = new float[n * corr_size];
        lib.tp_proc_get_corr2d_combo(proc, flat);
        float[][] corrs = new float[n][corr_size];
        for (int i = 0; i < n; i++) System.arraycopy(flat, i * corr_size, corrs[i], 0, corr_size);
        return corrs;
    }

    // TpProc is full-frame fixed-size -> no per-WOI realloc; report the full frame.
    @Override public int[] handleWH(int[] wh, boolean ref_scene) {
        return (wh != null) ? wh : new int[]{img_width, img_height};
    }

    // ---- oracle TD-correlation readback / erase (read gpu_corr_indices / gpu_corrs_td / combo_indices) ----
    @Override public int[] getCorrIndices() {
        int n = lib.tp_proc_num_corr_tiles(proc);
        num_corr_tiles = n; // propagate the native count to the base field (base getCorrIndices does this too;
                            // getFromGpu/getCorrTilesTd read num_corr_tiles -> without this they see 0). By Claude 06/26/2026.
        int[] out = new int[n];
        lib.tp_proc_get_corr_indices(proc, out, n);
        return out;
    }
    @Override public float[] getCorrTdData() {
        int n = lib.tp_proc_num_corr_tiles(proc);
        int td = 4 * getDttSize() * getDttSize();
        float[] out = new float[n * td];
        lib.tp_proc_get_corr_td(proc, out);
        return out;
    }
    // Re-upload host-built TD per-pair corr indices + data (inverse of getCorrIndices/getCorrTdData): the
    // TD->PD re-conversion path (TDCorrTile.convertTDtoPD). Base does JCuda cuMemcpyHtoD into null buffers.
    // ensureRbgCorr() guarantees the native corr buffers exist (idempotent). By Claude on 06/26/2026.
    @Override public void setCorrIndicesTdData(int num_tiles, int[] corr_indices, float[] fdata) {
        ensureRbgCorr();
        num_corr_tiles = num_tiles; // NOT YET TESTED on this path — base setCorrIndicesTdData sets this; readers (getFromGpu) rely on it. By Claude 06/27/2026.
        lib.tp_proc_set_corr_indices_td(proc, corr_indices, fdata, num_tiles);
    }
    @Override public int[] getCorrComboIndices() {
        int n = lib.tp_proc_num_corr_combo(proc);
        num_corr_combo_tiles = n; // NOT YET TESTED (latent: no live consumer on validated path) — side-effect parity w/ base, cf. f6dcc90f. By Claude 06/27/2026.
        int[] out = new int[n];
        lib.tp_proc_get_corr_combo_indices(proc, out, n);
        return out;
    }
    @Override public void eraseGpuCorrs() { lib.tp_proc_erase_corrs(proc); }

    // ---- LPF filter coefficients -> native __constant__ memory (same as GpuQuad's cuModuleGetGlobal path) ----
    @Override public void setLpfRbg(float[][] lpf_rbg, boolean debug) {
        int l = lpf_rbg[0].length;
        float[] lpf_flat = new float[lpf_rbg.length * l];
        for (int i = 0; i < lpf_rbg.length; i++)
            for (int j = 0; j < l; j++) lpf_flat[j + i * l] = lpf_rbg[i][j];
        lib.tp_proc_set_const(proc, "lpf_data", lpf_flat, lpf_flat.length);
    }
    @Override public void setLpfCorr(String const_name, float[] lpf_flat, boolean debug) {
        lib.tp_proc_set_const(proc, const_name, lpf_flat, lpf_flat.length);
    }
    // Debug-only constant-memory read-back (base uses gpuTileProcessor.module). No-op natively to avoid
    // the NPE when debug printing is enabled; the values are uploaded via tp_proc_set_const regardless.
    @Override public void printConstMem(String const_name, boolean flt) {
        if (flt) System.out.println("GpuQuadJna.printConstMem(" + const_name + "): skipped (native backend)");
    }

    // ---- textures (oracle path): textures_nonoverlap (DP) + readback. Matches GpuQuad.execTextures_DP:
    // textures_nonoverlap<<<1,1>>> CDP-builds the index list (create_nonoverlap_list) then CDP-launches
    // textures_accumulate. Production uses linescan_order=0 (so getExtra de-interleaves via texture_indices).
    @Override public void execTextures(
            double [] color_weights, boolean is_lwir, double min_shot, double scale_shot,
            double diff_sigma, double diff_threshold, double min_agree, boolean dust_remove,
            int keep_weights, boolean calc_textures, boolean calc_extra, boolean linescan_order) {
        int nc = Math.min(color_weights.length, 3);
        float [] weights3 = new float[3];
        weights3[0] = (float) color_weights[0];
        weights3[1] = (nc > 1) ? (float) color_weights[1] : 0.0f;
        weights3[2] = (nc > 2) ? (float) color_weights[2] : 0.0f;
        float [] params5 = { (float) min_shot, (float) scale_shot, (float) diff_sigma,
                             (float) diff_threshold, (float) min_agree };
        lib.tp_proc_exec_textures(proc, params5, weights3,
                is_lwir ? 1 : 0, dust_remove ? 1 : 0, keep_weights,
                calc_textures ? 1 : 0, calc_extra ? 1 : 0, linescan_order ? 1 : 0);
    }
    // read pnum (kernel-built count) then that many packed indices (= GpuQuad.getTextureIndices, sans the
    // private num_texture_tiles side-effect, which our overridden getFlatTextures does not depend on).
    @Override public int[] getTextureIndices() {
        int tilesX = img_width / getDttSize(), tilesY = img_height / getDttSize();
        int [] tmp = new int [tilesX * tilesY];
        int cnt = lib.tp_proc_get_texture_indices(proc, tmp, tmp.length);
        if (cnt < 0) cnt = 0;
        num_texture_tiles = cnt; // NOT YET TESTED (texture path off validated CUAS path) — side-effect parity w/ base. By Claude 06/27/2026.
        int [] out = new int [cnt];
        System.arraycopy(tmp, 0, out, 0, cnt);
        return out;
    }
    // diff_rgb_combo "extra": reshape diff_rgb_combo (texture_indices order, linescan_order=0) into
    // [num_cams*(num_colors+1)][tilesX*tilesY] keyed by ntile. Identical to GpuQuad.getExtra.
    @Override public float[][] getExtra() {
        int [] texture_indices = getTextureIndices();
        int num_tile_extra = num_cams * (num_colors + 1);
        float [] diff_rgb_combo = new float [texture_indices.length * num_tile_extra];
        if (diff_rgb_combo.length > 0) lib.tp_proc_get_diff_rgb_combo(proc, diff_rgb_combo, diff_rgb_combo.length);
        int tilesX = img_width / getDttSize(), tilesY = img_height / getDttSize();
        float [][] extra = new float [num_tile_extra][tilesX * tilesY];
        for (int i = 0; i < texture_indices.length; i++) {
            if (((texture_indices[i] >> GPUTileProcessor.LIST_TEXTURE_BIT) & 1) != 0) {
                int ntile = (texture_indices[i] >> GPUTileProcessor.TEXT_NTILE_SHIFT);
                for (int l = 0; l < num_tile_extra; l++) {
                    extra[l][ntile] = diff_rgb_combo[i * num_tile_extra + l];
                }
            }
        }
        return extra;
    }
    // de-pitch gpu_textures (= GpuQuad.getFlatTextures): num_tiles * (slices * (2*DTT)^2) floats.
    @Override public float[] getFlatTextures(int num_tiles, int num_colors, boolean keep_weights) {
        int texture_slices = num_colors + 1 + (keep_weights ? (num_cams + num_colors + 1) : 0);
        int texture_slice_size = (2 * getDttSize()) * (2 * getDttSize());
        int texture_tile_size = texture_slices * texture_slice_size;
        float [] out = new float [texture_tile_size * num_tiles];
        lib.tp_proc_get_textures(proc, out, num_tiles, num_colors, keep_weights ? 1 : 0);
        return out;
    }
}
