/*  ------------------------ Projet BAORadio -------------------- 
    Classes to compute 3D power spectrum and noise power spectrum
    R. Ansari - Nov 2008 ... Dec 2010 
---------------------------------------------------------------  */

#include "specpk.h"
#include "randr48.h"      
#include "ctimer.h"      

//------------------------------------
// Class SpectralShape 
// -----------------------------------

double Pnu1(double nu) 
{
  return ( sqrt(sqrt(nu)) / ((nu+1.0)/0.2) * 
	   (1+0.2*cos(2*M_PI*(nu-2.)*0.15)*exp(-nu/50.)) );
}

double Pnu2(double nu) 
{
  if (nu < 1.e-9) return 0.;
  return ((1.-exp(-nu/0.5))/nu*(1+0.25*cos(2*M_PI*nu*0.1)*exp(-nu/20.)) );
}


double Pnu3(double nu) 
{
  return ( log(nu/100.+1)*(1+sin(2*M_PI*nu/300))*exp(-nu/4000) );
}


double Pnu4(double nu) 
{
  double x = (nu-0.5)/0.05;
  double rc = 2*exp(-x*x);
  x = (nu-3.1)/0.27;
  rc += exp(-x*x);
  x = (nu-7.6)/1.4;
  rc += 0.5*exp(-x*x);
  return ( rc+2.*exp(-x*x) );
}

//--------------------------------------------------
// -- SpectralShape class : test P(k) class
//--------------------------------------------------
// Constructor
SpectralShape::SpectralShape(int typ)
{
  typ_=typ;
  SetRenormFac();
}

// Return the spectral power for a given wave number wk 
double SpectralShape::operator() (double wk)
{
  wk/=DeuxPI;
  double retv=1.;
  switch (typ_) 
    {
    case 1:
      retv=Pnu1(wk);
      break;
    case 2:
      retv=Pnu2(wk);
      break;
    case 3:
      retv=Pnu3(wk);
      break;
    case 4:
      retv=Pnu4(wk);
      break;
    default :
      {
  // global shape
      double csp = pow( (2*sin(sqrt(sqrt(wk/7.)))),2.);
      if (csp < 0.) return 0.;
      
      // Adding some pics
      double picpos[5] = {75.,150.,225.,300.,375.,};
      
      for(int k=0; k<5; k++) {
  	double x0 = picpos[k];
  	if ( (wk > x0-25.) && (wk < x0+25.) ) {
	  double x = (wk-x0);
	  csp *= (1.+0.5*exp(-(x*x)/(2.*5*5)));
	  break;
  	}
      }
      retv=csp;
      }
      break;
    }
  return retv*renorm_fac;
}
// Return a vector representing the power spectrum (for checking) 
Histo SpectralShape::GetPk(int n)
{
  if (n<16) n = 256;
  Histo h(0.,1024.*DeuxPI,n);
  for(int k=0; k<h.NBins(); k++)   h(k) = Value((k+0.5)*h.BinWidth());
  return h;	
}

double SpectralShape::Sommek2Pk(double kmax, int n)
{
  double dk=kmax/(double)n;
  double s=0.;
  for(int i=1; i<=n; i++) {
    double ck=(double)i*dk;
    s += Value(ck)*ck*ck;
  }
  return s*dk*4.*M_PI;
}
//--------------------------------------------------
// -- Four2DResponse class : test P(k) class

//---------------------------------------------------------------
// -- Four3DPk class :  3D fourier amplitudes and power spectrum 
//---------------------------------------------------------------
// Constructeur avec Tableau des coeff. de Fourier en argument
Four3DPk::Four3DPk(TArray< complex<TF> > & fourcoedd, RandomGeneratorInterface& rg)
  : rg_(rg), fourAmp(fourcoedd)
{
  SetPrtLevel();
  SetCellSize();
}
// Constructor
Four3DPk::Four3DPk(RandomGeneratorInterface& rg, sa_size_t szx, sa_size_t szy, sa_size_t szz)
  : rg_(rg), fourAmp(szx, szy, szz) 
{
  SetPrtLevel();
  SetCellSize();
}


