/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2020 NKI/AVL, Netherlands Cancer Institute
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice, this
 *    list of conditions and the following disclaimer
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
 * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <algorithm>
#include <cassert>

#include <algorithm>
#include <cif++/row.hpp>
#include <fstream>
#include <mutex>

#include <cif++/cif++.hpp>

#include <pdb-redo/BondMap.hpp>
#include <pdb-redo/Compound.hpp>
#include <ranges>

namespace pdb_redo
{



// --------------------------------------------------------------------

struct CompoundBondInfo
{
	std::string mID;
	std::set<std::tuple<uint32_t, uint32_t>> mBonded;

	[[nodiscard]] bool bonded(uint32_t a1, uint32_t a2) const
	{
		return mBonded.count({ a1, a2 }) > 0;
	}
};

// --------------------------------------------------------------------

class CompoundBondMap
{
  public:
	static CompoundBondMap &instance()
	{
		return sInstance;
	}

	bool bonded(const std::string &compoundID, const std::string &atomID1, const std::string &atomID2);

  private:
	CompoundBondMap() = default;

	uint32_t getAtomID(const std::string &atomID)
	{
		std::string id(atomID);

		uint32_t result;

		auto i = mAtomIDIndex.find(id);
		if (i == mAtomIDIndex.end())
		{
			result = static_cast<uint32_t>(mAtomIDIndex.size());
			mAtomIDIndex[id] = result;
		}
		else
			result = i->second;

		return result;
	}

