package com.elphel.imagej.cuas;

import java.util.concurrent.atomic.AtomicInteger;

import com.elphel.imagej.tileprocessor.ImageDtt;
import com.elphel.imagej.tileprocessor.QuadCLT;

import Jama.Matrix;

/*
 * Use fixed-position relative to initial max gaussian window weight
 * fx = A*cos2(r/R0*pi/2)*cos(K*r/R0*pi/2)+C
 * K<1, K~=2
 * A,K,R0,C
 * RR0= pi/(R0*2)
 * r=sqrt((x-X0)^2+(y-Y0)^2)
 * 
 * fx = A*cos2(r*RR0)*cos(K*RR0*r)+C
 * 
 *  
 *   
 * 
 */



public class CuasMotionLMA {
	private static final double R_OFFS = 1e-6;
	public static final int INDX_A =   0;
	public static final int INDX_C =   1;
	public static final int INDX_RR0 = 2;
	public static final int INDX_K =   3;
	public static final int INDX_X0 =  4; // relative to left
	public static final int INDX_Y0 =  5; // relative to top
	public static final int INDX_LEN = INDX_Y0+1;
	
	// result vector indices
	public static final int RSLT_X =     0;
	public static final int RSLT_Y =     1;
	public static final int RSLT_A =     2;
	public static final int RSLT_R0 =    3;
	public static final int RSLT_K =     4;
	public static final int RSLT_C =     5;
	public static final int RSLT_RMS =   6;
	public static final int RSLT_ITERS = 7;
	public static final int RSLT_LEN = RSLT_ITERS+1;
	
	public static final String [] LMA_TITLES = {"X-OFFS","Y-OFFS", "AMPLITUDE","RADIUS","OVERSHOOT","OFFSET","RMSE","ITERATIONS"};
	
	
	private int          width;
	private double [][]  window;
	private double []    y_vector;
	private double []    weights;
	private double []    full_vector = new double[INDX_LEN];
	private int []       pindx; // full parameter index for vector
	private int []       rindx; // vector index or -1
	private double []    last_rms =        null; // {rms, rms_pure}, matching this.vector
	private double []    good_or_bad_rms = null; // just for diagnostics, to read last (failed) rms
	private double []    initial_rms =     null; // {rms, rms_pure}, first-calcualted rms
	private double []    last_ymfx =       null;	
	private double [][]  last_jt =         null;
	private int          iters =           -2; // never ran

	public CuasMotionLMA(
			int         width,
			double     sigma) {
		this.width = width;
		window = new double [width][width];
		double k = -0.5/(sigma*sigma);
		for (int i = 0; i < width; i++) {
			for (int j = 0; j < width; j++) {
				window[i][j] = Math.exp(k*(i*i+j*j));
			}
		}
	}
	
	
	public int prepareLMA(
			boolean [] param_select,
			double []  tile_data,
			double     xc, // relative to center =width/2
			double     yc, // relative to center =width/2
			double     r0,
			double     k,
			int        debugLevel) {
		y_vector = tile_data;
		double x0 = Math.min(Math.max(xc + width/2, 0), width-1);
		double y0 = Math.min(Math.max(yc + width/2, 0), width-1);

		int ix0 = (int) Math.round(x0);
		int iy0 = (int) Math.round(y0);
		full_vector[INDX_A] = tile_data[ix0+iy0*width];
		full_vector[INDX_C] = 0;
		full_vector[INDX_RR0] = Math.PI/(2* r0);
		full_vector[INDX_K] = k; // > 1 (~2.0)
		full_vector[INDX_X0] = x0;
		full_vector[INDX_Y0] = y0;
		weights = new double [width*width];
		double sw = 0;
		for (int y = 0; y < width; y++) {
			int ay = Math.abs(y-iy0);
			for (int x = 0; x < width; x++) {
				int ax = Math.abs(x-ix0);
				double w = window[ay][ax]; // window to the nearest integer x,y
				weights[x + y*width] = w;
				sw += w;
			}
		}
		for (int i = 0; i < weights.length; i++) {
			weights[i] /= sw;
		}
		int indx = 0;
		for (int i = 0; i < INDX_LEN; i++) if (param_select[i]){
			indx++;
		}
		rindx = new int [INDX_LEN];
		pindx = new int [indx];
		indx= 0;
		for (int i = 0; i < INDX_LEN; i++) {
			if (param_select[i]) {
				pindx[indx] = i;
				rindx[i] = indx++;
			} else {
				rindx[i] = -1;
			}
		}
   		last_jt = new double [pindx.length][];
		
		double [] fx = getFxDerivs(
				getParametersVector(), // double []         vector,
				last_jt,               // final double [][] jt, // should be null or initialized with [vector.length][]
				debugLevel);           // final int         debug_level)
		last_rms =        new double [2];
		last_ymfx = getYminusFxWeighted(
				fx, // final double []   fx,
				last_rms); // final double []   rms_fp // null or [2]
		
		initial_rms =     last_rms.clone();
		good_or_bad_rms = last_rms.clone();
		return 0;
	}	
	
