package com.elphel.imagej.cuas.rt;

/**
 * Layer 2 recurrent accumulation buffer for the CUAS real-time detector.
 * By Claude (Fable 5) on 06/09/2026, design decisions by Andrey Filippov (2026-06-09).
 *
 * Maintains a 4D state (px, py, vx, vy) over an ROI, same shape as the ROI output
 * of CuasRTUtils.convolve3d()/filterConv5dROI(): [roi_npix][num_sub][vel_dim*vel_dim].
 * The sub-pixel dimension is carried through unchanged (pixel_decimate is currently 1).
 *
 * Update rule (applied once per `period` time units):
 *   buf_new = leak * (1 - w_eff) * shift(buf_old, v) + w_eff * conditioned(measurement)
 *
 * - shift(buf, v): every 4D point moves by its own velocity vector. The pixel shift
 *   over `period` time units is (vx_off * period / velocity_decimate) - exactly
 *   integer (1 px per minimum velocity step) when period == velocity_decimate.
 *   Points predicted outside the ROI are dropped (target left the watched area).
 * - conditioned(measurement): adaptive 4D spreading of the filterConv5dROI survivors,
 *   see condition(). Each survivor contributes its value with total spread weight 1.
 * - w_eff: fixed `w` when wthresh <= 0; otherwise per-pixel adaptive
 *   w_eff = w * tanh(strength / wthresh), where strength is the strongest conditioned
 *   measurement at that pixel. Dark frames then give w_eff ~ 0 (pure prediction),
 *   and `leak` alone controls gap survival: leak^n_gap >= 0.2-0.3 recommended.
 *   With fixed w, the effective per-update gap decay is leak * (1 - w).
 *
 * This is intentionally a simplified fixed/adaptive-gain filter, NOT a Kalman filter -
 * the 4D map state is too large; a proper KF is reserved for the per-target worker stage.
 * Single-threaded for now (ROI is small); optimize later.
 */
public class CuasRecurrentBuffer {
	private final int roi_width;
	private final int roi_height;
	private final int num_sub;            // pixel_decimate^2
	private final int vel_dim;            // 2*velocity_radius + 1
	private final int velocity_radius;
	private final int velocity_decimate;
	private final int period;             // time units between updates
	private final double spread;          // weight of an orthogonal 4D neighbor when spreading (diagonals get products)
	private double [][][] buffer;         // [roi_width*roi_height][num_sub][vel_dim*vel_dim]
	private int num_updates = 0;

	public CuasRecurrentBuffer (
			int roi_width,
			int roi_height,
			int num_sub,
			int velocity_radius,
			int velocity_decimate,
			int period,
			double spread) {
		this.roi_width =         roi_width;
		this.roi_height =        roi_height;
		this.num_sub =           num_sub;
		this.velocity_radius =   velocity_radius;
		this.vel_dim =           2 * velocity_radius + 1;
		this.velocity_decimate = velocity_decimate;
		this.period =            period;
		this.spread =            spread;
		this.buffer = new double [roi_width*roi_height][num_sub][vel_dim*vel_dim];
	}

	public double [][][] getBuffer() {
		return buffer;
	}

	public double [][][] getBufferCopy() {
		double [][][] copy = new double [buffer.length][num_sub][vel_dim*vel_dim];
		for (int p = 0; p < buffer.length; p++) for (int s = 0; s < num_sub; s++) {
			System.arraycopy(buffer[p][s], 0, copy[p][s], 0, buffer[p][s].length);
		}
		return copy;
	}

	public int getNumUpdates() {
		return num_updates;
	}

