/* See copyright information at the end of the file */

/*! @file lattice_scheme_generator.c
 * @brief Prints a circular array of atoms arranged in an hexagonal lattice.
 */

#include <assert.h>
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define SQRT3_2 0.86602540378443864676
#define MAX_BONDS 3

/*! Possible output formats. */
enum GENERATE {
	LAMMPS,			/*!< Lammps input. */
	XYZ,			/*!< XYZ file. */
	NGENERATE,
};

/*! Properties of each atom. */
enum PROPERTIES {
	EXCLUDED = 1 << 0, /*!< Whether this atom is excluded from the grid */
	BORDER = 1 << 1,   /*!< Whether this atoms belongs to the border */
	NPROPERTIES,
};

/*! A single atom in the grid. */
struct Atom {
	uint8_t mask;		/*!< mask of properties. @see PROPERTIES */
	double x, y;
	int bonds[MAX_BONDS]; /*!< Index of connected atoms in grid */
};

/*! The global grid of atoms */
struct Atoms {
	struct Atom *arr;	/*!< Array of all of the atoms, stored
				 * in row-major order. */
	size_t len;		/*!< Length of arr. */
	size_t capacity;	/*!< Total capacity of arr. */
	size_t excluded; /*!< Number of excluded atoms, so that
			   included ones are len - excluded */
	long bonds;	 /*!< Total number of bonds. */
};

/*! Rectangular grid out of which we carve a circle */
static struct Atoms grid = { 0 };
static int width, height; /* Of the initial rectangular grid */

static inline double
norm2(const struct Atom atom)
{
	return atom.x * atom.x + atom.y * atom.y;
}

static inline int
excluded(const struct Atom atom)
{
	return atom.mask & EXCLUDED;
}

static inline int
border(const struct Atom atom)
{
	return atom.mask & BORDER;
}

/*! Cut out a circle from grid.
 *
 * All of the atoms between the internal and the external radius will
 * be marked as BORDER atoms, so that they can be handled differently
 * when printing the output.
 *
 * @param internal_radius Radius of the inner circle.
 * @param external_radius Total radius of the circle.
 */
void
cut_circle(const double internal_radius, const double external_radius)
{
	long i, bond;

	for (i = 0; i < grid.len; i++) {
		if (norm2(grid.arr[i]) < external_radius * external_radius) {
			for (bond = 0; bond < MAX_BONDS; bond++) {
				if (norm2(grid.arr[grid.arr[i].bonds[bond]])
					> external_radius * external_radius) {
					grid.arr[i].bonds[bond] = -1;
				}
			}
		} else {
			grid.arr[i].mask |= EXCLUDED;
			grid.excluded++;
		}

		if (norm2(grid.arr[i]) > internal_radius * internal_radius)
			grid.arr[i].mask |= BORDER;
	}
}

/*! Print grid as a LAMMPS output file.
 *
 * @param box_width Width of the LAMMPS bounding box
 * (https://docs.lammps.org/region.html).
 * @param box_height Height of the LAMMPS bounding box
 * (https://docs.lammps.org/region.html).
 */
void
generate_lammps(const double box_width, const double box_height)
{
	int i;
	size_t bond, bond_counter;
	time_t epoch;
	struct tm *timeinfo;

	time(&epoch);
	timeinfo = localtime(&epoch);
	printf("LAMMPS data file, generated on %s# n_atoms\n",
		asctime(timeinfo)); /* asctime includes a newline */
	printf("%zu atoms\n2 atom types\n%ld bonds\n1 bond types\n",
		grid.len - grid.excluded, grid.bonds);
	printf("#scatola\n%.16f %.16f xlo xhi\n%.16f %.16f ylo yhi\n-5.0 5.0 "
	       "zlo zhi\n\n",
		-box_width, box_width, -box_height, box_height);
	printf("Masses\n\n1 31.06\n2 31.06\n\n");

	printf("Atoms # bond\n\n");
	for (i = 0; i < grid.len; i++) {
		if (!excluded(grid.arr[i])) {
			if (border(grid.arr[i])) {
				printf("%d 1 2 ", i);
			} else {
				printf("%d 1 1 ", i);
			}
			printf("%f %f 0.0 0 0 0\n", grid.arr[i].x,
				grid.arr[i].y);
		}
	}

	printf("\nBonds #\n\n");
	/*
	 * The first bond has to be indexed 1, index 0 will make lammps
	 * segfault
	 */
	for (i = 0, bond_counter = 1; i < grid.len; i++) {
#ifdef DEBUG
		if (i != 0 && i % width == 0)
			fprintf(stderr, "\n");
#endif
		if (!excluded(grid.arr[i])) {
			for (bond = 0; bond < MAX_BONDS; bond++) {
				if (grid.arr[i].bonds[bond] >= 0
					&& !excluded(
						grid.arr[grid.arr[i].bonds
								 [bond]]))
					printf("%zu 1 %d %d\n",
						bond_counter++, i,
						grid.arr[i].bonds[bond]);
#ifdef DEBUG
				fprintf(stderr, "%d -> %d ", i,
					grid.arr[i].bonds[bond]);
#endif
			}
		}
	}
}