	public double getRMS() {
		return last_rms[0];
	}
	
	public double getInitialRMS() {
		return initial_rms[0];
	}
	
	public double [] getResult() {
		double rslt[] = new double [RSLT_LEN];
		rslt[RSLT_X] =     getCenter()[0];
		rslt[RSLT_Y] =     getCenter()[1];
		rslt[RSLT_A] =     getA();
		rslt[RSLT_R0] =    getR0();
		rslt[RSLT_K] =     getK();
		rslt[RSLT_C] =     getC();
		rslt[RSLT_RMS] =   getRMS();
		rslt[RSLT_ITERS] = getIters();
		return rslt;
	}
	
	public double [] getCenter(){
		return new double [] {full_vector[INDX_X0] - width/2, full_vector[INDX_Y0]-width/2};
	}
	
	public double getRR0() {
		return full_vector[INDX_RR0];
	}

	public double getR0() {
		return Math.PI/(2 * full_vector[INDX_RR0]);
	}

	public double getA() {
		return full_vector[INDX_A];
	}

	public double getC() {
		return full_vector[INDX_C];
	}
	
	public double getK() {
		return full_vector[INDX_K];
	}
	
	public int getIters() {
		return iters;
	}
	
	public double [] getFullParametersVector() {
		return full_vector;
	}
	
	public double [] getParametersVector() {
		double [] vector = new double [pindx.length];
		for (int i = 0; i < vector.length; i++) {
			vector[i] = full_vector[pindx[i]];
		}
		return vector;
	}
	public void setParametersVector(double [] vector) {
		for (int i = 0; i < vector.length; i++) {
			full_vector[pindx[i]] = vector[i];
		}
	}
	