	/**
	 * One recurrent update step.
	 * @param measurement filterConv5dROI() output for the current time step,
	 *                    [roi_npix][num_sub][vel_dim*vel_dim]; may be null (dark/skipped
	 *                    frame) - then the update is pure decayed prediction.
	 * @param w       measurement weight (fixed gain), 0..1
	 * @param leak    extra decay of the predicted (shifted) buffer per update
	 * @param wthresh if > 0, per-pixel adaptive gain w_eff = w*tanh(strength/wthresh)
	 */
	public void update (
			double [][][] measurement,
			double        w,
			double        leak,
			double        wthresh) {
		double [][][] shifted = shiftByVelocity(buffer);
		double [][][] cond = (measurement != null) ? condition(measurement) : null;
		for (int p = 0; p < buffer.length; p++) {
			double w_eff = w;
			if (cond == null) {
				w_eff = 0.0;
			} else if (wthresh > 0) { // adaptive gain from the strongest conditioned measurement at this pixel
				double strength = 0;
				for (int s = 0; s < num_sub; s++) for (int v = 0; v < cond[p][s].length; v++) {
					if (cond[p][s][v] > strength) strength = cond[p][s][v];
				}
				w_eff = w * Math.tanh(strength / wthresh);
			}
			for (int s = 0; s < num_sub; s++) {
				for (int v = 0; v < buffer[p][s].length; v++) {
					double pred = leak * (1.0 - w_eff) * shifted[p][s][v];
					double meas = (cond != null) ? (w_eff * cond[p][s][v]) : 0.0;
					buffer[p][s][v] = pred + meas;
				}
			}
		}
		num_updates++;
	}

	/**
	 * Blind prediction: every 4D point moves by its own velocity, pixel shift =
	 * velocity * period (integer-clean when period == velocity_decimate).
	 * Multiple sources landing on the same destination are summed.
	 */
	private double [][][] shiftByVelocity(double [][][] src) {
		double [][][] dst = new double [src.length][num_sub][vel_dim*vel_dim];
		for (int p = 0; p < src.length; p++) {
			int px = p % roi_width;
			int py = p / roi_width;
			for (int s = 0; s < num_sub; s++) {
				for (int v = 0; v < src[p][s].length; v++) {
					double val = src[p][s][v];
					if (val == 0) continue;
					int vx_off = (v % vel_dim) - velocity_radius;
					int vy_off = (v / vel_dim) - velocity_radius;
					int dx = (int) Math.round(((double) vx_off * period) / velocity_decimate);
					int dy = (int) Math.round(((double) vy_off * period) / velocity_decimate);
					int nx = px + dx;
					int ny = py + dy;
					if ((nx < 0) || (nx >= roi_width) || (ny < 0) || (ny >= roi_height)) {
						continue; // predicted out of the ROI - drop
					}
					dst[ny * roi_width + nx][s][v] += val;
				}
			}
		}
		return dst;
	}

