package com.elphel.imagej.gpu.jna;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.file.Files;
import java.nio.file.Paths;

import com.sun.jna.Native;
import com.sun.jna.Pointer;

/**
 * Stage-1 validation: drive the native geometry path entirely across the JNA boundary
 * (no JCuda), then check the result against the tile_processor_gpu reference data.
 *
 * Flow (mirrors GpuQuad.execCalcReverseDistortions / execRotDerivs):
 *   load aux.geometry_correction  -> tp_set_geometry_correction
 *   tp_exec_calc_reverse_distortions
 *   tp_get_rbyrdist               -> compare to aux.rbyrdist (gold reference, maxRbyRDistErr style)
 *   load aux.correction_vector    -> tp_set_correction_vector
 *   tp_exec_rot_derivs(16)
 *   tp_get_rot_deriv              -> sanity: finite + cam-0 rotation matrix det ~ 1
 *
 * Data files are raw little-endian float32 (x86 fread in the harness).
 * Run:
 *   java -Djna.library.path=&lt;dir with libtileproc.so&gt; -cp target/classes:&lt;jna.jar&gt; \
 *        com.elphel.imagej.gpu.jna.Stage1 [srcdir] [libcudadevrt.a] [cltdir] [main|aux]
 * By Claude on 2026-06-25.
 */
public class Stage1 {
    static final int RBYRDIST_LEN   = 5001;
    static final int ROT_DERIV_LEN  = 5 * 16 * 3 * 3; // 720
    static final int NUM_CAMS       = 16;

    static float[] readLEFloats(String path) throws IOException {
        byte[] b = Files.readAllBytes(Paths.get(path));
        if ((b.length & 3) != 0) throw new IOException("not a multiple of 4 bytes: " + path);
        ByteBuffer bb = ByteBuffer.wrap(b).order(ByteOrder.LITTLE_ENDIAN);
        float[] f = new float[b.length / 4];
        for (int i = 0; i < f.length; i++) f[i] = bb.getFloat();
        return f;
    }

    public static void main(String[] args) throws IOException {
        String srcdir = (args.length > 0) ? args[0] : "/home/elphel/git/tile_processor_gpu/src";
        String devrt  = (args.length > 1) ? args[1] : "/usr/local/cuda/targets/x86_64-linux/lib/libcudadevrt.a";
        String cltdir = (args.length > 2) ? args[2] : "/home/elphel/git/tile_processor_gpu/clt";
        String set    = (args.length > 3) ? args[3] : "aux"; // aux = lwir(16), main = rgb(4)

        float[] fgc = readLEFloats(cltdir + "/" + set + ".geometry_correction");
        float[] fcv = readLEFloats(cltdir + "/" + set + ".correction_vector");
        float[] ref = readLEFloats(cltdir + "/" + set + ".rbyrdist");
        System.out.printf("Stage1 (JNA geometry): set=%s  gc=%d cv=%d ref_rByRDist=%d floats%n",
                set, fgc.length, fcv.length, ref.length);

        TpJna lib = Native.load("tileproc", TpJna.class);
        Pointer m = lib.tp_create_module(srcdir, devrt);
        if (m == null) { System.out.println("FAIL: tp_create_module:\n" + lib.tp_last_error()); System.exit(1); }
        Pointer inst = lib.tp_create_instance(m);
        if (inst == null) { System.out.println("FAIL: tp_create_instance:\n" + lib.tp_last_error()); System.exit(1); }

        boolean ok = true;

        // --- reverse distortion table: gc -> rByRDist, compare to reference ---
        check(lib.tp_set_geometry_correction(inst, fgc, fgc.length),    "tp_set_geometry_correction", lib);
        check(lib.tp_exec_calc_reverse_distortions(inst),               "tp_exec_calc_reverse_distortions", lib);
        float[] rbr = new float[RBYRDIST_LEN];
        check(lib.tp_get_rbyrdist(inst, rbr),                           "tp_get_rbyrdist", lib);
        double maxErr = 0; int worst = -1;
        int n = Math.min(ref.length, RBYRDIST_LEN);
        for (int i = 0; i < n; i++) {
            double e = Math.abs((double) rbr[i] - (double) ref[i]);
            if (e > maxErr) { maxErr = e; worst = i; }
        }
        boolean rbrOk = maxErr <= 1e-3; // matches GpuQuad.maxRbyRDistErr() tolerance
        ok &= rbrOk;
        System.out.printf("  rByRDist: maxErr=%.3e at i=%d (n=%d)  %s%n",
                maxErr, worst, n, rbrOk ? "PASS" : "FAIL (tol 1e-3)");

        // --- rotation derivatives: cv -> rot_deriv, sanity check ---
        check(lib.tp_set_correction_vector(inst, fcv, fcv.length),      "tp_set_correction_vector", lib);
        check(lib.tp_exec_rot_derivs(inst, NUM_CAMS),                   "tp_exec_rot_derivs", lib);
        float[] rot = new float[ROT_DERIV_LEN];
        check(lib.tp_get_rot_deriv(inst, rot),                          "tp_get_rot_deriv", lib);
        boolean finite = true;
        for (float v : rot) if (Float.isNaN(v) || Float.isInfinite(v)) { finite = false; break; }
        // rots[NUM_CAMS][3][3] (group 0) is a SCALED rotation per camera (zoom*R), so det~zoom^3, NOT 1.
        // Structural sanity that needs no gold reference: rows mutually orthogonal (R*R^T off-diagonal ~0).
        int nGood = 0; double worstOrtho = 0;
        for (int c = 0; c < NUM_CAMS; c++) {
            int o = c * 9;
            double[] nrm = new double[3];
            for (int i = 0; i < 3; i++) {
                nrm[i] = Math.sqrt(rot[o+3*i]*rot[o+3*i] + rot[o+3*i+1]*rot[o+3*i+1] + rot[o+3*i+2]*rot[o+3*i+2]);
            }
            // normalized off-diagonal dot products (0,1),(0,2),(1,2)
            int[][] pr = {{0,1},{0,2},{1,2}};
            double wOff = 0;
            for (int k = 0; k < 3; k++) {
                int a = pr[k][0], b = pr[k][1];
                double d = rot[o+3*a]*rot[o+3*b] + rot[o+3*a+1]*rot[o+3*b+1] + rot[o+3*a+2]*rot[o+3*b+2];
                double nd = (nrm[a]>0 && nrm[b]>0) ? Math.abs(d)/(nrm[a]*nrm[b]) : (Math.abs(d)<1e-12 ? 0 : 1);
                if (nd > wOff) wOff = nd;
            }
            if (wOff < 1e-3) nGood++;
            if (wOff > worstOrtho) worstOrtho = wOff;
        }
        boolean rotOk = finite && (nGood == NUM_CAMS);
        ok &= rotOk;
        System.out.printf("  rot_deriv: finite=%b  orthogonal-rows for %d/%d cams (worst off-diag=%.2e)  %s (sanity; no gold ref)%n",
                finite, nGood, NUM_CAMS, worstOrtho, rotOk ? "PASS" : "FAIL");

        lib.tp_destroy_instance(inst);
        lib.tp_destroy_module(m);
        System.out.println(ok ? "RESULT: PASS (geometry path matches reference via JNA)"
                              : "RESULT: FAIL");
        System.exit(ok ? 0 : 2);
    }

    static void check(int rc, String what, TpJna lib) {
        if (rc != 0) { System.out.printf("FAIL: %s -> rc=%d: %s%n", what, rc, lib.tp_last_error()); System.exit(3); }
    }
}
