/**********************************************************************
*                                                                     *
*  Copyright 2017 Max Planck Institute                                *
*                 for Dynamics and Self-Organization                  *
*                                                                     *
*  This file is part of bfps.                                         *
*                                                                     *
*  bfps is free software: you can redistribute it and/or modify       *
*  it under the terms of the GNU General Public License as published  *
*  by the Free Software Foundation, either version 3 of the License,  *
*  or (at your option) any later version.                             *
*                                                                     *
*  bfps is distributed in the hope that it will be useful,            *
*  but WITHOUT ANY WARRANTY; without even the implied warranty of     *
*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the      *
*  GNU General Public License for more details.                       *
*                                                                     *
*  You should have received a copy of the GNU General Public License  *
*  along with bfps.  If not, see <http://www.gnu.org/licenses/>       *
*                                                                     *
* Contact: Cristian.Lalescu@ds.mpg.de                                 *
*                                                                     *
**********************************************************************/



#include <string>
#include <cmath>
#include "NSVEcomplex_particles.hpp"
#include "scope_timer.hpp"
#include "particles/particles_sampling.hpp"
#include "particles/p2p_computer.hpp"
#include "particles/particles_inner_computer.hpp"

template <typename rnumber>
int NSVEcomplex_particles<rnumber>::initialize(void)
{
    TIMEZONE("NSVEcomplex_particles::initialize");
    this->NSVE<rnumber>::initialize();

    p2p_computer<double, long long int> current_p2p_computer;
    // TODO: particle interactions are switched off manually for testing purposes.
    // this needs to be fixed once particle interactions can be properly resolved.
    current_p2p_computer.setEnable(enable_p2p);
    //current_p2p_computer.setEnable(false);

    particles_inner_computer<double, long long int> current_particles_inner_computer(inner_v0);
    current_particles_inner_computer.setEnable(enable_inner);
    this->cutoff = 1.0;

    this->ps = particles_system_builder_with_p2p(
                this->fs->cvelocity,              // (field object)
                this->fs->kk,                     // (kspace object, contains dkx, dky, dkz)
                tracers0_integration_steps, // to check coherency between parameters and hdf input file (nb rhs)
                (long long int)nparticles,  // to check coherency between parameters and hdf input file
                this->fs->get_current_fname(),    // particles input filename
                std::string("/tracers0/state/") + std::to_string(this->fs->iteration), // dataset name for initial input
                std::string("/tracers0/rhs/")  + std::to_string(this->fs->iteration),  // dataset name for initial input
                tracers0_neighbours,        // parameter (interpolation no neighbours)
                tracers0_smoothness,        // parameter
                this->comm,
                this->fs->iteration+1,
                std::move(current_p2p_computer),
                std::move(current_particles_inner_computer),
                cutoff);

    this->particles_output_writer_mpi = new particles_output_hdf5<
        long long int, double, 6>(
                MPI_COMM_WORLD,
                "tracers0",
                nparticles,
                tracers0_integration_steps);
    this->particles_sample_writer_mpi = new particles_output_sampling_hdf5<
        long long int, double, 3>(
                MPI_COMM_WORLD,
                this->ps->getGlobalNbParticles(),
                (this->simname + "_particles.h5"),
                "tracers0",
                "position/0");


    /// allocate grad vel field
    this->nabla_u = new field<rnumber, FFTW, THREExTHREE>(
            this->nx, this->ny, this->nz,
            this->comm,
            this->fs->cvorticity->fftw_plan_rigor);
    return EXIT_SUCCESS;
}

template <typename rnumber>
int NSVEcomplex_particles<rnumber>::step(void)
{
    TIMEZONE("NSVEcomplex_particles::step");
    this->fs->compute_velocity(this->fs->cvorticity);
    this->fs->cvelocity->ift();
    if(enable_vorticity_omega){
        *this->tmp_vec_field = this->fs->cvorticity->get_cdata();
        this->tmp_vec_field->ift();
        this->ps->completeLoopWithVorticity(this->dt, *this->tmp_vec_field);
    }
    else{
        this->ps->completeLoop(this->dt);
    }
    this->NSVE<rnumber>::step();
    return EXIT_SUCCESS;
}

template <typename rnumber>
int NSVEcomplex_particles<rnumber>::write_checkpoint(void)
{
    TIMEZONE("NSVEcomplex_particles::write_checkpoint");
    this->NSVE<rnumber>::write_checkpoint();
    this->particles_output_writer_mpi->open_file(this->fs->get_current_fname());
    // TODO P2P write particle data too
    this->particles_output_writer_mpi->template save<6>(
            this->ps->getParticlesState(),
            this->ps->getParticlesRhs(),
            this->ps->getParticlesIndexes(),
            this->ps->getLocalNbParticles(),
            this->fs->iteration);
    this->particles_output_writer_mpi->close_file();
    return EXIT_SUCCESS;
}

