package com.elphel.imagej.gpu.jna;

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.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 final TpJna   lib;
    private final Pointer module;   // TpModule* (NVRTC kernels)
    private final Pointer proc;      // TpProc*   (persistent device buffers)
    private boolean closed = false;

    /**
     * @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
        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.
    }

    // Switch this shared GpuQuad to a different scene's QuadCLT. Base also clears
    // gpuTileProcessor.bayer_set (null here) — N/A natively: bayer + geometry are re-uploaded on the next
    // exec, kernels are unchanged. resetGeometryCorrection*/() are flag-only (safe to call).
    @Override public void updateQuadCLT(final QuadCLT quadCLT) {
        this.quadCLT = quadCLT;
        resetGeometryCorrection();
        resetGeometryCorrectionVector();
    }
    // Native re-uploads bayer each convert_direct -> nothing to reset (base clears the null gpuTileProcessor flag).
    @Override public void resetBayer() { /* no-op for the native backend */ }

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

    /** 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};
        lib.tp_proc_setup_rbg_corr(proc, Correlation2d.getNumPairs(num_cams), 0,0,0,0, cw[0],cw[1],cw[2], corr_out_rad);
        jna_rbg_corr_ready = true;
    }

    // ---- geometry ----
    @Override public void setGeometryCorrection() { setGeometryCorrection(quadCLT.getGeometryCorrection(), false); }
    @Override public void setGeometryCorrection(GeometryCorrection gc, boolean use_java_rByRDist) {
        float[] fgc = gc.expandSensors(GPUTileProcessor.MAX_NUM_CAMS).toFloatArray();
        lib.tp_proc_set_geometry(proc, fgc, fgc.length);
    }
    @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;
        }
        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);
        }
    }
    // 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;
    }

    // ---- tasks ----
    @Override public void setTasks(TpTask[] tile_tasks, boolean use_aux, boolean verify) {
        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); }
        lib.tp_proc_set_tasks(proc, ftasks, num_task_tiles, task_size * num_task_tiles);
    }

    // Read tasks back after calculate_tiles_offsets (computed centerXY/disp_dist) -> rebuild TpTask[].
    @Override public void updateTasks(TpTask[] tile_tasks, boolean use_aux) {
        num_task_tiles = tile_tasks.length;
        int task_size = getTaskSize();
        float[] ftasks = new float[task_size * num_task_tiles];
        lib.tp_proc_get_tasks(proc, ftasks, task_size * num_task_tiles);
        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); // apply_lpf = 1
    }
    @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);
    }
    @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);
        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};
    }

    // ---- 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)");
    }
}