// Generate mass field Fourier Coefficient
void Four3DPk::ComputeFourierAmp(SpectralShape& pk)
{
  // We generate a random gaussian real field  
  // fourAmp represent 3-D fourier transform of a real input array. 
  // The second half of the array along Y and Z contain negative frequencies
  //  double fnorm = 1./sqrt(2.*fourAmp.Size()); 
  double fnorm = 1.; 
  double kxx, kyy, kzz;
  // sa_size_t is large integer type  
  for(sa_size_t kz=0; kz<fourAmp.SizeZ(); kz++) {
    kzz =  (kz>fourAmp.SizeZ()/2) ? (double)(fourAmp.SizeZ()-kz)*dkz_ : (double)kz*dkz_; 
    for(sa_size_t ky=0; ky<fourAmp.SizeY(); ky++) {
      kyy =  (ky>fourAmp.SizeY()/2) ? (double)(fourAmp.SizeY()-ky)*dky_ : (double)ky*dky_; 
      for(sa_size_t kx=0; kx<fourAmp.SizeX(); kx++) {
	double kxx=(double)kx*dkx_;
	double wk = sqrt(kxx*kxx+kyy*kyy+kzz*kzz);
	double amp = sqrt(pk(wk)*fnorm/2.);      
	fourAmp(kx, ky, kz) = complex<TF>(rg_.Gaussian(amp), rg_.Gaussian(amp));   // renormalize fourier coeff usin 
      }
    }
  }
  if (prtlev_>2)
    cout << " Four3DPk::ComputeFourierAmp() done ..." << endl;
}

// Generate mass field Fourier Coefficient
void Four3DPk::ComputeNoiseFourierAmp(Four2DResponse& resp, bool crmask)
{
  TMatrix<r_4> mask(fourAmp.SizeY(), fourAmp.SizeX());
  // fourAmp represent 3-D fourier transform of a real input array. 
  // The second half of the array along Y and Z contain negative frequencies
  double kxx, kyy, kzz, rep, amp;
  // sa_size_t is large integer type  
  for(sa_size_t kz=0; kz<fourAmp.SizeZ(); kz++) {
    kzz =  (kz>fourAmp.SizeZ()/2) ? -(double)(fourAmp.SizeZ()-kz)*dkz_ : (double)kz*dkz_; 
    for(sa_size_t ky=0; ky<fourAmp.SizeY(); ky++) {
      kyy =  (ky>fourAmp.SizeY()/2) ? -(double)(fourAmp.SizeY()-ky)*dky_ : (double)ky*dky_; 
      for(sa_size_t kx=0; kx<fourAmp.SizeX(); kx++) {
	kxx=(double)kx*dkx_;
	rep = resp(kxx, kyy);
	if (crmask&&(kz==0))  mask(ky,kx)=((rep<1.e-8)?9.e9:(1./rep));
	if (rep<1.e-8)  fourAmp(kx, ky, kz) = complex<TF>(9.e9,0.);
	else {
	  amp = 1./sqrt(rep)/sqrt(2.);
	  fourAmp(kx, ky, kz) = complex<TF>(rg_.Gaussian(amp), rg_.Gaussian(amp));   
	}
      }
    }
  }
  if (prtlev_>2)  fourAmp.Show();
  if (crmask) {
    POutPersist po("mask.ppf");
    po << mask;
  }
  if (prtlev_>0)
    cout << " Four3DPk::ComputeNoiseFourierAmp() done ..." << endl;
}

// Compute mass field from its Fourier Coefficient
TArray<TF>  Four3DPk::ComputeMassDens()
{
  TArray<TF> massdens;
// Backward fourier transform of the fourierAmp array   
  FFTWServer ffts(true);                     
  ffts.setNormalize(true); 
  ffts.FFTBackward(fourAmp, massdens, true);
  //  cout << " Four3DPk::ComputeMassDens() done NbNeg=" << npbz << " / NPix=" <<  massDens.Size() << endl;
  cout << " Four3DPk::ComputeMassDens() done NPix=" <<  massdens.Size() << endl;
  return massdens;
}

