package com.elphel.imagej.vegetation;

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

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

public class VegetationSegment {
	public Rectangle   woi_veg;
	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_veg,
			Rectangle   woi,
			double terrain_offset,
			double []   scene_offsets, // has NaNs
			double [][]  tva,
			double []   confidence,
			boolean []  overlaid) {
		this.path = path;
		this.woi = woi;
		this.woi_veg = woi_veg;
		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_veg.x) + (wy + woi.y - woi_veg.y) * woi_veg.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_veg : 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,
			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","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_veg;
			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: // vegetation
					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: // 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;
		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)
		
		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);

		
		// Combine offsets per scene
/*
		double [][] combined_offsets = new double [num_scenes][full_length];
		String [] titles_scenes = new String [num_scenes];
		for (int ithread = 0; ithread < threads.length; ithread++) {
			threads[ithread] = new Thread() {
				public void run() {
					double [] sw =  new double [full_length];
					for (int nscene = ai.getAndIncrement(); nscene < num_scenes; nscene = ai.getAndIncrement()) {
						titles_scenes[nscene] = "scene "+nscene;
						double [] swd = combined_offsets[nscene];
						Arrays.fill(sw, 0);
						for (int ns = 0; ns < segments.length; ns++) {
							Rectangle woi_veg = segments[ns].woi;
							int woi_length = woi_veg.width*woi_veg.height;
							double scene_offset = segments[ns].scene_offsets[nscene];
							if (!Double.isNaN(scene_offset)) {
								double [] wnd = overlapCosineWindow (woi, width);
								for (int windx = 0; windx < woi_length; windx++) {
									int x = windx % woi_veg.width + woi_veg.x;
									int y = windx / woi_veg.width + woi_veg.y;
									int indx = x + y * full.width;
									for (int t = 0; t < preview_data.length; t++) {
										double w = wnd[windx];
										sw [indx] += w;
										swd[indx] += w * scene_offset ;
									}				
								}
							}
						}
						(new DoubleGaussianBlur()).blurDouble(
								sw,          //
								full.width,
								full.height,
								ksigma*width,              // double sigmaX,
								ksigma*width,              // double sigmaY,
								0.01);                 // double accuracy)
						(new DoubleGaussianBlur()).blurDouble(
								swd,          //
								full.width,
								full.height,
								ksigma*width,              // double sigmaX,
								ksigma*width,              // double sigmaY,
								0.01);                 // double accuracy)

						for (int i = 0; i < full_length; i++) {
							if (sw[i] < min_sw) {
								swd[i] = Double.NaN;
							} else {
								swd[i] /= sw[i];
							}
						}
					}
				}
			};
		}		      
		ImageDtt.startAndJoin(threads);
		
		double [][] tva = {preview_data[0][accum_indx],preview_data[1][accum_indx],preview_data[2][accum_indx]};
		vegetationModel.setTVA(tva);

	//				new String[] {"full average", "filtered average"});
	
//		String ref_scene = vegetationModel.reference_scene;
		String render_title = "render-ro_"+render_open+"-rna_"+render_no_alpha+"-amn"+alpha_min+"-amx"+alpha_max+"-wo"+weight_opaque+"-bp"+boost_parallax+".tiff"; 
		
		
		double [] render = renderWithAlpha(
				preview_data[2][accum_indx], // final double []   alpha,
				combined_offsets,            // final double [][] scene_offsets, // subtract if not null
				render_open,                 // final boolean     render_open,      // render open areas (no vegetation offset)
				render_no_alpha,             // final boolean     render_no_alpha,  // render where no opacity is available
				alpha_min,                   // final double      alpha_min,        // below - completely transparent vegetation
				alpha_max,                   // final double      alpha_max,        // above - completely opaque
				weight_opaque,               // final double      weight_opaque,    // render through completely opaque vegetation
				boost_parallax,              // final double      boost_parallax) {    // increase weights of scenes with high parallax relative to the reference one
				max_parallax);               // final double      max_parallax)
		
		// num_exaggerate
		double [][] result_renders = new double [1+num_exaggerate][];
		String []   result_titles = new String [result_renders.length];
		result_titles[0] = "full average";
		result_titles[1] = "filtered average";
		result_renders[0] = terrain_average;
		result_renders[1] = render;
		for (int n = 2; n <=  num_exaggerate; n++) {
			result_titles[n]="exaggerate x "+n;
			result_renders[n] = new double [full_length];
			for (int i = 0; i < full_length; i++) {
				result_renders[n][i] = terrain_average[i] + n * (render[i] - terrain_average[i]);
			}
		}
		
		ShowDoubleFloatArrays.showArrays(
				combined_offsets,
				full.width,
				full.height,
				true,
				ref_scene+"-scene_offsets",
				titles_scenes);
		
		
		
		ShowDoubleFloatArrays.showArrays(
				result_renders,
				full.width,
				full.height,
				true,
				ref_scene+"-"+render_title,
				result_titles);
		String [] synt_titles = {"render","synthetic", "synt_terrain", "synt vegetation","masked render", "render"};
		final double [][][] synt_data = new double [synt_titles.length][num_scenes][];
		final double      terrain_max = alpha_min;
		final double      vegetation_min = 0.5; // alpha_max;
		String synt_name = "synthetic-terr"+terrain_max+"-veg"+vegetation_min;
		synt_data[0] = terrain_rendered;
		synt_data[5] = terrain_rendered;
		double [][][] synt_tv= renderSynthetic(
				preview_data[2][accum_indx], // final double []   alpha,
				combined_offsets,            // final double [][] scene_offsets, // subtract if not null
				terrain_max,                 // final double      terrain_max,      // below - terrain is visible
				vegetation_min,              // final double      vegetation_min,   // above - vegetation is visible
				boost_parallax,              // final double      boost_parallax,   // increase weights of scenes with high parallax relative to the reference one
				max_parallax);               //final double      max_parallax) {
		synt_data[2] = synt_tv[0];
		synt_data[3] = synt_tv[1];
		
//		for (int nscene = 0; nscene < num_scenes; nscene++) { // use threads
//			synt_data[1][nscene] = synt_tv[0][nscene].clone(); 
//			synt_data[4][nscene] = terrain_rendered[nscene].clone();
//			double [] veg_scene = synt_tv[1][nscene];
//			for (int i = 0; i < veg_scene.length; i++) {
//				if (!Double.isNaN(veg_scene[i])) {
//					synt_data[1][nscene][i] = veg_scene[i];
//					synt_data[4][nscene][i] = Double.NaN;
//				}
//			}
//		}

		ai.set(0);
		for (int ithread = 0; ithread < threads.length; ithread++) {
			threads[ithread] = new Thread() {
				public void run() {
					for (int nscene = ai.getAndIncrement(); nscene < num_scenes; nscene = ai.getAndIncrement()) {
						synt_data[1][nscene] = synt_tv[0][nscene].clone(); 
						synt_data[4][nscene] = terrain_rendered[nscene].clone();
						double [] veg_scene = synt_tv[1][nscene];
						for (int i = 0; i < veg_scene.length; i++) {
							if (!Double.isNaN(veg_scene[i])) {
								synt_data[1][nscene][i] = veg_scene[i];
								synt_data[4][nscene][i] = Double.NaN;
							}
						}
						
					}
				}
			};
		}		      
		ImageDtt.startAndJoin(threads);
		
		
		
		
		ShowDoubleFloatArrays.showArraysHyperstack(
				synt_data,              // double[][][] pixels, 
				full.width,                // int          width, 
				ref_scene+"-"+synt_name,        // String       title, "time_derivs-rt"+diff_time_rt+"-rxy"+diff_time_rxy,
				titles_scenes,                    // String []    titles, // all slices*frames titles or just slice titles or null
				synt_titles,                // String []    frame_titles, // frame titles or null
				true);                     // boolean      show)
		
		*/
		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;
	}
	
}
