package com.elphel.imagej.common;
/**
** CholeskyBlock - multithreaded Chelesky decomposition
**
** Copyright (C) 2025 Elphel, Inc.
**
** -----------------------------------------------------------------------------**
**
**  CholeskyBlock.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/>.
** -----------------------------------------------------------------------------**
*/

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

import com.elphel.imagej.tileprocessor.ImageDtt;

import Jama.Matrix;

public class CholeskyBlock {
	public final int m;  // tile size
	public final int np; // number of elements in row/col
	public final int n;  // number of tiles in  row/col
	public final int nf; // number of full rows/cols
	public final double [] A;
	public final double [] L;
	
	public CholeskyBlock (
			double [][] A_in,
			int size) {
		m = size;
		np = A_in.length;
		nf = np / m;
		n = ((nf * m) < np)? (nf + 1) : nf; 
		A = new double [np*np]; //  [n*n*m*m];
		L = new double [A.length];
		setup_ATriangle(A_in);
	}
	
	private void setup_ATriangle(double [][] A_in) {
		for (int tile_row = 0; tile_row < nf; tile_row++) {
			for (int tile_col = 0; tile_col < tile_row; tile_col++) {
				int indx = indx_IJ(tile_row, tile_col);
				for (int k = 0; k < m; k++) {
					System.arraycopy(
							A_in[tile_row*m +k],
							tile_col * m,
							A,
							indx + m * k,
							m);
				}
			}
			// copy diagonal
			int indx = indx_IJ(tile_row, tile_row);
			for (int k = 0; k < m; k++) {
				System.arraycopy(
						A_in[tile_row*m +k],
						tile_row * m,
						A,
						indx + m * k,
						k+1);
			}
		}
		if (n > nf) { // if there are small tiles below and to the right
			int tile_row = nf;
			int h = np - m * nf;
			for (int tile_col = 0; tile_col < nf; tile_col++) {
				int indx = indx_IJ(tile_row, tile_col);
				for (int k = 0; k < h; k++) {
					System.arraycopy(
							A_in[tile_row*m +k],
							tile_col * m,
							A,
							indx + m * k,
							m);
				}
				
			}
			// copy diagonal
			int indx = indx_IJ(tile_row, tile_row);
			for (int k = 0; k < h; k++) {
				System.arraycopy(
						A_in[tile_row*m +k],
						tile_row * m,
						A,
						indx + h * k,
						k+1);
			}
		}
	}
	
	public Matrix getL() {
		return new Matrix(get_LTriangle(),np,np);
	}
	
	
	public double [][] get_LTriangle() {
		double [][] L_out = new double[np][np];
		for (int tile_row = 0; tile_row < nf; tile_row++) {
			for (int tile_col = 0; tile_col < tile_row; tile_col++) {
				int indx = indx_IJ(tile_row, tile_col);
				for (int k = 0; k < m; k++) {
					System.arraycopy(
							L,
							indx + m * k,
							L_out[tile_row*m +k],
							tile_col * m,
							m);
				}
			}
			// copy diagonal
			int indx = indx_IJ(tile_row, tile_row);
			for (int k = 0; k < m; k++) {
				System.arraycopy(
						L,
						indx + m * k,
						L_out[tile_row*m +k],
						tile_row * m,
						k+1);
			}
		}
		if (n > nf) { // if there are small tiles below and to the right
			int tile_row = nf;
			int h = np - m * nf;
			for (int tile_col = 0; tile_col < nf; tile_col++) {
				int indx = indx_IJ(tile_row, tile_col);
				for (int k = 0; k < h; k++) {
					System.arraycopy(
							L,
							indx + m * k,
							L_out[tile_row*m +k],
							tile_col * m,
							m);
				}
				
			}
			// copy diagonal
			int indx = indx_IJ(tile_row, tile_row);
			for (int k = 0; k < h; k++) {
				System.arraycopy(
						L,
						indx + h * k,
						L_out[tile_row*m +k],
						tile_row * m,
						k+1);
			}
		}
		return L_out;
		
	}
	
	/**
	 * Get index of the top-left tile corner
	 * @param i tile row
	 * @param j tile column
	 * @return index in A and L arrays
	 */
	public int indx_IJ(int i, int j) {
		return j * (m * np) + ((j >= nf) ? (np-nf*m): m) * m * i;
	}

	/**
	 * Get a new diagonal square submatrix 
	 * @param arr A or L packed array (line-scan order for each tile column)
	 * @param i tile index on the diagonal
	 * @return new array with tile data (m x m, or smaller for the bottom right) 
	 */
	public double [][] getDiagSquare(
			double [] arr,
			int i){
		int l =  (i >= nf) ? (np-nf*m): m;
		double [][] a_diag = new double[l][l];
		return getDiagSquare (arr, a_diag, i);
	}

