#ifndef __ANALYZE_H__
#define __ANALYZE_H__

/*! @file analyze.h
 * @brief Definitions needed in analyze.c
 */

#define MINIMA_CAP 1024
#define MAX_ANGLES 601

/*!Extract energies from input (a LAMMPS log file) and write them out
 * to output, returning the minimum value in fJ.
 *
 * In the eventuality of errors log to the console and put errno(3)
 * into errnum, for handling at the call site.
 * @param[in] r Global Region.
 * @param[in] input Name of the log file.
 * @param[in] output Name of the output file.
 * @param[out] frigid_minimum Pointer to a double in which the rigid energy is
 * stored.
 * @param[out] try Current repetition of the simulation.
 * @param[out] errnum Store errno in this memory location in case of error.
 * @return Minimum energy value found.
 */
double analyze_logfire(Region *r, const char input[static restrict 1],
	const char output[static restrict 1], double frigid_minimum[static 1],
	int try, int errnum[static 1]);
/*! Check LAMMPS data file (in the format of the write_data command)
 * for out of plane z components.
 *
 * @param tmp Memory Region.
 * @param filename LAMMPS data file to check.
 * @return true if no out of plane z components are found, false otherwise.
 */
bool ensure_2d_data(Region *tmp, const char filename[static 1]);
/*! Analyze angles simulation in directory nsdirname.
 *
 * - For each file named 'logfire_r\(.*\)\.lammps' create a file named
 *   'logfire_r\1.energies' extracting the timestep and energy columns
 *   from lammps outputs, also store each of the minimum energies in
 *   fminima
 *
 * - fminima[idx] represents the minimum energy at angle
 *   (double)idx/10.0 in fJ (LAMMPS default for micro units)
 *
 * - frigid_minima[idx] represents the initial inital energy at angle
 * (double)idx/10.0 in fJ.
 *
 * - Write out the energy differences to nsdirname/NORMALIZED_FILENAME
 *   and nsdirname/ANGLES_FILENAME, which contain per-atom normalized
 *   and non-normalized energy differences in zJ.
 *
 * @param r Global Region.
 * @param d Open directory where files will be searched.
 * @param nsdirname Directory name (without trailing slash).
 * @param m Metadata for the simulation.
 * @param ensure Whether or not to ensure that no off-plane z
 * components are found in results.
 * @param[out] frigid_minima Array in which rigid energies are stored for each
 * step.
 * @param[out] fminima Array in which energies are stored for each step.
 * @return true if no errors occurred, false otherwise.
 */
bool analyze_angles(Region *r, DIR *d,
	const char nsdirname[static restrict 1],
	Metadata m[static restrict 1], bool ensure,
	double frigid_minima[static restrict MAX_ANGLES],
	double fminima[static restrict MAX_ANGLES]);
/*! Analyze translational simulation in directory nsdirname:
 *
 * - Perform the same data extraction done in analyze_angles
 *
 * - Write out energy values as a function of the translation
 *   vector in two files named nsdirname/AVG_FILENAME and
 *   nsdirname/AVG_NORMALIZED_FILENAME, names should be explicative.
 *
 * @param[in] r Global Region.
 * @param[in] d Open directory where files will be searched.
 * @param[in] nsdirname Directory name (without trailing slash).
 * @param ensure Whether or not to ensure that no off-plane z
 * components are found in results.
 * @param[out] frigid_minima Array in which rigid energies are stored for each
 * step.
 * @param[out] fminima Array in which energies are stored for each step.
 * @param[in] m Metadata for the simulation.
 * @return true if no errors occurred, false otherwise.
 */
bool analyze_translate(Region *r, DIR *d,
	const char nsdirname[static restrict 1],
	Metadata m[static restrict 1], bool ensure,
	double frigid_minima[static restrict 1],
	double fminima[static restrict 1]);

#ifdef ANALYZE_IMPLEMENTATION

/*! Convert fJ to zJ. */
static inline double
z_from_f(double p)
{
	return p * 1e6;
}

