package com.elphel.imagej.cuas.rt;
/**
 **
 ** CuasConditioning.java - lean per-sensor image conditioning for the CUAS RT / tile-processor path.
 **
 ** Copyright (C) 2026 Elphel, Inc.
 **
 ** -----------------------------------------------------------------------------**
 **
 **  CuasConditioning.java is free software: you can redistribute it and/or modify
 **  it under the terms of the GNU General Public License as published by
 **  the Free Software Foundation, either version 3 of the License, or
 **  (at your option) any later version.
 **
 **  This program is distributed in the hope that it will be useful,
 **  but WITHOUT ANY WARRANTY; without even the implied warranty of
 **  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 **  GNU General Public License for more details.
 **
 **  You should have received a copy of the GNU General Public License
 **  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 ** -----------------------------------------------------------------------------**
 **
 */

import com.elphel.imagej.cuas.CuasMotion;
import com.elphel.imagej.gpu.GpuQuad;
import com.elphel.imagej.tileprocessor.QuadCLT;

/**
 * Lean, self-contained per-sensor conditioning for the CUAS real-time / tile-processor path.
 * By Claude on 06/30/2026.
 *
 * Deliberately bypasses the heavy QuadCLT "conditioning over conditioning" instance-array path:
 * it operates directly on raw jp4-domain per-sensor pixel arrays and the (borrowed) calibration
 * values pulled from a QuadCLT (lwir_scales / lwir_offsets / lwir_scales2 and the per-pixel FPN).
 *
 * Pipeline (per sensor), applied in place:
 *   1) Row/Col denoise (optional): estimate a per-row and a per-column average, optionally
 *      high-pass the 1-D profile (preserve the low frequencies for the downstream LoG), subtract.
 *      The HPF is applied to the length-H / length-W average vectors, NOT the full image
 *      (mathematically ~equivalent, far fewer operations).
 *   2) Photometric + fixed-pattern (2-map model, per-pixel-ready):
 *         corrected = scales2*(raw-offset)^2 + scale*(raw-offset)    - FPN(x,y)
 *      matching the existing model (QuadCLT.calibratePhotometric2 forward form) so that with
 *      scales2==0 (the abandoned 2nd-order term) it is bit-identical to today's additive path.
 *      scale = global per-sensor (uniform) for now; the per-pixel *scale* map (multiplicative
 *      dust) is the future improvement, hence the 2-image (scale, offset) framing.
 */
public class CuasConditioning {

	public static class Config {
		public boolean rowcol_enable        = true;   // subtract per-row and per-col noise
		public boolean rowcol_hpf           = true;   // high-pass the 1-D avg profile before subtracting
		public double  rowcol_hpf_sigma     = 8.0;    // px; ~ the LoG scale (KR3-ish) - noise band the LoG would pass
		public double  rowcol_max_abs       = Double.POSITIVE_INFINITY; // clip |pixel| for the average (outlier reject)
		public double  rowcol_weight_outlier= 0.01;   // small weight for out-of-range pixels
		public boolean photometric          = true;   // apply scale*(raw-offset) [+ scales2 quadratic]
		public boolean apply_fpn            = true;    // subtract per-pixel FPN (the offset map's per-pixel part)
		public Config() {}
	}

	/**
	 * Condition all sensors in place.
	 * @param image_data [sensor][0][pixel] mono LWIR, row-major (idx = y*width + x). Modified in place.
	 * @param width      image width in pixels (height = pixels/width).
	 * @param scales     per-sensor global scale  (lwir_scales;  null -> 1.0).
	 * @param offsets    per-sensor global offset (lwir_offsets; null -> 0.0), subtracted before scaling.
	 * @param scales2    per-sensor 2nd-order term (lwir_scales2; null -> 0.0; normally unused).
	 * @param fpn        [sensor][0][pixel] per-pixel additive FPN, subtracted after photometric (null -> none).
	 * @param cfg        configuration (null -> defaults).
	 */
	public static void condition(
			double [][][] image_data,
			int           width,
			double []     scales,
			double []     offsets,
			double []     scales2,
			double [][][] fpn,
			Config        cfg) {
		if (cfg == null) cfg = new Config();
		final int num_sens = image_data.length;
		for (int s = 0; s < num_sens; s++) if (image_data[s] != null && image_data[s][0] != null) {
			final double [] px = image_data[s][0];
			if (cfg.rowcol_enable) {
				subtractRowColSensor(px, width, cfg);
			}
			if (cfg.photometric || cfg.apply_fpn) {
				final double B0 = (scales  != null) ? scales [s] : 1.0;
				final double C0 = (offsets != null) ? offsets[s] : 0.0;
				final double A0 = (scales2 != null) ? scales2[s] : 0.0;
				final double [] fp = (cfg.apply_fpn && fpn != null && fpn[s] != null) ? fpn[s][0] : null;
				for (int i = 0; i < px.length; i++) {
					double p = px[i];
					if (cfg.photometric) {
						final double d = p - C0;
						p = A0 * d * d + B0 * d;   // == B0*(raw-C0) when A0==0 (linear, matches current additive path)
					}
					if (fp != null) p -= fp[i];
					px[i] = p;
				}
			}
		}
	}