	/**
	 * Get a new diagonal square submatrix 
	 * @param arr A or L packed array (line-scan order for each tile column)
	 * @param a_diag - preallocated array (should be smaller if needed) 
	 * @param i tile index on the diagonal
	 * @return new array with tile data (m x m, or smaller for the bottom right) 
	 */
	public double [][] getDiagSquare(
			double []   arr,
			double [][] a_diag,
			int         i){
		int indx = indx_IJ(i,i);
		int l = a_diag.length;
		for (int k = 0; k < l; k++) {
			int kl = k*l;
			System.arraycopy( arr, indx+kl, a_diag[k], 0, l);
		}
		return a_diag;
	}


	/**
	 * Get a new diagonal square submatrix, copy only lower triangle including diagonal  
	 * @param arr A or L packed array (line-scan order for each tile column)
	 * @param i tile index on the diagonal
	 * @return new array with tile data (m x m, or smaller for the bottom right) 
	 */
	public double [][] getDiagTriangle(
			double [] arr,
			int i){
		int l =  (i >= nf) ? (np-nf*m): m;
		double [][] a_diag = new double[l][l];
		return getDiagLTriangle (arr, a_diag, i);
	}

	/**
	 * Get a new diagonal square submatrix, copy only lower triangle including diagonal 
	 * @param arr A or L packed array (line-scan order for each tile column)
	 * @param a_diag - preallocated array (should be smaller if needed) 
	 * @param i tile index on the diagonal
	 * @return new array with tile data (m x m, or smaller for the bottom right) 
	 */
	public double [][] getDiagLTriangle(
			double [] arr,
			double [][] a_diag,
			int i){
		int indx = indx_IJ(i,i);
		int l = a_diag.length;
		for (int k = 0; k < l; k++) {
			int kl = k*l;
			System.arraycopy( arr, indx+kl, a_diag[k], 0, k+1);
		}
		return a_diag;
	}
	
	/**
	 * Save calculated tile L lower diagonal matrix to a packed array
	 * @param arr    A or L packed array (line-scan order for each tile column)
	 * @param l_diag lower triangular array with Cholesky decomposition
	 * @param i tile index on a diagonal
	 */
	public void putDiagLTriangle(
			double []   arr,
			double [][] l_diag,
			int         i) {
		int indx = indx_IJ(i,i);
		int l = l_diag.length;
		for (int k = 0; k < l; k++) {
			int kl = k*l;
			System.arraycopy(l_diag[k], 0, arr, indx+kl,  k+1);
		}
	}
	
	
	public void setL21(
			int       i,   // i > j,
			int       j) { // j <nf
		int indx_diag = indx_IJ(j,j);
		int indx_ij =   indx_IJ(i,j);
		int h = (i < nf) ? m : (np-nf*m);
		// prepare solving Lx = b, copy tile A -> L
		System.arraycopy(A, indx_ij, L, indx_ij, m * h);
		for (int l_row = 0; l_row < m; l_row++) { // was <m <h!
			for (int x_col= 0; x_col < h; x_col++) { // b-vector
				int lindx = indx_ij + m * x_col + l_row;
				double ls = L[lindx];
				for (int l_col = 0; l_col < l_row; l_col++) {
					ls -= L[indx_ij + m * x_col + l_col] * L[indx_diag + m* l_row+l_col];
				}
				L[lindx] = ls/L[indx_diag + (m + 1)* l_row];
			}
		}
		return;
	}
	
	public void setA22(
			int diag,// < col,  < row
			int row, // >= col
			int col) {
		int h = (row < nf) ? m : (np-nf*m);
		int indx_a =    indx_IJ(row,col);
		int indx_lrow = indx_IJ(row,diag);
		if (row == col) {
			for (int i = 0; i < h; i++) {
				for (int j = 0; j <= i; j++) {
					for (int k = 0; k < m; k++) { // was h
						A[indx_a + i * h + j] -= L[indx_lrow + i * m + k] * L[indx_lrow + j * m + k];
					}
				}
			}
		} else {
			int indx_lcol = indx_IJ(col,diag);
			for (int i = 0; i < h; i++) {
				for (int j = 0; j < m; j++) {
					for (int k = 0; k < m; k++) {
						A[indx_a + i * m + j] -= L[indx_lrow + i * m + k] * L[indx_lcol + j * m + k];
					}
				}
			}
		}
		return;
	}
	