double
analyze_logfire(Region *r, const char input[static restrict 1],
	const char output[static restrict 1],
	double frigid_minimum[restrict static 1], int try,
	int errnum[static 1])
{
	FILE *in, *out;
	String in_str = { 0 };
	String line = { 0 };
	double minimum = 1e6;
	bool collecting = false, first = true;
	*errnum = 0;

	out = in = NULL;
	if (!(in = fopen(input, "r"))) {
		DBG("Error opening %s for reading: %s\n", input,
			strerror(errno));
		*errnum = errno;
		goto cleanup;
	}
	if (!(out = fopen(output, "w"))) {
		DBG("Error opening %s for writing: %s\n", output,
			strerror(errno));
		*errnum = errno;
		goto cleanup;
	}
	if (!string_read(r, in, &in_str)) {
		DBG("Error reading from %s: %s\n", input, strerror(errno));
		*errnum = errno;
		goto cleanup;
	}

	line = string_next_line(&in_str);
	while (!string_is_null(line)) {
		string_trim(&line);
		if (string_starts_with(STRING_LITERAL("Step"), line)) {
			collecting = true;
			fprintf(out, "# ");
			fprintf(out, String_Fmt "\n", String_Arg(line));
			goto next;
		} else if (string_starts_with(
				   STRING_LITERAL("Loop time of"), line)) {
			collecting = false;
			goto next;
		}

		if (collecting) {
			String tok;
			double energy;

			if (string_starts_with(
				    STRING_LITERAL("WARNING"), line)) {
				DBG("Lammps Warning: " String_Fmt "\n",
					String_Arg(line));
				goto next;
			}

			fprintf(out, String_Fmt "\n", String_Arg(line));

			tok = string_next_tok(&line, ' ');
			tok = string_next_tok(&line, ' ');
			if (!string_strtod(tok, &energy)) {
				DBG("Expected floating "
				    "point "
				    "number, found: " String_Fmt "\n",
					String_Arg(tok));
				goto next;
			}
			if (first && try == 0) {
				*frigid_minimum = energy;
				first = false;
			}
			minimum = energy;
		}

	next:
		line = string_next_line(&in_str);
	}

	fclose(in);
	fclose(out);
	return minimum;
cleanup:
	if (in)
		fclose(in);
	if (out)
		fclose(out);
	return -1;
}

bool
ensure_2d_data(Region *tmp, const char filename[static 1])
{
	String contents = { 0 }, line = { 0 };
	double z = 0;

	if (!string_read_file(tmp, filename, &contents)) {
		DBG("Error reading from %s: %s\n", filename, strerror(errno));
		return false;
	}

	line = string_next_line(&contents);
	while (line.str
		&& !string_starts_with(STRING_LITERAL("Atoms"), line)) {
		line = string_next_line(&contents);
	}

	/* Skip blank line */
	line = string_next_line(&contents);
	line = string_next_line(&contents);
	while (line.str
		&& !string_starts_with(STRING_LITERAL("Velocities"), line)) {
		String last_line = line;
		String tok = { 0 };
		size_t count = 0;
		tok = string_next_tok(&line, ' ');
		count += tok.count + 1;
		tok = string_next_tok(&line, ' ');
		count += tok.count + 1;
		tok = string_next_tok(&line, ' ');
		count += tok.count + 1;
		tok = string_next_tok(&line, ' ');
		count += tok.count + 1;
		tok = string_next_tok(&line, ' ');
		count += tok.count + 1;
		tok = string_next_tok(&line, ' ');

		if (!string_strtod(tok, &z)) {
			DBG("In %s: expected z component, found: " String_Fmt
			    "\n",
				filename, String_Arg(tok));
		}

		if (z < -EPS || z > EPS) {
			DBG("%s: non zero z component\n", filename);
			fprintf(stderr, String_Fmt "\n",
				String_Arg(last_line));
			for (size_t i = 0; i < count - 4; i++) {
				fprintf(stderr, " ");
			}
			fprintf(stderr, "~~~~^\n");
			return false;
		}

		line = string_next_line(&contents);
	}

	return true;
}

