/* -*- mode: c++; c-basic-offset: 3 -*-
 *
 * PANDORE (PANTHEON Project)
 *
 * GREYC IMAGE
 * 6 Boulevard Maréchal Juin
 * F-14050 Caen Cedex France
 *
 * This file is free software. You can use it, distribute it
 * and/or modify it. However, the entire risk to the quality
 * and performance of this program is with you.
 *
 * Please email any bugs, comments, and/or additions to this file to:
 * pandore@greyc.ensicaen.fr
 * http://www.greyc.ensicaen.fr/EquipeImage/Pandore/
 */

/**
 * @author Jalal Fadili - 18-05-2005
 * @author Yannick Remy - 04-01-2006
 * @author Nael Ouedraogo - 04-01-2006
 */

/**
 * @file mca.ccpp
 *
 * Morphological Component Analysis of a 2D images (a matrix) using highly
 * redundant dictionaries.
 * The optimization pb is solved using a modified version of the BCR algorithm.
 * MCA2_bcr solves the following optimization problem
 *	(part_a,part_b) = argmin Sum_i || \Phi^+_i part_i ||_p + lambda * ||img - Sum_i part_i||_2^2
 *			p = 1 (l_1 norm: Soft thesholding as a solution).
 *			p = 0 (l_0 norm: difficult but approximated with a
 *                                Hard thresholding).
 *  Each component is supposed to be sparsely described in its corresponding
 * dictionary \Phi.
 *
 * La MCA ne traite que le cas des images reelles
 * en niveaux de gris a 2 dimensions.
 */ 

#define NB_ANGLES_COARSE 16

#ifdef _MSC_VER
#include <float.h>
#define isnormal(a) (!(_isnan((double)a) && (_finite((double)a))))
#endif

#undef MAIN
#include "dctutils.cpp"
#include "wutils.cpp"

namespace Utility{
#include "utility/pextractsubimage.cpp"
#include "utility/pinsertsubimage.cpp"
}
namespace Dwt{
#include "dwt.cpp"
#include "idwt.cpp"
#include "getsubband.cpp"
#include "setsubband.cpp"
}
namespace Udwt{
#include "udwt.cpp"
#include "iudwt.cpp"
}
namespace Isoudwt{
#include "isoudwt.cpp"
#include "isoiudwt.cpp"
}
namespace Dct{
#include "dct.cpp"
#include "idct.cpp"
#include "ldct.cpp"
#include "ildct.cpp"
}
namespace Dst{
#include "dst.cpp"
#include "idst.cpp"
#include "ldst.cpp"
#include "ildst.cpp"
}
namespace Dctiv{
#include "dctiv.cpp"
#include "ldctiv.cpp"
#include "ildctiv.cpp"
}
namespace Fct{
#include "fct.cpp"
#include "fict.cpp"
}

namespace Arithm{
#include "arithmetic/padd.cpp"
#include "arithmetic/psub.cpp"
}

namespace Stat{
#include "imagefeatureextraction/pmedianvalue.cpp"
}

namespace Fft{
#include "frequency/pfft.cpp"
}

#include "fat.cpp"
#include "fst.cpp"
#include "sumbands.cpp"
#include "mca_threshold.cpp"

/*****************************************************************
* Compare function for qsort().
*/
int compare( const float* val1, const float* val2 ) {
	if (*val1>*val2) return(1);
	if (*val1<*val2) return(-1);
	if (*val1==*val2) return(0);

	return FAILURE;
}

/**
 * Estimate the noise std by the MAD robust estimator.
 */
Float mad( Img2dsf & ims ) {
	Long x,y;
	Float median;
	median = Stat::PMedianValue(ims);
	for (y = 0; y< ims.Height() ;y++)
		for (x = 0; x< ims.Width() ;x++)
			ims[y][x] = (Float)fabs(ims[y][x] - median);
	return Stat::PMedianValue(ims);
}

