/**
 **
 ** CuasPoseRT.java - RT pose-adjustment prototype: per-scene 3-angle pose
 ** fitting against the persistent virtual-center reference.
 **
 ** Copyright (C) 2026 Elphel, Inc.
 **
 ** -----------------------------------------------------------------------------**
 **
 **  CuasPoseRT.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/>.
 ** -----------------------------------------------------------------------------**
 **
 */
package com.elphel.imagej.cuas.rt;

import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;

import com.elphel.imagej.cameras.CLTParameters;
import com.elphel.imagej.gpu.TpTask;
import com.elphel.imagej.tileprocessor.ErsCorrection;
import com.elphel.imagej.tileprocessor.Interscene;
import com.elphel.imagej.tileprocessor.OpticalFlow;
import com.elphel.imagej.tileprocessor.QuadCLT;
import com.elphel.imagej.tileprocessor.QuadCLTCPU;
import com.elphel.imagej.tileprocessor.TwoQuadCLT;

import ij.ImagePlus;
import ij.ImageStack;
import ij.Prefs;
import ij.process.FloatProcessor;

/**
 * Phase A of the RT pose-adjustment package (design: ANDREY_CONTINUE 2026-07-03,
 * details in imagej-elphel-internal handoffs). A top-level scene iterator that
 * re-generates the per-scene 3-angle (az, tilt, roll) poses of a CUAS sequence
 * by fitting each scene directly against the persistent VIRTUAL-CENTER reference
 * (center CLT), in ascending time order, seeded by prediction — the RT processing
 * pattern. Uses the proven offline measurement engine
 * (Interscene.adjustPairsLMAInterscene(), which assumes the reference GPU data is
 * set once for all scenes); Phase B will swap that engine for the lean
 * TD-average x virtual-center correlation with the GPU argmax+eigen kernels,
 * keeping this iterator and its comparison output unchanged as the oracle.
 *
 * Differences from the offline setInitialOrientationsCuas():
 *  - ascending scene order (RT: "now" keeps arriving) instead of descending;
 *  - pose seed = previous fitted pose (zero-order prediction; the fit is anchored
 *    to the virtual center, so prediction only warm-starts the LMA and never
 *    accumulates into the geometry) instead of IMS integration;
 *  - single pass with the FINAL combo DSI (no refinement pass needed - the offline
 *    refinement only existed because disparity arrived after initial orientations);
 *  - tile selection by simple strength threshold (curt_pose_str) of the combo DSI
 *    (equivalent to old getReliableTiles: >1.0 gives 1074 vs 1055 tiles on the
 *    reference sequence); the optimized minimal selection is a future slow-loop task.
 *
 * Output: per-scene CSV (predicted/fitted/stored poses + deltas + rms) in the model
 * directory and a console summary of fitted-vs-stored statistics. Stored poses
 * (scenes_poses of the center ErsCorrection, restored from INTERFRAME corr-xml)
 * are the comparison truth - the re-adjusted offline set fitted from the same
 * DSI vintage.
 *
 * By Claude on 07/03/2026, from Andrey's design.
 */
public class CuasPoseRT {
	public static final double [] ZERO3 = {0.0, 0.0, 0.0};
	// Per-scene measurement characterization slices of the -POSE-RT-HYPER stack
	// (z = scenes, t = components, matching the renderSceneSequence make_hyper convention).
	// dx,dy,strength from coord_motion[1] (vector_XYS); the eigen four from coord_motion[2]
	// {eig_x,eig_y,l0,l1} (filled when imp.eig_use is on; NaN otherwise): sqrt_l0/sqrt_l1 =
	// correlation-peak ellipse half-axes (pix), elong = sqrt(l1/l0) (large = linear feature,
	// e.g. straight horizon: precise only perpendicular), eig0_ang = direction of the
	// precise axis (eigenvector of the smaller eigenvalue), radians [0,PI).
	public static final String [] HYPER_COMPONENTS =
		{"dx", "dy", "strength", "dxy", "sqrt_l0", "sqrt_l1", "elong", "eig0_ang"};