	/**
	 * Adaptive 4D spreading of the filtered measurement ("conditioning").
	 * Each survivor (value > 0 after filterConv5dROI) is spread into its 4D
	 * (px, py, vx, vy) neighbourhood, except along "resolved" dimensions:
	 * a dimension is resolved when another survivor exists among the 3^4-1
	 * immediate 4D neighbors offset along that dimension - the survivors already
	 * span it, so expanding would only blur real resolution.
	 * - 1 isolated survivor: full 3x3x3x3 hypercube (orthogonal neighbors weight
	 *   `spread`, diagonals `spread^2` etc.).
	 * - 2 survivors adjacent in velocity: that velocity axis not expanded (2x3x3x3).
	 * Total contributed weight per survivor is 1 regardless of geometry (weights are
	 * normalized over the in-bounds part of the adaptive kernel).
	 */
	private double [][][] condition(double [][][] meas) {
		double [][][] out = new double [meas.length][num_sub][vel_dim*vel_dim];
		for (int p = 0; p < meas.length; p++) {
			int px = p % roi_width;
			int py = p / roi_width;
			for (int s = 0; s < num_sub; s++) {
				for (int v = 0; v < meas[p][s].length; v++) {
					double m = meas[p][s][v];
					if (m <= 0) continue; // survivors only
					int vx = v % vel_dim;
					int vy = v / vel_dim;
					// Find resolved dimensions: scan the 3^4-1 immediate 4D neighbors for other survivors
					boolean res_x = false, res_y = false, res_vx = false, res_vy = false;
					for (int dpy = -1; dpy <= 1; dpy++) {
						int ny = py + dpy;
						if ((ny < 0) || (ny >= roi_height)) continue;
						for (int dpx = -1; dpx <= 1; dpx++) {
							int nx = px + dpx;
							if ((nx < 0) || (nx >= roi_width)) continue;
							int np = ny * roi_width + nx;
							for (int dvy = -1; dvy <= 1; dvy++) {
								int nvy = vy + dvy;
								if ((nvy < 0) || (nvy >= vel_dim)) continue;
								for (int dvx = -1; dvx <= 1; dvx++) {
									int nvx = vx + dvx;
									if ((nvx < 0) || (nvx >= vel_dim)) continue;
									if ((dpx == 0) && (dpy == 0) && (dvx == 0) && (dvy == 0)) continue;
									if (meas[np][s][nvy * vel_dim + nvx] > 0) {
										if (dpx != 0) res_x =  true;
										if (dpy != 0) res_y =  true;
										if (dvx != 0) res_vx = true;
										if (dvy != 0) res_vy = true;
									}
								}
							}
						}
					}
					int rx0 =  res_x  ? 0 : -1, rx1 =  res_x  ? 0 : 1;
					int ry0 =  res_y  ? 0 : -1, ry1 =  res_y  ? 0 : 1;
					int rvx0 = res_vx ? 0 : -1, rvx1 = res_vx ? 0 : 1;
					int rvy0 = res_vy ? 0 : -1, rvy1 = res_vy ? 0 : 1;
					// First pass: total in-bounds weight of the adaptive kernel
					double sum_w = 0;
					for (int dpy = ry0; dpy <= ry1; dpy++) {
						int ny = py + dpy;
						if ((ny < 0) || (ny >= roi_height)) continue;
						for (int dpx = rx0; dpx <= rx1; dpx++) {
							int nx = px + dpx;
							if ((nx < 0) || (nx >= roi_width)) continue;
							for (int dvy = rvy0; dvy <= rvy1; dvy++) {
								int nvy = vy + dvy;
								if ((nvy < 0) || (nvy >= vel_dim)) continue;
								for (int dvx = rvx0; dvx <= rvx1; dvx++) {
									int nvx = vx + dvx;
									if ((nvx < 0) || (nvx >= vel_dim)) continue;
									sum_w += kernelWeight(dpx, dpy, dvx, dvy);
								}
							}
						}
					}
					if (sum_w <= 0) continue;
					// Second pass: spread m with total weight 1
					double scale = m / sum_w;
					for (int dpy = ry0; dpy <= ry1; dpy++) {
						int ny = py + dpy;
						if ((ny < 0) || (ny >= roi_height)) continue;
						for (int dpx = rx0; dpx <= rx1; dpx++) {
							int nx = px + dpx;
							if ((nx < 0) || (nx >= roi_width)) continue;
							int np = ny * roi_width + nx;
							for (int dvy = rvy0; dvy <= rvy1; dvy++) {
								int nvy = vy + dvy;
								if ((nvy < 0) || (nvy >= vel_dim)) continue;
								for (int dvx = rvx0; dvx <= rvx1; dvx++) {
									int nvx = vx + dvx;
									if ((nvx < 0) || (nvx >= vel_dim)) continue;
									out[np][s][nvy * vel_dim + nvx] += scale * kernelWeight(dpx, dpy, dvx, dvy);
								}
							}
						}
					}
				}
			}
		}
		return out;
	}

	// Unnormalized spreading weight: center 1, each non-zero offset dimension multiplies by `spread`
	// (orthogonal neighbor = spread, 2D diagonal = spread^2, ..., full 4D diagonal = spread^4).
	private double kernelWeight(int dpx, int dpy, int dvx, int dvy) {
		double w = 1.0;
		if (dpx != 0) w *= spread;
		if (dpy != 0) w *= spread;
		if (dvx != 0) w *= spread;
		if (dvy != 0) w *= spread;
		return w;
	}
}
