/** \brief   A small tool to read Gadget files and extract a given set
 *           of particles.
 *
 * (C)  2009, S.R. Knollmann (steffen.knollmann@uam.es)
 */


/***********************************************************************
 *    Includes                                                         *
 ***********************************************************************/
#include "../version.h"
#include "../config.h"
#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <math.h>
#include <string.h>
#include "libutil/xmem.h"
#include "libutil/xstring.h"
#include "libutil/xfile.h"
#include "libutil/diediedie.h"
#include "libio/io.h"
#include "datatypes.h"


/***********************************************************************
 *    Global variables, used for the command line arguments            *
 ***********************************************************************/

/** The filename of the input data */
char *fInName;

/** The type of the input data */
io_file_type_t fInType;

/** The filename of the input data */
char *fInPartName;


/***********************************************************************\
 *    Definitions of local functions                                   *
 ***********************************************************************/
static void
local_parse_cmdline(int argc, char **argv);

static void
local_get_ids(io_logging_t log,
              const char   *fInPartName,
              uint32_t     *numObjects,
              uint32_t     **ids,
              uint32_t     **numIds,
              uint32_t     **offsetIds,
              uint32_t     **haloIds);

static void
local_set_scaling(io_logging_t log,
                  io_file_t    f,
                  double       *r_fac,
                  double       *v_fac,
                  double       *m_fac);

static io_file_t
local_open_snapshot(io_logging_t   log,
                    char           *fInName,
                    io_file_type_t fInType,
                    io_file_t      *f,
                    uint64_t       *ptotal,
                    uint64_t       *pskip,
                    uint64_t       *pread);

static void
local_get_memory(io_logging_t   log,
                 uint32_t       numObjects,
                 const uint32_t *numIds,
                 part_t         *pbuf,
                 part_t         *pselect,
                 uint32_t       **numSelected);

static void
local_select_particles_by_id(io_logging_t   log,
                             const part_t   pbuf,
                             uint64_t       numParts,
                             part_t         pselect,
                             uint32_t       *numSelected,
                             const uint32_t *ids,
                             uint32_t       numIds);

static int
cmp_ids(const void *p1, const void *p2);

static uint64_t
local_get_new_bunch(io_logging_t log,
                    uint64_t     ptotal,
                    uint64_t     pskip,
                    uint64_t     pread,
                    part_t       pbuf,
                    io_file_t    f);

static uint64_t
local_fill_particles(part_t       parts,
                     io_file_t    f,
                     io_logging_t log,
                     uint64_t     pskip,
                     uint64_t     pread);

static void
local_output(io_logging_t log,
             uint32_t     haloID,
             part_t       pselect,
             uint32_t     numSelected,
             double       r_fac,
             double       v_fac,
             double       m_fac);


/***********************************************************************\
 *    MAIN                                                             *
 ***********************************************************************/