/*! Print grid as an XYZ file. */
void
generate_xyz(void)
{
	int i;

	printf("%zu\n\n", grid.len - grid.excluded);
	for (i = 0; i < grid.len; i++) {
		if (!excluded(grid.arr[i])) {
			if (border(grid.arr[i])) {
				printf("N ");
			} else {
				printf("C ");
			}
			printf("%f %f 0.0\n", grid.arr[i].x, grid.arr[i].y);
		}
	}
}

static void
usage(const char *program)
{
	fprintf(stderr, "Usage: %s [-lx] internal-radius external-radius\n\n",
		program);
	fprintf(stderr, "Options:\n");
	fprintf(stderr, "\t-l: Output lammps data\n");
	fprintf(stderr, "\t-x: Ouput xyz data\n");
	exit(EXIT_FAILURE);
}

int
main(int argc, char **argv)
{
	char *endptr;
	double internal_radius, external_radius;
	long x, y;
	int i;
	size_t bond;
	static double spacing = 5.8;
	double box_width, box_height;
	enum GENERATE mode;

	if (argc < 4) {
		usage(argv[0]);
	}

	if (strlen(argv[1]) == 2) {
		switch (argv[1][1]) {
		case 'l':
			mode = LAMMPS;
			break;
		case 'x':
			mode = XYZ;
			break;
		default:
			usage(argv[0]);
		}
	} else {
		usage(argv[0]);
	}

	errno = 0;
	internal_radius = strtod(argv[2], &endptr);
	internal_radius /= 2.;
	if (errno) {
		fprintf(stderr, "%s", strerror(errno));
		exit(EXIT_FAILURE);
	}

	external_radius = strtod(argv[3], &endptr);
	external_radius /= 2.;
	if (errno) {
		fprintf(stderr, "%s", strerror(errno));
		exit(EXIT_FAILURE);
	}

	assert(external_radius > 0 && internal_radius > 0
		&& "Radii must be positive");
	height = width = (int)external_radius * 3;
	if (width % 2 != 0)
		width = height += 1;
	box_height = box_width = (double)width * spacing;
	grid.capacity = (size_t)(width * height);
	grid.arr = calloc(grid.capacity, sizeof(*grid.arr));
	grid.len = 0;
#ifdef DEBUG
	fprintf(stderr, "w: %d h: %d\n", width, height);
#endif

	for (i = 0, y = -height / 2 + 1; y <= height / 2; y++) {
#ifdef DEBUG
		fprintf(stderr, "r%03ld ", y);
		if (y % 2 == 0)
			fprintf(stderr, "   ");
#endif
		for (x = -width / 2 + 1; x <= width / 2; x++, i++) {
			if (y % 2 == 0) {
				struct Atom atom = {
					.mask = 0,
					.x = spacing / 2. + (double)(x - 1) * spacing,
					.y = (double)(y - 1) * SQRT3_2 * spacing,
					/* (i+width) (i+width+1)
					 *   x       x
					 *    \     /
					 *     \   /
					 *      \ /
					 *       x-------x
					 *       i       (i+1)
					 */
					.bonds = {
						x == width / 2 ? -1 : i + 1,
						y == height / 2 ? -1 : i + width,
						y == width / 2 || x == width / 2 ? -1 : i + width + 1,
					},
				};
				grid.arr[grid.len++] = atom;
			} else {
				struct Atom atom = {
					.mask = 0,
					.x = (double)(x - 1) * spacing,
					.y = (double)(y - 1) * SQRT3_2 * spacing,
					/* (i+width-1) (i+width)
					 *  x         x
					 *   \       /
					 *    \     /
					 *     \   /
					 *      \ /
					 *       x---------x
					 *       i         (i+1)
					 */
					.bonds = {
						x == width / 2 ? -1 : i + 1,
						y == height / 2 || x == - width / 2 + 1 ? -1 : i + width - 1,
						y == height / 2 ? - 1 : i + width,
					},
				};
				grid.arr[grid.len++] = atom;
			}
#ifdef DEBUG
			fprintf(stderr, "%03d ", i);
#endif
		}
#ifdef DEBUG
		fprintf(stderr, "\n");
#endif
	}

	cut_circle((double)internal_radius * spacing,
		(double)external_radius * spacing);

	for (i = 0; i < grid.len; i++) {
		if (!excluded(grid.arr[i])) {
			for (bond = 0; bond < MAX_BONDS; bond++) {
				if (grid.arr[i].bonds[bond] >= 0
					&& !excluded(
						grid.arr[grid.arr[i].bonds
								 [bond]]))
					grid.bonds++;
			}
		}
	}

	if (mode == LAMMPS) {
		generate_lammps(box_width, box_height);
	} else if (mode == XYZ) {
		generate_xyz();
	}

	free(grid.arr);
	return 0;
}

/*
 * Copyright ©️ 2023 Mario Forzanini <mf@marioforzanini.com>
 *
 * This file is part of my bachelor thesis.
 *
 * This file 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.
 *
 * This file 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 this file. If not, see <https://www.gnu.org/licenses/>.
 *
 */