//
//  Olaf Ronneberger
//  Sample Solution for Exercise 6 (Lecture "3D Image Analysis" Summer 2014) 
//  Spherical Harmonics
//
//  compile with:
//  g++ -Wall -O2 -g compute_basis.cc -lgsl -lgslcblas -lblitz -o compute_basis -lhdf5



#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <vector>
#include <complex>
#include <gsl/gsl_sf_legendre.h>

// HDF5
#include "BlitzHDF5Helper.hh"
// Blitz++
#include <blitz/array.h>

int main( int argc, char** argv)
{
  /*-----------------------------------------------------------------------
   *  Define some constants
   *-----------------------------------------------------------------------*/
  std::complex<double> com_i(0,1);
  int ellMax = 6;
  int nTheta = 180;
  int nPhi = 360;
  int nBasisFuncs = (ellMax+1)*(ellMax+1);
  
  /*-----------------------------------------------------------------------
   *  Compute the Spherical Harmonics basis functions
   *-----------------------------------------------------------------------*/
  blitz::Array<std::complex<double>,3> basis( nBasisFuncs, nTheta, nPhi);
  
  for( int ell = 0; ell <= ellMax; ++ell)
  {
    for( int m = -ell; m <= ell; ++m)
    {
      for( int thetaQuant = 0; thetaQuant < nTheta; ++thetaQuant)
      {
        double theta = double( thetaQuant)/nTheta * M_PI;
        for( int phiQuant = 0; phiQuant < nPhi; ++phiQuant)
        {
          double phi = double(phiQuant)/nPhi * 2*M_PI;
          int lev = ell * (ell+1) + m;
          
          if( m >= 0)
          {
            basis(lev, thetaQuant, phiQuant) = 
                gsl_sf_legendre_sphPlm(ell,m,cos(theta)) 
                *  exp(  m * phi * com_i);
          }
          else
          {
            basis(lev, thetaQuant, phiQuant) = 
                pow(-1,m) * gsl_sf_legendre_sphPlm(ell,-m,cos(theta)) 
                *  exp(  m * phi * com_i);
          }
          
        }
      }
    }
  }

  /*-----------------------------------------------------------------------
   *  Save real and imaginary part
   *-----------------------------------------------------------------------*/
  //BlitzH5File outFile( "basis.h5", BlitzH5File::Replace);
  blitz::Array<double,3> basis_real( basis.shape());
  blitz::Array<double,3> basis_imag( basis.shape());
  basis_real = real(basis);
  basis_imag = imag(basis);
  
  std::string resultFileName = "basis.h5";
  writeBlitzToHDF5( basis_real, "basis_real", resultFileName);
  writeBlitzToHDF5( basis_imag, "basis_imag", resultFileName);
  //outFile.writeDataSetSimple( basis_real, "basis_real");
  //outFile.writeDataSetSimple( basis_imag, "basis_imag");
  
  
  /*-----------------------------------------------------------------------
   *  check orhtogonality
   *-----------------------------------------------------------------------*/
  for( int row = 0; row < nBasisFuncs; ++row)
  {
    for( int col = 0; col < nBasisFuncs; ++col)
    {
      std::complex<double> sumval = 0;
      
      for( int thetaQuant = 0; thetaQuant < nTheta; ++thetaQuant)
      {
        double theta = double( thetaQuant)/nTheta * M_PI;
        double ds =  sin(theta) * (2*M_PI/nPhi) * (M_PI/nTheta);
        
        for( int phiQuant = 0; phiQuant < nPhi; ++phiQuant)
        {
          sumval += basis( row, thetaQuant, phiQuant) 
              * conj(basis( col, thetaQuant, phiQuant))
              * ds;
        }
      }
      std::cout << std::setw(2) << floor(sumval.real()*1000.0+0.5)/1000 << " ";
      
    } 
    std::cout << std::endl;
    
  }
  

  return 0;
}