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(). Survivors are grouped into 4D-connected clusters; each compact
 *   cluster (a single target) injects total energy equal to the cluster maximum,
 *   independent of the survivor count - this compensates the filter, which does not
 *   conserve energy (it zeroes weaker components), and keeps the output continuous
 *   when a sub-pixel position change switches between 1 and 2 surviving cells.
 *   Large clusters (weak signal passing the filter linearly) bypass the
 *   renormalization and spreading - identity.
 * - 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 final int max_clust;          // maximal 4D cluster size treated as a compact target; larger - linear mode
	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,
			int max_clust) {
		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.max_clust =         max_clust;
		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").
	 * Survivors (value > 0 after filterConv5dROI) are grouped into 4D-connected
	 * clusters (3^4-1 connectivity in (px, py, vx, vy)). Three regimes:
	 * - 1 isolated survivor: full 3x3x3x3 spread with fading peripherals
	 *   (orthogonal neighbors weight `spread`, diagonals products of it).
	 * - 2..max_clust near-equal survivors (one target straddling cell boundaries;
	 *   near-equality is guaranteed by the filter): each survivor's spread is
	 *   limited along the "resolved" dimensions spanned by its neighbors
	 *   (e.g. 2x3x3x3 for a velocity-adjacent pair - no fading added there).
	 * - more than max_clust survivors in one cluster: weak signal that passed the
	 *   filter linearly - identity pass-through, no spreading, no renormalization.
	 * Energy normalization (compact clusters): the filter does not conserve energy -
	 * it zeroes weaker components - so a sub-pixel target shift switching between
	 * 1 and 2 surviving cells must not change the injected energy. Each cluster
	 * contributes a total equal to its maximum M, shared between survivors
	 * proportionally to their values: survivor i contributes M*m_i/sum(m_j).
	 * (1 survivor -> total m1; a near-equal pair -> ~M/2 each, total M = m1 -
	 * continuous across the transition.)
	 */
	private double [][][] condition(double [][][] meas) {
		double [][][] out = new double [meas.length][num_sub][vel_dim*vel_dim];
		final int nvel = vel_dim * vel_dim;
		final int [] cell_stack = new int [meas.length * nvel];
		final int [] label =      new int [meas.length * nvel];
		for (int s = 0; s < num_sub; s++) {
			// Pass 1: label 4D-connected clusters of survivors, collect per-cluster size/max/sum
			java.util.Arrays.fill(label, -1);
			java.util.ArrayList<double[]> clusters = new java.util.ArrayList<double[]>(); // {size, max, sum}
			for (int cell0 = 0; cell0 < label.length; cell0++) {
				if ((label[cell0] >= 0) || (meas[cell0 / nvel][s][cell0 % nvel] <= 0)) continue;
				int ncl = clusters.size();
				double cl_size = 0, cl_max = 0, cl_sum = 0;
				int sp = 0;
				cell_stack[sp++] = cell0;
				label[cell0] = ncl;
				while (sp > 0) {
					int cell = cell_stack[--sp];
					int p = cell / nvel, v = cell % nvel;
					double val = meas[p][s][v];
					cl_size += 1;
					cl_sum +=  val;
					if (val > cl_max) cl_max = val;
					int px = p % roi_width,  py = p / roi_width;
					int vx = v % vel_dim,    vy = v / vel_dim;
					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;
									int ncell = np * nvel + nvy * vel_dim + nvx;
									if ((label[ncell] < 0) && (meas[np][s][nvy * vel_dim + nvx] > 0)) {
										label[ncell] = ncl;
										cell_stack[sp++] = ncell;
									}
								}
							}
						}
					}
				}
				clusters.add(new double[] {cl_size, cl_max, cl_sum});
			}
			// Pass 2: spread each survivor according to its cluster regime
			for (int p = 0; p < meas.length; p++) {
				int px = p % roi_width;
				int py = p / roi_width;
				for (int v = 0; v < meas[p][s].length; v++) {
					double m = meas[p][s][v];
					if (m <= 0) continue; // survivors only
					double [] cl = clusters.get(label[p * nvel + v]);
					if (cl[0] > max_clust) { // weak signal, linear mode - identity, no renormalization
						out[p][s][v] += m;
						continue;
					}
					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;
					// 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;
					// Spread this survivor's share of the cluster energy: total = M*m_i/sum(m_j),
					// so the whole cluster injects exactly its maximum M (continuous when a
					// sub-pixel target shift changes the number of surviving cells)
					double scale = (cl[1] * m / cl[2]) / 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;
	}
}
