/*
 * LAMMPSDataFileExporter.cpp
 *
 *  Created on: Mar 1, 2020
 *      Author: lorenzo
 */

#include "LAMMPSDataFileExporter.h"

#include "../System.h"
#include "../topology/Topology.h"
#include "../utils/strings.h"

#include <fstream>

namespace ba {

LAMMPSDataFileExporter::LAMMPSDataFileExporter(std::string atom_style) :
		_atom_style(atom_style) {
	if(std::find(_supported_atom_styles.begin(), _supported_atom_styles.end(), atom_style) == _supported_atom_styles.end()) {
		std::string error = fmt::format("Exporting LAMMPS atom_style '{}' is not supported", atom_style);
		throw std::runtime_error(error);
	}
}

LAMMPSDataFileExporter::~LAMMPSDataFileExporter() {

}

void LAMMPSDataFileExporter::_write_system_to_stream(std::shared_ptr<System> system, std::ostream &output) {
	auto type_map = _ba_LAMMPS_type_map(system);
	_set_id_shift(system);

	Topology topology(system.get());
	topology.apply(system.get());
	uint N_bonds = topology.bonds().size();
	uint N_angles = topology.angles().size();
	uint N_dihedrals = topology.dihedrals().size();

	output << "Generated by baggianalysis" << std::endl << std::endl;

	output << fmt::format("{} atoms", system->N()) << std::endl;
	output << fmt::format("{} atom types", type_map.size()) << std::endl;
	if(N_bonds > 0) {
		output << fmt::format("{} bonds", N_bonds) << std::endl;
		output << fmt::format("{} bond types", topology.bond_types().size()) << std::endl;
	}

	if(N_angles > 0) {
		output << fmt::format("{} angles", N_angles) << std::endl;
		output << fmt::format("{} angle types", topology.angle_types().size()) << std::endl;
	}

	if(N_dihedrals > 0) {
		output << fmt::format("{} dihedrals", N_dihedrals) << std::endl;
		output << fmt::format("{} dihedral types", topology.dihedral_types().size()) << std::endl;
	}

	output << std::endl;

	// print the box
	auto box_low = -system->box / 2.;
	auto box_high = box_low + system->box;
	output << fmt::format("{} {} xlo xhi", box_low[0], box_high[0]) << std::endl;
	output << fmt::format("{} {} ylo yhi", box_low[1], box_high[1]) << std::endl;
	output << fmt::format("{} {} zlo zhi", box_low[2], box_high[2]) << std::endl;

	output << std::endl;
	output << "Masses" << std::endl << std::endl;

	std::map<particle_type, bool> printed;
	for(auto type : type_map) {
		printed[type.first] = false;
	}

	for(auto p : system->particles()) {
		auto &type = p->type();
		if(!printed[type]) {
			output << fmt::format("{} {}", type_map[type], p->mass()) << std::endl;
			printed[type] = true;
		}
	}

	output << std::endl;
	output << "Atoms" << std::endl << std::endl;

	for(auto p : system->particles()) {
		output << _particle_line(p, type_map[p->type()]) << std::endl;
	}

	if(N_bonds > 0) {
		output << std::endl;
		output << "Bonds" << std::endl << std::endl;

		// LAMMPS's ids start from 1
		uint i = 1;
		for(auto bond : topology.bonds()) {
			output << fmt::format("{} {} {} {}", i, bond.type,
					_shifted(bond[0]), _shifted(bond[1])) << std::endl;
			i++;
		}
	}

	if(N_angles > 0) {
		output << std::endl;
		output << "Angles" << std::endl << std::endl;

		// LAMMPS's ids start from 1
		uint i = 1;
		for(auto angle : topology.angles()) {
			output << fmt::format("{} {} {} {} {}", i, angle.type,
					_shifted(angle[0]), _shifted(angle[1]), _shifted(angle[2])) << std::endl;
			i++;
		}
	}

	if(N_dihedrals > 0) {
		output << std::endl;
		output << "Dihedrals" << std::endl << std::endl;

		// LAMMPS's ids start from 1
		uint i = 1;
		for(auto dihedral : topology.dihedrals()) {
			output << fmt::format("{} {} {} {} {} {}", i, dihedral.type,
					_shifted(dihedral[0]), _shifted(dihedral[1]), _shifted(dihedral[2]), _shifted(dihedral[3])) << std::endl;
			i++;
		}
	}
}

void LAMMPSDataFileExporter::_set_id_shift(std::shared_ptr<System> system) {
	auto indexes = system->indexes();
	std::sort(indexes.begin(), indexes.end());
	_id_shift = 1 - indexes[0];
}

int LAMMPSDataFileExporter::_shifted(int idx) {
	return idx + _id_shift;
}

std::map<particle_type, int> LAMMPSDataFileExporter::_ba_LAMMPS_type_map(std::shared_ptr<System> system) {
	std::map<particle_type, int> type_map;

	auto types = system->types();
	// sorting and removal is faster than other options (see https://stackoverflow.com/a/56034014/5140209)
	// TODO run some performance tests
	std::sort(types.begin(), types.end());
	types.erase(std::unique(types.begin(), types.end()), types.end());

	bool non_positive_found = false;
	bool not_integer_found = false;
	int smaller_type = 0;
	for(auto type: types) {
		try {
			int int_type = utils::lexical_cast<int>(type);
			if(int_type < 1) {
				non_positive_found = true;
			}
			if(int_type < smaller_type) {
				smaller_type = int_type;
			}

			type_map[type] = int_type;
		}
		catch(utils::bad_lexical_cast &) {
			not_integer_found = true;
		}
	}

	if(not_integer_found) {
		BA_WARNING("Non-integer particle type(s) found: particles will be assigned integer types, please check that this is not an issue");

		int curr_type = 1;
		for(auto type: types) {
			type_map[type] = curr_type;
			curr_type++;
		}
	}
	else if(non_positive_found) {
		int type_shift = 1 - smaller_type;
		BA_WARNING("Non-positive particle type(s) found: all particle types will be shifted by {} so that the smallest one will be 1 as requested by LAMMPS", type_shift);

		for(auto type: types) {
			type_map[type] += type_shift;
		}
	}

	return type_map;
}

std::string LAMMPSDataFileExporter::_particle_line(std::shared_ptr<Particle> p, int int_type) {
	int p_id = _shifted(p->index());
	particle_type p_type = p->type();
	std::string line = fmt::format("{}", p_id);
	if(_atom_style == "bond" || _atom_style == "full" || _atom_style == "molecular") {
		int mol_id = ba::utils::lexical_cast<int>(ba::utils::split(p->molecule()->name(), "_")[1]) + 1;
		line += fmt::format(" {}", mol_id);
	}

	// print the type
	line += fmt::format(" {}", int_type);

	if(_atom_style == "full") {
		line += fmt::format(" {}", p->charge());
	}

	// print the position
	line += fmt::format(" {} {} {}", p->position()[0], p->position()[1], p->position()[2]);

	return line;
}

#ifdef PYTHON_BINDINGS

void export_LAMMPSDataFileExporter(py::module &m) {
	py::class_<LAMMPSDataFileExporter, BaseExporter, std::shared_ptr<LAMMPSDataFileExporter>> exporter(m, "LAMMPSDataFileExporter",
			R"pbdoc(
Export configurations to the LAMMPS data file format.

Note that LAMMPS requires particle indexes and types to be positive-definite integers. By contrast, baggianalysis uses
zero-based indexes and string types. If the system to be printed does not contain LAMMPS-valid values (which is common if 
baggianalysis is being used to convert between different formats) then the exporter will try to fix them either
by shifting values so they start from 1 or, in the case of string types, by using 1-based sequences.
	)pbdoc");

	exporter.def(py::init<std::string>(),
			R"pbdoc(
The constructor takes one mandatory argument.

Parameters
----------
	atom_style : str
		The atom style to be used in printing the data file (see `here <https://lammps.sandia.gov/doc/atom_style.html>`_ for details). As of now, only the "atomic" and "bond" styles are supported.
	)pbdoc");
}

#endif

} /* namespace ba */
