#include "ornstein_uhlenbeck_process.hpp"
#include <cmath>
#include <cstring>
#include <cassert>
#include "scope_timer.hpp"
#include <algorithm>
#include <chrono>


template <class rnumber,field_backend be>
ornstein_uhlenbeck_process<rnumber,be>::ornstein_uhlenbeck_process(
    const char *NAME,
    int nx,
    int ny,
    int nz,
    double ou_kmin,
    double ou_kmax,
    double ou_energy_amplitude,
    double ou_gamma_factor,
    double DKX,
    double DKY,
    double DKZ,
    unsigned FFTW_PLAN_RIGOR)
{
    TIMEZONE("ornstein_uhlenbeck_process::ornstein_uhlenbeck_process");
    strncpy(this->name, NAME, 256);
    this->name[255] = '\0';
    this->iteration = 0;
    this->ou_field = new field<rnumber,be,THREE>(
        nx,ny,nz, MPI_COMM_WORLD, FFTW_PLAN_RIGOR);
    *this->ou_field = 0.0;
    this->ou_field->dft();

    this->ou_field_vort = new field<rnumber,be,THREE>(
        nx,ny,nz, MPI_COMM_WORLD, FFTW_PLAN_RIGOR);
    *this->ou_field_vort = 0.0;
    this->ou_field_vort->dft();

    this->B = new field<rnumber,be,THREExTHREE>(
        nx,ny,nz, MPI_COMM_WORLD, FFTW_PLAN_RIGOR);

    this->kk = new kspace<be,SMOOTH>(
        this->ou_field->clayout, DKX, DKY, DKZ);

    this->ou_kmin_squ = pow(ou_kmin,2);
    this->ou_kmax_squ = pow(ou_kmax,2);
    this->ou_energy_amplitude = ou_energy_amplitude;
    this->ou_gamma_factor = ou_gamma_factor;
    this->epsilon = pow((this->ou_energy_amplitude/this->kolmogorov_constant), 3./2.);

    assert(this->kk->kM2 >= this->ou_kmax_squ);

    gen.resize(omp_get_max_threads());

    long now;

    if (this->ou_field->clayout->myrank == 0){
	now = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
    }
    MPI_Bcast(&now,1,MPI_LONG,0,this->ou_field->comm);

    for(int thread=0;thread<omp_get_max_threads();thread++)
    {
            long current_seed =
                    this->ou_field->clayout->myrank*omp_get_max_threads() +
                    thread+now;
            gen[thread].seed(current_seed);
    }


    this->initialize_B();

}

template <class rnumber,field_backend be>
ornstein_uhlenbeck_process<rnumber,be>::~ornstein_uhlenbeck_process()
{
    TIMEZONE("ornstein_uhlenbeck_process::~ornstein_uhlenbeck_process");
    delete this->kk;
    delete this->ou_field;
    delete this->ou_field_vort;
    delete this->B;
}

template <class rnumber,field_backend be>
void ornstein_uhlenbeck_process<rnumber,be>::step(
    double dt)
{
    // put in "CFL"-criterium TODO!!!
    TIMEZONE("ornstein_uhlenbeck_process::step");
    this->kk->CLOOP_K2(
                [&](ptrdiff_t cindex,
                    ptrdiff_t xindex,
                    ptrdiff_t yindex,
                    ptrdiff_t zindex,
                    double k2){
        if (k2 <= this->ou_kmax_squ && k2 >= this->ou_kmin_squ)
        {
            double g = this->gamma(sqrt(k2));
            int tr = omp_get_thread_num();
            double random[6] =
            {
                this->d(this->gen[tr]),this->d(this->gen[tr]),this->d(this->gen[tr]),
                this->d(this->gen[tr]),this->d(this->gen[tr]),this->d(this->gen[tr])
            };
            for (int cc=0; cc<3; cc++) for (int i=0; i<2; i++)
                this->ou_field->cval(cindex,cc,i) =
                    (1-dt*g) * this->ou_field->cval(cindex,cc,i) +
                    sqrt(dt) * (
                        this->B->cval(cindex,0,cc,0) * random[0+3*i] +
                        this->B->cval(cindex,1,cc,0) * random[1+3*i] +
                        this->B->cval(cindex,2,cc,0) * random[2+3*i] );

        }
    });

    this->ou_field->symmetrize();
    this->calc_ou_vorticity();

}