int
main(int argc, char **argv)
{
	uint32_t     numObjects;
	uint32_t     *ids;
	uint32_t     *numIds;
	uint32_t     *offsetIds;
	uint32_t     *haloIds;
	uint32_t     i;
	uint64_t     ptotal;
	part_t       pbuf, pselect;
	uint32_t     *numSelected;
	uint64_t     actual_read = UINT64_C(0);
	io_logging_t log;
	io_file_t    f;
	uint64_t     pskip = UINT64_C(0);
	uint64_t     pread = UINT64_MAX;
	double       r_fac, v_fac, m_fac;

	/* Do the command line stuff */
	local_parse_cmdline(argc, argv);

	/* Start logging */
	log = io_logging_start("-", 10,
	                       IO_LOGGING_FLAGS_DUPLICATE_CRITICAL
	                       | IO_LOGGING_FLAGS_USE_STDOUT_FOR_CRITICAL);

	/* Get the particle IDs to select for */
	io_logging_part(log, "Set-up");
	io_logging_section(log, "Particle IDs");
	local_get_ids(log, fInPartName, &numObjects, &ids,
	              &numIds, &offsetIds, &haloIds);

	/* Open the snapshot file */
	io_logging_section(log, "Snapshot file");
	local_open_snapshot(log, fInName, fInType, &f, &ptotal,
	                    &pskip, &pread);

	/* Get memory for particle storages */
	io_logging_section(log, "Getting memory for particles");
	local_get_memory(log, numObjects, numIds, &pbuf, &pselect, &numSelected);

	/* Actually do something */
	io_logging_part(log, "Run");
	pread = BUFFER_SIZE;
	for (pskip = UINT64_C(0); pskip < ptotal; pskip += pread) {
		/* Start a new iteration */
		io_logging_section(log, "New bunch");

		/* Get a new bunch of particles */
		io_logging_subsection(log, "Get particles");
		actual_read = local_get_new_bunch(log, ptotal, pskip,
		                                  pread, pbuf, f);

		/* Check which particles to take */
		io_logging_subsection(log, "Select particles");
		for (i = 0; i < numObjects; i++) {
			io_logging_subsection(log, "Halo %" PRIu32, haloIds[i]);
			local_select_particles_by_id(log, pbuf, actual_read,
			                             pselect + offsetIds[i],
			                             numSelected + i,
			                             ids + offsetIds[i],
			                             numIds[i]);
		}
	}

	/* Write out */
	io_logging_part(log, "Output");
	local_set_scaling(log, f, &r_fac, &v_fac, &m_fac);

	for (i = 0; i < numObjects; i++) {
		local_output(log, haloIds[i],
		             pselect + offsetIds[i], numSelected[i],
		             r_fac, v_fac, m_fac);
	}

	/* Clean up */
	io_logging_part(log, "Clean up");
	xfree(pbuf);
	xfree(ids);
	xfree(pselect);
	xfree(numIds);
	xfree(offsetIds);
	xfree(haloIds);
	xfree(numSelected);
	xfree(fInPartName);
	xfree(fInName);
	io_file_close(log, &f);
	io_logging_stop(&log);

	/* Done */
#ifdef XMEM_TRACK_MEM
	xmem_info(stdout);
#endif
	return EXIT_SUCCESS;
} /* main */

/***********************************************************************\
 *    Implementation of local functions                                *
 ***********************************************************************/
static void
local_parse_cmdline(int argc, char **argv)
{
	/* Check for sufficient parameters */
	if (argc != 4) {
		fprintf(stderr,
		        "Usage: %s <snapshot name> "
		        "<snapshot type> <ID files name>\n\n",
		        argv[0]);
		exit(EXIT_FAILURE);
	}

	/* Get the input file name */
	fInName = xstrdup(argv[1]);

	/* Get the input file type */
	fInType = (io_file_type_t)atoi(argv[2]);

	/* Get the filename of the ID file */
	fInPartName = xstrdup(argv[3]);

	/* Done */
	return;
} /* local_parse_cmdline */