	/**
	 * The A2/RT ingest: read a scene's RAW /jp4/ source, condition it with the CURRENT
	 * calibration (curt_calib-updated lwir scales/offsets/scales2 + per-pixel FPN from the
	 * scene QuadCLT), and feed the GPU DIRECTLY - completely bypassing the QuadCLT
	 * image_data mechanics (per Andrey: this is the permanent RT shape, especially as more
	 * data becomes GPU-memory-resident; QuadCLT stays for geometry/poses only).
	 *
	 * GPU handoff contract (validated by the bayer-guard fix, ebef0b23): bind the scene
	 * FIRST (saveQuadClt - conditional, resets the guard only on instance switch), clear
	 * hasNewImageData, then FORCE-upload - setBayerImages(data,true) sets bayer_set, so the
	 * pipeline's later pull-variant setBayerImages(false) skips its re-pull and the
	 * conditioned pixels survive through convert_direct. TELL if this ever breaks: the
	 * downstream results become identical to the prepared-data path (silent re-pull).
	 *
	 * By Claude on 07/05/2026, from Andrey's design.
	 * @param scene      physical scene: /jp4/ source, calibration globals, FPN, geometry
	 * @param cfg        conditioning configuration (null - defaults: rowcol+photometric+FPN)
	 * @param threadsMax max read threads
	 * @param debugLevel debug verbosity
	 * @return true on success; false if raw source missing (GPU state untouched)
	 */
	public static boolean conditionSceneToGpu(
			final QuadCLT scene,
			final Config  cfg,
			final int     threadsMax,
			final int     debugLevel) {
		return conditionSceneToGpuCuas(scene, cfg, null, null, threadsMax, debugLevel);
	}

	/**
	 * CUAS wrapper: the ONE place where per-scene MOTION enters the RT ingest (Andrey's
	 * design 2026-07-05). Sets the scene's ERS rates (when given) before the ingest so the
	 * downstream grid transform (transformToScenePxPyD -> ErsCorrection, the fully-preserved
	 * RGB-proven path) applies the rolling-shutter correction; omegas==null &&
	 * velocities==null leaves the scene's current ERS state untouched (legacy behavior).
	 * For CUAS: omegas from the wobble radius + rotation RPM (analytic, per scene phase);
	 * velocities - PROVISION only, pass {0,0,0} (or null) for the rotation-only camera.
	 * FUTURE (GPU migration): the same omegas/velocities feed the GPU corr_vector
	 * imu_rot/imu_move (get_tiles_offsets has per-sensor ERS support, currently dormant) -
	 * to be enabled after verifying no double-application against the CPU pXpYD ERS.
	 * By Claude on 07/05/2026, from Andrey's design.
	 * @param omegas     angular rates {az, tilt, roll} rad/s, scene-vs-virtual-center
	 *                   relative (wobble-scale; the constant-rotation part cancels against
	 *                   the co-rotating center) - or null (with velocities null: keep state)
	 * @param velocities linear rates {vx, vy, vz} m/s (provision; null/{0,0,0} for CUAS)
	 */
	public static boolean conditionSceneToGpuCuas(
			final QuadCLT   scene,
			final Config    cfg,
			final double [] omegas,
			final double [] velocities,
			final int       threadsMax,
			final int       debugLevel) {
		if ((omegas != null) || (velocities != null)) {
			scene.getErsCorrection().setErsDt(
					(velocities != null) ? velocities : new double [3],
					(omegas != null)     ? omegas     : new double [3]);
		}
		return conditionSceneToGpuInner(scene, cfg, threadsMax, debugLevel);
	}

	private static boolean conditionSceneToGpuInner(
			final QuadCLT scene,
			final Config  cfg,
			final int     threadsMax,
			final int     debugLevel) {
		final double [][][] data = CuasMotion.readRawImageData(scene, threadsMax, debugLevel);
		if (data == null) {
			return false;
		}
		final int width = scene.getGeometryCorrection().getSensorWH()[0];
		condition(
				data,                    // double [][][] image_data (in place)
				width,                   // int           width
				scene.getLwirScales(),   // double []     scales
				scene.getLwirOffsets(),  // double []     offsets
				scene.getLwirScales2(),  // double []     scales2
				scene.getFPN(),          // double [][][] fpn (null - none)
				cfg);                    // Config        (null - defaults)
		final GpuQuad gpuQuad = scene.getGPUQuad();
		scene.saveQuadClt();                 // bind this scene (conditional; resets bayer guard on switch)
		scene.setHasNewImageData(false);     // the guard's other re-pull trigger - must be clear
		gpuQuad.setBayerImages(data, true);  // force H2D; sets bayer_set -> pipeline skips its pull
		return true;
	}