/**
 * Desallocate a FCT collection
 */
Errc desallocFCTCollec(Collection & collec) {
   Collection ** img = collec.GETPARRAY("curvelets",Collection);
   Img2dsl ** curv_size = collec.GETPARRAY("curv_size",Img2dsl);
   Long nbs = collec.GETVALUE("nbscales",Long);
   Img1dsl * nbangles = (Img1dsl *)collec.GETPOBJECT("nbangles",Img1dsl);
   Img2dsf** curvdata_tmp=NULL;
   
   for (int it = 0; it < nbs; it++) {
     curvdata_tmp=img[it]->GETPARRAY("curvangles",Img2dsf);
     for (int j=0; j < (*nbangles)[it]; j++) curvdata_tmp[j]->Delete();
     delete[] curvdata_tmp;
     curv_size[it]->Delete();
   }	      			   
   delete[] img;
   delete[] curv_size;
   nbangles->Delete();
   collec.Delete();   
   return SUCCESS;
}

/**
 * Estimate the norm of curvelets
 */
Errc computeL2Norm( Img2dsf &ims, Dictionary dict, Float** & E ) {
   Long nbscales =  (Long)dict.GetPar1();
   Long nbangles_coarse = NB_ANGLES_COARSE;

   Img1dsl *nbangles = new Img1dsl(nbscales);
   // Evaluate the number of angles for each scale
   for (Long i=0; i<nbscales-1; i++) {
      (*nbangles)[i]=nbangles_coarse*(Long)pow(2,(ceil((nbscales-(i+2))/2.0)));
   }
   (*nbangles)[nbscales-1]=1;  

   Long Height = ims.Height();
   Long Width = ims.Width();   

   Img2dsf * im1 = new Img2dsf(Height, Width);
   Img2dsf * im2 = new Img2dsf(Height, Width);
   Collection *res = new Collection();
   Collection ** les_curvelets;
   Img2dsl ** curv_size;
   Img2dsf ** curvdata_tmp=NULL;
   
   E = new Float*[nbscales];
   for(Long i = 0; i < nbscales; i++) {
      E[i] = new Float[(*nbangles)[i]];
   }

   (*im1) = (float)(1.0F / sqrt((Float)(Height * Width)));
   (*im2) = 0.0F;
   
   // im1 is a Dirac
   Fft::PFFT(*im1,*im2,*im1,*im2); 
   Fct::FCT(*im1, *res, nbscales, nbangles_coarse);
   
   im1->Delete();
   im2->Delete();

   les_curvelets = res->GETPARRAY("curvelets",Collection);
   curv_size = res->GETPARRAY("curv_size", Img2dsl);
   
   for (Long i = 0; i < nbscales-1 ; i++) {
      curvdata_tmp=les_curvelets[i]->GETPARRAY("curvangles",Img2dsf);
      for (Long j=0; j < (*nbangles)[i]; j++) {
	 E[i][j] = 0; 
	 for (Long k=0; k < (*curv_size[i])[j][0]; k++) {
	    for (Long l=0; l < (*curv_size[i])[j][1]; l++) {
	       E[i][j] += pow((*(curvdata_tmp[j]))[k][l], 2);
	    }
	 }
	 E[i][j] /= (*curv_size[i])[j][1] * (*curv_size[i])[j][0];
	 E[i][j] = sqrt(E[i][j]);
	 if (!isnormal(E[i][j]))
	 	std::cout<<"Something wrong happening, NaN, Inf or 0 l_2 norm"<<std::endl;
      }      
   }
   desallocFCTCollec(*res);
   return SUCCESS;
}

/**
 * Get the starting point.
 */
