/*  ------------------------ Projet BAORadio -------------------- 
    Classe  ForegroundCleaner
    R. Ansari , C. Magneville - Juin 2010 
---------------------------------------------------------------  */

#include "fgndsub.h"
#include "lobe.h"

#include "cubedef.h"
#include "matharr.h"
#include "poly.h"

#include "ctimer.h"

/* --Methode-- */
ForegroundCleaner::ForegroundCleaner(Four2DResponse& arrep, Four2DResponse& tbeam, TArray< TF >& skycube)
  : arrep_(arrep) , tbeam_(tbeam), skycube_(skycube)
{
  double dxdeg = ThetaSizeDegre/(double)NTheta;
  double dydeg = PhiSizeDegre/(double)NPhi;
  dx_ = DegreeToRadian(dxdeg);
  dy_ = DegreeToRadian(dydeg);
  freq0_ = Freq0MHz;
  dfreq_ = FreqSizeMHz/(double)NFreq;
  cout << " ForegroundCleaner: " << " dx=" << dxdeg << " dy=" << dydeg << " degres ( dx_rad=" << dx_ << " dy_rad=" << dy_ << ")" 
       << " Freq0=" << freq0_ << " deltaFreq=" << dfreq_ << " MHz" << endl;
  skycube.Show();
}

/* --Methode-- */
void ForegroundCleaner::BeamCorrections()
{
  BeamEffect beam(arrep_);
  beam.Correct2RefLobe(tbeam_, skycube_, dx_, dy_, freq0_, dfreq_);
  cout << " ForegroundCleaner::BeamCorrections() done " << endl;
}

/* --Methode-- */
int ForegroundCleaner::CleanNegatives(TF seuil)
{
  sa_size_t nneg = 0.;
  for(sa_size_t kz=0; kz<skycube_.SizeZ(); kz++) 
    for(sa_size_t ky=0; ky<skycube_.SizeY(); ky++) 
      for(sa_size_t kx=0; kx<skycube_.SizeX(); kx++) 
	if (skycube_(kx, ky, kz) < seuil)  {
	  nneg++; skycube_(kx, ky, kz)=seuil;
	}
  cout << " ForegroundCleaner::CleanNegatives " << nneg << " sky-pixels <" << seuil << " changed to" << seuil << endl;
  return (int)nneg;
}

/* --Methode-- */
int ForegroundCleaner::CleanPointSources(double nsigmas)
{
  Timer tm("ForegroundCleaner::CleanPointSources");
  TArray< TF > sky2d(skycube_.SizeX(), skycube_.SizeY());
  for(sa_size_t ky=0; ky<sky2d.SizeY(); ky++) 
    for(sa_size_t kx=0; kx<sky2d.SizeX(); kx++) 
      sky2d(kx, ky) = skycube_(Range(kx), Range(ky), Range::all()).Sum();


  double mean, sigma;
  

  TArray< TF > amz(1,1,skycube_.SizeZ()), asz(1,1,skycube_.SizeZ());
  for(sa_size_t kz=0; kz<skycube_.SizeZ(); kz++)  {
    TArray< TF > slice = skycube_(Range::all() , Range::all(), Range(kz));
    MeanSigma(slice, mean, sigma);
    amz(0,0,kz)=mean;  asz(0,0,kz)=sigma; 
  }

  MeanSigma(sky2d, mean, sigma);
  cout << " ForegroundCleaner::CleanPointSources 2D Sky projection, mean=" << mean << " sigma=" << sigma << endl;  
  
  TF seuil = (TF)(mean+nsigmas*sigma);
  
  sa_size_t srccnt=0;
  for(sa_size_t ky=0; ky<skycube_.SizeY(); ky++) 
    for(sa_size_t kx=0; kx<skycube_.SizeX(); kx++) {
      if (sky2d(kx,ky)>seuil) {
	srccnt++; 
	skycube_(Range(kx), Range(ky), Range::all()) = amz;
      }
    }
  
  cout << " Cleaned NSrc= " << srccnt << " 2D source/pixels (TotNPix=" << sky2d.Size() 
       << ")-> " << 100.*srccnt/sky2d.Size() <<  "% with S>" << seuil << " NSigmas=" << nsigmas << endl;
  return (int)srccnt;
}


