//  Usuall mathematical functions and operations on arrays
//                     R. Ansari, C.Magneville   03/2000

#include "machdefs.h"
#include <math.h>
#include "matharr.h"

// ----------------------------------------------------
//          Application d'une fonction
// ----------------------------------------------------

/*!
  \class SOPHYA::MathArray
  \ingroup TArray
  Class for simple mathematical operation on arrays 
  \warning Instanciated only for \b real and \b double (r_4, r_8) type arrays
*/

//! Apply Function In Place (function double version)
/*!
  \param a : array to be replaced in place
  \param f : function for replacement
  \return Return an array \b a filled with function f(a(i,j))
*/
template <class T>
TArray<T>& MathArray<T>::ApplyFunctionInPlace(TArray<T> & a, Arr_DoubleFunctionOfX f)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("MathArray<T>::ApplyFunctionInPlace(TArray<T> & a..) Not Allocated Array a !");
  T * pe;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    for(k=0; k<maxx; k+=step )  pe[k] = (T)(f((double)pe[k])); 
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  pe[k] = (T)(f((double)pe[k]));
    }
  }
  return(a);
}

//! Apply Function In Place (function float version)
/*!
  \param a : array to be replaced in place
  \param f : function for replacement
  \return Return an array \b a filled with function f(a(i,j))
*/
template <class T>
TArray<T>& MathArray<T>::ApplyFunctionInPlaceF(TArray<T> & a, Arr_FloatFunctionOfX f)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("MathArray<T>::ApplyFunctionInPlaceF(TArray<T> & a..) Not Allocated Array a !");
  T * pe;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    for(k=0; k<maxx; k+=step )  pe[k] = (T)(f((float)pe[k])); 
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  pe[k] = (T)(f((float)pe[k]));
    }
  }
  return(a);
}


//! Apply Function (function double version)
/*!
  \param a : argument array of the function
  \param f : function for replacement
  \return Return a new array filled with function f(a(i,j))
*/
template <class T>
TArray<T> MathArray<T>::ApplyFunction(TArray<T> const & a, Arr_DoubleFunctionOfX f)
{
  TArray<T> ra;
  ra = a;
  ApplyFunctionInPlace(ra, f);
  return(ra);
}

//! Apply Function (function float version)
/*!
  \param a : argument array of the function
  \param f : function for replacement
  \return Return a new array filled with function f(a(i,j))
*/
template <class T>
TArray<T> MathArray<T>::ApplyFunctionF(TArray<T> const & a, Arr_FloatFunctionOfX f)
{
  TArray<T> ra;
  ra = a;
  ApplyFunctionInPlaceF(ra, f);
  return(ra);
}

//! Compute \b mean and \b sigma of elements of array \b a, return \b mean
template <class T>
double MathArray<T>::MeanSigma(TArray<T> const & a, double & mean, double & sig)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("MathArray<T>::MeanSigma(TArray<T> const & a..) Not Allocated Array a !");
  const T * pe;
  sa_size_t j,k;
  mean=0.;
  sig = 0.;
  double valok;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    for(k=0; k<maxx; k+=step )  { 
      valok = (double) pe[k]; 
      mean += valok;  sig += valok*valok;
    }
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      for(k=0; k<gpas; k+=step) { 
	valok = (double) pe[k]; 
	mean += valok;  sig += valok*valok;
      }
    }
  }
  double dsz = (double)(a.Size());
  mean /= dsz;
  sig = sig/dsz - mean*mean;
  if (sig >= 0.) sig = sqrt(sig);
  return(mean);
}


//-------------------------------------------------------------------------------
//      Definition utilitaire d'application de fonction 
inline complex<r_8> ApplyComplexDoubleFunction(complex<r_8> z, 
					       Arr_ComplexDoubleFunctionOfX f)
{
  return(f(z));
}

inline complex<r_4> ApplyComplexDoubleFunction(complex<r_4> z, 
					       Arr_ComplexDoubleFunctionOfX f)
{
  complex<r_8> zd((r_8)z.real(), (r_8)z.imag());
  zd = f(zd);
  complex<r_4> zr((r_4)zd.real(), (r_4)zd.imag());
  return(zr);
}

//-------------------------------------------------------------------------------

