package com.elphel.imagej.cuas;

import java.util.concurrent.atomic.AtomicInteger;

import com.elphel.imagej.common.ShowDoubleFloatArrays;
import com.elphel.imagej.tileprocessor.ImageDtt;
import com.elphel.imagej.tileprocessor.OpticalFlow;
import com.elphel.imagej.tileprocessor.QuadCLT;
import com.elphel.imagej.tileprocessor.QuadCLTCPU;

import ij.ImagePlus;

public class CorrectionFPN {
		final QuadCLT quadCLT;
		public CorrectionFPN(QuadCLTCPU quadCLT) {
			this.quadCLT = (QuadCLT) quadCLT;
		}
	
		  public void setApplyFPN(
				  QuadCLT []    quadCLTs,
				  double [][][] fpn) { // can be null to reset
			  quadCLT.setImageData(fpn);
			  quadCLT.setFPN(fpn);
			  for (int nscene = 0; nscene < quadCLTs.length; nscene++) if (quadCLTs[nscene] != null){
				  quadCLTs[nscene].setFPN(fpn);
				  quadCLTs[nscene].applyFPN();
			  }
			  return;
		  }
		  
		  
		  
		  public void applyFPN(
				  double [][][] fpn) {
			  if (quadCLT.getImageFpnApplied() != fpn) {
				  if (quadCLT.getImageData() == null) {
					  System.out.println("applyFPN (): image_data==null");
					  return;
				  }
				  final double [][][] image_data = quadCLT.getImageData();
				  final double [][][] image_fpn_applied = quadCLT.getImageFpnApplied();
				  final Thread[] threads = ImageDtt.newThreadArray(ImageDtt.THREADS_MAX);
				  final AtomicInteger ai = new AtomicInteger(0);
				  for (int ithread = 0; ithread < threads.length; ithread++) {
					  threads[ithread] = new Thread() {
						  public void run() {
							  for (int nChn = ai.getAndIncrement(); nChn < image_data.length; nChn = ai.getAndIncrement()) {
								  for (int ncol = 0; ncol < image_data[nChn].length; ncol++) {
									  for (int npix = 0; npix < image_data[nChn][ncol].length; npix++) {
										  if (image_fpn_applied != null) {
											  image_data[nChn][ncol][npix] += image_fpn_applied[nChn][ncol][npix];
										  }
										  if (fpn != null) {
											  image_data[nChn][ncol][npix] -= fpn[nChn][ncol][npix]; 
										  }
									  }
								  }
							  }
						  }
					  };
				  }
				  ImageDtt.startAndJoin(threads);
				  quadCLT.setFpnApplied(fpn);
				  quadCLT.setHasNewImageData(true);
			  }
			  return;
		  }

		  public ImagePlus saveShowFPN(
				  double [][][] fpn,
				  int           width,
				  boolean       save,
				  boolean       show) {
			  ImagePlus imp =  CorrectionFPN.showFPN(
					  fpn, // double [][][] fpn,
					  width, // int     width,
					  show); // boolean show)
			  if (save && (imp != null)) {
				  quadCLT.saveImagePlusInModelDirectory(
						  QuadCLT.CENTER_FPN_SUFFIX, // String      suffix, // null - use title from the imp
							imp); // ImagePlus   imp)
			  }
			  return imp;
		  }

		
		
	  public double [][][] readImageFPN (
			  int sens_mask){

		  int [] wh = new int [2];
		  double [][] fpn_data2 = quadCLT.readDoubleArrayFromModelDirectory(
				  QuadCLT.CENTER_FPN_SUFFIX, // String      suffix,
				  0,                 // int         num_slices, // (0 - all)
				  wh); // int []      wh)
		  if (fpn_data2 == null) {
			  return null;
		  }
		  int num_sensors = quadCLT.getNumSensors();
		  int used_sensors = 0;
		  for (int nsens = 0; nsens < num_sensors; nsens++) if ( (sens_mask & (1 << nsens)) != 0) {
			  used_sensors++;
		  }
		  int num_colors = fpn_data2.length/used_sensors;
		  double [][][] fpn = new double [num_sensors][num_colors][];
		  int isens = 0;
		  for (int nsens = 0; nsens < num_sensors; nsens++) if ( (sens_mask & (1 << nsens)) != 0) {
			  for (int ncol = 0; ncol < num_colors; ncol++) {
				  fpn[nsens][ncol] = fpn_data2[ncol * used_sensors + isens];
			  }
			  isens++;
		  }
		  return fpn;
	  }

