#include "PBC_Util.h"
#include "Nedelec3D_Util.h"

#include <algorithm>
#include <unordered_map>

namespace {

int findMemberIndex(const std::vector<int>& values, int key)
{
	for (size_t i = 0; i < values.size(); i++)
	{
		if (values[i] == key)
			return static_cast<int>(i);
	}
	return -1;
}

std::vector<int> findDuplicateValues(const std::vector<int>& values)
{
	std::unordered_map<int, int> counts;
	for (int v : values)
		counts[v]++;
	std::vector<int> duplicates;
	for (const auto& entry : counts)
	{
		if (entry.second > 1)
			duplicates.push_back(entry.first);
	}
	std::sort(duplicates.begin(), duplicates.end());
	return duplicates;
}

} // namespace

PbcEdgePairs collectPbcEdgePairs(
	Mesh_3D* mesh,
	const std::vector<Eigen::VectorXi>& srcDomains,
	const std::vector<Eigen::VectorXi>& dstDomains,
	const Eigen::VectorXd& pbcAngles,
	const Eigen::MatrixXd& pbcWaveVecs,
	bool useFaceStyleEdgeTol)
{
	const int nbrGroups = static_cast<int>(srcDomains.size());
	Eigen::VectorXcd unitPhi = Eigen::VectorXcd::Ones(nbrGroups);
	Eigen::VectorXi srcEdgeIndex;
	Eigen::VectorXi dstEdgeIndex;
	Eigen::VectorXcd edgePhi;
	mesh->GetIndexOfPBC(srcDomains, dstDomains, unitPhi, pbcAngles, pbcWaveVecs,
		srcEdgeIndex, dstEdgeIndex, edgePhi, useFaceStyleEdgeTol);

	sortUniquePbcPairs(srcEdgeIndex, dstEdgeIndex, edgePhi);

	PbcEdgePairs pairs;
	const int n = static_cast<int>(srcEdgeIndex.size());
	pairs.src.resize(n);
	pairs.dst.resize(n);
	pairs.sign.resize(n);
	for (int i = 0; i < n; i++)
	{
		pairs.src[i] = srcEdgeIndex(i);
		pairs.dst[i] = dstEdgeIndex(i);
		pairs.sign[i] = edgePhi(i);
	}
	return pairs;
}

PbcEdgePairs collectPbcFacePairs(
	Mesh_3D* mesh,
	const std::vector<Eigen::VectorXi>& srcDomains,
	const std::vector<Eigen::VectorXi>& dstDomains,
	const Eigen::VectorXd& pbcAngles,
	const Eigen::MatrixXd& pbcWaveVecs)
{
	Eigen::VectorXi srcFaceIndex;
	Eigen::VectorXi dstFaceIndex;
	Eigen::VectorXcd facePhi;
	mesh->GetIndexOfPBCFaces(srcDomains, dstDomains, pbcAngles, pbcWaveVecs,
		srcFaceIndex, dstFaceIndex, facePhi);

	PbcEdgePairs pairs;
	const int n = static_cast<int>(srcFaceIndex.size());
	pairs.src.resize(n);
	pairs.dst.resize(n);
	pairs.sign.resize(n);
	for (int i = 0; i < n; i++)
	{
		pairs.src[i] = srcFaceIndex(i);
		pairs.dst[i] = dstFaceIndex(i);
		pairs.sign[i] = facePhi(i);
	}
	return pairs;
}

