/* compute projection (or slice) of a grid variable
 * (to be called from FLASH)
 * by Christoph Federrath, 2013
 */

#include "HDFIO.h"
#include "GRID3D.h"
#include "mangle_names.h"
#include <iostream>

GRID3D *grid;


// === allocate/create local c++ grids
extern "C" void FTOC(create_grids_c)(int* ngrids)
{
	grid = new GRID3D[*ngrids];
}

// === initialize local c++ grid index *ng-1, set dims and ranges
extern "C" void FTOC(init_grid_c)(int* ng, int* px, int* py, int* pz, 
                                  double* xl, double* xr, 
                                  double* yl, double* yr, 
                                  double* zl, double* zr)
{
	grid[*ng-1] = GRID3D(*px, *py, *pz);
	grid[*ng-1].set_bnds(*xl, *xr, *yl, *yr, *zl, *zr);
	//grid[*ng-1].print_bnds();
	grid[*ng-1].clear();
}

// === calls the GRID3D.h function to add data o the grid
extern "C" void FTOC(add_to_grid_c)(int* ng, double* ccx, double* ccy, double* ccz, 
                                    double* dx, double* dy, double* dz, double* dat)
{
	grid[*ng-1].add_coord_fields(*ccx, *ccy, *ccz, *dx, *dy, *dz, *dat);
}

// === copies local c++ grid to fortran grid
extern "C" void FTOC(retrieve_grid_c)(int* ng, double* grid_out)
{
	int ntot = grid[*ng-1].get_ntot();
	for (int n=0; n<ntot; n++)
		grid_out[n] = grid[*ng-1].field[n];
}

// === writes out HDF5 file
extern "C" void FTOC(writehdf5_grid_c)(int* ng, char* filename, char* writeflag, char* datasetname, 
                                       double* outgrid, double* sim_time, char* outdir)
{
	/// write HDF5 output
	HDFIO HDFOutput = HDFIO();
	//printf("filename in c++: %s\n",filename);
	std::string hdf5filename = filename;
	std::string hdf5writeflag = writeflag;
	std::string hdf5datasetname = datasetname;
	if (hdf5writeflag == "new") {
		HDFOutput.create(hdf5filename);
	}
	if (hdf5writeflag == "add") {
		HDFOutput.open(hdf5filename, 'w');
	}
	std::vector<int> grid_dims(2);
    std::string direction_string = outdir;
    if (direction_string == "xy") {
        grid_dims[0] = grid[*ng-1].get_dimy();
        grid_dims[1] = grid[*ng-1].get_dimx();
    }
    if (direction_string == "xz") {
        grid_dims[0] = grid[*ng-1].get_dimz();
        grid_dims[1] = grid[*ng-1].get_dimx();
    }
    if (direction_string == "yz") {
        grid_dims[0] = grid[*ng-1].get_dimz();
        grid_dims[1] = grid[*ng-1].get_dimy();
    }
	HDFOutput.write(outgrid, hdf5datasetname, grid_dims, H5T_NATIVE_DOUBLE);

	if (hdf5writeflag == "new") {
		std::vector<int> hdf5dims(1); hdf5dims[0] = 1;
		HDFOutput.write(sim_time, "time", hdf5dims, H5T_NATIVE_DOUBLE);
		/// write direction string
		std::string direction_string = outdir;
		if(direction_string == "xy") direction_string = "z";
		if(direction_string == "xz") direction_string = "y";
		if(direction_string == "yz") direction_string = "x";
		HDFOutput.write(direction_string.c_str(), "direction", hdf5dims, H5T_C_S1);
		/// write minmax_xyz
		hdf5dims.resize(2); hdf5dims[0] = 3; hdf5dims[1] = 2;
		double minmax_xyz[6];
		minmax_xyz[0] = grid[*ng-1].get_xmin(); minmax_xyz[1] = grid[*ng-1].get_xmax(); //min, max, x
		minmax_xyz[2] = grid[*ng-1].get_ymin(); minmax_xyz[3] = grid[*ng-1].get_ymax(); //min, max, y
		minmax_xyz[4] = grid[*ng-1].get_zmin(); minmax_xyz[5] = grid[*ng-1].get_zmax(); //min, max, z
		HDFOutput.write(minmax_xyz, "minmax_xyz", hdf5dims, H5T_NATIVE_DOUBLE);
	}

	HDFOutput.close();
	std::cout << " *** Wrote movie dataset " << hdf5datasetname << " to " << hdf5filename << " ****" << std::endl;
}

// === delete c++ grids
extern "C" void FTOC(finalize_grid_c)()
{
	delete [] grid;
}

