/* -*- 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 - 02-02-2005
 * @author Jalal Fadili - 02-07-2006
 */

/**
 * @file emipaint.ccpp
 *
 * EM_Inpaint: Bayesian Inpainting of 2D images (a matrix) using redundant dictionaries.
 *	   The optimization pb is solved using the EM algorithm.
 *	   EM_Inpaint solves the following (MAP) optimization problem
 *		(part) = argmin lambda * Psi(\Phi^+ part) +  ||img - M \Phi coeff||_2^2/(2 sigma^2)
 *				M is the mask.
 *				Phi is the dictionary.
 *				Psi is any penalty on the dictionary coefficients
 *                              (hopefully strictly convex function).
 *	   The image is supposed to be sparsely described in the dictionary \Phi.
 *  
 *  Inputs:
 *    img	     	2D image matrix, nxn, n = 2^J
 *    dict		Dictionary collection (see directory for available transforms).
 *    lambda		The regularisation parameter is either fixed or decreasing (deterministic annealing EM).
 *    sigma		Standard deviation of the noise (can be estimated and updated automatically).
 *    itermax		Nb of iterations.
 *    emtype		Type of EM algorithm (0: regular, 1: deterministic annealing, 2: Starck heuristic). 
 *    mask		The binary mask (image of ones and zeros).
 *  Outputs:
 *    part 		Estimated inpainted image (nxn image)
 *
 *  Description
 *    The dictionaries and their parameters can be built using the dictionary operator.
 * L'inpainting 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 Arithm{
#include "arithmetic/padd.cpp"
#include "arithmetic/psub.cpp"
#include "transformation/pzoom.cpp"
#include "arithmetic/pmult.cpp"
}
namespace Stat{
#include "imagefeatureextraction/psumvalue.cpp"
#include "imagefeatureextraction/penergyvalue.cpp"
#include "imagefeatureextraction/pvaluenumber.cpp"
#include "imagefeatureextraction/pmedianvalue.cpp"
}
namespace Dwt{
#include "dwt.cpp"
#include "idwt.cpp"
#include "getsubband.cpp"
#include "setsubband.cpp"
}
namespace Utility{
#include "utility/pextractsubimage.cpp"
#include "utility/pinsertsubimage.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 Fft{
#include "frequency/pfft.cpp"
}
namespace Fct{
#include "fct.cpp"
#include "fict.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.
 * Only non-occluded pixels are taken into account.
 */
float mad(float* data,int n,float* mask,int nmask) {
   int x,i;
   float median;
   float* tmp = new float[nmask];
      
   for(x=0,i=0;x<n;x++){
      if(mask[x]) tmp[i++] = data[x];
   }

   data=tmp;
   n=nmask;
   
   qsort((void*)data,n,sizeof(float),(int (*)(const void *,const void*))compare);
      
   if(n%2) // Odd number of elements.
      median=(data[(n+1)/2-1]);
   else // Even number of elements.
      median=((data[n/2-1] + data[n/2])/2);
   
   for(x=0;x<n;x++)
      data[x]=(float)fabs((double)(data[x]-median));
	
   qsort(data,n,sizeof(float),(int (*)(const void *,const void*))compare);
   
   if(n%2) // Odd number of elements.
      return(data[(n+1)/2-1]);
   else // Even number of elements.
      return((data[n/2-1] + data[n/2])/2);
   
}
/**
 * Desallocate a FCT collection
 */