template <typename rnumber>
int NSVEcomplex_particles<rnumber>::finalize(void)
{
    TIMEZONE("NSVEcomplex_particles::finalize");
    delete this->nabla_u;
    delete this->particles_output_writer_mpi;
    delete this->particles_sample_writer_mpi;
    this->NSVE<rnumber>::finalize();
    return EXIT_SUCCESS;
}

/** \brief Compute fluid stats and sample particle data.
 */

template <typename rnumber>
int NSVEcomplex_particles<rnumber>::do_stats()
{
    TIMEZONE("NSVEcomplex_particles::do_stats");
    /// perform fluid stats
    this->NSVE<rnumber>::do_stats();

    /// check if particle stats should be performed now;
    /// if not, exit method.
    if (!(this->iteration % this->niter_part == 0))
        return EXIT_SUCCESS;

    /// allocate temporary data array
    std::unique_ptr<double[]> pdata0 = this->ps->extractParticlesState(0, 3);
    std::unique_ptr<double[]> pdata1 = this->ps->extractParticlesState(3, 6);
    std::unique_ptr<double[]> pdata2(new double[9*this->ps->getLocalNbParticles()]);

    /// sample position
    this->particles_sample_writer_mpi->template save_dataset<3>(
            "tracers0",
            "position",
            pdata0.get(),
            &pdata0,
            this->ps->getParticlesIndexes(),
            this->ps->getLocalNbParticles(),
            this->ps->get_step_idx()-1);

    /// sample orientation
    this->particles_sample_writer_mpi->template save_dataset<3>(
            "tracers0",
            "orientation",
            pdata0.get(),
            &pdata1,
            this->ps->getParticlesIndexes(),
            this->ps->getLocalNbParticles(),
            this->ps->get_step_idx()-1);

    /// sample velocity
    this->ps->sample_compute_field(*this->tmp_vec_field, pdata1.get());
    this->particles_sample_writer_mpi->template save_dataset<3>(
            "tracers0",
            "velocity",
            pdata0.get(),
            &pdata1,
            this->ps->getParticlesIndexes(),
            this->ps->getLocalNbParticles(),
            this->ps->get_step_idx()-1);

    /// sample velocity gradient
    /// fs->cvelocity should contain the velocity in Fourier space
    this->fs->compute_velocity(this->fs->cvorticity);
    compute_gradient(
            this->fs->kk,
            this->fs->cvelocity,
            this->nabla_u);
    this->nabla_u->ift();
    this->ps->sample_compute_field(*this->nabla_u, pdata2.get());
    this->particles_sample_writer_mpi->template save_dataset<9>(
            "tracers0",
            "velocity_gradient",
            pdata0.get(),
            &pdata2,
            this->ps->getParticlesIndexes(),
            this->ps->getLocalNbParticles(),
            this->ps->get_step_idx()-1);

    /// compute acceleration and sample it
    this->fs->compute_Lagrangian_acceleration(this->tmp_vec_field);
    this->tmp_vec_field->ift();
    this->ps->sample_compute_field(*this->tmp_vec_field, pdata1.get());
    this->particles_sample_writer_mpi->template save_dataset<3>(
            "tracers0",
            "acceleration",
            pdata0.get(),
            &pdata1,
            this->ps->getParticlesIndexes(),
            this->ps->getLocalNbParticles(),
            this->ps->get_step_idx()-1);

    // deallocate temporary data array
    // TODO: is it required/safe to call the release method here?
    //pdata.release();

    return EXIT_SUCCESS;
}

template <typename rnumber>
int NSVEcomplex_particles<rnumber>::read_parameters(void)
{
    TIMEZONE("NSVEcomplex_particles::read_parameters");
    this->NSVE<rnumber>::read_parameters();
    hid_t parameter_file = H5Fopen((this->simname + ".h5").c_str(), H5F_ACC_RDONLY, H5P_DEFAULT);
    this->niter_part = hdf5_tools::read_value<int>(parameter_file, "parameters/niter_part");
    this->nparticles = hdf5_tools::read_value<int>(parameter_file, "parameters/nparticles");
    this->tracers0_integration_steps = hdf5_tools::read_value<int>(parameter_file, "parameters/tracers0_integration_steps");
    this->tracers0_neighbours = hdf5_tools::read_value<int>(parameter_file, "parameters/tracers0_neighbours");
    this->tracers0_smoothness = hdf5_tools::read_value<int>(parameter_file, "parameters/tracers0_smoothness");
    this->enable_p2p = hdf5_tools::read_value<int>(parameter_file, "parameters/tracers0_enable_p2p");
    this->enable_inner = hdf5_tools::read_value<int>(parameter_file, "parameters/tracers0_enable_inner");
    this->enable_vorticity_omega = hdf5_tools::read_value<int>(parameter_file, "parameters/tracers0_enable_vorticity_omega");
    this->cutoff = hdf5_tools::read_value<double>(parameter_file, "parameters/tracers0_cutoff");
    this->inner_v0 = hdf5_tools::read_value<double>(parameter_file, "parameters/tracers0_inner_v0");
    H5Fclose(parameter_file);
    return EXIT_SUCCESS;
}

template class NSVEcomplex_particles<float>;
template class NSVEcomplex_particles<double>;