	  public ImagePlus debugFPN(
			  QuadCLT []     quadCLTs,
			  double [][][]  fpn,
			  int []         range,
			  int            nsens,
			  double         um_sigma,
			  boolean show) {
		  int fpn_width =quadCLT.getTilesX()*quadCLT.getTileSize(); // see if center_CLT can be used
		  String dbg_title = quadCLT.getImageName()+"-DEBUG_FPN_ROWCOL_SENS_"+nsens;
		  ImagePlus imp = debugFPN(
				  quadCLTs,   // final QuadCLT []     quadCLTs,
				  fpn,        // final double [][][]  fpn,
				  range,      // final int []         range, // required
				  nsens,      // final int            nsens,
				  fpn_width,  // final int            width,
				  um_sigma,   // final double         um_sigma,
				  1.0,        // final double         um_weight,
				  dbg_title); // final String         title)
		  if (imp != null) {
			  quadCLT.saveImagePlusInModelDirectory(
					  dbg_title, // String      suffix, // null - use title from the imp
					  imp); // ImagePlus   imp)
			  if (show) {
				  imp.show();
			  }
		  }
		  return imp; // ************** restart from here!
	  }
	
	
	  public static ImagePlus showFPN(
			  double [][][] fpn,
			  int     width,
			  boolean show) {
		  int num_sens = fpn.length;
		  int num_colors = -1; // fpn[0].length;
		  for (int i = 0; i < num_sens; i++) {
			  if (fpn[i] != null) {
				  num_colors = fpn[i].length;
				  break;
			  }
		  }
		  String [] titles = new String[num_sens];
		  String [] top_titles = new String[num_colors];
		  for (int ncol = 0; ncol < num_colors;ncol++) {
			  top_titles[ncol] = "Color-"+ncol;
		  }

		  double [][][] img_data = new double [num_colors][num_sens][];
		  for (int nsens = 0; nsens < num_sens; nsens++) if (fpn[nsens] != null) {
			  titles[nsens]= "Sens-"+nsens;
			  for (int ncol = 0; ncol < num_colors;ncol++) {
				  img_data [ncol][nsens] = fpn[nsens][ncol];
			  }
		  }
		  String fpn_title = "FPN_data";
		  ImagePlus imp = ShowDoubleFloatArrays.showArraysHyperstack(
				  img_data,     // double[][][] pixels, 
				  width,    // int          width, 
				  fpn_title,        // String       title, "time_derivs-rt"+diff_time_rt+"-rxy"+diff_time_rxy,
				  titles,           // String []    titles, // all slices*frames titles or just slice titles or null
				  top_titles,       // String []    frame_titles, // frame titles or null
				  show);            // boolean      show)
		  return imp;
	  }

	  /**
	   * Calculate sensors FPN in a simple way average each channel/color for a scene sequence. Intended to be
	   * used in CUAS mode averaging several (normally 2 full periods of rotation.
	   * TODO: make more accurate and universal by back-propagating the corrected image and subtracting per-sensor
	   * versions before averaging.
	   * @param quadCLTs sequence of scenes with conditioned images (getImageData() != null), without FPN applied
	   * @param range    a first_last scene index pair. May be adjusted to include an integer number of rotations.
	   * @param sensor_mask bitmask which sensors to process (normally -1 - all)
	   * @param debugLevel debug level
	   * @return [sensor][color][pixel] average FPN image. May be saved as image_data to a virtual (center) scene
	   * (a QuadCLT instance). 
	   */
	  public static double [][][] calculateFPN(
			  final QuadCLT []     quadCLTs,
			  final int []         range, // required
			  final int            sensor_mask,
			  final int            debugLevel){
		  QuadCLT first_scene =  quadCLTs[range[0]];
		  final int num_sens =   first_scene.getNumSensors();    	
		  final int num_colors = first_scene.getNumColors();
		  final int width =      first_scene.getTilesX()*first_scene.getTileSize();
		  final int height =     first_scene.getTilesY()*first_scene.getTileSize();
		  final int num_pix =    width*height;
		  final int num_scenes = range[1]-range[0]+1;
		  final double [][][] fpn = new double [num_sens][num_colors][num_pix];
		  final Thread[] threads = ImageDtt.newThreadArray();
		  final AtomicInteger ai = new AtomicInteger(0);
		  for (int ithread = 0; ithread < threads.length; ithread++) {
			  threads[ithread] = new Thread() {
				  public void run() {
					  for (int nChn = ai.getAndIncrement(); nChn < num_sens; nChn = ai.getAndIncrement()) if (((sensor_mask >> nChn) & 1) != 0) {
						  for (int nscene = range[0]; nscene <= range[1]; nscene++) {
							  for (int ncol = 0; ncol < num_colors; ncol++) {
								  double [] img_slice = quadCLTs[nscene].getImageData()[nChn][ncol];
								  for (int npix = 0; npix < num_pix; npix++) {
									  fpn[nChn][ncol][npix] += img_slice[npix];
								  }
							  }
						  }
						  double scale = 1.0/num_scenes;
						  for (int ncol = 0; ncol < num_colors; ncol++) {
							  for (int npix = 0; npix < num_pix; npix++) {
								  fpn[nChn][ncol][npix] *= scale;
							  }
						  }
					  }
				  }
			  };
		  }		      
		  ImageDtt.startAndJoin(threads);
		  return fpn;
	  }
	  