void mergeDoublePbcPairs(
	const PbcEdgePairs& pairs1,
	std::complex<double> blochPhi1,
	const PbcEdgePairs& pairs2,
	std::complex<double> blochPhi2,
	PbcEdgePairs& merged,
	std::vector<std::complex<double>>& mergedPhi)
{
	const int nbr2 = static_cast<int>(pairs1.src.size());
	const int nbr3 = static_cast<int>(pairs2.src.size());

	std::vector<int> allDst;
	allDst.reserve(nbr2 + nbr3);
	allDst.insert(allDst.end(), pairs1.dst.begin(), pairs1.dst.end());
	allDst.insert(allDst.end(), pairs2.dst.begin(), pairs2.dst.end());
	const std::vector<int> ovDstIndex = findDuplicateValues(allDst);
	const int nbr1 = static_cast<int>(ovDstIndex.size());

	merged.src.clear();
	merged.dst.clear();
	merged.sign.clear();
	mergedPhi.clear();
	merged.src.reserve(nbr2 + nbr3 - nbr1);
	merged.dst.reserve(nbr2 + nbr3 - nbr1);
	merged.sign.reserve(nbr2 + nbr3 - nbr1);
	mergedPhi.reserve(nbr2 + nbr3 - nbr1);

	std::vector<char> remove1(nbr2, 0);
	std::vector<char> remove2(nbr3, 0);

	for (int ovDst : ovDstIndex)
	{
		const int index1 = findMemberIndex(pairs2.dst, ovDst);
		if (index1 < 0)
			continue;

		const int src2 = pairs2.src[index1];
		const int index3 = findMemberIndex(pairs1.dst, src2);
		if (index3 < 0)
			continue;

		const int src1 = pairs1.src[index3];
		const std::complex<double> sign =
			pairs2.sign[index1] * pairs1.sign[index3];

		merged.src.push_back(src1);
		merged.dst.push_back(ovDst);
		merged.sign.push_back(sign);
		mergedPhi.push_back(blochPhi1 * blochPhi2 * sign);

		remove1[index3] = 1;
		remove2[index1] = 1;
		for (int i = 0; i < nbr2; i++)
		{
			if (pairs1.dst[i] == ovDst)
				remove1[i] = 1;
		}
		for (int i = 0; i < nbr3; i++)
		{
			if (pairs2.dst[i] == ovDst)
				remove2[i] = 1;
		}
	}

	for (int i = 0; i < nbr2; i++)
	{
		if (remove1[i])
			continue;
		merged.src.push_back(pairs1.src[i]);
		merged.dst.push_back(pairs1.dst[i]);
		merged.sign.push_back(pairs1.sign[i]);
		mergedPhi.push_back(blochPhi1 * pairs1.sign[i]);
	}

	for (int i = 0; i < nbr3; i++)
	{
		if (remove2[i])
			continue;
		merged.src.push_back(pairs2.src[i]);
		merged.dst.push_back(pairs2.dst[i]);
		merged.sign.push_back(pairs2.sign[i]);
		mergedPhi.push_back(blochPhi2 * pairs2.sign[i]);
	}
}

void sortUniquePbcPairs(
	Eigen::VectorXi& srcEdgeIndex,
	Eigen::VectorXi& dstEdgeIndex,
	Eigen::VectorXcd& edgePhi)
{
	if (dstEdgeIndex.size() == 0)
		return;

	Eigen::VectorXi tempIndex = Eigen::VectorXi::LinSpaced(dstEdgeIndex.size(), 0, dstEdgeIndex.size() - 1);
	QuickSort(dstEdgeIndex, tempIndex, 0, static_cast<int>(dstEdgeIndex.size()) - 1);
	Unique(dstEdgeIndex, tempIndex);

	Eigen::VectorXcd sortedPhi = Eigen::VectorXcd::Zero(dstEdgeIndex.size());
	Eigen::VectorXi sortedSrc = Eigen::VectorXi::Zero(dstEdgeIndex.size());
	for (int i = 0; i < dstEdgeIndex.size(); i++)
	{
		sortedPhi(i) = edgePhi(tempIndex(i));
		sortedSrc(i) = srcEdgeIndex(tempIndex(i));
	}
	edgePhi = sortedPhi;
	srcEdgeIndex = sortedSrc;
}

void fillPbcGroupFromPhy(
	const Phy_WaveOpticsModel* phy,
	int group,
	std::vector<Eigen::VectorXi>& srcDomains,
	std::vector<Eigen::VectorXi>& dstDomains,
	Eigen::VectorXd& pbcAngles,
	Eigen::MatrixXd& pbcWaveVecs)
{
	const int nbrPBC = phy->GetNbrPBCInGroup(group);
	srcDomains.resize(nbrPBC);
	dstDomains.resize(nbrPBC);
	pbcAngles = Eigen::VectorXd(nbrPBC);
	pbcWaveVecs = Eigen::MatrixXd::Zero(3, nbrPBC);
	for (int i = 0; i < nbrPBC; i++)
	{
		phy->GetSrcDomainList(group, i, srcDomains[i]);
		phy->GetDstDomainList(group, i, dstDomains[i]);
		pbcAngles(i) = phy->GetPBCAngle(group, i);
		Eigen::Vector3d waveVec;
		phy->GetPBCWaveVec(group, i, waveVec);
		pbcWaveVecs.col(i) = waveVec;
	}
}