/*!
  \class SOPHYA::ComplexMathArray
  \ingroup TArray
  Class for simple mathematical operation on arrays 
  \warning Instanciated only for \b real and \b double (r_4, r_8) complex arrays
*/

//! Apply Function In Place (complex arrays)
/*!
  \param a : complex array to be replaced in place
  \param f : function for replacement
  \return Return an array \b a filled with function f(a(i,j))
*/
template <class T>
TArray< complex<T> >& ComplexMathArray<T>::ApplyFunctionInPlace(TArray< complex<T> > & a, Arr_ComplexDoubleFunctionOfX f)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray< complex<T> >::ApplyFunctionInPlace(TArray< complex<T> > & a..) Not Allocated Array a !");
  complex<T> * pe;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    for(k=0; k<maxx; k+=step )  pe[k] = ApplyComplexDoubleFunction(pe[k],f); 
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  pe[k] = ApplyComplexDoubleFunction(pe[k],f);
    }
  }
  return(a);
}



//! Apply Function (complex arrays)
/*!
  \param a : argument array of the function
  \param f : function for replacement
  \return Return a new array filled with function f(a(i,j))
*/
template <class T>
TArray< complex<T> > ComplexMathArray<T>::ApplyFunction(TArray< complex<T> > const & a, Arr_ComplexDoubleFunctionOfX f)
{
  TArray< complex<T> > ra;
  ra = a;
  ApplyFunctionInPlace(ra, f);
  return(ra);
}

//! Create a complex array, from a real and an imaginary arrays
/*!
  \param p_real : array containing the real part of the complex output array
  \param p_imag : array containing the imaginary part of the complex output array
  \return Return a new complex array build from \b p_real and \b p_imag
*/
template <class T>
TArray< complex<T> > ComplexMathArray<T>::FillFrom(TArray<T> const & p_real,
						   TArray<T> const & p_imag)
{
  if (p_real.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray<T>::FillFrom() - Not Allocated Array ! ");
  bool smo;
  if (!p_real.CompareSizes(p_imag, smo)) 
    throw(SzMismatchError("ComplexMathArray<T>::FillFrom() SizeMismatch")) ;

  TArray< complex<T> > ra;
  ra.ReSize(p_real);

  complex<T> * pe;
  const T * per;
  const T * pei;
  sa_size_t j,k,ka;
  if (smo && (p_real.AvgStep() > 0) && (p_imag.AvgStep() > 0))   {  // regularly spaced elements
    sa_size_t step = p_real.AvgStep();
    sa_size_t stepa = p_imag.AvgStep();
    sa_size_t maxx = p_real.Size()*step;
    per = p_real.Data();
    pei = p_imag.Data();
    pe = ra.Data();
    for(k=0, ka=0;  k<maxx;  k+=step, ka+=stepa )  
      pe[k] = complex<T>(per[k], pei[ka]) ;
  }
  else {    // Non regular data spacing ...
    int_4 ax,axa;
    sa_size_t step, stepa;
    sa_size_t gpas, naxa;
    p_real.GetOpeParams(p_imag, smo, ax, axa, step, stepa, gpas, naxa);
    for(j=0; j<naxa; j++)  {
      per = p_real.Data()+p_real.Offset(ax,j);
      pei = p_imag.Data()+p_imag.Offset(axa,j);
      pe = ra.Data()+ra.Offset(ax,j);
      for(k=0, ka=0;  k<gpas;  k+=step, ka+=stepa)  
	pe[k] = complex<T>(per[k], pei[ka]) ;
    }
  }
  return(ra);
}


//! Returns the real part of the complex input array.
/*!
  \param a : input complex array
  \return Return a new array filled with the real part of the input complex array elements
*/

template <class T>
TArray<T> ComplexMathArray<T>::real(TArray< complex<T> > const & a)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray< complex<T> >::real(TArray< complex<T> >& a) Not Allocated Array a !");
  TArray<T> ra;
  ra.ReSize(a);

  const complex<T> * pe;
  T * po;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    po = ra.Data();
    for(k=0; k<maxx; k+=step )  po[k] = pe[k].real();
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      po = ra.DataBlock().Begin()+ra.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  po[k] = pe[k].real();
    }
  }
  return(ra);
}

//! Returns the imaginary part of the complex input array.
/*!
  \param a : input complex array
  \return Return a new array filled with the imaginary part of the input complex array elements
*/