// Compute power spectrum as a function of wave number k 
// cells with amp^2=re^2+im^2>s2cut are ignored
// Output : power spectrum (profile histogram)
HProf Four3DPk::ComputePk(double s2cut, int nbin, double kmin, double kmax)
{
  // The second half of the array along Y (matrix rows) contain
  // negative frequencies
  //  int nbh = sqrt(fourAmp.SizeX()*fourAmp.SizeX()+fourAmp.SizeY()*fourAmp.SizeY()/4.+fourAmp.SizeZ()*fourAmp.SizeY()/4.);
  // The profile histogram will contain the mean value of FFT amplitude
  // as a function of wave-number k = sqrt((double)(kx*kx+ky*ky))
  //  if (nbin < 1) nbin = nbh/2;
  if ((kmax<0.)||(kmax<kmin)) {
    kmin=0.;
    double maxx=fourAmp.SizeX()*dkx_;
    double maxy=fourAmp.SizeY()*dky_/2;
    double maxz=fourAmp.SizeZ()*dkz_/2;
    kmax=sqrt(maxx*maxx+maxy*maxy+maxz*maxz);
  }
  if (nbin<2) nbin=128; 
  HProf hp(kmin, kmax, nbin);
  hp.SetErrOpt(false);
  ComputePkCumul(hp, s2cut);
  return hp;
}

// Compute power spectrum as a function of wave number k 
// Cumul dans hp - cells with amp^2=re^2+im^2>s2cut are ignored
void Four3DPk::ComputePkCumul(HProf& hp, double s2cut)
{
  uint_8 nmodeok=0;
  // fourAmp represent 3-D fourier transform of a real input array. 
  // The second half of the array along Y and Z contain negative frequencies
  double kxx, kyy, kzz;
  // sa_size_t is large integer type  
  // We ignore 0th term in all frequency directions ...
  for(sa_size_t kz=1; kz<fourAmp.SizeZ(); kz++) {
    kzz =  (kz > fourAmp.SizeZ()/2) ? (double)(fourAmp.SizeZ()-kz)*dkz_ : (double)kz*dkz_; 
    for(sa_size_t ky=1; ky<fourAmp.SizeY(); ky++) {
      kyy =  (ky > fourAmp.SizeY()/2) ? (double)(fourAmp.SizeY()-ky)*dky_ : (double)ky*dky_; 
      for(sa_size_t kx=1; kx<fourAmp.SizeX(); kx++) {  // ignore the 0th coefficient (constant term)
	double kxx=(double)kx*dkx_;
	complex<TF> za = fourAmp(kx, ky, kz);
	if (za.real()>8.e9) continue;
	double wk = sqrt(kxx*kxx+kyy*kyy+kzz*kzz);
	double amp2 = za.real()*za.real()+za.imag()*za.imag();
	if ((s2cut>1.e-9)&&(amp2>s2cut))  continue;
	hp.Add(wk, amp2);
	nmodeok++;
      }
    }
  }
  if ((prtlev_>1)||((prtlev_>0)&&(s2cut>1.e-9))) {
    cout << " Four3DPk::ComputePkCumul/Info : NModeOK=" << nmodeok << " / NMode=" << fourAmp.Size() 
	 << " -> " << 100.*(double)nmodeok/(double)fourAmp.Size() << "%" << endl;
  }
  return;
}

//-----------------------------------------------------
// -- MassDist2D class :  2D mass distribution 
// --- PkNoiseCalculator : Classe de calcul du spectre de bruit PNoise(k) 
// determine par une reponse 2D de l'instrument
//-----------------------------------------------------
PkNoiseCalculator::PkNoiseCalculator(Four3DPk& pk3, Four2DResponse& rep, double s2cut, int ngen, 
				     const char* tit)
  : pkn3d(pk3), frep(rep), S2CUT(s2cut), NGEN(ngen), title(tit) 
{
  SetPrtLevel();
}

HProf PkNoiseCalculator::Compute()
{
  Timer tm(title.c_str());
  tm.Nop();
  HProf hnd;
  cout << "PkNoiseCalculator::Compute() " << title << "  NGEN=" << NGEN << " S2CUT=" << S2CUT << endl;
  for(int igen=0; igen<NGEN; igen++) {
    pkn3d.ComputeNoiseFourierAmp(frep);
    if (igen==0) hnd = pkn3d.ComputePk(S2CUT);
    else pkn3d.ComputePkCumul(hnd,S2CUT);
    if ((prtlev_>0)&&(igen>0)&&(((igen-1)%prtmodulo_)==0)) 
      cout << " PkNoiseCalculator::Compute() - done igen=" << igen << " / MaxNGen=" << NGEN << endl;
  }
  return hnd;
}


//-----------------------------------------------------
// -- MassDist2D class :  2D mass distribution 
//-----------------------------------------------------
// Constructor
MassDist2D::MassDist2D(GenericFunc& pk, int size, double meandens) 
: pkSpec(pk) , sizeA((size>16)?size:16) ,  massDens(sizeA, sizeA), 
  meanRho(meandens) , fg_fourAmp(false) , fg_massDens(false)
{
}