namespace {

void expandOrder2PbcPairs(
	Mesh_3D* mesh,
	const PbcEdgePairs& edgePairs,
	const std::vector<std::complex<double>>& edgePhi,
	const PbcEdgePairs& facePairs,
	const std::vector<std::complex<double>>& facePhi,
	Eigen::VectorXi& srcDofIndex,
	Eigen::VectorXi& dstDofIndex,
	Eigen::VectorXcd& dofPhi)
{
	// Order-2 tangential PBC = order-1 edge pairing/signs, with each edge enriched
	// to two DOFs (e, e+nE) sharing the same Bloch×orient φ.
	//
	// Face DOFs are intentionally NOT constrained here:
	//   - getBF face modes (BF 13–20) have vanishing tangential trace on their face,
	//     so Et continuity on PBC walls is carried entirely by edge DOFs;
	//   - MATLAB-style identity face expand (f→f, f+nF→f+nF) is wrong on this
	//     translation mesh (Face vertex order usually does not match under +dis)
	//     and actively breaks L/R symmetry (A/B: with faces L/R~0.67, without ~0.98).
	(void)facePairs;
	(void)facePhi;
	const int nE = mesh->GetNbrEdge();
	std::vector<int> srcAll, dstAll;
	std::vector<std::complex<double>> phiAll;
	srcAll.reserve(edgePairs.src.size() * 2);
	dstAll.reserve(edgePairs.dst.size() * 2);
	phiAll.reserve(srcAll.capacity());

	for (size_t i = 0; i < edgePairs.src.size(); i++)
	{
		const std::complex<double> phi = edgePhi[i];
		srcAll.push_back(Nedelec3D::edgeGlobalDof(edgePairs.src[i], 0, nE));
		dstAll.push_back(Nedelec3D::edgeGlobalDof(edgePairs.dst[i], 0, nE));
		phiAll.push_back(phi);
		srcAll.push_back(Nedelec3D::edgeGlobalDof(edgePairs.src[i], 1, nE));
		dstAll.push_back(Nedelec3D::edgeGlobalDof(edgePairs.dst[i], 1, nE));
		phiAll.push_back(phi);
	}

	std::vector<int> filteredSrc, filteredDst;
	std::vector<std::complex<double>> filteredPhi;
	for (size_t i = 0; i < srcAll.size(); i++)
	{
		if (srcAll[i] == dstAll[i])
			continue;
		filteredSrc.push_back(srcAll[i]);
		filteredDst.push_back(dstAll[i]);
		filteredPhi.push_back(phiAll[i]);
	}

	const int n = static_cast<int>(filteredDst.size());
	srcDofIndex = Eigen::VectorXi(n);
	dstDofIndex = Eigen::VectorXi(n);
	dofPhi = Eigen::VectorXcd(n);
	for (int i = 0; i < n; i++)
	{
		srcDofIndex(i) = filteredSrc[static_cast<size_t>(i)];
		dstDofIndex(i) = filteredDst[static_cast<size_t>(i)];
		dofPhi(i) = filteredPhi[static_cast<size_t>(i)];
	}
	sortUniquePbcPairs(srcDofIndex, dstDofIndex, dofPhi);
}

bool collectMergedPbcPairGroups(
	Mesh_3D* mesh,
	const Phy_WaveOpticsModel* phy,
	int elementOrder,
	PbcEdgePairs& mergedEdges,
	std::vector<std::complex<double>>& mergedEdgePhi,
	PbcEdgePairs& mergedFaces,
	std::vector<std::complex<double>>& mergedFacePhi)
{
	const int nbrGroups = phy->GetNbrPBCGroups();
	if (nbrGroups <= 0)
		return false;

	// Do NOT loosen edge tol for order-2: it wrecked the validated top (z=+5e-6) face.
	const bool looseEdgeTol = false;

	std::vector<Eigen::VectorXi> srcDomains;
	std::vector<Eigen::VectorXi> dstDomains;
	Eigen::VectorXd pbcAngles;
	Eigen::MatrixXd pbcWaveVecs;
	fillPbcGroupFromPhy(phy, 0, srcDomains, dstDomains, pbcAngles, pbcWaveVecs);
	const PbcEdgePairs edgePairs1 = collectPbcEdgePairs(mesh, srcDomains, dstDomains, pbcAngles, pbcWaveVecs, looseEdgeTol);
	const PbcEdgePairs facePairs1 = collectPbcFacePairs(mesh, srcDomains, dstDomains, pbcAngles, pbcWaveVecs);

	if (nbrGroups >= 2)
	{
		std::vector<Eigen::VectorXi> srcDomains2;
		std::vector<Eigen::VectorXi> dstDomains2;
		Eigen::VectorXd pbcAngles2;
		Eigen::MatrixXd pbcWaveVecs2;
		fillPbcGroupFromPhy(phy, 1, srcDomains2, dstDomains2, pbcAngles2, pbcWaveVecs2);
		const PbcEdgePairs edgePairs2 = collectPbcEdgePairs(mesh, srcDomains2, dstDomains2, pbcAngles2, pbcWaveVecs2, looseEdgeTol);
		const PbcEdgePairs facePairs2 = collectPbcFacePairs(mesh, srcDomains2, dstDomains2, pbcAngles2, pbcWaveVecs2);

		const std::complex<double> phi1 = phy->GetPBCPhi(0, 0);
		const std::complex<double> phi2 = phy->GetPBCPhi(1, 0);

		mergeDoublePbcPairs(edgePairs1, phi1, edgePairs2, phi2, mergedEdges, mergedEdgePhi);
		mergeDoublePbcPairs(facePairs1, phi1, facePairs2, phi2, mergedFaces, mergedFacePhi);
	}
	else
	{
		const std::complex<double> phi1 = phy->GetPBCPhi(0, 0);
		mergedEdges = edgePairs1;
		mergedFaces = facePairs1;
		mergedEdgePhi.resize(edgePairs1.src.size());
		mergedFacePhi.resize(facePairs1.src.size());
		for (size_t i = 0; i < edgePairs1.src.size(); i++)
			mergedEdgePhi[i] = phi1 * edgePairs1.sign[i];
		for (size_t i = 0; i < facePairs1.src.size(); i++)
			mergedFacePhi[i] = phi1 * facePairs1.sign[i];
	}

	return mergedEdges.dst.size() > 0 || mergedFaces.dst.size() > 0;
}

} // namespace