/* Used to sort in descending order with qsort(3). */
static int
int_cmp(const void *a, const void *b)
{
	return *(int *)a - *(int *)b;
}

bool
analyze_angles(Region *r, DIR *d, const char nsdirname[static restrict 1],
	Metadata m[static restrict 1], bool ensure,
	double frigid_minima[static restrict MAX_ANGLES],
	double fminima[static restrict MAX_ANGLES])
{
	FILE *angles = NULL, *normalized = NULL, *rigid = NULL,
	     *normalized_movable = NULL;
	struct dirent *dir;
	char *angles_path, *normalized_path, *rigid_path,
		*normalized_movable_path;
	size_t dirlen, angles_len, normalized_len, rigid_len,
		normalized_movable_len;
	bool already_checked = false;
	Region *tmp = region_alloc(32 * MB);
	static int indexes[MAX_ANGLES] = { 0 };
	int indexes_len = 0;
	memset(indexes, 0, LEN(indexes) * sizeof(int));

	dirlen = strlen(nsdirname);

	while ((dir = readdir(d))) {
		if (dir->d_type != DT_REG) {
			continue;
		}
		String sname = make_string(strlen(dir->d_name), dir->d_name);
		if (string_starts_with(STRING_LITERAL("logfire_r"), sname)
			&& string_ends_with(
				STRING_LITERAL(".lammps"), sname)) {
			char *path, *outpath;
			int errnum = 0, idx, try;
			double fminimum;
			/* '/' and '\0' add 2 chars */
			size_t path_len, outpath_len, d_name_len;

			d_name_len = strlen(dir->d_name);
			path_len = (dirlen + d_name_len + 2) * sizeof(char);
			path = region_malloc(tmp, path_len);

			/* len(.lammps) = 6, len(.energies) = 8 */
			outpath_len = path_len + 2;
			outpath = region_malloc(tmp, outpath_len);
			String tok;

			if (snprintf(path, path_len, "%s/%s", nsdirname,
				    dir->d_name)
				< 0) {
				DBG("Error encoding.%s", "\n");
				goto err;
			}
			strcpy(outpath, path);
			strcpy(&outpath[path_len - 1 - strlen("lammps")],
				"energies");
			outpath[outpath_len - 1] = '\0';

			LOG("Analyzing %s\n", path);
			/* Why is this here? */
			sname = make_string(d_name_len + 1, dir->d_name);
			tok = string_next_tok(&sname, '_');
			tok = string_next_tok(&sname, '_');
			assert(tok.str[0] == 'r');
			string_take(&tok, 1);
			if (!string_strtoi(tok, &idx)) {
				DBG("Error extracting "
				    "index from pathname %s, "
				    "expected an integer, "
				    "found " String_Fmt "\n",
					path, String_Arg(tok));
				goto err;
			}

			tok = string_next_tok(&sname, '_');
			assert(tok.str[0] == 't');
			string_take(&tok, 1);
			if (!string_strtoi(tok, &try)) {
				DBG("Error extracting try index from "
				    "pathname %s, expected an "
				    "integer, found " String_Fmt "\n",
					path, String_Arg(tok));
				goto err;
			}

			VERBOSE("idx: %d, try: %d\n", idx, try);

			assert(idx >= 0 && idx < MAX_ANGLES);
			fminimum = analyze_logfire(tmp, path, outpath,
				&frigid_minima[idx], try, &errnum);
			if (errnum != 0) {
				DBG("Errors occurred while analyzing "
				    "%s: %s\n",
					path, strerror(errnum));
				goto err;
			}
			indexes[indexes_len++] = idx;
			VERBOSE("%s: fminimum: %.15e, "
				"frigid_minimum: %.15e\n",
				dir->d_name, fminimum, frigid_minima[idx]);
			fminima[idx] = fminimum;
			region_reset(tmp);
		} else if (!already_checked && ensure
			&& string_starts_with(
				STRING_LITERAL("conf_finale"), sname)) {
			char *path;
			size_t path_len;

			path_len = dirlen + strlen(dir->d_name) + 2;
			path = region_malloc(tmp, path_len);
			snprintf(path, path_len, "%s/%s", nsdirname,
				dir->d_name);
			VERBOSE("Checking %s for out of plane z "
				"components.\n",
				path);
			if (!ensure_2d_data(tmp, path)) {
				goto err;
			} else {
				already_checked = true;
				region_reset(tmp);
			}
		}
	}

	/* strlen(angles.energies) = 15 */
	angles_len = (dirlen + ANGLES_FILENAME_LEN + 3) * sizeof(char);
	angles_path = region_malloc(tmp, angles_len);
	if (snprintf(
		    angles_path, angles_len, "%s/" ANGLES_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}

	/* strlen(normalized.energies) = 19 */
	normalized_len
		= (dirlen + NORMALIZED_FILENAME_LEN + 3) * sizeof(char);
	normalized_path = region_malloc(tmp, normalized_len);
	if (snprintf(normalized_path, normalized_len,
		    "%s/" NORMALIZED_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}

	rigid_len = (dirlen + RIGID_FILENAME_LEN + 3) * sizeof(char);
	rigid_path = region_malloc(tmp, rigid_len);
	if (snprintf(rigid_path, rigid_len, "%s/" RIGID_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}

	normalized_movable_len
		= (dirlen + NORMALIZED_MOVABLE_FILENAME_LEN + 3)
		* sizeof(char);
	normalized_movable_path = region_malloc(tmp, normalized_movable_len);
	if (snprintf(normalized_movable_path, normalized_movable_len,
		    "%s/" NORMALIZED_MOVABLE_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}

	if (!(angles = fopen(angles_path, "w"))) {
		DBG("Error opening %s: %s\n", angles_path, strerror(errno));
		goto err;
	}
	if (!(normalized = fopen(normalized_path, "w"))) {
		DBG("Error opening %s: %s\n", normalized_path,
			strerror(errno));
		goto err;
	}
	if (!(rigid = fopen(rigid_path, "w"))) {
		DBG("Error opening %s: %s\n", rigid_path, strerror(errno));
		goto err;
	}

	if (!(normalized_movable = fopen(normalized_movable_path, "w"))) {
		DBG("Error opening %s: %s\n", normalized_movable_path,
			strerror(errno));
		goto err;
	}

	qsort(indexes, (size_t)indexes_len, sizeof(int), int_cmp);
	fprintf(angles, "# angle [deg]\ttotal energy [zJ]\n");
	fprintf(normalized, "# angle [deg]\tenergy/N\n [zJ]");
	fprintf(normalized_movable, "# angle [deg]\tenergy/N [zJ]\n");
	fprintf(rigid, "# angle\tenergy/N/V0\n");
	for (int i = 0; i < indexes_len; i++) {
		int index = indexes[i];
		double zenergy = z_from_f(fminima[index]);
		double zrigid_energy = z_from_f(frigid_minima[index]);

		VERBOSE("Energy @ %.1f deg: %.15e zJ, rigid energy: %.15e "
			"zJ\n",
			(double)index * m->step, zenergy, zrigid_energy);
		fprintf(angles, "%.2f %.15e\n", (double)index * m->step,
			zenergy - zrigid_energy);
		fprintf(normalized, "%.2f %.15e\n", (double)index * m->step,
			(zenergy - zrigid_energy) / m->tot_natoms);
		fprintf(normalized_movable, "%.2f %.15e\n",
			(double)index * m->step,
			(zenergy - zrigid_energy) / m->natoms);
		fprintf(rigid, "%.2f %.15e\n", (double)index * m->step,
			frigid_minima[index] / m->natoms / m->v0);
	}

	fclose(angles);
	fclose(normalized);
	fclose(rigid);
	fclose(normalized_movable);
	region_free(tmp);
	return true;
err:
	if (angles)
		fclose(angles);
	if (normalized)
		fclose(normalized);
	if (rigid)
		fclose(rigid);
	if (normalized_movable)
		fclose(normalized_movable);
	region_free(tmp);
	return false;
}

bool
analyze_translate(Region *r, DIR *d, const char nsdirname[static restrict 1],
	Metadata m[static restrict 1], bool ensure,
	double frigid_minima[static restrict 1],
	double fminima[static restrict 1])
{
	struct dirent *dir;
	size_t dirlen, normalized_len, normalized_movable_len, rigid_len,
		translation_len;
	char *normalized_path = NULL, *normalized_movable_path = NULL,
	     *rigid_path = NULL, *translation_path = NULL;
	bool ok = true, already_checked = false;
	double norm = 0.0;
	FILE *normalized = NULL, *normalized_movable = NULL, *rigid = NULL,
	     *translation = NULL;
	Region *tmp = region_alloc(8 * MB);
	int *indexes = NULL;
	int indexes_len = 0;
	indexes = ecalloc((size_t)m->nstep + 1, sizeof(int));

	dirlen = strlen(nsdirname);

	while ((dir = readdir(d))) {
		if (dir->d_type != DT_REG) {
			continue;
		}
		String sname = make_string(strlen(dir->d_name), dir->d_name);
		if (string_starts_with(STRING_LITERAL("logfire_r"), sname)
			&& string_ends_with(
				STRING_LITERAL(".lammps"), sname)) {
			char *path, *outpath;
			int errnum = 0, idx;
			double fminimum;
			size_t path_len, outpath_len;
			path_len = (sname.count + dirlen + 2) * sizeof(char);
			path = region_malloc(tmp, path_len);

			outpath_len = path_len + 2;
			outpath = region_malloc(tmp, outpath_len);
			String tok;

			if (snprintf(path, path_len, "%s/%s", nsdirname,
				    dir->d_name)
				< 0) {
				DBG("Error encoding.%s", "\n");
				goto err;
			}
			strcpy(outpath, path);
			strcpy(&outpath[path_len - 1 - strlen("lammps")],
				"energies");
			outpath[outpath_len - 1] = '\0';

			LOG("Analyzing %s\n", path);
			tok = string_next_tok(&sname, '_');
			tok = string_next_tok(&sname, '_');
			assert(tok.str[0] == 'r');
			string_take(&tok, 1);
			if (!string_strtoi(tok, &idx)) {
				DBG("Error extracting index from "
				    "pathname %s, expected an "
				    "integer, found " String_Fmt "\n",
					path, String_Arg(tok));
				goto err;
			}
			VERBOSE("idx: %d\n", idx);
			fminimum = analyze_logfire(tmp, path, outpath,
				&frigid_minima[idx], 0, &errnum);
			if (errnum != 0) {
				DBG("Errors occurred while analyzing "
				    "%s: %s\n",
					path, strerror(errnum));
				goto err;
			}
			assert(indexes_len < m->nstep + 1);
			indexes[indexes_len++] = idx;
			VERBOSE("%s: fminimum: %.15e, "
				"frigid_minimum: %.15e\n",
				dir->d_name, fminimum, frigid_minima[idx]);
			fminima[idx] = fminimum;
			region_reset(tmp);
		} else if (!already_checked && ensure
			&& string_starts_with(
				STRING_LITERAL("conf_finale"), sname)) {
			char *path;
			size_t path_len;

			path_len = dirlen + strlen(dir->d_name) + 2;
			path = region_malloc(tmp, path_len);
			snprintf(path, path_len, "%s/%s", nsdirname,
				dir->d_name);
			VERBOSE("Checking %s for out of plane z "
				"components.\n",
				path);
			if (!ensure_2d_data(tmp, path)) {
				goto err;
			} else {
				already_checked = true;
				region_reset(tmp);
			}
		}
	}

	/* strlen(normalized.energies) = 19 */
	normalized_len
		= (dirlen + NORMALIZED_FILENAME_LEN + 3) * sizeof(char);
	normalized_path = region_malloc(tmp, normalized_len);
	if (snprintf(normalized_path, normalized_len,
		    "%s/" NORMALIZED_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}
	normalized_movable_len
		= (dirlen + NORMALIZED_MOVABLE_FILENAME_LEN + 3)
		* sizeof(char);
	normalized_movable_path = region_malloc(tmp, normalized_movable_len);
	if (snprintf(normalized_movable_path, normalized_movable_len,
		    "%s/" NORMALIZED_MOVABLE_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}
	rigid_len = (dirlen + RIGID_FILENAME_LEN + 3) * sizeof(char);
	rigid_path = region_malloc(tmp, rigid_len);
	if (snprintf(rigid_path, rigid_len, "%s/" RIGID_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}
	/* strlen(translation.energies) = 20 */
	translation_len
		= (dirlen + TRANSLATION_FILENAME_LEN + 3) * sizeof(char);
	translation_path = region_malloc(tmp, translation_len);
	if (snprintf(translation_path, translation_len,
		    "%s/" TRANSLATION_FILENAME, nsdirname)
		< 0) {
		DBG("Encoding error.%s", "\n");
		goto err;
	}

	if (!(normalized = fopen(normalized_path, "w"))) {
		DBG("Error opening %s: %s\n", normalized_path,
			strerror(errno));
		goto err;
	}
	if (!(normalized_movable = fopen(normalized_movable_path, "w"))) {
		DBG("Error opening %s: %s\n", normalized_movable_path,
			strerror(errno));
		goto err;
	}
	if (!(rigid = fopen(rigid_path, "w"))) {
		DBG("Error opening %s: %s\n", rigid_path, strerror(errno));
		goto err;
	}
	if (!(translation = fopen(translation_path, "w"))) {
		DBG("Error opening %s: %s\n", translation_path,
			strerror(errno));
		goto err;
	}

	norm = sqrt(m->translation[0] * m->translation[0]
		       + m->translation[1] * m->translation[1])
		* m->apot;
	qsort(indexes, (size_t)indexes_len, sizeof(int), int_cmp);
	fprintf(translation, "# distance [um]\ttotal energy [zJ]\n");
	fprintf(normalized, "# distance [um]\tenergy/N [zJ]\n");
	fprintf(normalized_movable, "# distance [um]\tenergy/N [zJ]\n");
	fprintf(rigid, "# distance [um]\trigid energy/N [zJ]\n");
	for (int i = 0; i < indexes_len; i++) {
		int index = indexes[i];
		double zenergy = z_from_f(fminima[index]);
		double zrigid_energy = z_from_f(frigid_minima[index]);
		double distance = norm * (double)i / (double)m->nstep;
		VERBOSE("Energy @ %.5f um: %.15e zJ, rigid energy: %.15e "
			"zJ\n",
			distance, zenergy, zrigid_energy);
		fprintf(rigid, "%.5f %.15e\n", distance,
			zrigid_energy / m->natoms);
		fprintf(translation, "%.5f %.15e\n", distance,
			zenergy - zrigid_energy);
		fprintf(normalized, "%.5f %.15e\n", distance,
			(zenergy - zrigid_energy) / m->tot_natoms);
		fprintf(normalized_movable, "%.5f %.15e\n", distance,
			(zenergy - zrigid_energy) / m->natoms);
	}

	free(indexes);
	fclose(normalized);
	fclose(normalized_movable);
	fclose(translation);
	fclose(rigid);
	region_free(tmp);
	return ok;
err:
	if (indexes)
		free(indexes);
	if (normalized)
		fclose(normalized);
	if (normalized_movable)
		fclose(normalized_movable);
	if (rigid)
		fclose(rigid);
	if (translation)
		fclose(translation);
	region_free(tmp);
	return false;
}

#endif /* ANALYZE_IMPLEMENTATION */
#endif /* __ANALYZE_H__ */