	public void choleskyBlockMulti() {
		final Thread[] threads =       ImageDtt.newThreadArray();
		final AtomicInteger ai =       new AtomicInteger(0);
		final double [][] Am = new double[m][m];
		final double [][] Ah = (n > nf) ? new double[(np-nf*m)][(np-nf*m)] : null;
		final double [][] A1 = (np < m) ? Ah : Am; // smaller than a tile
		final double [][] Lm = new double[m][m];
		final double [][] Lh = (n > nf) ? new double[(np-nf*m)][(np-nf*m)] : null;
		final double [][] L1 = (np < m) ? Lh : Lm;  // smaller than a tile
		// Extract top-left tile (only lower triangle)
		getDiagLTriangle(
				A,   // double [] arr,
				A1,  // double [][] a_diag,
				0);  // int i)
		// Decompose, get L (lower triangle)
		cholesky_single (
				A1,  // double [][] a,
				L1); // double [][] l)
		// Save it to L-array:
		putDiagLTriangle(
				L,   // double []   arr,
				L1,  // double [][] l_diag,
				0);  // int         i)
		// Calculate first column under diagonal (L21) - maybe use m
		ai.set(1); // start from the second tile row
		for (int ithread = 0; ithread < threads.length; ithread++) { // first sum for pairs
			threads[ithread] = new Thread() {
				public void run() {
					for (int tile_row = ai.getAndIncrement(); tile_row < n; tile_row = ai.getAndIncrement()) {
						setL21(tile_row, 0);
					}
				}
			};
		}		      
		ImageDtt.startAndJoin(threads);

		for (int tile_diag = 1; tile_diag < n; tile_diag++) {
			final int ftile_diag = tile_diag;
			// Calculate A in one tile column of the remaining A2'
			// start with diagonal (top) tile, and calculate its Cholesky
			// In parallel, calculate A for all tiles in that column below diagonal
			if (tile_diag == (n-1)) {
				System.out.println("choleskyBlockMulti() last pass n= "+n+", tile_diag="+tile_diag);
			}
			ai.set(ftile_diag);
			for (int ithread = 0; ithread < threads.length; ithread++) { // first sum for pairs
				threads[ithread] = new Thread() {
					public void run() {
						for (int nRow= ai.getAndIncrement(); nRow < n; nRow = ai.getAndIncrement()) {
//							if ((nRow == (n-1)) && (ftile_diag == (n-1))) {
//								System.out.println("calling setA22() first: diag="+(ftile_diag-1)+", nRow="+nRow);
//							}

							setA22(
									ftile_diag-1, // int diag,// < col,  < row
									nRow,         // int row, // >= col
									ftile_diag);  // int col) 							
							if (nRow == ftile_diag) {
								double [][] At = (ftile_diag < nf) ? Am: Ah; // full or reduced array
								double [][] Lt = (ftile_diag < nf) ? Lm: Lh; // full or reduced array
								// Extract top-left tile (only lower triangle)
								getDiagLTriangle(
										A,            // double [] arr,
										At,           // double [][] a_diag,
										ftile_diag);  // int i)
								// Decompose, get L (lower triangle)
								cholesky_single (
										At,           // double [][] a,
										Lt);          // double [][] l)
								// Save it to L-array:
								putDiagLTriangle(
										L,            // double []   arr,
										Lt,           // double [][] l_diag,
										ftile_diag);  // int         i)
								/*
								// Calculate first column under diagonal (L21) - maybe use m
								for (int tr = ftile_diag; tr < n; tr++) {
									setL21(tr, ftile_diag);
								}
								*/
							}
						}
					}
				};
			}		      
			ImageDtt.startAndJoin(threads);
			if (ftile_diag < (n-1)) {
				// Now in parallel calculate L in column ftile_diag under diagonal and
				// finish A2' to the right of ftile_diag column
				final int left_rows = n - ftile_diag - 1;
				final int num_tiles = left_rows * (left_rows + 1) / 2 + 1;

				ai.set(0);
				for (int ithread = 0; ithread < threads.length; ithread++) { // first sum for pairs
					threads[ithread] = new Thread() {
						public void run() {
							for (int ntile = ai.getAndIncrement(); ntile < num_tiles; ntile = ai.getAndIncrement()) {
								if (ntile == 0) {
									// Calculate first column of L under diagonal (L21) - maybe use m

									for (int tr = ftile_diag + 1; tr < n; tr++) {
										setL21(
												tr,          // row > column 
												ftile_diag); // column
									}
								} else  {
									int nrow = (int) Math.floor(-0.5 + 0.5* Math.sqrt(1 + 8 * (ntile-1)));
									int ncol = (ntile-1) - (nrow * (nrow + 1) /2);
									int row = ftile_diag + nrow + 1;
									int col = ftile_diag + ncol + 1;
//									if ((row == (n-1)) && (col == (n-1))) {
//										System.out.println("calling setA22(): diag="+(ftile_diag-1)+", row="+row);
//									}
									setA22(
											ftile_diag-1, // int diag,// < col,  < row
											row, // int row, // >= col
											col); // int col) 							
								}
							}
						}
					};
				}		      
				ImageDtt.startAndJoin(threads);
			}
		}
	}
	
	
	// Cholesky-Banachiewicz Algorithm ?
	// Single-threaded
	public static double [][] cholesky_single (
			double [][] a,
			double [][] l) {
		  int n = a.length;
//		  double [][] l = new double[n][n];
		  for (int j = 0; j < n; j++) {
			  Arrays.fill(l[j], 0.0);
		  }
	      // Main loop.
	      for (int j = 0; j < n; j++) {
	         double[] Lrowj = l[j];
	         double d = 0.0;
	         for (int k = 0; k < j; k++) {
	            double[] Lrowk = l[k];
	            double s = 0.0;
	            for (int i = 0; i < k; i++) {
	               s += Lrowk[i]*Lrowj[i];
	            }
	            Lrowj[k] = s = (a[j][k] - s)/l[k][k];
	            d = d + s*s;
	         }
	         d = a[j][j] - d;
	         l[j][j] = Math.sqrt(Math.max(d,0.0));
	      }
	      return l;
	}
	
	
	
	

}