static void
local_get_ids(io_logging_t log,
              const char   *fInPartName,
              uint32_t     *numObjects,
              uint32_t     **ids,
              uint32_t     **numIds,
              uint32_t     **offsetIds,
              uint32_t     **haloIds)
{
	FILE     *f;
	uint32_t i, j;

	io_logging_subsection(log, "Opening");
	f = xfopen(fInPartName, "r");

	io_logging_msg(log, INT32_C(0),
	               "%s opened successfully", fInPartName);

	io_logging_subsection(log, "Getting number of objects in the file");
	if (fscanf(f, "%" SCNu32, numObjects) != 1) {
		io_logging_fatal(log, "Could not aquire object count");
		diediedie(EXIT_FAILURE);
	}
	io_logging_msg(log, INT32_C(0),
	               "%" PRIu32 " objects in file", *numObjects);

	io_logging_subsection(log, "Getting memory");
	*numIds = xmalloc(sizeof(uint32_t) * (*numObjects));
	io_logging_msg(log, INT32_C(0), "Aquired %fkB for ID count storage",
	               (*numObjects) * sizeof(uint32_t) / 1024.);
	*offsetIds = xmalloc(sizeof(uint32_t) * (*numObjects));
	io_logging_msg(log, INT32_C(0), "Aquired %fkB for ID offset storage",
	               (*numObjects) * sizeof(uint32_t) / 1024.);
	*haloIds = xmalloc(sizeof(uint32_t) * (*numObjects));
	io_logging_msg(log, INT32_C(0), "Aquired %fkB for Halo ID storage",
	               (*numObjects) * sizeof(uint32_t) / 1024.);

	*ids        = NULL;
	**offsetIds = 0;
	for (j = 0; j < *numObjects; j++) {
		io_logging_subsection(log, "Getting IDs for halo %" PRIu32, j);
		if (fscanf(f, "%" SCNu32 " %" SCNu32,
		           (*numIds) + j, (*haloIds) + j) != 2) {
			io_logging_fatal(log, "Could not aquire particle count");
			diediedie(EXIT_FAILURE);
		}
		io_logging_msg(log, INT32_C(0),
		               "%" PRIu32 " has %" PRIu32 " particles",
		               (*haloIds)[j], (*numIds)[j]);
		if (j > 0)
			(*offsetIds)[j] = (*offsetIds)[j - 1] + (*numIds)[j-1];

		*ids = xrealloc(*ids,
		                sizeof(uint32_t)
		                * ((*numIds)[j] + (*offsetIds)[j]));
		io_logging_msg(log, INT32_C(0), "Aquired %fkB for ID storage",
		               (*numIds)[j] * sizeof(uint32_t) / 1024.);

		for (i = 0; i < (*numIds)[j]; i++) {
			if (fscanf(f, "%" SCNu32, (*ids) + (*offsetIds)[j] + i) != 1) {
				io_logging_fatal(log,
				                 "Could not aquire particle ID %" PRIu32,
				                 i);
				diediedie(EXIT_FAILURE);
			}
		}
		io_logging_msg(log, 0, "Read IDs of %" PRIu32 " particles", *numIds);

		qsort((*ids) + (*offsetIds)[j], (*numIds)[j],
		      sizeof(uint32_t), &cmp_ids);
	}

	return;
} /* local_get_ids */

static io_file_t
local_open_snapshot(io_logging_t   log,
                    char           *fInName,
                    io_file_type_t fInType,
                    io_file_t      *f,
                    uint64_t       *ptotal,
                    uint64_t       *pskip,
                    uint64_t       *pread)
{
	io_file_t dummy;

	io_logging_subsection(log, "Opening");
	*f = io_file_open(log, fInName,
	                  fInType,
	                  IO_FILE_UNKOWN_SWAPPING,
	                  IO_FILE_READ, 1);
	if (*f == NULL)
		diediedie(EXIT_FAILURE);

	io_logging_subsection(log, "Initializing");
	io_file_init(log, *f);

	io_logging_subsection(log, "Get particle numbers");
	*ptotal = io_file_get_numpart(log, *f, pskip, pread);
	io_logging_msg(log, INT32_C(0),
	               "Total number of particles: %" PRIu64,
	               *ptotal);

	io_logging_subsection(log, "File information");
	io_file_log(log, *f);

	return dummy;
} /* local_open_snapshot */

static void
local_set_scaling(io_logging_t log,
                  io_file_t    f,
                  double       *r_fac,
                  double       *v_fac,
                  double       *m_fac)
{
	double boxsize, expansion, pmass;

	io_logging_subsection(log, "Setting scaling values");
	io_file_get(log, f, IO_FILE_GET_BOXSIZE, &boxsize);
//	if (boxsize > 4000.)
//		boxsize /= 1000.;
	io_file_get(log, f, IO_FILE_GET_A, &expansion);
	io_file_get(log, f, IO_FILE_GET_PMASS, &pmass);
	*v_fac = boxsize * 100. / expansion;
	*r_fac = boxsize;
	*m_fac = pmass;
	io_logging_msg(log, INT32_C(0), "r_fac: %f", *r_fac);
	io_logging_msg(log, INT32_C(0), "v_fac: %f", *v_fac);
	io_logging_msg(log, INT32_C(0), "m_fac: %f", *m_fac);

	return;
} /* local_open_snapshot */