Float getstartpoint( Img2dsf & ims, Dictionary & dict, Float ** & E ) {
	Long i;
	Long height   = ims.Height();
	Long width    = ims.Width();
	Long nbofdict = dict.GetNbTransforms();
	
	Float *maxcoeffs = new Float[nbofdict];//maxcoeff[i] is the max coeff of the transform i
	Float mincoeffs;
	Char  *nameofdict;
	Imx3dsf *coeffs;
	Collection* coeffsCol = new Collection;
	for (i = 0, mincoeffs = MAXFLOAT; i < nbofdict; i++) {
		nameofdict = dict[i].GetNameofTransform();
		FAT(ims,dict[i],*coeffsCol);
		if (!coeffsCol->Exists("Collection")) {
			// FAT returns a collection which containd an Image
			coeffs = (Imx3dsf *)coeffsCol->GETPOBJECT("Image",Imx3dsf);
			if (!strcmp(nameofdict,"UDWT") || !strcmp(nameofdict,"ISOUDWT")) {
				maxcoeffs[i] = maxabs((Float*)(coeffs->Vector(0))+height*width,coeffs->VectorSize()-height*width);
			} else {
				maxcoeffs[i] = maxabs(coeffs->Vector(0),coeffs->VectorSize());	       
			}
			coeffs->Delete();
		} else {
			// FAT returns a collection which contains a Collection
			if (!strcmp(nameofdict,"FCT")) {
				Collection ** les_curvelets;
				Collection * collec = (Collection *)(coeffsCol->GETPOBJECT("Collection",Collection));
				Long nbscales = collec->GETVALUE("nbscales", Long);
				Img1dsl* nbangles = (Img1dsl*)(collec->GETPOBJECT("nbangles",Img1dsl));
				Float max =MINFLOAT; 
				les_curvelets = collec->GETPARRAY("curvelets",Collection);
				Img2dsf**  curvdata_tmp=NULL;
				for (int j=0; j<nbscales-1; j++) {			    
					// Calculer le max des coeffs
					curvdata_tmp=les_curvelets[j]->GETPARRAY("curvangles",Img2dsf);
					for (int k=0;k<(*nbangles)[j]; k++) {
						//max =MINFLOAT;
						for (int m=0; m < curvdata_tmp[k]->Width(); m++) {
							for (int n=0; n < curvdata_tmp[k]->Height(); n++) {
								Float tmp = fabs((*(curvdata_tmp[k]))[n][m] / E[j][k]);
								max = max < tmp ? tmp : max;
							}
						}
						//std::cout<<"max["<<j<<"]["<<k<<"]= "<<max<<std::endl;
					}
				}
				maxcoeffs[i] = max;			
				// Desallocation
				desallocFCTCollec(*collec);
			}		      
		}
		coeffsCol->Delete();
		std::cout<<"Max : "<<maxcoeffs[i]<< std::endl;
	}
	//Le coeff choisi est la deuxieme plus grande valeur dans maxcoeffs
	if (nbofdict==1)
		mincoeffs = maxcoeffs[0];
	else {
		qsort((void*)maxcoeffs,nbofdict,sizeof(float),(int (*)(const void*,const void*))compare);
		mincoeffs = maxcoeffs[nbofdict-2];
	}
	delete []maxcoeffs;
	return mincoeffs;	
}

/*
 * Apply the total-variation constraint step 
 * with \gamma as the regularization parameter.
 * In 1D the TV norm is equivalent to the l_1 norm of the haar-wavelet coefficients.
 * This approximation is used in 2D although not rigorous. 
 * This step is then almost equivalent to soft-thresholding the haar detail coefficients
 */
