package com.elphel.imagej.vegetation;

import java.awt.Rectangle;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicInteger;

import com.elphel.imagej.common.GenericJTabbedDialog;
import com.elphel.imagej.common.ShowDoubleFloatArrays;
import com.elphel.imagej.tileprocessor.ImageDtt;
import com.elphel.imagej.tileprocessor.QuadCLT;

import ij.ImagePlus;
import ij.gui.GenericDialog;

public class VegetationSegment {
	public Rectangle   woi_max;
	public Rectangle   woi;
	public double []   scene_offsets;
	public double [][] tva; // woi_veg
	public double      terrain_offset;
	public double []   confidence; // woi
	public boolean []  overlaid;
	public String      path;
	public VegetationSegment(
			String      path, 
			Rectangle   woi_max,
			Rectangle   woi,
			double terrain_offset,
			double []   scene_offsets, // has NaNs
			double [][]  tva,
			double []   confidence,
			boolean []  overlaid) {
		this.path = path;
		this.woi = woi;
		this.woi_max = woi_max;
		this.scene_offsets = scene_offsets;
		this.tva =  tva;
		this.terrain_offset = terrain_offset;
		this.confidence = confidence;
		this.overlaid = overlaid;
		correctTerrainOffset();
	}
	
	public void correctTerrainOffset() {
		if (!Double.isNaN(terrain_offset) && (overlaid != null)) {
			int woi_length = woi.width * woi.height;
			for (int windx = 0; windx < woi_length; windx++) {
				int wx = windx % woi.width;
				int wy = windx / woi.height;
				int wvindx = (wx + woi.x - woi_max.x) + (wy + woi.y - woi_max.y) * woi_max.width;
				if (overlaid[windx]) {
					tva[0][wvindx] += terrain_offset;
				}
			}
		}
	}
	