	public int runLma( // <0 - failed, >=0 iteration number (1 - immediately)
			double lambda,           // 0.1
			double lambda_scale_good,// 0.5
			double lambda_scale_bad, // 8.0
			double lambda_max,       // 100
			double rms_diff,         // 0.001
			int    num_iter,         // 20
			int    debug_level)
	{
		boolean last_run = true;
		boolean [] rslt = {false,false};
		this.last_rms = null; // remove?
		int iter = 0;
		for (iter = 0; iter < num_iter; iter++) {
			rslt =  lmaStep(
					lambda,
					rms_diff,
					debug_level);
			if (rslt == null) {
				iters = -1;
				return iters;
			}
			if (debug_level > 1) {
				System.out.println("LMA step "+iter+": {"+rslt[0]+","+rslt[1]+"} full RMS= "+good_or_bad_rms[0]+
						" ("+initial_rms[0]+"), pure RMS="+good_or_bad_rms[1]+" ("+initial_rms[1]+") + lambda="+lambda);
			}
			if (rslt[1]) {
				break;
			}
			if (rslt[0]) { // good
				lambda *= lambda_scale_good;
			} else {
				lambda *= lambda_scale_bad;
				if (lambda > lambda_max) {
					break; // not used in lwir
				}
			}
//			if (dbg_prefix != null) {
//				showDebugImage(dbg_prefix+"-"+iter+(rslt[0]?"-GOOD":"-BAD"));
//			}

		}
		if (rslt[0]) { // better
			if (iter >= num_iter) { // better, but num tries exceeded
				if (debug_level > 1) System.out.println("Step "+iter+": Improved, but number of steps exceeded maximal");
			} else {
				if (debug_level > 1) System.out.println("Step "+iter+": LMA: Success");
			}

		} else { // improved over initial ?
			if (last_rms[0] < initial_rms[0]) { // NaN
				rslt[0] = true;
				if (debug_level > 1) System.out.println("Step "+iter+": Failed to converge, but result improved over initial");
			} else {
				if (debug_level > 1) System.out.println("Step "+iter+": Failed to converge");
			}
		}
//		if (dbg_prefix != null) {
//			showDebugImage(dbg_prefix+"-FINAL");
//		}
		boolean show_intermediate = true;
		if (show_intermediate && (debug_level > 0)) {
			System.out.println("LMA: full RMS="+last_rms[0]+" ("+initial_rms[0]+"), pure RMS="+last_rms[1]+" ("+initial_rms[1]+") + lambda="+lambda);
		}
		if (debug_level > 2){ 
//			String [] lines1 = printOldNew(false); // boolean allvectors)
//			System.out.println("iteration="+iter);
//			for (String line : lines1) {
//				System.out.println(line);
//			}
		}
		if (debug_level > 0) {
			if ((debug_level > 1) ||  last_run) { // (iter == 1) || last_run) {
				if (!show_intermediate) {
					System.out.println("LMA: iter="+iter+",   full RMS="+last_rms[0]+" ("+initial_rms[0]+"), pure RMS="+last_rms[1]+" ("+initial_rms[1]+") + lambda="+lambda);
				}
//				String [] lines = printOldNew(false); // boolean allvectors)
//				for (String line : lines) {
//					System.out.println(line);
//				}
			}
		}
		if ((debug_level > -2) && !rslt[0]) { // failed
			if ((debug_level > 1) || (iter == 1) || last_run) {
				System.out.println("LMA failed on iteration = "+iter);
//				String [] lines = printOldNew(true); // boolean allvectors)
//				for (String line : lines) {
//					System.out.println(line);
//				}
			}
			System.out.println();
		}
		// No need to updateFullParameters() in this implementation - they are already updated
		/*
		if (rslt[0]) {
			updateFullParameters();
		}
		*/
		iters = rslt[0]? iter : -1;
		return iters;
	}

	
	