Errc TVCorrection(Img2dsf &ims,Dictionary & dict,Float gamma) {
	Long n  =  ims.Width();
	Long m  =  ims.Height();
	Long d;
	Img3dsf wc(4,m,n);
	Img2dsf temp;
	Char* nameofdict = dict.GetNameofTransform();
	
	// Apply only to BV components.
	if (!strcmp(nameofdict,"DWT") || !strcmp(nameofdict,"UDWT") || !strcmp(nameofdict,"ISOUDWT") || !strcmp(nameofdict,"FCT")){
		// Create the haar qmf.
		Qmf qmf("haar",1);
		Udwt::UDWT(ims, qmf, wc, 1);
		for (d = 1;d <= 3; d++) {
			temp.New(m,n,&wc[d][0][0]);	
			Thresh(temp,temp,gamma,Soft);
			temp.Delete();
		}
		Udwt::IUDWT(wc, qmf, ims, 1);
	}
	
	wc.Delete();
	return SUCCESS;
}

/**
 * MCA function.
 */
Errc MCA(Img2dsf & ims, Dictionary& dict, Ulong itermax, int cartoon_part, Float gamma, int expdecrease, Float stop, Img3dsf & imd) {
   Long i,iter,nb;
   Long n = ims.Width();
   Long m = ims.Height();
   Float sigma, stopcriterion, deltamax, delta, lambda;
	
   Img2dsf imest(ims.Props());
   Img2dsf residual(ims.Props());
   Img2dsf Rnb(ims.Props());
   Img2dsf partnb;
	
   //For Fat,fst,mca_threshold operator.
   Collection* tempCol = new Collection;
   // norm of Curvelets. Only use for FCT
   Float ** E = 0;

	
   //We display the dictionary.
   std::cout<<"MCA with "<<dict.GetNbTransforms()<<" dictionaries:"<<std::endl<<dict;
	
   //Initializations.
   int J = nextpow2(MIN(m,n));

   // To estimate the WGN standard deviation using the MAD.
   // The qmf is quite arbitrary, sufficiently regular to approach a good band-pass.
   Qmf qmf("daubechies",4);
   Img2dsf wc(m,n),hh(m>>1,n>>1);
   Dwt::DWT(ims,qmf,wc,1);
   Dwt::GetSubband(wc,hh,m>>1,n>>1);
   sigma = (Float)(mad(hh)/.6745F);
   stopcriterion = stop*sigma;
   wc.Delete();
   hh.Delete();
	
   // If FCT is in the dictionary we compute the L2norm of curvelets
   for(Long nbtrans=0; nbtrans<(Long)dict.GetNbTransforms(); nbtrans++) {
      if(!strcmp(dict[nbtrans].GetNameofTransform(),"FCT")) {
	 //dict[nbtrans].Par1++; // To account for number of scales compared to CurveLab enumeration.
	 computeL2Norm(ims, dict[nbtrans], E);
      }
   }

   // Calculate the starting thd, which is the minimum of maximal coefficients 
   // of the image in each transform of the dictionary.
   deltamax = getstartpoint(ims,dict, E);

   delta    = deltamax;
	
   if(expdecrease)	
      // Exponential decrease.
      lambda = (Float)pow((double)deltamax/stopcriterion,1./(1.0-itermax)); 
   else	
      // Slope of the linear decrease.
      lambda = (deltamax-stopcriterion)/(itermax-1);
	
   std::cout<<stopcriterion<<" "<<deltamax<<" "<<lambda<<std::endl;

   // Start the modified Block Relaxation Algorithm.
   for(iter=0;iter<(int)itermax;iter++, std::cout<<"\r          \r"<<(100*(Float)iter/itermax)<<"%"<<std::flush){
      for( i = 1 ; i <= J+1; i++){
	 // Calculate the residual image.
	 SumBands(imd,imest);
	 Arithm::PSub(ims,imest,residual);
	 for(nb = 0; nb < (Long)dict.GetNbTransforms(); nb++){

	    // Update Part nb assuming other parts fixed.
	    // Solve for Part nb the marginal penalized minimization problem 
	    // (l_0 penalization -> Hard thesholding, l_1 -> Soft).
	    partnb.New(m,n,&imd[nb][0][0]);
	    Arithm::PAdd(partnb,residual,Rnb);
	    FAT(Rnb,dict[nb],*tempCol);
	    NormThresh(*tempCol,dict[nb],delta,E,Hard);
	    FST(*tempCol,dict[nb],partnb);
	    if(gamma && cartoon_part && nb==cartoon_part-1)
	       TVCorrection(partnb,dict[nb],gamma);
	    partnb.Delete();
	    
	    /* Desallocation */
	    if(tempCol->Exists("Collection")) {
	       /* Desallocation of the Colletion in tempCol*/
	       Collection * collec = (Collection *)tempCol->GETPOBJECT("Collection",Collection);
	       desallocFCTCollec(*collec);
	    }
	    else {
	       /* Desallocation of the Image in tempCol */
	       if(!strcmp(dict[nb].GetNameofTransform(),"UDWT") || !strcmp(dict[nb].GetNameofTransform(),"ISOUDWT")) {
		  Img3dsf * img = (Img3dsf *)tempCol->GETPOBJECT("Image", Img3dsf);
		  img->Delete();
	       }			   
	       else {
		  Img2dsf * img = (Img2dsf *)tempCol->GETPOBJECT("Image", Img2dsf);
		  img->Delete();
	       }
	    }
	    tempCol->Delete();
	 }
      }
	  
      // Update the regularization parameter delta.
      if(expdecrease)	
	 delta *= lambda; // Exponential decrease.
      else
	 delta -= lambda; // Linear decrease.
   }

   /*for(Long nbtrans=0; nbtrans<(Long)dict.GetNbTransforms(); nbtrans++) {
     if(!strcmp(dict[nbtrans].GetNameofTransform(),"FCT")) {
     dict[nbtrans].Par1--; // To account for number of scales compared to CurveLab enumeration.
     }
     }*/

   imest.Delete();
   residual.Delete();
   Rnb.Delete();
   return SUCCESS;
}
  