	public static Rectangle getBounds( // [0] - woi, [1] - woi_weg
			VegetationSegment [] segments,
			boolean use_veg) {
		Rectangle bounds = null;
		for (VegetationSegment segment:segments) {
			Rectangle selected_woi = use_veg ?segment.woi_max : segment.woi;
			if (bounds == null) {
				bounds = new Rectangle(selected_woi);
			} else {
				bounds.add(selected_woi);
			}
		}
		return bounds;
	}
	public static void combineSegments(
			VegetationLMA        vegetationLMA,
			VegetationSegment [] segments,
			boolean              crop_combo,
			boolean              keep_partial,
			int                  width,
			final double         um_sigma,
			final double         um_weight,
			final boolean        render_open,      // render open areas (no vegetation offset)
			final boolean        render_no_alpha,  // render where no opacity is available
			final double         alpha_min,        // below - completely transparent vegetation
			final double         alpha_max,        // above - completely opaque
			final double         weight_opaque,    // render through completely opaque vegetation
			final double         boost_parallax,   // increase weights of scenes with high parallax relative to the reference one
			final double         max_parallax,			
			final int            num_exaggerate) { // show amplified difference from filtering
		//FIXME: Handle woi separately!
		final Rectangle combo_woi =     crop_combo? VegetationSegment.getBounds(segments, false):vegetationLMA.getFull();
		final Rectangle combo_woi_veg = crop_combo? VegetationSegment.getBounds(segments, true): vegetationLMA.getFull();
		final Rectangle out_woi = new Rectangle(combo_woi);
		out_woi.add(combo_woi_veg);
		
		final Thread[] threads = ImageDtt.newThreadArray(QuadCLT.THREADS_MAX);
		final AtomicInteger ai = new AtomicInteger(0);
		double ksigma = 1.0; 
		double min_sw = 1e-4;
		int out_length = out_woi.width*out_woi.height;
		int accum_indx = segments.length;
		String [] top_titles = {"terrain","vegetation","alpha","elevation","terrain_elevation","confidence"};
		double [][][] preview_data = new double [top_titles.length][accum_indx+1][out_length];
		String [] titles = new String[accum_indx+1];
		for (int ns = 0; ns < segments.length; ns++) {
			Rectangle woi =     segments[ns].woi; // for confidence and limits for terrain
			Rectangle woi_veg = segments[ns].woi_max;
			titles[ns] = woi.toString();
			int woi_length =     woi.width * woi.height;  
			int woi_veg_length = woi_veg.width * woi_veg.height;  
			double [][] woi_tva = segments[ns].tva;
			double []   confidence = segments[ns].confidence;
			for (int t = 0; t < top_titles.length; t++) {
				Arrays.fill(preview_data[t][ns], Double.NaN); // java.lang.ArrayIndexOutOfBoundsException: Index 3 out of bounds for length 3
				switch (t) {
				case 0: // terrain cut out from woi_tva
					for (int windx = 0; windx < woi_length; windx++) {
						int wx = windx % woi.width;
						int wy = windx / woi.width;
						int wxsrc = wx +  woi.x-woi_veg.x; // >= 0
						int wysrc = wy +  woi.y-woi_veg.y; // >= 0
						int wsrc = wxsrc + wysrc *  woi_veg.width; // index in woi_tva[t][windx];
						int x = wx +  woi.x - out_woi.x;
						int y = wy +  woi.y - out_woi.y;
						int indx = x + y * out_woi.width;
						preview_data[t][ns][indx] = woi_tva[t][wsrc]; // java.lang.ArrayIndexOutOfBoundsException: Index -43 out of bounds for length 550
					}
					break;
				case 1: // vegetation
				case 2: // alpha
				case 3: // elevation
					for (int windx = 0; windx < woi_veg_length; windx++) {
						int x = windx % woi_veg.width + woi_veg.x - out_woi.x;
						int y = windx / woi_veg.width + woi_veg.y - out_woi.y;
						int indx = x + y * out_woi.width;
						preview_data[t][ns][indx] = woi_tva[t][windx];
					}
					break;
				case 4: // terrain elevation
					for (int windx = 0; windx < woi_veg_length; windx++) {
						int x = windx % woi_veg.width + woi_veg.x - out_woi.x;
						int y = windx / woi_veg.width + woi_veg.y - out_woi.y;
						int indx = x + y * out_woi.width;
						preview_data[t][ns][indx] = woi_tva[t][windx]; // skipping terrain elevation average
					}
					break;
				case 5: // confidence
					for (int windx = 0; windx < woi_length; windx++) {
						int x = windx % woi.width + woi_veg.x - out_woi.x;
						int y = windx / woi.width + woi_veg.y - out_woi.y;
						int indx = x + y * out_woi.width;
						preview_data[t][ns][indx] = confidence[windx];
					}
					break;
				}
			}
		}
		titles[accum_indx] = "accum";
		double [][] sum_w =  new double [preview_data.length][out_length];
		double [][] sum_wd = new double [preview_data.length][out_length];
		for (int ns = 0; ns < segments.length; ns++) {
			Rectangle woi = segments[ns].woi;
			int woi_length = woi.width*woi.height;
//			double [][] woi_tva = segments[ns].tva;			
			double [] wnd = overlapCosineWindow (woi, width);
			for (int windx = 0; windx < woi_length; windx++) {
				int x = windx % woi.width + woi.x - out_woi.x;
				int y = windx / woi.width + woi.y - out_woi.y;
				int indx = x + y * out_woi.width;
				for (int t = 0; t < preview_data.length; t++) {
					double d = preview_data[t][ns][indx];
					if (!Double.isNaN(d)) {
						double w = wnd[windx];
						sum_w [t][indx] += w;
						sum_wd[t][indx] += w * d ;
					}
				}				
			}
		}
		
		for (int t = 0; t < preview_data.length; t++) {
			for (int i = 0; i < out_length; i++) {
				if (sum_w[t][i] == 0) {
					preview_data[t][accum_indx][i] = Double.NaN;
				} else {
					preview_data[t][accum_indx][i] = sum_wd[t][i]/sum_w[t][i];
				}
			}
		}
		VegetationModel vegetationModel = vegetationLMA.getModel();
		String ref_scene = vegetationModel.reference_scene;
		String pre_title = ref_scene+"-preview_segments-"+out_woi.x+"-"+out_woi.y+"-"+out_woi.width+"-"+out_woi.height;
		if (keep_partial) {
		ShowDoubleFloatArrays.showArraysHyperstack(
				preview_data,     // double[][][] pixels, 
				out_woi.width,    // int          width, 
				pre_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
				true);            // boolean      show)
		} else {
			double [][] accum_only = new double [preview_data.length][];
			for (int t = 0; t < preview_data.length; t++) {
				accum_only[t] = preview_data[t][accum_indx];
			}
			ShowDoubleFloatArrays.showArrays(
					accum_only,
					out_woi.width,
					out_woi.height,
					true,
					pre_title+"-ACCUM",
					top_titles);
		}
		
		double [] terrain_reference =   crop(
				vegetationModel.terrain_scenes_render[vegetationModel.reference_index],
				vegetationLMA.getFull().width,
				out_woi);
		double [] terrain_accumulated = crop(
				vegetationModel.terrain_average_render,
				vegetationLMA.getFull().width,
				out_woi);
		double [] terrain_um = 	terrain_accumulated.clone();
		for (int i = 0; i < terrain_um.length; i++) {
			if (Double.isNaN(terrain_um[i])) {
				terrain_um[i] = 0;
			}
		}
		VegetationModel.unsharpMask(
				terrain_um,    // final double [] data,
				out_woi.width, // final int       width,
				um_sigma,      // final double    um_sigma,
				um_weight);    // final double    um_weight)
		String [] result_titles = {"ref_render", "terrain_accum",String.format("terrain_accum_UM%5.2f-%4.2f",um_sigma,um_weight),"terrain"};
		double [][] result_img = {terrain_reference,terrain_accumulated,terrain_um,preview_data[0][accum_indx]};
		boolean normalize = true;
		int norm_index = 1;
		if (normalize) {
		result_img =  normalize (
				result_img, // final double [][] data,
				norm_index); // final int         norm_index){
		}
		String res_title = ref_scene+"-result_terrain-"+out_woi.x+"-"+out_woi.y+"-"+out_woi.width+"-"+out_woi.height;
		if (normalize) {
			res_title+= "-normalized"+norm_index;
		}
		ShowDoubleFloatArrays.showArrays(
				result_img,
				out_woi.width,
				out_woi.height,
				true,
				res_title,
				result_titles);

		
		return;
	}
	
	
	public static double [] overlapCosineWindow(
			Rectangle woi,
			int       width){
		if (width > woi.width/2) {
			width = woi.width/2;
		}
		if (width > woi.height/2) {
			width = woi.height/2;
		}
		double [] window = new double [woi.width*woi.height];
		double [] wnd_x = new double  [woi.width];
		double [] wnd_y = new double  [woi.height];
		Arrays.fill(wnd_x, 1.0);
		Arrays.fill(wnd_y, 1.0);
		for (int i = 0; i < width; i++) {
			double w = 0.5* (1 - Math.cos(Math.PI * (i + 0.5) / width));
			wnd_x[i] = w;
			wnd_x[woi.width - i - 1] = w;
			wnd_y[i] = wnd_x[i];
			wnd_y[woi.height - i - 1] = w;
		}
		int indx = 0;
		for (int y = 0; y < woi.height; y++) {
			for (int x = 0; x < woi.width; x++) {
				window[indx++] = wnd_x[x] * wnd_y[y];
			}
		}
		return window;
	}
	