	/** Row/Col denoise for one sensor's flat pixel array (in place): subtract per-row, then per-col. */
	public static void subtractRowColSensor(final double [] px, final int width, final Config cfg) {
		final int height = px.length / width;
		// --- per-row profile (average along x for each row y) ---
		double [] rowProf = lineAverage(px, width, height, true, cfg.rowcol_max_abs, cfg.rowcol_weight_outlier);
		if (cfg.rowcol_hpf) rowProf = highpass(rowProf, cfg.rowcol_hpf_sigma);
		for (int y = 0; y < height; y++) {
			final double sub = rowProf[y];
			if (sub != 0.0) {
				final int row0 = y * width;
				for (int x = 0; x < width; x++) px[row0 + x] -= sub;
			}
		}
		// --- per-col profile (average along y for each col x), recomputed on row-corrected data ---
		double [] colProf = lineAverage(px, width, height, false, cfg.rowcol_max_abs, cfg.rowcol_weight_outlier);
		if (cfg.rowcol_hpf) colProf = highpass(colProf, cfg.rowcol_hpf_sigma);
		for (int y = 0; y < height; y++) {
			final int row0 = y * width;
			for (int x = 0; x < width; x++) px[row0 + x] -= colProf[x];
		}
	}

	/**
	 * Outlier-weighted average producing a 1-D profile.
	 * @param perRow true  -> length-height profile, averaging along x (per-row offset);
	 *               false -> length-width  profile, averaging along y (per-col offset).
	 */
	private static double [] lineAverage(
			final double [] px, final int width, final int height,
			final boolean perRow, final double max_abs, final double weight_outlier) {
		final int n = perRow ? height : width;
		final double [] sum = new double [n];
		final double [] wgt = new double [n];
		int idx = 0;
		for (int y = 0; y < height; y++) {
			for (int x = 0; x < width; x++) {
				final double d = px[idx++];
				if (Double.isNaN(d)) continue;
				final double w = (Math.abs(d) <= max_abs) ? 1.0 : weight_outlier;
				final int k = perRow ? y : x;
				sum[k] += w * d;
				wgt[k] += w;
			}
		}
		for (int k = 0; k < n; k++) sum[k] = (wgt[k] > 0.0) ? (sum[k] / wgt[k]) : 0.0;
		return sum;
	}

	/** 1-D Gaussian high-pass of a profile: x - smooth(x, sigma), reflect-padded. Returns a new array. */
	public static double [] highpass(final double [] x, final double sigma) {
		if (sigma <= 0.0) return x.clone();
		final int r = Math.max(1, (int) Math.ceil(3.0 * sigma));
		final double [] k = new double [2 * r + 1];
		double ksum = 0.0;
		for (int i = -r; i <= r; i++) { final double v = Math.exp(-0.5 * (i * i) / (sigma * sigma)); k[i + r] = v; ksum += v; }
		for (int i = 0; i < k.length; i++) k[i] /= ksum;
		final int n = x.length;
		final double [] hp = new double [n];
		for (int i = 0; i < n; i++) {
			double sm = 0.0;
			for (int j = -r; j <= r; j++) {
				int ii = i + j;
				if (ii < 0)      ii = -ii - 1;          // reflect
				else if (ii >= n) ii = 2 * n - ii - 1;  // reflect
				if (ii < 0) ii = 0; else if (ii >= n) ii = n - 1;
				sm += k[j + r] * x[ii];
			}
			hp[i] = x[i] - sm;
		}
		return hp;
	}

	/** Build the per-sensor per-pixel OFFSET map (post-scale, subtractive) = scale*offset + FPN,
	 *  so that (scale*raw - offset_map) == scale*(raw-offset) - FPN. For the future 2-map API /
	 *  visualization; not required by {@link #condition}. */
	public static double [][] buildOffsetMaps(
			final double [][][] fpn, final double [] scales, final double [] offsets, final int width) {
		final int num_sens = fpn.length;
		final double [][] out = new double [num_sens][];
		for (int s = 0; s < num_sens; s++) if (fpn[s] != null && fpn[s][0] != null) {
			final double B0 = (scales  != null) ? scales [s] : 1.0;
			final double C0 = (offsets != null) ? offsets[s] : 0.0;
			final double [] fp = fpn[s][0];
			final double [] m = new double [fp.length];
			final double base = B0 * C0;
			for (int i = 0; i < fp.length; i++) m[i] = base + fp[i];
			out[s] = m;
		}
		return out;
	}
}