static void
local_get_memory(io_logging_t   log,
                 uint32_t       numObjects,
                 const uint32_t *numIds,
                 part_t         *pbuf,
                 part_t         *pselect,
                 uint32_t       **numSelected)
{
	uint32_t totNumIds, i;

	io_logging_subsection(log, "Memory for particle buffer");
	*pbuf = xmalloc(sizeof(part_struct_t) * BUFFER_SIZE);
	io_logging_msg(log, INT32_C(0),
	               "Got %fkB for storing %i particles",
	               sizeof(part_struct_t) * BUFFER_SIZE / 1024.,
	               BUFFER_SIZE);

	io_logging_subsection(log, "Memory for selected particle count");
	*numSelected = xmalloc(sizeof(uint32_t) * numObjects);
	io_logging_msg(log, INT32_C(0),
	               "Aquired %fkB for particle select count",
	               numObjects * sizeof(uint32_t) / 1024.);

	io_logging_subsection(log, "Memory for selected particle");
	totNumIds = 0;
	for (i = 0; i < numObjects; i++) {
		totNumIds     += numIds[i];
		(*numSelected)[i] = 0;
	}
	*pselect = xmalloc(sizeof(part_struct_t) * totNumIds);
	io_logging_msg(log, INT32_C(0),
	               "Got %fkB for storing %i particles",
	               sizeof(part_struct_t) * totNumIds / 1024.,
	               totNumIds);

	return;
} /* local_get_memory */

static uint64_t
local_get_new_bunch(io_logging_t log,
                    uint64_t     ptotal,
                    uint64_t     pskip,
                    uint64_t     pread,
                    part_t       pbuf,
                    io_file_t    f)
{
	uint64_t actual_read;

	pread = (ptotal - pskip < pread) ? ptotal - pskip : pread;
	io_logging_msg(log, INT32_C(0),
	               "pskip = %" PRIu64 "\tpread = %" PRIu64,
	               pskip, pread);

	actual_read = local_fill_particles(pbuf, f, log, pskip, pread);
	io_logging_msg(log, INT32_C(0),
	               "Received %" PRIu64 " particles",
	               actual_read);

	return actual_read;
} /* local_get_new_bunch */

static void
local_select_particles_by_id(io_logging_t   log,
                             const part_t   pbuf,
                             uint64_t       numParts,
                             part_t         pselect,
                             uint32_t       *numSelected,
                             const uint32_t *ids,
                             uint32_t       numIds)
{
	uint64_t i;
	uint32_t targetID;
	uint32_t *foundID;
	uint32_t numFound = 0;

	/*
	 * Loop over all particles in the buffer and check for each if it is
	 * a selected particle by looking up the ID
	 */
	for (i = UINT64_C(0); i < numParts; i++) {
		targetID = (uint32_t)((pbuf + i)->id);
		foundID  = bsearch(&targetID, ids,
		                   numIds, sizeof(uint32_t),
		                   &cmp_ids);
		if (foundID != NULL) {
			memcpy(pselect + *numSelected, pbuf + i, sizeof(part_struct_t));
			(*numSelected)++;
			numFound++;
			if (*numSelected == numIds)
				break;
		}
	}
	io_logging_msg(log, INT32_C(0), "Found %" PRIu32 " IDs", numFound);

	return;
} /* local_select_particles */