	/**
	 * Re-generate per-scene poses against the virtual-center reference, RT-style.
	 *
	 * @param clt_parameters processing parameters (uses imp.curt_pose_str threshold,
	 *                       imp.lma_use_R, LMA/offset limits as the offline flow)
	 * @param center_CLT     virtual-center reference: must have CenterClt (fclt) and
	 *                       the final combo DSI restored
	 * @param quadCLTs       scene sequence (time-ascending indices), instances built
	 *                       with image data (e.g. by CuasRanging.prepareFpixels())
	 * @param debugLevel     debug level
	 * @return fitted poses [nscene][2][3] ({xyz,atr}, xyz locked to zero) or null
	 */
	public static double [][][] testPoseSequence(
			final CLTParameters clt_parameters,
			final QuadCLT       center_CLT,
			final QuadCLT []    quadCLTs,
			final int           debugLevel) {
		final double  min_str =     clt_parameters.imp.curt_pose_str;   // e.g. 1.0
		final boolean use_lma_dsi = clt_parameters.imp.use_lma_dsi;
		final int     margin =      clt_parameters.imp.margin;
		final double  mb_tau =      clt_parameters.imp.mb_tau;
		final double  mb_max_gain = clt_parameters.imp.mb_max_gain_inter;
		final double  min_offset =  clt_parameters.imp.min_offset;
		final double  max_offset =  clt_parameters.imp.max_rel_offset *
				center_CLT.getTilesX() * center_CLT.getTileSize();
		final double  max_rms =     clt_parameters.imp.eig_use ?
				clt_parameters.imp.eig_max_rms : clt_parameters.imp.max_rms;
		final double [] reg_weights = clt_parameters.ilp.ilma_regularization_weights;
		final double [] lma_rms = new double [clt_parameters.imp.eig_use ? 6 : 4];

		// 3-angle parameter selection: azimuth, tilt (+ optionally roll), XYZ locked
		// (rotation-only pose: displacements are good enough from the mechanical
		// construction; same as offline CUAS lock_position=true)
		final boolean [] param_select = ErsCorrection.getParamSelect(
				false,                          // use_Z (lock_position)
				clt_parameters.imp.lma_use_R,   // use_R
				false,                          // use_XY (lock_position)
				true,                           // use_AT
				false, false, false);           // ERS az/tilt/roll - off
		if (debugLevel > -3) {
			System.out.print("CuasPoseRT.testPoseSequence(): param_select=[");
			for (boolean b : param_select) System.out.print(b ? "+" : "-");
			System.out.println("], min_str="+min_str+", max_offset="+max_offset);
		}

		// ---- Reference prep (ONCE) - the persistent virtual-center TD reference ----
		final boolean is_aux = center_CLT.isAux();
		final double [][] center_dsi = center_CLT.dsi;
		final double [] disparity =     center_dsi[is_aux ? TwoQuadCLT.DSI_DISPARITY_AUX     : TwoQuadCLT.DSI_DISPARITY_MAIN];
		final double [] disparity_lma = center_dsi[is_aux ? TwoQuadCLT.DSI_DISPARITY_AUX_LMA : TwoQuadCLT.DSI_DISPARITY_MAIN_LMA];
		final double [] strength =      center_dsi[is_aux ? TwoQuadCLT.DSI_STRENGTH_AUX      : TwoQuadCLT.DSI_STRENGTH_MAIN];
		final double [] center_disparity = (use_lma_dsi && (disparity_lma != null)) ? disparity_lma : disparity;
		final boolean [] reliable_ref = new boolean [strength.length];
		int num_reliable = 0;
		for (int i = 0; i < strength.length; i++) {
			reliable_ref[i] = (strength[i] > min_str) && !Double.isNaN(center_disparity[i]);
			if (reliable_ref[i]) num_reliable++;
		}
		if (debugLevel > -3) {
			System.out.println("CuasPoseRT.testPoseSequence(): num_reliable="+num_reliable+
					" (strength > "+min_str+")");
		}
		final double [][] pXpYD_center = OpticalFlow.transformToScenePxPyD(
				null,             // final Rectangle [] extra_woi,
				center_disparity, // final double []    disparity_ref,
				ZERO3,            // final double []    scene_xyz,
				ZERO3,            // final double []    scene_atr,
				center_CLT,       // final QuadCLT      scene_QuadClt,
				center_CLT);      // final QuadCLT      reference_QuadClt
		final float [][] fclt = new float [][] {center_CLT.getCenterClt()};
		final TpTask [][] tp_tasks_center = Interscene.setReferenceGPU(
				clt_parameters,   // CLTParameters      clt_parameters,
				center_CLT,       // QuadCLT            ref_scene,
				center_disparity, // double []          ref_disparity,
				pXpYD_center,     // double [][]        ref_pXpYD,
				fclt,             // final float [][]   fclt,
				reliable_ref,     // final boolean []   selection,
				margin,           // final int          margin,
				mb_tau,           // double             mb_tau,
				mb_max_gain,      // double             mb_max_gain,
				null,             // double [][]        mb_vectors (center does not move)
				debugLevel);      // int                debug_level

		// ---- Stored poses (comparison truth + seed) ----
		final ErsCorrection ers_center = center_CLT.getErsCorrection();
		final double [][][] stored = new double [quadCLTs.length][][];
		for (int i = 0; i < quadCLTs.length; i++) if (quadCLTs[i] != null) {
			stored[i] = ers_center.getSceneXYZATR(quadCLTs[i].getImageName()); // may be null
		}
		int earliest = -1;
		for (int i = 0; i < quadCLTs.length; i++) if (quadCLTs[i] != null) { earliest = i; break; }
		if (earliest < 0) {
			System.out.println("CuasPoseRT.testPoseSequence(): no scenes");
			return null;
		}
		// Seed: earliest stored pose ("cheating" bootstrap); identity if unavailable.
		// Any scene works as the entry point - every fit is anchored to the virtual
		// center, not chained to neighbors.
		double [][] predicted = (stored[earliest] != null) ?
				new double [][] {stored[earliest][0].clone(), stored[earliest][1].clone()} :
				new double [][] {ZERO3.clone(), ZERO3.clone()};
		if (stored[earliest] == null) {
			System.out.println("CuasPoseRT.testPoseSequence(): WARNING - no stored pose for seed scene "+
					quadCLTs[earliest].getImageName()+", seeding with identity");
		}

		// Angle->pixel conversion, same scales the LMA uses (GeometryCorrection par_scales):
		// az/tilt ~ focal length, roll ~ distortion radius (edge-of-field displacement)
		final double px_per_rad_at = 1000.0 * center_CLT.getGeometryCorrection().focalLength /
				center_CLT.getGeometryCorrection().pixelSize;
		final double px_per_rad_r =  1000.0 * center_CLT.getGeometryCorrection().distortionRadius /
				center_CLT.getGeometryCorrection().pixelSize;
		final double [] px_per_rad = {px_per_rad_at, px_per_rad_at, px_per_rad_r};

		// ---- Ascending per-scene loop (the RT iterator) ----
		final int tilesX = center_CLT.getTilesX();
		final int tilesY = center_CLT.getTilesY();
		final int num_pix = tilesX * tilesY;
		final ArrayList<float [][]> hyper_scenes = new ArrayList<float [][]>(); // per used scene: [component][tile]
		final ArrayList<String>     hyper_ts =     new ArrayList<String>();
		if (!clt_parameters.imp.eig_use && (debugLevel > -4)) {
			System.out.println("CuasPoseRT.testPoseSequence(): imp.eig_use is OFF - the 4 eigen slices of "+
					"-POSE-RT-HYPER (sqrt_l0, sqrt_l1, elong, eig0_ang) will be all NaN");
		}
		final double [][][] fitted = new double [quadCLTs.length][][];
		final int [] fail_reason = new int [1];
		final StringBuffer sb = new StringBuffer();
		sb.append("index, timestamp, pred_az, pred_tilt, pred_roll, fit_az, fit_tilt, fit_roll,"+
				" stored_az, stored_tilt, stored_roll, dfit_az, dfit_tilt, dfit_roll, rms, weight, num_tiles, fail\n");
		int num_ok = 0, num_fail = 0;
		double [] sum2_d = new double [3];
		double [] max_d =  new double [3];
		int num_cmp = 0;
		double prev_ts = Double.NaN;
		double [][] prev_pose = null;
		for (int nscene = earliest; nscene < quadCLTs.length; nscene++) {
			if (quadCLTs[nscene] == null) continue;
			final String ts_name = quadCLTs[nscene].getImageName();
			final double ts = Double.parseDouble(ts_name.replace("_", "."));
			// ERS velocities from finite differences of the pose sequence (substitute
			// for the offline PIMU rates; only the residual/wobble rates ~0.05 rad/s -
			// ERS params are not fitted, disable_ers=true below)
			double [][] dxyzatr_dt = new double [][] {ZERO3.clone(), ZERO3.clone()};
			if ((prev_pose != null) && !Double.isNaN(prev_ts) && (ts > prev_ts)) {
				for (int j = 0; j < 3; j++) {
					dxyzatr_dt[1][j] = (predicted[1][j] - prev_pose[1][j]) / (ts - prev_ts);
				}
			}
			quadCLTs[nscene].getErsCorrection().setErsDt(
					QuadCLTCPU.scaleDtToErs(clt_parameters, dxyzatr_dt));
			fail_reason[0] = 0;
			final double [][][] coord_motion_rslt = new double [3][][]; // [centers, vector_XYS, eigen] of the last LMA cycle
			double [][] pose = Interscene.adjustPairsLMAInterscene(
					clt_parameters,      // CLTParameters  clt_parameters,
					true,                // boolean        initial_adjust,
					false,               // boolean        fpn_disable,
					true,                // boolean        disable_ers,
					new double [] {min_offset, max_offset, 0.0}, // double [] min_max,
					fail_reason,         // int []         fail_reason,
					center_CLT,          // QuadCLT        ref_QuadClt (virtual center)
					center_disparity,    // double []      ref_disparity,
					center_CLT,          // QuadCLT        first_QuadClt (virtual center)
					pXpYD_center,        // double [][]    pXpYD_ref,
					clt_parameters.imp.cuas_min_confidence, // double min_confidence (as re-adjust: final-DSI vintage)
					reliable_ref,        // boolean []     reliable_ref,
					null,                // boolean []     reliable_all,
					tp_tasks_center[0],  // TpTask[]       tp_tasks_ref,
					quadCLTs[nscene],    // QuadCLT        scene_QuadClt,
					predicted,           // double [][]    camera_xyzatr (initial),
					predicted,           // double [][]    scene_xyzatr_pull,
					new double [2][3],   // double [][]    ref_xyzatr,
					param_select,        // boolean[]      param_select,
					reg_weights,         // double []      param_regweights,
					lma_rms,             // double []      rms_out,
					max_rms,             // double         max_rms,
					mb_tau,              // double         mb_tau,
					mb_max_gain,         // double         mb_max_gain,
					null,                // double [][]    mb_vectors (MB off in phase A),
					coord_motion_rslt,   // double [][][]  coord_motion_rslt (last-cycle measurement, for -POSE-RT-HYPER),
					debugLevel - 2);     // int            debug_level
			final boolean ok = (pose != null);
			if (ok) {
				fitted[nscene] = pose;
				num_ok++;
			} else {
				// coast: keep the prediction as this scene's pose, continue the sequence
				fitted[nscene] = new double [][] {predicted[0].clone(), predicted[1].clone()};
				num_fail++;
				if (debugLevel > -4) {
					System.out.println("CuasPoseRT: scene "+nscene+" ("+ts_name+") FAILED, reason="+
							fail_reason[0]+" ("+Interscene.getFailReason(fail_reason[0])+"), coasting prediction");
				}
			}
			// Collect measurement characterization slices for -POSE-RT-HYPER
			{
				final float [][] comps = new float [HYPER_COMPONENTS.length][num_pix];
				for (float [] c : comps) Arrays.fill(c, Float.NaN);
				final double [][] vxys = coord_motion_rslt[1]; // {dx,dy,strength} per tile or null
				final double [][] eig =  (coord_motion_rslt.length > 2) ? coord_motion_rslt[2] : null;
				if (vxys != null) {
					for (int i = 0; i < num_pix; i++) {
						if (vxys[i] != null) {
							comps[0][i] = (float) vxys[i][0];
							comps[1][i] = (float) vxys[i][1];
							comps[2][i] = (float) vxys[i][2];
							comps[3][i] = (float) Math.sqrt(vxys[i][0]*vxys[i][0] + vxys[i][1]*vxys[i][1]);
						}
						if ((eig != null) && (eig[i] != null)) { // {eig_x, eig_y, l0, l1}, l0 <= l1
							final double l0 = eig[i][2], l1 = eig[i][3];
							comps[4][i] = (float) Math.sqrt(l0);
							comps[5][i] = (float) Math.sqrt(l1);
							comps[6][i] = (l0 > 0) ? (float) Math.sqrt(l1/l0) : Float.NaN;
							double ang = Math.atan2(eig[i][1], eig[i][0]); // precise-axis direction
							if (ang < 0) ang += Math.PI;                   // fold to [0, PI)
							comps[7][i] = (float) ang;
						}
					}
				}
				hyper_scenes.add(comps);
				hyper_ts.add(ts_name);
			}
			// CSV row + comparison stats
			final double [] st = (stored[nscene] != null) ? stored[nscene][1] : null;
			sb.append(nscene+", "+ts_name);
			for (int j = 0; j < 3; j++) sb.append(String.format(", %.6f", predicted[1][j]));
			for (int j = 0; j < 3; j++) sb.append(String.format(", %.6f", fitted[nscene][1][j]));
			for (int j = 0; j < 3; j++) sb.append((st != null) ? String.format(", %.6f", st[j]) : ", NaN");
			for (int j = 0; j < 3; j++) {
				double d = (st != null) ? (fitted[nscene][1][j] - st[j]) : Double.NaN;
				sb.append((st != null) ? String.format(", %.6f", d) : ", NaN");
				if (ok && (st != null)) {
					sum2_d[j] += d * d;
					if (Math.abs(d) > max_d[j]) max_d[j] = Math.abs(d);
				}
			}
			if (ok && (stored[nscene] != null)) num_cmp++;
			sb.append(String.format(", %.6f, %.1f, %.0f, %d\n",
					ok ? lma_rms[0] : Double.NaN, lma_rms[2], lma_rms[3], fail_reason[0]));
			// Scene identification line - attributes the preceding (unlabeled) LMA iteration
			// prints from adjustPairsLMAInterscene to this scene (index of N + timestamp)
			if (debugLevel > -4) {
				StringBuffer sl = new StringBuffer();
				sl.append("CuasPoseRT scene "+nscene+" (of "+quadCLTs.length+") "+ts_name+
						(ok ? " Done." : " FAILED (coasted).")+
						String.format(" RMS=%.4f, weight=%.1f, num=%.0f", lma_rms[0], lma_rms[2], lma_rms[3]));
				if (st != null) {
					sl.append(" dstored[pix]=[");
					for (int j = 0; j < 3; j++) {
						sl.append(String.format((j > 0 ? ", %7.4f" : "%7.4f"),
								(fitted[nscene][1][j] - st[j]) * px_per_rad[j]));
					}
					sl.append("]");
				}
				System.out.println(sl.toString());
			}
			// zero-order prediction for the next scene
			prev_pose = predicted;
			prev_ts = ts;
			predicted = new double [][] {fitted[nscene][0].clone(), fitted[nscene][1].clone()};
		}

		// ---- Output ----
		final String csv_path = center_CLT.getX3dDirectory(true) + Prefs.getFileSeparator() +
				center_CLT.getImageName() + "-POSE-RT-TEST.csv";
		try (FileWriter fw = new FileWriter(csv_path)) {
			fw.write(sb.toString());
			System.out.println("CuasPoseRT.testPoseSequence(): saved "+csv_path);
		} catch (IOException e) {
			System.out.println("CuasPoseRT.testPoseSequence(): FAILED to save "+csv_path+": "+e.getMessage());
		}
		// reliable_tiles selection mask (1=selected), same grid as the DSI
		{
			final float [] frel = new float [num_pix];
			for (int i = 0; i < num_pix; i++) frel[i] = reliable_ref[i] ? 1.0f : 0.0f;
			final ImagePlus imp_rel = new ImagePlus(center_CLT.getImageName()+"-POSE-RT-RELIABLE",
					new FloatProcessor(tilesX, tilesY, frel));
			center_CLT.saveImagePlusInModelDirectory(null, imp_rel); // title as filename
		}
		// -POSE-RT-HYPER: z = scenes (top slider), t = components (bottom slider) -
		// same layout as renderSceneSequence make_hyper (component is the outer slice order)
		if (!hyper_scenes.isEmpty()) {
			final ImageStack stack_hyper = new ImageStack(tilesX, tilesY);
			for (int nc = 0; nc < HYPER_COMPONENTS.length; nc++) {
				for (int ns = 0; ns < hyper_scenes.size(); ns++) {
					stack_hyper.addSlice(HYPER_COMPONENTS[nc]+":"+hyper_ts.get(ns), hyper_scenes.get(ns)[nc]);
				}
			}
			final ImagePlus imp_hyper = new ImagePlus(center_CLT.getImageName()+"-POSE-RT-HYPER", stack_hyper);
			imp_hyper.setDimensions(1, hyper_scenes.size(), HYPER_COMPONENTS.length);
			imp_hyper.setOpenAsHyperStack(true);
			imp_hyper.getProcessor().resetMinAndMax();
			center_CLT.saveImagePlusInModelDirectory(null, imp_hyper); // title as filename
		}
		System.out.println("CuasPoseRT.testPoseSequence(): scenes OK="+num_ok+", failed(coasted)="+num_fail+
				", compared="+num_cmp);
		if (num_cmp > 0) {
			final String [] names = {"azimuth", "tilt", "roll"};
			for (int j = 0; j < 3; j++) {
				final double rms_rad = Math.sqrt(sum2_d[j]/num_cmp);
				System.out.println(String.format(
						"  fitted-stored %8s: RMS=%7.4f pix, max=%7.4f pix   (%8.4f mrad, max %8.4f mrad)",
						names[j], rms_rad*px_per_rad[j], max_d[j]*px_per_rad[j],
						1000*rms_rad, 1000*max_d[j]));
			}
			System.out.println(String.format(
					"  (pixel scales: az/tilt %.1f pix/rad = focal/pixelSize, roll %.1f pix/rad = distortionRadius/pixelSize)",
					px_per_rad_at, px_per_rad_r));
		}
		return fitted;
	}
}