/* --Methode-- */
TArray< TF >  ForegroundCleaner::extractLSSCubeP1(TArray< TF >& synctemp, TArray< TF >& specidx)
{
  Timer tm("ForegroundCleaner::extractLSSCubeP1");
// Inputs : maplss, mapsyc, freq0, dfreq
// Outputs : synctemp, specidx  (reconstructed foreground temperature and spectral index
// Return_Array : foreground subtracted LSS signal 
  sa_size_t sz[5];   sz[0]=skycube_.SizeX();  sz[1]=skycube_.SizeY();
  synctemp.SetSize(2, sz); 
  specidx.SetSize(2, sz);
  TArray<r_4> omap; 
  omap.SetSize(skycube_, true);
  Vector vlnf(skycube_.SizeZ());
  int nprt = 0;
  // double freq0 : Frequence premier index en k (MHz)
  // double dfreq :   // largeur en frequence de chaque plan (Mhz)  
  for(sa_size_t k=0; k<skycube_.SizeZ(); k++)  
    vlnf(k)=log((double)k*dfreq_+freq0_);

  sa_size_t nbinfini=0;
  sa_size_t nbbad=0;
  sa_size_t imodprt=skycube_.SizeX()/6;
  sa_size_t jmodprt=skycube_.SizeY()/6;
  for(sa_size_t i=0; i<skycube_.SizeX(); i++) 
    for(sa_size_t j=0; j<skycube_.SizeY(); j++)  {
      r_8 s1, sx, sx2, sy, sxy;
      s1=sx=sx2=sy=sxy=0.;
      for(sa_size_t k=0; k<skycube_.SizeZ(); k++)  {
	double lnf = vlnf(k);
	double ttot=(r_8)(skycube_(i,j,k));
	if (ttot < 1.e-5) continue;
	double lntt=log(ttot);
	s1+=1.;  sx+=lnf;  sx2+=(lnf*lnf);
	sy+=lntt;    sxy+=(lnf*lntt);
      }
      double beta = (sx*sxy-sx2*sy)/(sx*sx-s1*sx2);
      double alpha = (s1*sxy-sx*sy)/(s1*sx2-sx*sx);
      double T0 = exp(beta+alpha*vlnf(0));

      bool fgnan = false;
      if (!isfinite(alpha)||(!isfinite(beta))) {
	cout << "extractLSSCubeP1[" << i << "," << j << "]/ Not finite alpha, beta - (mapsync="
	     << skycube_(i,j,0) << " ... " << skycube_(i,j,skycube_.SizeZ()-1) << ")" << endl; 
	alpha=beta=-999.;
	fgnan = true;  nbinfini++;
      }
      else {
	double axp1 = beta+alpha*vlnf(0);
	double axp2 = beta+alpha*vlnf(vlnf.Size()-1);
	
	if ((axp1<-70.)||(axp1>70.)||(axp2<-70.)||(axp2>70.)) {
	cout << "extractLSSCubeP1[" << i << "," << j << "] BAD alpha=" << alpha 
	     << " beta=" << beta << " T0=" << T0 << " - (mapsync="
	     << skycube_(i,j,0) << " ... " << skycube_(i,j,skycube_.SizeZ()-1) << ")" << endl; 
	fgnan = true;  nbbad++;
	}
      }
      if ((i%imodprt==0)&&(j%jmodprt==0)) 
	cout << "extractLSSCubeP1[" << i << "," << j << "]: T0=" << T0 << " alpha=" << alpha 
	     << " (mapsync=" << skycube_(i,j,0) << " ... " << skycube_(i,j,skycube_.SizeZ()-1) << ")" << endl; 
      synctemp(i,j) = T0;
      specidx(i,j) = alpha;
      if (fgnan) {
	for(sa_size_t k=0; k<skycube_.SizeZ(); k++)   omap(i,j,k) = 0.;
      }
      else {
	for(sa_size_t k=0; k<skycube_.SizeZ(); k++) {
	  r_4 fittedtemp = (r_4)(exp(beta+alpha*vlnf(k)));
	  omap(i,j,k) = skycube_(i,j,k)-fittedtemp;
	}
      }
    }
  cout << " ForegroundCleaner::extractLSSCubeP1() - NbNan alpha/beta=" << nbinfini << " NbBAD =" << nbbad << endl;
  return omap;
}

/*
static inline val_polyn2(double alpha, double beta, double gamma, double x)
{
  return (beta+alpha*x+gamma*x*x);
}
*/