	private boolean [] lmaStep(
			double lambda,
			double rms_diff,
			int debug_level) {
		boolean [] rslt = {false,false};
		double [] parameters_vector = getParametersVector();
		// maybe the following if() branch is not needed - already done in prepareLMA !
		if (this.last_rms == null) { //first time, need to calculate all (vector is valid)
			last_rms = new double[2];
			if (debug_level > 1) {
				System.out.println("lmaStep(): first step");
			}
			double [] fx = getFxDerivs(
					parameters_vector, // parameters_vector, // double []         vector,
					last_jt,           // final double [][] jt, // should be null or initialized with [vector.length][]
					debug_level);      // final int         debug_level)
			last_ymfx = getYminusFxWeighted(
					fx, // final double []   fx,
					last_rms); // final double []   rms_fp // null or [2]
			this.initial_rms = this.last_rms.clone();
			this.good_or_bad_rms = this.last_rms.clone();
			if (last_ymfx == null) {
				return null; // need to re-init/restart LMA
			}
			// TODO: Restore/implement
			if (debug_level > 3) {
				double    delta = 1E-5;
			 	double delta_err=compareJT(
			 			parameters_vector, // double [] vector,
						delta,             // double    delta,
						false); // last3only);        // boolean   last3only); // do not process samples - they are tested before
				System.out.println("\nMaximal error = "+delta_err);
				
				/*
				 dbgJacobians(
							corr_vector, // GeometryCorrection.CorrVector corr_vector,
							1E-5, // double delta,
							true); //boolean graphic)
				*/
			}
		}
		
		if (debug_level > 3) {
			double    delta = 1E-5;
		 	double delta_err=compareJT(
		 			parameters_vector, // double [] vector,
					delta,             // double    delta,
					false); // last3only);        // boolean   last3only); // do not process samples - they are tested before
			System.out.println("\nMaximal error = "+delta_err);
			
			/*
			 dbgJacobians(
						corr_vector, // GeometryCorrection.CorrVector corr_vector,
						1E-5, // double delta,
						true); //boolean graphic)
			*/
		}
		
		Matrix y_minus_fx_weighted = new Matrix(this.last_ymfx, this.last_ymfx.length);

		Matrix wjtjlambda = new Matrix(getWJtJlambda(
				lambda, // *10, // temporary
				this.last_jt)); // double [][] jt)
		
		if (debug_level>2) {
			System.out.println("JtJ + lambda*diag(JtJ");
			wjtjlambda.print(18, 6);
		}
		Matrix jtjl_inv = null;
		try {
			jtjl_inv = wjtjlambda.inverse(); // check for errors
		} catch (RuntimeException e) {
			rslt[1] = true;
			if (debug_level > 0) {
				System.out.println("Singular Matrix!");
			}
			return rslt;
		}
		if (debug_level>2) {
			System.out.println("(JtJ + lambda*diag(JtJ).inv()");
			jtjl_inv.print(18, 6);
		}
//last_jt has NaNs
		Matrix jty = (new Matrix(this.last_jt)).times(y_minus_fx_weighted);
		if (debug_level>2) {
			System.out.println("Jt * (y-fx)");
			jty.print(18, 6);
		}
		
		Matrix mdelta = jtjl_inv.times(jty);
		if (debug_level>2) {
			System.out.println("mdelta");
			mdelta.print(18, 10);
		}

		double scale = 1.0;
		double []  delta =      mdelta.getColumnPackedCopy();
		double []  new_vector = parameters_vector.clone();
		for (int i = 0; i < parameters_vector.length; i++) {
			new_vector[i] += scale * delta[i];
		}
		
		double [] fx = getFxDerivs(
				new_vector, // double []         vector,
				last_jt,           // final double [][] jt, // should be null or initialized with [vector.length][]
				debug_level);      // final int         debug_level)
		double [] rms = new double[2];
		last_ymfx = getYminusFxWeighted(
				fx, // final double []   fx,
				rms); // final double []   rms_fp // null or [2]
		if (debug_level > 2) {
			/*
			dbgYminusFx(this.last_ymfx, "next y-fX");
			dbgXY(new_vector, "XY-correction");
			*/
		}

		if (last_ymfx == null) {
			return null; // need to re-init/restart LMA
		}

		this.good_or_bad_rms = rms.clone();
		if (rms[0] < this.last_rms[0]) { // improved
			rslt[0] = true;
			rslt[1] = rms[0] >=(this.last_rms[0] * (1.0 - rms_diff));
			this.last_rms = rms.clone();
			parameters_vector = new_vector.clone();
			setParametersVector(new_vector);
			if (debug_level > 2) {
				// print vectors in some format
				/*
				System.out.print("delta: "+corr_delta.toString()+"\n");
				System.out.print("New vector: "+new_vector.toString()+"\n");
				System.out.println();
				*/
			}
		} else { // worsened
			rslt[0] = false;
			rslt[1] = false; // do not know, caller will decide
			// restore state
			fx = getFxDerivs(
					parameters_vector, // double []         vector,
					last_jt,           // final double [][] jt, // should be null or initialized with [vector.length][]
					debug_level);      // final int         debug_level)
			last_ymfx = getYminusFxWeighted(
					fx, // final double []   fx,
					this.last_rms); // final double []   rms_fp // null or [2]
			if (last_ymfx == null) {
				return null; // need to re-init/restart LMA
			}
			if (debug_level > 2) {
				/*
				 dbgJacobians(
							corr_vector, // GeometryCorrection.CorrVector corr_vector,
							1E-5, // double delta,
							true); //boolean graphic)
							*/
			}
		}
		return rslt;
	}

	
	private double [][] getWJtJlambda( // USED in lwir
			final double      lambda,
			final double [][] jt)
	{
		final int num_pars = jt.length;
		final int nup_points = jt[0].length;
		final double [][] wjtjl = new double [num_pars][num_pars];
		for (int i = 0; i < num_pars; i++) {
			for (int j = i; j < num_pars; j++) {
				double d = 0.0;
				for (int k = 0; k < nup_points; k++) {
//					if (jt[i][k] != 0) {
//						d+=0; // ???
//					}
					d += weights[k]*jt[i][k]*jt[j][k];
					if (Double.isNaN(d)) {
						System.out.println("CuasMotionLMA.getWJtJlambda(): NAN i="+i+", j="+j+", k="+k);
					}
				}
				wjtjl[i][j] = d;
				if (i == j) {
					wjtjl[i][j] += d * lambda;
				} else {
					wjtjl[j][i] = d;
				}
			}
		}
		return wjtjl;
	}
	
	
	