Errc desalocFCTCollec(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) = 1.0F / sqrt((Float)(Height * Width));
   (*im2) = 0;
   
   // 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;
      }      
   }
   desalocFCTCollec(*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
	       desalocFCTCollec(*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;
}

/**
 * Error Change Measure in HS-norm sense.
 */
Double ErrorChange(Img3dsf& imparts,Img2dsf& imest) {
   Img2dsf imsum(imest.Props());
   Img2dsf imnum(imest.Props());

   SumBands(imparts,imsum);

   Arithm::PSub(imest,imsum,imnum);

   //printf("%.20f\n",Stat::PEnergyValue(imnum)/Stat::Energie(imest));
	
   return(Stat::PEnergyValue(imnum));
}

/**
 * MCA function.
 */
Errc EM_Inpaint(Img2dsf& ims, Dictionary& dict, Img2dsf& imsmask, Ulong itermax, Float epschange, Float lambda, Float sigma, ThdType thdtype, int emtype, Img2dsf& imd) {
   int n = ims.Width();
   int m = ims.Height();
   int nbofdict,t,iter,nb;
   float deltamax;
   Img1dsf delta(itermax);
   Img1dsf T(itermax);
   Img2dsf residual(ims.Props());
   Img2dsf Rnb(ims.Props());
   Img2dsf imtmp(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.
   nbofdict = dict.GetNbTransforms();
   Img3dsf imparts(nbofdict,m,n);
   memcpy(imparts.Vector(),ims.Vector(),m*n*sizeof(Float));
   std::cout<<"Inpaint with "<<nbofdict<<" transforms:"<<std::endl<<dict;	
	
   //Initializations.
   //int J=nextpow2(MIN(m,n));
   int m0 = n*m - Stat::PValueNumber(imsmask); // Number of missing samples.
	
   if(!sigma){
      // 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);
      Img2dsf maskred(m>>1,n>>1);
      Dwt::DWT(imsmask, qmf, wc, 1);
      Dwt::GetSubband(wc, maskred, m>>1,n>>1);
      sigma = (Float)(mad(hh.Vector(),hh.VectorSize(),maskred.Vector(),Stat::PValueNumber(maskred))/.6745F);
      wc.Delete();
      hh.Delete();
      maskred.Delete();
   }
	
   // If FCT is in the dictionary we compute the L2norm of curvelets
   for(Long nbtrans=0; nbtrans<(Long)nbofdict; nbtrans++) {
      if(!strcmp(dict[nbtrans].GetNameofTransform(),"FCT")) {
	 dict[nbtrans].Par1++;
	 computeL2Norm(ims, dict[nbtrans], E);
      }
   }
	
   switch(emtype){
   case 0:
      std::cout<<"ECM algorithm with std "<<sigma<<std::endl;
      delta=lambda;
      //for(iter=0;iter<(int)itermax;iter++) 
      //std::cout<<lambda<<" "<<sigma<<" "<<delta[i]<<" ";
      break;
   case 1:
      std::cout<<"ECM algorithm with std estimation"<<std::endl;
      // Calculate detreministic annealing cooling schedule (exponential).
      deltamax = getstartpoint(ims,dict,E);
      std::cout<<"deltmax="<<deltamax<<" "<<"lambda="<<lambda<<" "<<"stdinit="<<sigma<<std::endl;
      delta=lambda;
      for(iter=0;iter<(int)itermax;iter++){
	 T[iter] = pow(deltamax*deltamax/(lambda*lambda*sigma*sigma),1-(float)iter/(float)(itermax-1)); //Exponential decrease.
	 //std::cout<<delta[iter]<<" "<<sigma<<" "<<T[iter]<<" ";
      }
      sigma = sqrt(T[0])*sigma;
      break;
   case 2:
      std::cout<<"Deterministic Annealing ECM algorithm"<<std::endl;
      // Calculate the starting thd, which is the minimum of maximal coefficients in magnitude
      // of the image in each transform of the dictionary.
      deltamax = getstartpoint(ims,dict,E);
      std::cout<<"deltmax="<<deltamax<<" "<<"lambda="<<lambda<<" "<<"std="<<sigma<<std::endl;
      for(iter=0;iter<(int)itermax;iter++){
	 delta[iter] = (-(deltamax-lambda*sigma)*(float)iter/(float)(itermax-1) + deltamax)/sigma; //Linear decrease.
	 //std::cout<<lambda<<" "<<sigma<<" "<<delta[iter]<<" ";
      }
      break;
   }
		
   // Start the Cyclic EM (ECM) Algorithm.
   for(iter=0,itermax=(emtype==2 ? itermax:1);iter<(int)itermax;iter++, std::cout<<"\r          \r"<<(100*(float)iter/(float)itermax)<<"%"<<std::flush){
      t=1;
      do{
	 //sigma = (emtype==1 ? sqrt(T[iter])*sigma:sigma);
	 //for(t=1;t<(int)itermax;t++){
	 for(nb=0;nb<nbofdict;nb++){
	    // E step: estimate the missing observations: 
	    // Y^(t)_i =   Yobs_i	if observation i is available 
	    // or Ymiss_i  		if observation is missing; Ymiss(t)=(1-M) Phi coeff^(t)
	    // Calculate the reconstructed residual image.
	    SumBands(imparts,imd);
	    Arithm::PMult(imd,imsmask,imtmp);
	    Arithm::PSub(ims,imtmp,residual);
		   
	      
	    // M step: solve for part the penalised least-squares problem conditionally on Y^(t), Phi and sigma^2(t)
	    // min lambda * ||coeff||_p +  ||Y^(t) - \Phi coeff||_2^2/(2 sigma(t)^2) -> p=1 ==> Soft thresholding.
	    partnb.New(m,n,&imparts[nb][0][0]);
	    Arithm::PAdd(partnb,residual,Rnb);
	    FAT(Rnb,dict[nb],*tempCol);
	    NormThresh(*tempCol,dict[nb],delta[iter]*sigma,E,thdtype);
	    FST(*tempCol,dict[nb],partnb);
	    partnb.Delete();
		   
	    /* Desallocation */
	    if(tempCol->Exists("Collection")) {
	       /* Desallocation of the Colletion in tempCol*/
	       Collection * collec = (Collection *)tempCol->GETPOBJECT("Collection",Collection);
	       desalocFCTCollec(*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();
	 }
	 // M step: update the variance for ECM
	 // sigma  = sqrt(m0*sigma^2 + sum(mask(:).*(img(:)-parts(:)).^2))/n;
	 if(emtype==1){
	    Arithm::PSub(ims,imd,imtmp);
	    Arithm::PMult(imtmp,imtmp,imtmp);
	    Arithm::PMult(imtmp,imsmask,imtmp);
	    sigma = (Float)(sqrt(m0*sigma*sigma + Stat::PSumValue(imtmp))/n);
	    //std::cout<<sigma<<" "<<std::flush;
	 }
		  	
	 //std::cout<<delta[iter]<<std::endl;
	 //getchar();
	  
      }
      while((++t)<=1000 && ErrorChange(imparts,imd)>=epschange);
   }
	
   imtmp.Delete();
   imparts.Delete();
   residual.Delete();
   Rnb.Delete();
		
   return SUCCESS;
}

#define USAGE   "USAGE : %s itermax epsilon lambda sigma thdtype emtype [dict_in] [im_in] [im_mask] [im_out]\n"
#define PARC    6
#define FINC    3
#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[0]->Type()==Po_Collection && objs[1]->Type()==Po_Img2dsf &&  objs[2]->Type()==Po_Img2dsf){
      Collection* const colin=(Collection *) objs[0];
      Dictionary dict(*colin);
      Img2dsf* const ims = (Img2dsf*)objs[1];	
      Img2dsf* const imsmask = (Img2dsf*)objs[2];
      
      if ((ims->Height() == imsmask->Height()) && (ims->Width() == imsmask->Width())) {
	 Check(*ims);
	 objd[0]=new Img2dsf(ims->Size());
	 Img2dsf* const imd = (Img2dsf*)objd[0];
	 
	 ThdType thdtype=Hard;
	 
	 if (atoi(parv[4])==1) thdtype=Hard;
	 else{
	    if (atoi(parv[4])==0) thdtype=Soft;
	    else {
	       fprintf(stderr,"ERROR: threshold type must be either 1 (hard) or 0 (soft)\n");
	       result =FAILURE;
	       Exit(result);
	    }
	 }
	 
	 result=EM_Inpaint(*ims, dict, *imsmask, atoi(parv[0]), (float)atof(parv[1]), (float)atof(parv[2]), (float)atof(parv[3]), thdtype, atoi(parv[5]), *imd);
      }	
      else{
	 fprintf(stderr,"ERROR: Input images must have the same size\n");
	 result =FAILURE;
      }
      
   }
   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;
}