template <class rnumber, field_backend be>
void ornstein_uhlenbeck_process<rnumber,be>::initialize_B()
{
    TIMEZONE("ornstein_uhlenbeck_process::initialize_B");
    *this->B = 0.0;
    this->kk->CLOOP(
        [&](ptrdiff_t cindex,
            ptrdiff_t xindex,
            ptrdiff_t yindex,
            ptrdiff_t zindex){
            double ksqu = pow(this->kk->kx[xindex],2)+
                          pow(this->kk->ky[yindex],2)+
                          pow(this->kk->kz[zindex],2);
            if (ksqu <= this->ou_kmax_squ && ksqu >= this->ou_kmin_squ)
            {
                double kabs = sqrt(ksqu);
                double sigma;
                if (ksqu == 0 || ksqu == this->ou_kmax_squ)
                {
                    ksqu = 1; kabs = 1; sigma = 0;
                }
                else
                {
                    sigma =
                        sqrt(4*this->gamma(kabs)*this->energy(kabs)
                        /(4*M_PI*ksqu));
                }

                for(int i=0;i<3;i++) {
                    for(int j=0;j<3;j++) {

                        if (i+j == 0)
                            this->B->cval(cindex,i,j,0) =
                                sigma/2. * (1-pow(this->kk->kx[xindex],2)/ksqu);

                        if (i+j == 4)
                            this->B->cval(cindex,i,j,0) =
                                sigma/2. * (1-pow(this->kk->kz[zindex],2)/ksqu);

                        if (i+j == 1)
                            this->B->cval(cindex,i,j,0) =
                                sigma/2. * (0-this->kk->kx[xindex]*this->kk->ky[yindex]/ksqu);

                        if (i+j == 3)
                            this->B->cval(cindex,i,j,0) =
                                sigma/2. * (0-this->kk->kz[zindex]*this->kk->ky[yindex]/ksqu);

                        if (i+j == 2) {

                            if(i==j)
                                this->B->cval(cindex,i,j,0) =
                                    sigma/2. * (1-pow(this->kk->ky[yindex],2)/ksqu);

                            if(i!=j)
                                this->B->cval(cindex,i,j,0) =
                                    sigma/2. * (0-this->kk->kx[xindex]*this->kk->kz[zindex]/ksqu);
                        }
                    }
                }
            }

        });
}


template <class rnumber, field_backend be>
void ornstein_uhlenbeck_process<rnumber, be>::let_converge(void)
{
  double ou_kmin = sqrt(this->ou_kmin_squ);
  double tau = 1.0/this->gamma(ou_kmin);
  double dt = tau/1000.;

  for(int i=0; i<2000; i++)
  {
    this->step(dt);
  }
}

template <class rnumber, field_backend be>
void ornstein_uhlenbeck_process<rnumber, be>::strip_from_field(
  field<rnumber, be, THREE> *src)
{
  assert(src->real_space_representation==false);

  this->kk->CLOOP_K2(
          [&](ptrdiff_t cindex,
            ptrdiff_t xindex,
            ptrdiff_t yindex,
            ptrdiff_t zindex,
            double k2){

              if (k2 <= this->ou_kmax_squ && k2 >= this->ou_kmin_squ){

                for(int cc=0; cc < 3; cc++){
                  for(int imag=0; imag < 2; imag++){
                    src->cval(cindex,cc,imag) = 0;
                  }
                }
              }

      }

  );

}

template <class rnumber, field_backend be>
void ornstein_uhlenbeck_process<rnumber,be>::add_to_field_replace(
  field<rnumber, be, THREE> *src, std::string uv)
{
  assert(src->real_space_representation==false);
  assert((uv == "vorticity") || (uv == "velocity"));

  field<rnumber, be, THREE> *field_to_replace;

  if (uv == "vorticity") field_to_replace = this->ou_field_vort;
  else field_to_replace = this->ou_field;

  this->kk->CLOOP_K2(
          [&](ptrdiff_t cindex,
            ptrdiff_t xindex,
            ptrdiff_t yindex,
            ptrdiff_t zindex,
            double k2){

              if (k2 <= this->ou_kmax_squ && k2 >= this->ou_kmin_squ){

                rnumber tmp;

                for(int cc=0; cc < 3; cc++){
                  for(int imag=0; imag < 2; imag++){
                    tmp = field_to_replace->cval(cindex,cc,imag);
                    src->cval(cindex,cc,imag) = tmp;
                  }
                }
              }

      }

  );
}


template <class rnumber, field_backend be>
void ornstein_uhlenbeck_process<rnumber,be>::calc_ou_vorticity(void)
{
  this->kk->CLOOP_K2(
              [&](ptrdiff_t cindex,
                  ptrdiff_t xindex,
                  ptrdiff_t yindex,
                  ptrdiff_t zindex,
                  double k2){
      if (k2 <= this->kk->kM2)
      {
          this->ou_field_vort->cval(cindex,0,0) = -(this->kk->ky[yindex]*this->ou_field->cval(cindex,2,1) - this->kk->kz[zindex]*this->ou_field->cval(cindex,1,1));
          this->ou_field_vort->cval(cindex,0,1) =  (this->kk->ky[yindex]*this->ou_field->cval(cindex,2,0) - this->kk->kz[zindex]*this->ou_field->cval(cindex,1,0));
          this->ou_field_vort->cval(cindex,1,0) = -(this->kk->kz[zindex]*this->ou_field->cval(cindex,0,1) - this->kk->kx[xindex]*this->ou_field->cval(cindex,2,1));
          this->ou_field_vort->cval(cindex,1,1) =  (this->kk->kz[zindex]*this->ou_field->cval(cindex,0,0) - this->kk->kx[xindex]*this->ou_field->cval(cindex,2,0));
          this->ou_field_vort->cval(cindex,2,0) = -(this->kk->kx[xindex]*this->ou_field->cval(cindex,1,1) - this->kk->ky[yindex]*this->ou_field->cval(cindex,0,1));
          this->ou_field_vort->cval(cindex,2,1) =  (this->kk->kx[xindex]*this->ou_field->cval(cindex,1,0) - this->kk->ky[yindex]*this->ou_field->cval(cindex,0,0));
      }
      else
          std::fill_n((rnumber*)(this->ou_field_vort->get_cdata()+3*cindex), 6, 0.0);
  }
  );
  this->ou_field_vort->symmetrize();
}

template class ornstein_uhlenbeck_process<float,FFTW>;
template class ornstein_uhlenbeck_process<double,FFTW>;