// To the computation job
void MassDist2D::Compute()
{
  ComputeFourierAmp();
  ComputeMassDens(); 	
}

// Generate mass field Fourier Coefficient
void MassDist2D::ComputeFourierAmp()
{
  if (fg_fourAmp) return; // job already done
  // We generate a random gaussian real field  
  double sigma = 1.;
// The following line fills the array by gaussian random numbers  
//--Replaced--  massDens = RandomSequence(RandomSequence::Gaussian, 0., sigma);
// Can be replaced by 
  DR48RandGen rg;
  for(sa_size_t ir=0; ir<massDens.NRows(); ir++) {
  	for(sa_size_t jc=0; jc<massDens.NCols(); jc++) {
      massDens(ir, jc) = rg.Gaussian(sigma); 
  	}
  }
// --- End of random filling

  // Compute fourier transform of the random gaussian field -> white noise 
  FFTWServer ffts(true);                     
  ffts.setNormalize(true); 
  ffts.FFTForward(massDens, fourAmp);
    
  // fourAmp represent 2-D fourier transform of a real input array. 
  // The second half of the array along Y (matrix rows) contain
  // negative frequencies
//  double fnorm = 1./sqrt(2.*fourAmp.Size()); 
// PUT smaller value for fnorm and check number of zeros
  double fnorm = 1.; 
  // sa_size_t is large integer type  
  for(sa_size_t ky=0; ky<fourAmp.NRows(); ky++) {
    double kyy = ky;
    if (ky > fourAmp.NRows()/2) kyy = fourAmp.NRows()-ky;  // negative frequencies 
    for(sa_size_t kx=0; kx<fourAmp.NCols(); kx++) {
      double wk = sqrt((double)(kx*kx+kyy*kyy));
      double amp = pkSpec(wk)*fnorm;      
      fourAmp(ky, kx) *= amp;   // renormalize fourier coeff using 
    }
  }
  fg_fourAmp = true;
  cout << " MassDist2D::ComputeFourierAmp() done ..." << endl;
}

// Compute mass field from its Fourier Coefficient
void MassDist2D::ComputeMassDens()
{
  if (fg_massDens) return; // job already done
  if (!fg_fourAmp) ComputeFourierAmp();   // Check fourier amp generation

// Backward fourier transform of the fourierAmp array   
  FFTWServer ffts(true);                     
  ffts.setNormalize(true); 
  ffts.FFTBackward(fourAmp, massDens, true);
// We consider that massDens represents delta rho/rho 
// rho = (delta rho/rho + 1) * MeanDensity 
  massDens += 1.;
// We remove negative values 
  sa_size_t npbz = 0;
  for (sa_size_t i=0; i<massDens.NRows(); i++) 
    for (sa_size_t j=0; j<massDens.NCols(); j++) 
      if (massDens(i,j) < 0.) { npbz++; massDens(i,j) = 0.; }
  massDens *= meanRho;
  cout << " MassDist2D::ComputeMassDens() done NbNeg=" << npbz << " / NPix=" <<  massDens.Size() << endl;
}

// Compute power spectrum as a function of wave number k 
// Output : power spectrum (profile histogram)
HProf MassDist2D::ReconstructPk(int nbin)
{
  // The second half of the array along Y (matrix rows) contain
  // negative frequencies
  int nbh = sqrt(2.0)*fourAmp.NCols();
  // The profile histogram will contain the mean value of FFT amplitude
  // as a function of wave-number k = sqrt((double)(kx*kx+ky*ky))
  if (nbin < 1) nbin = nbh/2;
  HProf hp(-0.5, nbh-0.5, nbin);
  hp.SetErrOpt(false);

  for(int ky=0; ky<fourAmp.NRows(); ky++) {
    double kyy = ky;
    if (ky > fourAmp.NRows()/2)  kyy = fourAmp.NRows()-ky;  // negative frequencies
    for(int kx=0; kx<fourAmp.NCols(); kx++) {
      double wk = sqrt((double)(kx*kx+kyy*kyy));
      complex<r_8> za = fourAmp(ky, kx);
      double amp = sqrt(za.real()*za.real()+za.imag()*za.imag());
      hp.Add(wk, amp);
    }
  }
  return hp;
}