	public static double [] crop(
			double [] data,
			int       width,
			Rectangle woi) {
		int height = data.length/width;
		if ((woi.x==0) && (woi.y==0) && (woi.width==width) && (woi.height==height)) {
			return data.clone();
		}
		double [] cropped = new double [woi.width*woi.height];
		Arrays.fill (cropped, Double.NaN);
		for (int row = 0; row < woi.height; row ++) {
			int src_row = row + woi.y;
			if ((src_row >0) && (src_row < height)) {
				int col0 = 0;
				int col1 = woi.width;
				int src_col0 = col0+woi.x;
				int src_col1 = src_col0+ woi.width;
				if (src_col0 < 0) {
					col0 -= src_col0;
					src_col0 = 0;
				}
				if (src_col1 > width) {
					col1 -= src_col1 - width;
					src_col1 = width;
				}
				int row_len = col1 - col0; 
				if (row_len > 0) {
					System.arraycopy(
							data,
							src_row * width + src_col0,
							cropped,
							row*woi.width + col0,
							row_len);
				}
			}
		}
		return cropped;
	}
	/**
	 * Normalize slices to match the selected one by average and standard deviation
	 * @param data       array of image slices
	 * @param norm_index index (0-based) of the reference slice
	 * @return normalized image
	 */
	public static double [][] normalize (
			final double [][] data,
			final int         norm_index){
		final int num_slices = data.length;
		final int num_pix = data[0].length; 
		double [][] normalized = new double [num_slices][];
		for (int i = 0; i < num_slices; i++) {
			normalized[i] = data[i].clone();
		}
		final double [] mean = new double[num_slices];
		final double [] std = new double[num_slices];
		final Thread[] threads = ImageDtt.newThreadArray(QuadCLT.THREADS_MAX);
		final AtomicInteger ai = new AtomicInteger(0);
		for (int ithread = 0; ithread < threads.length; ithread++) { // first sum for pairs
			threads[ithread] = new Thread() {
				public void run() {
					for (int nSlice = ai.getAndIncrement(); nSlice < num_slices; nSlice = ai.getAndIncrement()) {
						double [] slice = normalized[nSlice];
						double s0=0, s1=0, s2=0;
						for (int i = 0; i < slice.length; i++) if (!Double.isNaN(slice[i])) {
							s0 += 1;
							s1 += slice[i];
							s2 += slice[i]*slice[i];
						}
						s1 /= s0;
						s2 /= s0;
						mean[nSlice] = s1;
						std[nSlice] = Math.sqrt(s2 - s1 * s1);
					}
				}
			};
		}		      
		ImageDtt.startAndJoin(threads);
		final double [] a = new double [num_slices], b = new double[num_slices];
		Arrays.fill(a, 1);
		for (int n = 0; n < num_slices; n++) if (n != norm_index){
			a[n] = std[norm_index]/std[n];
			b[n] = mean[norm_index] - a[n]*mean[n]; 
		}
		ai.set(0);
		for (int ithread = 0; ithread < threads.length; ithread++) { // first sum for pairs
			threads[ithread] = new Thread() {
				public void run() {
					for (int nPix = ai.getAndIncrement(); nPix < num_pix; nPix = ai.getAndIncrement()) {
						for (int n = 0; n < num_slices; n++) if (n != norm_index){
							double d = normalized[n][nPix];
							if (!Double.isNaN(d)) {
								normalized[n][nPix] = a[n]*d + b[n]; 
							}
						}						
					}
				}
			};
		}		      
		ImageDtt.startAndJoin(threads);
		return normalized;
	}