bool collectMergedPbcConstraints(
	Mesh_3D* mesh,
	const Phy_WaveOpticsModel* phy,
	int elementOrder,
	Eigen::VectorXi& srcDofIndex,
	Eigen::VectorXi& dstDofIndex,
	Eigen::VectorXcd& dofPhi)
{
	PbcEdgePairs mergedEdges;
	std::vector<std::complex<double>> mergedEdgePhi;
	PbcEdgePairs mergedFaces;
	std::vector<std::complex<double>> mergedFacePhi;
	if (!collectMergedPbcPairGroups(mesh, phy, elementOrder, mergedEdges, mergedEdgePhi, mergedFaces, mergedFacePhi))
		return false;

	if (elementOrder == 2)
	{
		expandOrder2PbcPairs(mesh, mergedEdges, mergedEdgePhi, mergedFaces, mergedFacePhi,
			srcDofIndex, dstDofIndex, dofPhi);
		return dstDofIndex.size() > 0;
	}

	const int n = static_cast<int>(mergedEdges.src.size());
	srcDofIndex = Eigen::VectorXi(n);
	dstDofIndex = Eigen::VectorXi(n);
	dofPhi = Eigen::VectorXcd(n);
	for (int i = 0; i < n; i++)
	{
		srcDofIndex(i) = mergedEdges.src[static_cast<size_t>(i)];
		dstDofIndex(i) = mergedEdges.dst[static_cast<size_t>(i)];
		dofPhi(i) = mergedEdgePhi[static_cast<size_t>(i)];
	}
	sortUniquePbcPairs(srcDofIndex, dstDofIndex, dofPhi);
	return dstDofIndex.size() > 0;
}