	  public static int getImageHeight(
			  double [][][] image_data,
			  int width) {
		  for (int nChn = 0; nChn < image_data.length; nChn++) {
			  if (image_data[nChn] != null) {
				  return image_data[nChn][0].length/width;
			  }
		  }
		  return 0;
		  
	  }
	  
	  public static double [][]   getRowAvg(
			  final double [][][] image_data,
			  final int           width,
			  final double        max_abs,         // only average within +/- max_abs
			  final double        weight_outlier){ // small weight of outliers to avoid instability
		  final int num_sens = image_data.length;
		  final int height = getImageHeight(
				  image_data, // final double [][][] image_data,
				  width);    // final int width){
		  final double [][] image_row_avg =  new double [num_sens][width];
		  final double [][] weight_row_avg = new double [num_sens][width];
		  final Thread[] threads = ImageDtt.newThreadArray();
		  final AtomicInteger ai = new AtomicInteger(0);
		  for (int ithread = 0; ithread < threads.length; ithread++) {
			  threads[ithread] = new Thread() {
				  public void run() {
					  for (int nChn = ai.getAndIncrement(); nChn < num_sens; nChn = ai.getAndIncrement()) if (image_data[nChn] != null){
						  int indx = 0;
						  for (int y = 0; y < height; y++) {
							  for (int x = 0; x < width; x++) {
								  double d = image_data[nChn][0][indx++];
								  double w = (Math.abs(d) <= max_abs) ? 1.0 : weight_outlier;
								  image_row_avg[nChn][x] +=  w*d;
								  weight_row_avg[nChn][x] += w;
							  }
						  }
						  for (int x = 0; x < width; x++) {
							  image_row_avg[nChn][x] /= weight_row_avg[nChn][x];
						  }
					  }
				  }
			  };
		  }		      
		  ImageDtt.startAndJoin(threads);
		  return image_row_avg;
	  }

	  
	  

	  public static double [][] getColAvg(
			  final double [][][] image_data,
			  final int           width,
			  final double        max_abs,         // only average within +/- max_abs
			  final double        weight_outlier){ // small weight of outliers to avoid instability
		  final int num_sens = image_data.length;
		  final int height = getImageHeight(
				  image_data, // final double [][][] image_data,
				  width);    // final int width){
		  final double [][] image_col_avg =  new double [num_sens][height];
		  final double [][] weight_col_avg = new double [num_sens][height];
		  final Thread[] threads = ImageDtt.newThreadArray();
		  final AtomicInteger ai = new AtomicInteger(0);
		  for (int ithread = 0; ithread < threads.length; ithread++) {
			  threads[ithread] = new Thread() {
				  public void run() {
					  for (int nChn = ai.getAndIncrement(); nChn < num_sens; nChn = ai.getAndIncrement()) if (image_data[nChn] != null){
						  int indx = 0;
						  for (int y = 0; y < height; y++) {
							  for (int x = 0; x < width; x++) {
								  double d = image_data[nChn][0][indx++];
								  double w = (Math.abs(d) <= max_abs) ? 1.0 : weight_outlier;
								  image_col_avg[nChn][y] +=  w*d;
								  weight_col_avg[nChn][y] += w;
							  }
						  }
						  for (int y = 0; y < height; y++) {
							  image_col_avg[nChn][y] /= weight_col_avg[nChn][y];
						  }
					  }
				  }
			  };
		  }		      
		  ImageDtt.startAndJoin(threads);
		  return image_col_avg;
	  }
	  