	private double [] getYminusFxWeighted( // problems. at least with eigen?
			final double []   fx,
			final double []   rms_fp) { // null or [2]
		final double []     wymfw =       new double [fx.length];
		double s_rms = 0.0;
		for (int i = 0; i < fx.length; i++) {
			double d = y_vector[i] - fx[i];
			double wd = d * weights[i];
			wymfw[i] = wd;
			s_rms += d * wd;
		}
		double rms = Math.sqrt(s_rms); // assuming sum_weights == 1.0; /pure_weight); they should be re-normalized after adding regularization
		if (rms_fp != null) {
			rms_fp[0] = rms;
			if (rms_fp.length > 1) {
				rms_fp[1] = rms; // _pure;
			}
		}
		return wymfw;
	}
	
 	private double compareJT(
			double [] vector,
			double    delta,
			boolean   last3only) { // do not process samples - they are tested before
		double []  errors=new double [vector.length];
		double [][] jt =  new double [vector.length][];
		System.out.print("Parameters vector = [");
		for (int i = 0; i < vector.length; i++) {
			System.out.print(vector[i]);
			if (i < (vector.length -1)) System.out.print(", ");
		}
		System.out.println("]");
		getFxDerivs(
				vector,
				jt, // final double [][] jt, // should be null or initialized with [vector.length][]
				1); // debug_level);
		double [][] jt_delta =  getFxDerivsDelta(
				vector, // double []         vector,
				delta, // final double      delta,
				-1); // final int         debug_level)
		int start_index = last3only? (weights.length-3) : 0;
		for (int n = start_index; n < weights.length; n++) if (weights[n] > 0) {
			System.out.print(String.format("%3d",n));
			for (int i = 0; i < vector.length; i++) {
				System.out.print(String.format("\t%12.9f",jt[i][n]));
			}			
			for (int i = 0; i < vector.length; i++) {
				System.out.print(String.format("\t%12.9f",jt_delta[i][n]));
			}			
			for (int i = 0; i < vector.length; i++) {
				System.out.print(String.format("\t%12.9f",jt[i][n]-jt_delta[i][n]));
			}			
			System.out.println();
			/*
			System.out.println(String.format(
					"%3d\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f\t%12.9f",
					n, jt[0][n], jt[1][n], jt[2][n], jt[3][n],
					jt_delta[0][n], jt_delta[1][n], jt_delta[2][n], jt_delta[3][n],
					jt[0][n]-jt_delta[0][n],jt[1][n]-jt_delta[1][n],jt[2][n]-jt_delta[2][n],jt[3][n]-jt_delta[3][n]));
					*/
			for (int i = 0; i < vector.length; i++) {
				errors[i] = Math.max(errors[i], jt[i][n]-jt_delta[i][n]);
			}
		}
		for (int i = 0; i < vector.length; i++) {
			System.out.print("\t\t");
		}			
		for (int i = 0; i < vector.length; i++) {
			System.out.print(String.format("\t%12.9f",errors[i]));
		}			
        /*
		System.out.println(String.format(
				"-\t-\t-\t-\t-\t-\t-\t-\t-\t%12.9f\t%12.9f\t%12.9f\t%12.9f",
				errors[0], errors[1], errors[2], errors[3]));
				*/
		double err=0;
		for (int i = 0; i < vector.length; i++) {
			err = Math.max(errors[i], err);
		}
		return err;
	}
	