template <class T>
TArray<T> ComplexMathArray<T>::imag(TArray< complex<T> > const & a)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray< complex<T> >::imag(TArray< complex<T> >& a) Not Allocated Array a !");
  TArray<T> ra;
  ra.ReSize(a);

  const complex<T> * pe;
  T * po;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    po = ra.Data();
    for(k=0; k<maxx; k+=step )  po[k] = pe[k].imag();
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      po = ra.DataBlock().Begin()+ra.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  po[k] = pe[k].imag();
    }
  }
  return(ra);
}

//! Returns the module squared of the complex input array.
/*!
  \param a : input complex array
  \return Return a new array filled with the module squared of the input complex array elements 
*/

template <class T>
TArray<T> ComplexMathArray<T>::module2(TArray< complex<T> > const & a)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray< complex<T> >::module2(TArray< complex<T> >& a) Not Allocated Array a !");
  TArray<T> ra;
  ra.ReSize(a);

  const complex<T> * pe;
  T * po;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    po = ra.Data();
    for(k=0; k<maxx; k+=step )  
      po[k] = (pe[k].real()*pe[k].real()+pe[k].imag()*pe[k].imag());
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      po = ra.DataBlock().Begin()+ra.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  
	po[k] = (pe[k].real()*pe[k].real()+pe[k].imag()*pe[k].imag());
    }
  }
  return(ra);
}

//! Returns the module of the complex input array.
/*!
  \param a : input complex array
  \return Return a new array filled with the module of the input complex array elements 
*/

template <class T>
TArray<T> ComplexMathArray<T>::module(TArray< complex<T> > const & a)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray< complex<T> >::module(TArray< complex<T> >& a) Not Allocated Array a !");
  TArray<T> ra;
  ra.ReSize(a);

  const complex<T> * pe;
  T * po;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    po = ra.Data();
    for(k=0; k<maxx; k+=step )  
      po[k] = sqrt((double)(pe[k].real()*pe[k].real()+pe[k].imag()*pe[k].imag()));
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      po = ra.DataBlock().Begin()+ra.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  
	po[k] = sqrt((double)(pe[k].real()*pe[k].real()+pe[k].imag()*pe[k].imag()));
    }
  }
  return(ra);
}


//! Returns the phase of the complex input array.
/*!
  \param a : input complex array
  \return Return a new array filled with the phase of the input complex array elements 
*/

template <class T>
TArray<T> ComplexMathArray<T>::phase(TArray< complex<T> > const & a)
{
  if (a.NbDimensions() < 1) 
    throw RangeCheckError("ComplexMathArray< complex<T> >::phase(TArray< complex<T> >& a) Not Allocated Array a !");
  TArray<T> ra;
  ra.ReSize(a);

  const complex<T> * pe;
  T * po;
  sa_size_t j,k;
  if (a.AvgStep() > 0)   {  // regularly spaced elements
    sa_size_t step = a.AvgStep(); 
    sa_size_t maxx = a.Size()*step;
    pe = a.Data();
    po = ra.Data();
    for(k=0; k<maxx; k+=step )  
      po[k] = atan2((double)pe[k].imag(), (double)pe[k].real());
  }
  else {    // Non regular data spacing ...
    int_4 ka = a.MaxSizeKA();
    sa_size_t step = a.Step(ka);
    sa_size_t gpas = a.Size(ka)*step;
    sa_size_t naxa = a.Size()/a.Size(ka);
    for(j=0; j<naxa; j++)  {
      pe = a.DataBlock().Begin()+a.Offset(ka,j);
      po = ra.DataBlock().Begin()+ra.Offset(ka,j);
      for(k=0; k<gpas; k+=step)  
	po[k] = atan2((double)pe[k].imag(), (double)pe[k].real());
    }
  }
  return(ra);
}


///////////////////////////////////////////////////////////////
#ifdef __CXX_PRAGMA_TEMPLATES__
#pragma define_template MathArray<r_4>
#pragma define_template MathArray<r_8>
#pragma define_template ComplexMathArray<r_4>
#pragma define_template ComplexMathArray<r_8>
#endif

#if defined(ANSI_TEMPLATES) || defined(GNU_TEMPLATES)
template class MathArray<r_4>;
template class MathArray<r_8>;
template class ComplexMathArray<r_4>;
template class ComplexMathArray<r_8>;
#endif