	  public static double [][][] applyRowCol(
			  final double [][][] image_data,
			  final double [][]   image_row_avg,
			  final double [][]   image_col_avg,
			  final boolean       inplace ){
		  final int num_sens = image_data.length;
		  final double [][][] out_data = new double [num_sens][1][];
		  int w=0,h=0;
		  for (int nchn = 0; nchn < num_sens; nchn ++) if (image_data[nchn] != null) {
			  out_data[nchn][0] = inplace ? image_data[nchn][0] : image_data[nchn][0].clone();
			  w = image_row_avg[nchn].length;
			  h = image_col_avg[nchn].length;
		  }
		  final int width =  w;
		  final int height = h;
		  final Thread[] threads = ImageDtt.newThreadArray();
		  final AtomicInteger ai = new AtomicInteger(0);
		  for (int ithread = 0; ithread < threads.length; ithread++) {
			  threads[ithread] = new Thread() {
				  public void run() {
					  for (int nChn = ai.getAndIncrement(); nChn < num_sens; nChn = ai.getAndIncrement()) if (image_data[nChn] != null){
						  int indx = 0;
						  for (int y = 0; y < height; y++) {
							  for (int x = 0; x < width; x++) {
								  out_data[nChn][0][indx++] -=  image_row_avg[nChn][x] + image_col_avg[nChn][y];
							  }
						  }
					  }
				  }
			  };
		  }		      
		  ImageDtt.startAndJoin(threads);
		  return out_data;
	  }
	  
	  
	  
	  
	  public static ImagePlus debugFPN(
			  final QuadCLT []     quadCLTs,
			  final double [][][]  fpn,
			  final int []         range, // required
			  final int            nsens,
			  final int            width,
			  final double         um_sigma,
			  final double         um_weight,
			  final String         title){
		  int                     ncol = 0;
		  int num_sens = quadCLTs[range[0]].getNumSensors();
		  double max_abs =        100.0;
		  double weight_outlier = 0.001;
		  int           num_pix = quadCLTs[range[0]].getImageData()[nsens][ncol].length;
		  String [] titles_top = {"src","src-fpn","UM"+um_sigma, "row_col"};
		  double [][][] data = new double [titles_top.length][range[1]-range[0]+1][num_pix];
		  String [] titles = new String [range[1]-range[0]+1];
		  for (int nscene = range[0]; nscene <= range[1]; nscene++) {
			  int iscene = nscene - range[0];
			  titles [iscene] = quadCLTs[nscene].getImageName();
			  data[0][iscene] = quadCLTs[nscene].getImageData()[nsens][ncol].clone();
			  data[1][iscene] = data[0][iscene].clone();  
			  for (int i = 0; i < num_pix; i++) {
				  data[1][iscene][i] -= fpn[nsens][ncol][i];
			  }
			  data[2][iscene] = data[1][iscene].clone();
			  OpticalFlow.applyUMDouble(
					  data[2][iscene], // final double [] data,
					  width,         // final int       width,
					  um_sigma,      // final double um_sigma,
					  um_weight);    // final double um_weight)
			  
			  double [][][] data_chn = new double [num_sens][][];
			  data_chn[nsens] = new double [][] {data[2][iscene]};
			  double [][] image_row_avg = getRowAvg(
					  data_chn, // final double [][][] image_data,
					  width,             // final int           width,
					  max_abs,           // final double        max_abs,         // only average within +/- max_abs
					  weight_outlier);   // final double        weight_outlier)
			  double [][] image_col_avg = getColAvg(
					  data_chn, // final double [][][] image_data,
					  width,             // final int           width,
					  max_abs,           // final double        max_abs,         // only average within +/- max_abs
					  weight_outlier);   // final double        weight_outlier)
			  double [][][] data_out = applyRowCol(
					  data_chn,          // final double [][][] image_data,
					  image_row_avg,     // final double [][]   image_row_avg,
					  image_col_avg,     // final double [][]   image_col_avg,
					  false);            // final boolean       inplace )
			  data[3][iscene] = data_out[nsens][0];
		  }
		  
		  ImagePlus imp = ShowDoubleFloatArrays.showArraysHyperstack(
				  data,        // double[][][] pixels, 
				  width,       // int          width, 
				  title,       // String       title, "time_derivs-rt"+diff_time_rt+"-rxy"+diff_time_rxy,
				  titles,      // String []    titles, // all slices*frames titles or just slice titles or null
				  titles_top,  // String []    frame_titles, // frame titles or null
				  false);      // boolean      show)
		  return imp;
	  }	    
	  
	  
}