	private double [][] getFxDerivsDelta(
			double []         vector,
			final double      delta,
			final int         debug_level) {
		double [][] jt =  new double [vector.length][weights.length];
		for (int nv = 0; nv < vector.length; nv++) {
			double [] vpm = vector.clone();
			vpm[nv]+= 0.5*delta;
			double [] fx_p =  getFxDerivs(
					vpm,
					null, // final double [][] jt, // should be null or initialized with [vector.length][]
					debug_level);
			vpm[nv]-= delta;
			double [] fx_m =  getFxDerivs(
					vpm,
					null, // final double [][] jt, // should be null or initialized with [vector.length][]
					debug_level);
			for (int i = 0; i < weights.length; i++) if (weights[i] > 0) {
				jt[nv][i] = (fx_p[i]-fx_m[i])/delta;
			}
		}
		return jt;
	}

	
	
	private double [] getFxDerivs(
			final double []   vector,
			final double [][] jt, // should be null or initialized with [vector.length][]
			final int         debug_level) {
		int height = y_vector.length/width;
		final double [] fx = new double [y_vector.length];
		if (jt != null) {
			for (int i = 0; i < jt.length; i++) {
				jt[i] = new double [y_vector.length]; // weights.length];
			}
		}
		double A=    (rindx[INDX_A] >=0)?   vector[rindx[INDX_A]] :   full_vector[INDX_A] ; 
		double C=    (rindx[INDX_C] >=0)?   vector[rindx[INDX_C]] :   full_vector[INDX_C] ;
		double RR0 = (rindx[INDX_RR0] >=0)? vector[rindx[INDX_RR0]] : full_vector[INDX_RR0] ;
		double K =   (rindx[INDX_K] >=0)?   vector[rindx[INDX_K]] :   full_vector[INDX_K] ; // >1.0
		double X0 =  (rindx[INDX_X0] >=0)?  vector[rindx[INDX_X0]] :  full_vector[INDX_X0] ;
		double Y0 =  (rindx[INDX_Y0] >=0)?  vector[rindx[INDX_Y0]] :  full_vector[INDX_Y0] ;
		for (int y = 0; y < height; y++) {
			double dy = y-Y0;
			double dy2 = dy*dy;
			for (int x = 0; x < width; x++) {
				int indx = y*width+x;
				double dx = x-X0;
				double dx2 = dx*dx;
				double r2 = dx2+dy2+R_OFFS;
				double r = Math.sqrt(r2);
				if (r*RR0 > Math.PI/2) { // outside circle where function is defined
					fx[indx] = C;
					if (jt != null) {
						double df_dC = 1;
						if (rindx[INDX_C] >=0) {
							jt[rindx[INDX_C]][indx] = df_dC;
						}
					}					
				} else {
					double W1 = Math.cos(r*RR0);
					double W2 = Math.cos(r*RR0*K);
					fx[indx] = A * (W1 * W1 * W2) + C;
					if (jt != null) {
						double dr2_dX0 = -2 * dx;
						double dr2_dY0 = -2 * dy;
						double dr_dr2 = 0.5 / r; // r !=0
						double sin_rRR0 = Math.sin(r * RR0);
						double dW1_dRR0 = -r * sin_rRR0;     // 
						double dW1_dr = -RR0 *sin_rRR0;
						double sin_rKRR0 = Math.sin(r * RR0 * K);
						double dW2_dRR0 =  -K * r * sin_rKRR0;
						double dW2_dK =    -r * RR0 * sin_rKRR0;
						double dW2_dr =    -K * RR0 * sin_rKRR0; 
						double df_dA = W1*W1*W2; // +
						double df_dC = 1;        // +  
						double df_dW1 = 2 * A* W1 * W2;
						double df_dW2 = A * W1 * W1;
						double df_dK =  df_dW2 * dW2_dK; 
						double df_dRR0 = df_dW1*dW1_dRR0+df_dW2*dW2_dRR0;
						double df_dr = df_dW1 * dW1_dr + df_dW2 * dW2_dr;  
						double df_dr2 = df_dr * dr_dr2;
						double df_dX0 = df_dr2 * dr2_dX0;
						double df_dY0 = df_dr2 * dr2_dY0;
						double [] df_dfp = {df_dA,df_dC,df_dRR0, df_dK, df_dX0, df_dY0};
						for (int i = 0; i < vector.length; i++) {
							jt[i][indx] = df_dfp[pindx[i]]; 
						}
					}
				}
			}
		}
		return fx;
	}

	
	

}