bool assemblePbcProjectionMatrix(
	Mesh_3D* mesh,
	const Phy_WaveOpticsModel* phy,
	int elementOrder,
	int dof,
	bool isReal,
	Eigen::SparseMatrix<double, Eigen::RowMajor>& P_real,
	Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor>& P_complex)
{
	Eigen::VectorXi srcDofIndex;
	Eigen::VectorXi dstDofIndex;
	Eigen::VectorXcd dofPhi;
	if (!collectMergedPbcConstraints(mesh, phy, elementOrder, srcDofIndex, dstDofIndex, dofPhi))
		return false;

	if (isReal)
		buildPeriodicProjectionReal(dof, dstDofIndex, srcDofIndex, dofPhi, P_real);
	else
		buildPeriodicProjectionComplex(dof, dstDofIndex, srcDofIndex, dofPhi, P_complex);
	return true;
}

void buildPeriodicProjectionReal(
	int dof,
	const Eigen::VectorXi& dstIndex,
	const Eigen::VectorXi& srcIndex,
	const Eigen::VectorXcd& indexPhi,
	Eigen::SparseMatrix<double, Eigen::RowMajor>& P)
{
	std::vector<Eigen::Triplet<double>> tempP;
	int num = 0;
	for (int i = 0; i < dof; i++)
	{
		tempP.emplace_back(i, i, 1.0);
		if (num < dstIndex.size() && i == dstIndex(num))
		{
			tempP.emplace_back(dstIndex(num), srcIndex(num), indexPhi(num).real());
			num++;
		}
	}
	std::sort(tempP.begin(), tempP.end(),
		[](const Eigen::Triplet<double>& a, const Eigen::Triplet<double>& b)
		{
			return a.col() < b.col();
		});

	std::vector<Eigen::Triplet<double>> tripleP;
	num = 0;
	for (int i = 0; i < dof; i++)
	{
		bool deleteCol = false;
		for (int j = 0; j < dstIndex.size(); j++)
		{
			if (i == dstIndex(j))
			{
				deleteCol = true;
				break;
			}
		}
		if (!deleteCol)
		{
			for (const auto& triplet : tempP)
			{
				if (triplet.col() == i)
					tripleP.emplace_back(triplet.row(), triplet.col() - num, triplet.value());
			}
		}
		else
		{
			num++;
		}
	}
	P = Eigen::SparseMatrix<double, Eigen::RowMajor>(dof, dof - dstIndex.size());
	P.setFromTriplets(tripleP.begin(), tripleP.end());
}

void buildPeriodicProjectionComplex(
	int dof,
	const Eigen::VectorXi& dstIndex,
	const Eigen::VectorXi& srcIndex,
	const Eigen::VectorXcd& indexPhi,
	Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor>& P)
{
	std::vector<Eigen::Triplet<std::complex<double>>> tempP;
	int num = 0;
	for (int i = 0; i < dof; i++)
	{
		tempP.emplace_back(i, i, 1.0);
		if (num < dstIndex.size() && i == dstIndex(num))
		{
			tempP.emplace_back(dstIndex(num), srcIndex(num), indexPhi(num));
			num++;
		}
	}
	std::sort(tempP.begin(), tempP.end(),
		[](const Eigen::Triplet<std::complex<double>>& a, const Eigen::Triplet<std::complex<double>>& b)
		{
			return a.col() < b.col();
		});

	std::vector<Eigen::Triplet<std::complex<double>>> tripleP;
	num = 0;
	for (int i = 0; i < dof; i++)
	{
		bool deleteCol = false;
		for (int j = 0; j < dstIndex.size(); j++)
		{
			if (i == dstIndex(j))
			{
				deleteCol = true;
				break;
			}
		}
		if (!deleteCol)
		{
			for (const auto& triplet : tempP)
			{
				if (triplet.col() == i)
					tripleP.emplace_back(triplet.row(), triplet.col() - num, triplet.value());
			}
		}
		else
		{
			num++;
		}
	}
	P = Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor>(dof, dof - dstIndex.size());
	P.setFromTriplets(tripleP.begin(), tripleP.end());
}