static uint64_t
local_fill_particles(part_t       parts,
                     io_file_t    f,
                     io_logging_t log,
                     uint64_t     pskip,
                     uint64_t     pread)
{
	io_file_strg_struct_t strg;
	uint64_t              actual_read = UINT64_C(0);

	/* Set the storage */
	io_logging_msg(log, INT32_C(0), "Setting the storage");
	strg.posx.val      = (void *)&(parts->pos.x);
	strg.posx.stride   = (char *)&((parts + 1)->pos.x)
	                     - (char *)&((parts)->pos.x);
	strg.posy.val      = (void *)&(parts->pos.y);
	strg.posy.stride   = (char *)&((parts + 1)->pos.y)
	                     - (char *)&((parts)->pos.y);
	strg.posz.val      = (void *)&(parts->pos.z);
	strg.posz.stride   = (char *)&((parts + 1)->pos.z)
	                     - (char *)&((parts)->pos.z);
	strg.momx.val      = (void *)&(parts->vel.x);
	strg.momx.stride   = (char *)&((parts + 1)->vel.x)
	                     - (char *)&((parts)->vel.x);
	strg.momy.val      = (void *)&(parts->vel.y);
	strg.momy.stride   = (char *)&((parts + 1)->vel.y)
	                     - (char *)&((parts)->vel.y);
	strg.momz.val      = (void *)&(parts->vel.z);
	strg.momz.stride   = (char *)&((parts + 1)->vel.z)
	                     - (char *)&((parts)->vel.z);
	strg.weight.val    = (void *)&(parts->weight);
	strg.weight.stride = (char *)&((parts + 1)->weight)
	                     - (char *)&((parts)->weight);
	strg.u.val         = (void *)&(parts + 1)->u;
	strg.u.stride      = (char *)&((parts + 1)->u)
	                      - (char *)&((parts)->u);
	strg.id.val        = (void *)&(parts->id);
	strg.id.stride     = (char *)&((parts + 1)->id)
	                     - (char *)&((parts)->id);
	strg.bytes_float   = sizeof(parts->pos.x);
	strg.bytes_int     = sizeof(parts->id);

	/* Log the storage */
	io_file_strg_log(log, strg);

	/* Do the read */
	io_logging_msg(log, INT32_C(0), "Start reading");
	actual_read = io_file_readpart(log, f, pskip, pread, strg);

	return actual_read;
} /* local_fill_particles */

static int
cmp_ids(const void *p1, const void *p2)
{
	if (*((uint32_t *)p1) < *((uint32_t *)p2))
		return -1;

	if (*((uint32_t *)p1) > *((uint32_t *)p2))
		return 1;

	return 0;
}

static void
local_output(io_logging_t log,
             uint32_t     haloId,
             part_t       pselect,
             uint32_t     numSelected,
             double       r_fac,
             double       v_fac,
             double       m_fac)
{
	uint32_t i;
	FILE     *f;
	char     fname[512];

	/* Generate the output name */
	sprintf(fname, "%07" PRIu32 ".dat", haloId);

	/* Open the file */
	f = xfopen(fname, "w");

	/* Dump all particles to the output file */
	for (i = 0; i < numSelected; i++) {
#ifdef STEREO2
#else
		fprintf(f,
		        "%" PRIu32
		        "\t%" PRIu32
		        "\t%" PRIfpv
		        "\t%" PRIfpv
		        "\t%" PRIfpv
		        "\t%" PRIfpv
		        "\t%" PRIfpv
		        "\t%" PRIfpv
		        "\t%" PRIfpv
		        "\n",
		        pselect[i].id,
		        (pselect[i].u >= 0.0) ? 0 : (int)(-(pselect[i].u)),
		        pselect[i].weight * m_fac * 1e10,
		        pselect[i].pos.x * r_fac,
		        pselect[i].pos.y * r_fac,
		        pselect[i].pos.z * r_fac,
		        pselect[i].vel.x * v_fac,
		        pselect[i].vel.y * v_fac,
		        pselect[i].vel.z * v_fac);
#endif
	}

	/* Clean */
	xfclose(&f);

	/* Done */
	return;
} /* local_output */