#define USAGE   "USAGE : %s itermax cartoon gamma expdecrease stop [dict_in|-] [im_in|-] [im_out|-]"
#define PARC    5
#define FINC    2
#define FOUTC   1
#define MASK    0

int main(int argc, char* argv[]) {
   Errc result;              // The result code of the execution.
   Pobject *mask;            // The mask.
   Pobject *objin[FINC+1];   // The input objects;
   Pobject *objs[FINC+1];    // The source objects masked by mask.
   Pobject *objout[FOUTC+1]; // The output object unmasked by mask.
   Pobject *objd[FOUTC+1];   // The result object.
   char *parv[PARC+1];       // The input parameters.

   ReadArgs(argc,argv,PARC,FINC,FOUTC,&mask,objin,objs,objout,objd,parv,USAGE,MASK);
   
   if(objs[1]->Type()== Po_Img2dsf && objs[0]->Type()==Po_Collection) {
      Img2dsf* const ims=(Img2dsf*)objs[1];	
      Collection* const colin=(Collection *) objs[0];
      Dictionary dict(*colin);
      Check(*ims);
      if((unsigned int)atoi(parv[1]) > dict.GetNbTransforms()){
      		fprintf(stderr,"ERROR: The cartoon part index must be >=1 and <= %d\n",dict.GetNbTransforms());
		Exit(FAILURE);
      }
      objd[0] = new Img3dsf( dict.GetNbTransforms() ,ims->Height(), ims->Width());
      Img3dsf* const imd = (Img3dsf*)objd[0];
      *imd = 0;
      result = MCA(*ims, dict, atoi(parv[0]), atoi(parv[1]), (Float)atof(parv[2]), atoi(parv[3]), (Float)atof(parv[4]), *imd);
   } else {
      fprintf(stderr,"ERROR: Pandore type not supported by this operator\n");
      result =FAILURE;
   }
   
   WriteArgs(argc,argv,PARC,FINC,FOUTC,&mask,objin,objs,objout,objd,MASK);
   Exit(result);
   return 0;
}