	std::map<std::string, uint32_t> mAtomIDIndex;
	std::vector<CompoundBondInfo> mCompounds;
	std::mutex mMutex;
    static CompoundBondMap sInstance;
};

CompoundBondMap CompoundBondMap::sInstance;

bool CompoundBondMap::bonded(const std::string &compoundID, const std::string &atomID1, const std::string &atomID2)
{
	std::scoped_lock lock(mMutex);

	using namespace std::literals;

	std::string id(compoundID);
	uint32_t a1 = getAtomID(atomID1);
	uint32_t a2 = getAtomID(atomID2);
	if (a1 > a2)
		std::swap(a1, a2);

	for (auto &bi : mCompounds)
	{
		if (bi.mID != id)
			continue;

		return bi.bonded(a1, a2);
	}

	bool result = false;

	// not found in our cache, calculate
	CompoundBondInfo bondInfo{ id };

	auto compound = cif::compound_factory::instance().create(compoundID);
	if (not compound)
	{
		if (cif::VERBOSE >= 0)
			std::cerr << "Missing compound bond info for " << compoundID << '\n';
	}
	else
	{
		for (auto &atom : compound->bonds())
		{
			uint32_t ca1 = getAtomID(atom.atom_id[0]);
			uint32_t ca2 = getAtomID(atom.atom_id[1]);
			if (ca1 > ca2)
				std::swap(ca1, ca2);

			bondInfo.mBonded.insert({ ca1, ca2 });
			result = result or (a1 == ca1 and a2 == ca2);
		}
	}

	mCompounds.push_back(bondInfo);

	return result;
}

// --------------------------------------------------------------------

BondMap::BondMap(BondMap &&bm)
	: dim(bm.dim)
	, index(std::move(bm.index))
	, bond(std::move(bm.bond))
	, bond_1_4(std::move(bm.bond_1_4))
	, link(std::move(bm.link))
{
}

BondMap &BondMap::operator=(BondMap &&bm)
{
	dim = bm.dim;
	index.swap(bm.index);
	bond.swap(bm.bond);
	bond_1_4.swap(bm.bond_1_4);
	link.swap(bm.link);

	return *this;
}

BondMap::BondMap(const cif::datablock &db, std::optional<std::tuple<cif::point, float>> around, int model_nr)
{
	using namespace cif::literals;

	auto &compoundBondInfo = CompoundBondMap::instance();

	// First collect the atoms from the datablock
	std::vector<cif::const_row_handle> atoms;

	cif::crystal crystal(db);

	for (auto rh : db["atom_site"].find("pdbx_PDB_model_num"_key == model_nr or "pdbx_PDB_model_num"_key == cif::null))
	{
		if (around)
		{
			const auto &[p, r] = *around;
			const auto &[x, y, z] = rh.get<float, float, float>("Cartn_x", "Cartn_y", "Cartn_z");

			const auto &[d, pt, op] = crystal.closest_symmetry_copy(p, { x, y, z });

			if (d <= r)
				atoms.push_back(rh);
		}
		else
			atoms.push_back(rh);
	}

	dim = static_cast<uint32_t>(atoms.size());

	for (auto atom : atoms)
	{
		auto id = atom["id"].template get<std::string>();
		index[id] = static_cast<uint32_t>(index.size());
		rIndex.emplace_back(id);
	}

	auto bindAtoms = [this, &atoms](const std::string &a, const std::string &b)
	{
		uint32_t ixa = index[a];
		uint32_t ixb = index[b];

		auto altA = atoms[ixa].get<std::optional<std::string>>("label_alt_id");
		auto altB = atoms[ixb].get<std::optional<std::string>>("label_alt_id");

		if (not altA.has_value() or not altB.has_value() or *altA == *altB)
			bond.insert(key(ixa, ixb));
	};

	auto linkAtoms = [this, &bindAtoms](const std::string &a, const std::string &b)
	{
		bindAtoms(a, b);

		link[a].insert(b);
		link[b].insert(a);
	};

	// collect all compounds first
	std::set<std::string> compounds;
	for (const auto &comp_id : db["chem_comp"].rows<std::string>("id"))
		compounds.insert(comp_id);

	// make sure we also have all residues in the polyseq
	for (const auto &mon_id : db["entity_poly_seq"].rows<std::string>("mon_id"))
	{
		if (compounds.count(mon_id))
			continue;

		if (cif::VERBOSE > 1)
			std::cerr << "Warning: mon_id " << mon_id << " is missing in the chem_comp category\n";
		compounds.insert(mon_id);
	}

	cif::progress_bar progress_bar(compounds.size(), "Creating bond map");

	// some helper indices to speed things up a bit
	using atom_map_key_type = std::tuple<std::string, int, std::string, std::string>;
	std::map<atom_map_key_type, std::string> atomMapByAsymSeqAndAtom;
	for (auto a : atoms)
	{
		atom_map_key_type key = a.get("label_asym_id", "label_seq_id", "label_atom_id", "auth_seq_id");
		atomMapByAsymSeqAndAtom[key] = a.get<std::string>("id");
	}

	// first link all residues in a polyseq

	std::string lastAsymID, lastAuthSeqID;
	int lastSeqID = 0;
	for (const auto &[asymID, seqID, authSeqID] : db["pdbx_poly_seq_scheme"].rows<std::string, int, std::string>("asym_id", "seq_id", "pdb_seq_num"))
	{
		if (asymID != lastAsymID) // first in a new sequece
		{
			lastAsymID = asymID;
			lastSeqID = seqID;
			lastAuthSeqID = authSeqID;
			continue;
		}

		auto kc = make_tuple(asymID, lastSeqID, "C", lastAuthSeqID);
		auto kn = make_tuple(asymID, seqID, "N", authSeqID);

		if (atomMapByAsymSeqAndAtom.count(kc) and atomMapByAsymSeqAndAtom.count(kn))
		{
			auto c = atomMapByAsymSeqAndAtom.at(kc);
			auto n = atomMapByAsymSeqAndAtom.at(kn);

			bindAtoms(c, n);
		}
		// if (not(c.empty() or n.empty()))

		lastSeqID = seqID;
		lastAuthSeqID = authSeqID;
	}

	for (auto l : db["struct_conn"])
	{
		atom_map_key_type ka = l.get("ptnr1_label_asym_id", "ptnr1_label_seq_id", "ptnr1_label_atom_id", "ptnr1_auth_seq_id");
		atom_map_key_type kb = l.get("ptnr2_label_asym_id", "ptnr2_label_seq_id", "ptnr2_label_atom_id", "ptnr2_auth_seq_id");

		if (atomMapByAsymSeqAndAtom.count(ka) and atomMapByAsymSeqAndAtom.count(kb))
		{
			auto a = atomMapByAsymSeqAndAtom.at(ka);
			auto b = atomMapByAsymSeqAndAtom.at(kb);

			linkAtoms(a, b);
		}
	}

	// then link all atoms in the compounds

	std::vector<cif::const_row_handle> rAtoms;
	std::string lastEntityID, lastCompID;

	lastAsymID.clear();
	lastAuthSeqID.clear();

	auto bindCompoundAtoms = [&]
	{
		for (uint32_t i = 0; i + 1 < rAtoms.size(); ++i)
		{
			auto id_i = rAtoms[i].get<std::string>("id");
			auto atom_id_i = rAtoms[i].get<std::string>("label_atom_id");

			for (uint32_t j = i + 1; j < rAtoms.size(); ++j)
			{
				auto atom_id_j = rAtoms[j].get<std::string>("label_atom_id");

				if (compoundBondInfo.bonded(lastCompID, atom_id_i, atom_id_j))
					bindAtoms(id_i, rAtoms[j].get<std::string>("id"));
			}
		}
	};

	for (auto atom : atoms)
	{
		const auto [asym_id, seq_id, auth_seq_id, entity_id, comp_id] = atom.get<std::string, int, std::string, std::string, std::string>(
			"label_asym_id", "label_seq_id", "auth_seq_id", "label_entity_id", "label_comp_id");

		if (asym_id == lastAsymID and entity_id == lastEntityID and lastCompID == comp_id and lastSeqID == seq_id and lastAuthSeqID == auth_seq_id)
		{
			rAtoms.emplace_back(atom);
			continue;
		}

		bindCompoundAtoms();

		rAtoms = { atom };
		lastAsymID = asym_id;
		lastEntityID = entity_id;
		lastCompID = comp_id;
		lastSeqID = seq_id;
		lastAuthSeqID = auth_seq_id;
	}

	bindCompoundAtoms();

	// start by creating an index for single bonds

	std::multimap<uint32_t, uint32_t> b1_2;
	for (auto [a, b] : bond)
	{
		b1_2.insert({ a, b });
		b1_2.insert({ b, a });
	}

	std::multimap<uint32_t, uint32_t> b1_3;
	for (uint32_t i = 0; i < dim; ++i)
	{
		auto a = b1_2.equal_range(i);

		std::vector<uint32_t> s;
		for (auto j = a.first; j != a.second; ++j)
			s.push_back(j->second);

		for (std::size_t si1 = 0; si1 + 1 < s.size(); ++si1)
		{
			for (std::size_t si2 = si1 + 1; si2 < s.size(); ++si2)
			{
				uint32_t x = s[si1];
				uint32_t y = s[si2];

				if (isBonded(x, y))
					continue;

				b1_3.insert({ x, y });
				b1_3.insert({ y, x });
			}
		}
	}

	for (uint32_t i = 0; i < dim; ++i)
	{
		auto a1 = b1_2.equal_range(i);
		auto a2 = b1_3.equal_range(i);

		for (auto ai1 = a1.first; ai1 != a1.second; ++ai1)
		{
			for (auto ai2 = a2.first; ai2 != a2.second; ++ai2)
			{
				uint32_t b1 = ai1->second;
				uint32_t b2 = ai2->second;

				if (isBonded(b1, b2))
					continue;

				bond_1_4.insert(key(b1, b2));
			}
		}
	}
}

std::vector<std::string> BondMap::linked(const std::string &atom_id) const
{
	auto i = link.find(atom_id);

	std::vector<std::string> result;

	if (i != link.end())
		result = std::vector<std::string>(i->second.begin(), i->second.end());

	return result;
}

std::vector<std::string> BondMap::atomIDsForCompound(const std::string &compoundID)
{
	std::vector<std::string> result;

	auto *compound = cif::compound_factory::instance().create(compoundID);

	if (compound == nullptr)
		throw BondMapException("Missing bond information for compound " + compoundID);

	for (auto &compAtom : compound->atoms())
		result.push_back(compAtom.id);

	return result;
}

} // namespace pdbx