/* --Methode-- */
TArray< TF >  ForegroundCleaner::extractLSSCubeP2(TArray< TF >& synctemp, TArray< TF >& specidx)
{
  Timer tm("ForegroundCleaner::extractLSSCubeP2");
// Inputs : maplss, mapsyc, freq0, dfreq
// Outputs : synctemp, specidx  (reconstructed foreground temperature and spectral index
// Return_Array : foreground subtracted LSS signal 
  sa_size_t sz[5];   sz[0]=skycube_.SizeX();  sz[1]=skycube_.SizeY();
  synctemp.SetSize(2, sz); 
  specidx.SetSize(2, sz);
  TArray<r_4> omap; 
  omap.SetSize(skycube_, true);
  Vector vlnf(skycube_.SizeZ());
  Vector vlnT(skycube_.SizeZ());
  int nprt = 0;
  // double freq0 : Frequence premier index en k (MHz)
  // double dfreq :   // largeur en frequence de chaque plan (Mhz)  
  for(sa_size_t k=0; k<skycube_.SizeZ(); k++)  
    vlnf(k)=log((double)k*dfreq_+freq0_);
  vlnf -= vlnf(0);
  //  cout << " DBG*extractLSSCubeP2 vlnf(0)=" << vlnf(0) << " vlnf(1)=" << vlnf(1) 
  //       << "vlnf(last)=" << vlnf(vlnf.Size()-1) << endl;
  sa_size_t nbinfini=0;
  sa_size_t nbbad=0;
  sa_size_t imodprt=skycube_.SizeX()/6;
  sa_size_t jmodprt=skycube_.SizeY()/6;
  for(sa_size_t i=0; i<skycube_.SizeX(); i++) 
    for(sa_size_t j=0; j<skycube_.SizeY(); j++)  {
      vlnT = -12.;
      Poly polyn;
      for(sa_size_t k=0; k<skycube_.SizeZ(); k++)  {
	double lnf = vlnf(k);
	double ttot=(r_8)(skycube_(i,j,k));
	if (ttot < 1.e-5) continue;
	vlnT(k)=log(ttot);
      }
      polyn.Fit(vlnf,vlnT,2);
      double beta = polyn[0];
      double alpha = polyn[1];
      double gamma = polyn[2];
      double T0 = exp(polyn(vlnf(0))); // exp( val_polyn2(alpha, beta, gamma, vlnf(0)) ); 

      bool fgnan = false;
      if (!isfinite(alpha)||(!isfinite(beta))||(!isfinite(gamma))) {
	cout << "extractLSSCubeP2[" << i << "," << j << "]/ Not finite alpha, beta - (mapsync="
	     << skycube_(i,j,0) << " ... " << skycube_(i,j,skycube_.SizeZ()-1) << ")" << endl; 
	alpha=beta=gamma=-999.;
	fgnan = true;  nbinfini++;
      }
      else {
	double axp1 = polyn(vlnf(0));
	double axp2 = polyn(vlnf(vlnf.Size()-1));
	
	if ((axp1<-70.)||(axp1>70.)||(axp2<-70.)||(axp2>70.)) {
	  cout << "extractLSSCubeP2[" << i << "," << j << "] BAD alpha=" << alpha 
	    //	       << " beta=" << beta << " gamma=" << gamma << " T0=" << T0 << " axp1=" << axp1 << " axp2=" << axp2 
	       << " beta=" << beta << " gamma=" << gamma << " T0=" << T0 << " - (mapsync=" << skycube_(i,j,0) 
	       << " ... " << skycube_(i,j,skycube_.SizeZ()-1) << ")" << endl;
	  fgnan = true;  nbbad++;
	}
      }
      if ((i%imodprt==0)&&(j%jmodprt==0)) 
	cout << "extractLSSCubeP2[" << i << "," << j << "]: T0=" << T0 << " alpha=" << alpha << " gamma=" << gamma  
	     << " (mapsync=" << skycube_(i,j,0) << " ... " << skycube_(i,j,skycube_.SizeZ()-1) << ")" << endl; 
      synctemp(i,j) = T0;
      specidx(i,j) = alpha;
      if (fgnan) {
	for(sa_size_t k=0; k<skycube_.SizeZ(); k++)   omap(i,j,k) = 0.;
      }
      else {
	for(sa_size_t k=0; k<skycube_.SizeZ(); k++) {
	  r_4 fittedtemp = (r_4)( exp(polyn(vlnf(k))) );
	  omap(i,j,k) = skycube_(i,j,k)-fittedtemp;
	}
      }
    }
  cout << " ForegroundCleaner::extractLSSCubeP2() - NbNan alpha/beta=" << nbinfini << " NbBAD =" << nbbad << endl;
  return omap;
}