	public static ImagePlus generateSineTarget() {
		int    width =       1024;
		int    height =      1024;
		double period_hor =     0.3333;  // fraction of width -  1.0 single horizontal period
		double period_vert =    0.3333; // fraction of height - 1.0 single vertical period
		double phase_hor =      0;   // phase at left edge as fraction of 2*PI (0.25 - 90 degree shift)
		double phase_vert=      0;// phase at top edge as fraction of 2*PI (0.25 - 90 degree shift)
		GenericJTabbedDialog gd = new GenericJTabbedDialog("Set CLT parameters",600,250);
		gd.addNumericField("Width",            width,      0,4,"pix","Image width." );
		gd.addNumericField("Height",           width,      0,4,"pix","Image height." );
		gd.addNumericField("Horizontal period",period_hor, 5,8,"",   "Horisontal period (1.0 - one full period, 0.2 - two full periods)." );
		gd.addNumericField("Vertical period",  period_vert,5,8,"",   "Vertical period (1.0 - one full period, 0.2 - two full periods)." );
		gd.addNumericField("Horizontal phase", phase_hor,  5,8,"",   "Horisontal phase at left edge (1.0 - 2*PI, 0.25 - 90 degrees)." );
		gd.addNumericField("Vertical phase",   phase_vert, 5,8,"",   "Vertical phase at top edge (1.0 - 2*PI, 0.25 - 90 degrees)." );
        gd.showDialog();
		if (gd.wasCanceled()) return null;
		width =       (int) gd.getNextNumber();
		height =      (int) gd.getNextNumber();
		period_hor =        gd.getNextNumber(); // fraction of width -  1.0 single horizontal period
		period_vert =       gd.getNextNumber(); // fraction of height - 1.0 single vertical period
		phase_hor =         gd.getNextNumber(); // phase at left edge as fraction of 2*PI (0.25 - 90 degree shift)
		phase_vert=         gd.getNextNumber(); // phase at top edge as fraction of 2*PI (0.25 - 90 degree shift)
		double [] data = generateSineTarget(
				width,       // int width,
				height,      // int height,
				period_hor,  // double period_hor,  // fraction of width -  1.0 single horizontal period
				period_vert, //double period_vert, // fraction of height - 1.0 single vertical period
				phase_hor,   // double phase_hor,   // phase at left edge as fraction of 2*PI (0.25 - 90 degree shift)
				phase_vert); // double phase_vert) {// phase at top edge as fraction of 2*PI (0.25 - 90 degree shift)
		String title = "pattern_perx"+period_hor+"_pery"+period_vert+"_phx"+phase_hor+"_phy"+phase_vert; 
		ImagePlus imp = ShowDoubleFloatArrays.makeArrays(
				data, // double[] pixels,
				width, // int width,
				height, // int height,
				title); // String title);
  	    imp.getProcessor().resetMinAndMax();
		imp.show();
		return imp;
	}

	
	public static double [] generateSineTarget(
			int width,
			int height,
			double period_hor,  // fraction of width -  1.0 single horizontal period
			double period_vert, // fraction of height - 1.0 single vertical period
			double phase_hor,   // phase at left edge as fraction of 2*PI (0.25 - 90 degree shift)
			double phase_vert) {// phase at top edge as fraction of 2*PI (0.25 - 90 degree shift)
		double [] data = new double[width*height];
		double [] vx= new double [width];
		for (int x = 0; x <width; x++) {
			vx[x] = Math.sin(2*Math.PI*((x / (period_hor*width)) + phase_hor));
		}
		for (int y = 0; y < height; y++) {
			double vy = Math.sin(2*Math.PI*((y / (period_vert * height)) + phase_vert));
			for (int x = 0; x <width; x++) {
				data [y*width+x] = vx[x] * vy;
			}
		}
		return data;
	}
	
}