// === writes out HDF5 file
extern "C" void FTOC(writehdf5_sinkparticles_c)(int* ng, char* filename, int* outnp, double* outracc, 
                                                double* outpx, double* outpy, double* outpz, double* outpm)
{
	/// write HDF5 output
	HDFIO HDFOutput = HDFIO();
	//printf("filename in c++: %s\n",filename);
	std::string hdf5filename = filename;
	HDFOutput.open(hdf5filename, 'w');

	int np = *outnp;
	
	/// write numpart
	std::vector<int> hdf5dims;
	hdf5dims.resize(1); hdf5dims[0] = 1;
	HDFOutput.write(&np, "numpart", hdf5dims, H5T_NATIVE_INT);
	/// write the particle accretion radius
	double r_accretion = 0.;
	if (np > 0)
	{
		r_accretion = *outracc;
	}
	hdf5dims.resize(1); hdf5dims[0] = 1;
	HDFOutput.write(&r_accretion, "r_accretion", hdf5dims, H5T_NATIVE_DOUBLE);
	if (np > 0)
	{
	  //std::cout << " Number of particles in file: " << np << std::endl;
	  std::vector<double> posx;
	  std::vector<double> posy;
	  std::vector<double> posz;
	  std::vector<double> mass;

	  for (int i = 0; i < np; i++)
	  {
		posx.push_back(outpx[i]);
		posy.push_back(outpy[i]);
		posz.push_back(outpz[i]);
		mass.push_back(outpm[i]);
	  }

	  // write particlepositions
	  hdf5dims.resize(2); hdf5dims[0] = 3; hdf5dims[1] = np;
	  double * ppos = new double[3*np];
	  for (int dir = 0; dir < 3; dir++)
	    for (int i = 0; i < np; i++)
	    {
	      long index = dir*np + i;
	      if (dir == 0) ppos[index] = posx[i];
	      if (dir == 1) ppos[index] = posy[i];
	      if (dir == 2) ppos[index] = posz[i];
	    }
	  HDFOutput.write(ppos, "particlepositions", hdf5dims, H5T_NATIVE_DOUBLE);
	  delete [] ppos;
	  // write particlemasses
	  hdf5dims.resize(1); hdf5dims[0] = np;
	  double * pmass = new double[np];
	  for (int i = 0; i < np; i++)
	    pmass[i] = mass[i];
	  HDFOutput.write(pmass, "particlemasses", hdf5dims, H5T_NATIVE_DOUBLE);
	  delete [] pmass;
	  /// rotate particles
	  std::vector<double> particle_pos(3);
	  std::vector<double> posx_proj;
	  std::vector<double> posy_proj;
	  std::vector<double> posz_proj;
	  std::vector<double> mass_proj;
	  for (int i = 0; i < np; i++)
	  {
	    particle_pos[0] = posx[i];
	    particle_pos[1] = posy[i];
	    particle_pos[2] = posz[i];
	    // check depth range (3rd dimension)
	    if( (particle_pos[0] >= grid[*ng-1].get_xmin()) && 
		    (particle_pos[0] <= grid[*ng-1].get_xmax()) &&
		    (particle_pos[1] >= grid[*ng-1].get_ymin()) && 
		    (particle_pos[1] <= grid[*ng-1].get_ymax()) &&
		    (particle_pos[2] >= grid[*ng-1].get_zmin()) && 
		    (particle_pos[2] <= grid[*ng-1].get_zmax()) )
	    {
	      posx_proj.push_back(particle_pos[0]);
	      posy_proj.push_back(particle_pos[1]);
	      posz_proj.push_back(particle_pos[2]);
	      mass_proj.push_back(mass[i]);
	    }
	  }
	  // write out projected particle positions
	  const int np_proj = posx_proj.size();
	  if (np_proj > 0)
	  {
	    double * ppos_proj = new double[3*np_proj];
	    double * pmass_proj = new double[np_proj];
	    for (int i = 0; i < np_proj; i++)
	    {
	      ppos_proj[0*np_proj+i] = posx_proj[i];
	      ppos_proj[1*np_proj+i] = posy_proj[i];
	      ppos_proj[2*np_proj+i] = posz_proj[i];
	      pmass_proj[i]          = mass_proj[i];
	    }
	    hdf5dims.resize(2); hdf5dims[0] = 3; hdf5dims[1] = np_proj;
	    HDFOutput.write(ppos_proj, "particlepositions_proj", hdf5dims, H5T_NATIVE_DOUBLE);
	    hdf5dims.resize(1); hdf5dims[0] = np_proj;
	    HDFOutput.write(pmass_proj, "particlemasses_proj", hdf5dims, H5T_NATIVE_DOUBLE);
	    delete [] ppos_proj;
	    delete [] pmass_proj;
	  }
	}

	HDFOutput.close();
	std::cout << " *** Wrote " << np << " sink particle(s) to " << hdf5filename << " ****" << std::endl;
}