#include "OpticsFEM_3D_SBC_Scatter.h"
#include "../function/BF.h"
#include "../function/Gauss.h"
#include "../common/define.h"
#include <Eigen/SparseLU>

#include <cmath>
#include <fstream>
#include <iostream>
#ifdef _WIN32
#include <direct.h>
#define MKDIR(p) _mkdir(p)
#else
#include <sys/stat.h>
#define MKDIR(p) mkdir(p, 0755)
#endif

using namespace Eigen;
using std::complex;

void OpticsFEM_3D_SBC_Scatter::GetMesh(Mesh_3D* mesh)
{
	_mMesh = mesh;
}

void OpticsFEM_3D_SBC_Scatter::SetPhysics(double lda0,
	const VectorXd& eps, const VectorXd& mur,
	const Vector3cd& Einc,
	const VectorXi& outDomains, const VectorXi& incDomains)
{
	_mLda0 = lda0;
	_mEps = eps;
	_mMur = mur;
	_mEinc = Einc;
	_mMesh->FindTriByDomains(outDomains, _mOutTri);
	_mMesh->FindTriByDomains(incDomains, _mIncTri);
}

Vector3d OpticsFEM_3D_SBC_Scatter::crossNormal(const Vector3d& n, const Vector3d& a)
{
	return n.cross(a.cross(n));
}

void OpticsFEM_3D_SBC_Scatter::faceMap(int numFace, Vector3d& x2, Vector3d& y2, Vector3d& z2,
	Vector3i& bfIndex, Vector3i& edgeSlot)
{
	switch (numFace)
	{
	case 1: x2 << 1, 0, 0; y2 << 0, 1, 0; z2 << 0, 0, 1; bfIndex << 1, 2, 4; edgeSlot << 0, 1, 3; break;
	case 2: x2 << 1, 0, 0; y2 << 0, 1, 0; z2 << 0, 0, 0; bfIndex << 1, 3, 5; edgeSlot << 0, 2, 4; break;
	case 3: x2 << 1, 0, 0; y2 << 0, 0, 0; z2 << 0, 1, 0; bfIndex << 2, 3, 6; edgeSlot << 1, 2, 5; break;
	case 4: x2 << 0, 0, 0; y2 << 1, 0, 0; z2 << 0, 1, 0; bfIndex << 4, 5, 6; edgeSlot << 3, 4, 5; break;
	default: break;
	}
}

void OpticsFEM_3D_SBC_Scatter::Assemble_WaveEquation()
{
	const double k0 = 2.0 * Pi / _mLda0;
	Gauss gauss;
	int nGP = gauss.GetNbrGaussPoints(THREEDIM, TETRAHEDRON, BF_LINEFUNC * 2);
	double* u = new double[nGP];
	double* v = new double[nGP];
	double* w = new double[nGP];
	double* wg = new double[nGP];
	gauss.GetGaussPoints(THREEDIM, TETRAHEDRON, u, v, w, wg);

	BF bfN, bfC;
	const int dof = bfN.GetNbrBF(THREEDIM, TETRAHEDRON, BF_NEDELEC, BF_LINEFUNC);
	bfC.GetNbrBF(THREEDIM, TETRAHEDRON, BF_CURL_NEDELEC, BF_LINEFUNC);
	Vector3d vertex[4];
	Matrix3d Jac, InvJac, TJac;

	for (int n = 0; n < _mMesh->GetNbrTet(); n++)
	{
		for (int i = 0; i < 4; i++)
			_mMesh->GetVertex(_mMesh->GetTet(n, i) - 1, vertex[i]);

		Jac(0, 0) = vertex[0](0) - vertex[3](0); Jac(0, 1) = vertex[0](1) - vertex[3](1); Jac(0, 2) = vertex[0](2) - vertex[3](2);
		Jac(1, 0) = vertex[1](0) - vertex[3](0); Jac(1, 1) = vertex[1](1) - vertex[3](1); Jac(1, 2) = vertex[1](2) - vertex[3](2);
		Jac(2, 0) = vertex[2](0) - vertex[3](0); Jac(2, 1) = vertex[2](1) - vertex[3](1); Jac(2, 2) = vertex[2](2) - vertex[3](2);
		const double detJ = Jac.determinant();
		if (std::abs(detJ) < 1e-30)
			continue;
		const double detJac = std::abs(detJ);
		InvJac = Jac.inverse();
		TJac = Jac.transpose() / detJ;

		int dom = _mMesh->GetDomainOfTet(n);
		const double eps = (dom > 0 && dom <= _mEps.size()) ? _mEps(dom - 1) : _mEps(0);

		Vector3d E[6], curlE[6];
		for (int k = 0; k < nGP; k++)
		{
			for (int j = 0; j < dof; j++)
			{
				bfN.GetValueBF(j + 1, u[k], v[k], w[k], E[j]);
				E[j] = InvJac * E[j];
				bfC.GetValueBF(j + 1, u[k], v[k], w[k], curlE[j]);
				curlE[j] = TJac * curlE[j];
			}
			for (int i = 0; i < dof; i++)
				for (int j = 0; j < dof; j++)
				{
					const complex<double> val =
						wg[k] * detJac * curlE[i].dot(curlE[j])
						- wg[k] * detJac * k0 * k0 * eps * E[i].dot(E[j]);
					const int ri = _mMesh->GetEdgeOfTet(n, i) - 1;
					const int cj = _mMesh->GetEdgeOfTet(n, j) - 1;
					_mTripleA.emplace_back(ri, cj, val);
				}
		}
	}
	delete[] u; delete[] v; delete[] w; delete[] wg;
}

void OpticsFEM_3D_SBC_Scatter::Assemble_SBC_Face(bool isInc)
{
	const double k0 = 2.0 * Pi / _mLda0;
	const VectorXi& triList = isInc ? _mIncTri : _mOutTri;

	Gauss gauss;
	int nGP = gauss.GetNbrGaussPoints(TWODIM, TRIANGLE, BF_LINEFUNC * 2);
	double* u = new double[nGP];
	double* v = new double[nGP];
	double* w = new double[nGP];
	double* wg = new double[nGP];
	gauss.GetGaussPoints(TWODIM, TRIANGLE, u, v, w, wg);

	BF bfN;
	const int nBf = bfN.GetNbrBF(THREEDIM, TETRAHEDRON, BF_NEDELEC, BF_LINEFUNC);
	(void)nBf;
	Vector3d vertex[4];
	Matrix3d Jac, InvJac;

	for (int n = 0; n < triList.size(); n++)
	{
		const int triIdx = triList(n);
		const int domain = _mMesh->GetDomainOfTri(triIdx);
		const int numTet = _mMesh->GetConnOfTri(triIdx, 0) - 1;
		const int numFace = _mMesh->GetConnOfTri(triIdx, 1);

		for (int i = 0; i < 4; i++)
			_mMesh->GetVertex(_mMesh->GetTet(numTet, i) - 1, vertex[i]);

		Vector3d x2, y2, z2;
		Vector3i bfIndex, edgeSlot;
		faceMap(numFace, x2, y2, z2, bfIndex, edgeSlot);

		Vector3d x3, y3, z3;
		if (numFace == 1) { x3 << vertex[0](0), vertex[1](0), vertex[2](0); y3 << vertex[0](1), vertex[1](1), vertex[2](1); z3 << vertex[0](2), vertex[1](2), vertex[2](2); }
		else if (numFace == 2) { x3 << vertex[0](0), vertex[1](0), vertex[3](0); y3 << vertex[0](1), vertex[1](1), vertex[3](1); z3 << vertex[0](2), vertex[1](2), vertex[3](2); }
		else if (numFace == 3) { x3 << vertex[0](0), vertex[2](0), vertex[3](0); y3 << vertex[0](1), vertex[2](1), vertex[3](1); z3 << vertex[0](2), vertex[2](2), vertex[3](2); }
		else { x3 << vertex[1](0), vertex[2](0), vertex[3](0); y3 << vertex[1](1), vertex[2](1), vertex[3](1); z3 << vertex[1](2), vertex[2](2), vertex[3](2); }

		Jac(0, 0) = vertex[0](0) - vertex[3](0); Jac(0, 1) = vertex[0](1) - vertex[3](1); Jac(0, 2) = vertex[0](2) - vertex[3](2);
		Jac(1, 0) = vertex[1](0) - vertex[3](0); Jac(1, 1) = vertex[1](1) - vertex[3](1); Jac(1, 2) = vertex[1](2) - vertex[3](2);
		Jac(2, 0) = vertex[2](0) - vertex[3](0); Jac(2, 1) = vertex[2](1) - vertex[3](1); Jac(2, 2) = vertex[2](2) - vertex[3](2);
		const double detJ = Jac.determinant();
		if (std::abs(detJ) < 1e-30)
			continue;
		InvJac = Jac.inverse();

		const double la = (Vector3d(x3(0), y3(0), z3(0)) - Vector3d(x3(1), y3(1), z3(1))).norm();
		const double lb = (Vector3d(x3(0), y3(0), z3(0)) - Vector3d(x3(2), y3(2), z3(2))).norm();
		const double lc = (Vector3d(x3(1), y3(1), z3(1)) - Vector3d(x3(2), y3(2), z3(2))).norm();
		const double a2 = la * la, b2 = lb * lb, c2 = lc * lc;
		double heron = (a2 + b2 + c2) * (a2 + b2 - c2) * (a2 - b2 + c2) * (b2 + c2 - a2);
		if (heron < 0.0) heron = 0.0;
		const double integCoe = 0.25 * std::sqrt(heron);

		Vector3d normal;
		_mMesh->GetNormOfFace(domain, normal);

		int tetDom = _mMesh->GetDomainOfTet(numTet);
		const double eps = (tetDom > 0 && tetDom <= _mEps.size()) ? _mEps(tetDom - 1) : _mEps(0);
		const complex<double> nn = std::sqrt(complex<double>(eps, 0.0));

		Vector3i mapIdx;
		for (int i = 0; i < 3; i++)
			mapIdx(i) = _mMesh->GetEdgeOfTet(numTet, edgeSlot(i)) - 1;

		Matrix3cd Ae = Matrix3cd::Zero();
		Vector3cd Be = Vector3cd::Zero();
		Vector3d Egp[3];

		for (int k = 0; k < nGP; k++)
		{
			const double ugp = u[k];
			const double vgp = v[k];
			const double wgp = 1.0 - u[k] - v[k];
			const double u2 = x2(0) * ugp + x2(1) * vgp + x2(2) * wgp;
			const double v2 = y2(0) * ugp + y2(1) * vgp + y2(2) * wgp;
			const double w2 = z2(0) * ugp + z2(1) * vgp + z2(2) * wgp;
			for (int j = 0; j < 3; j++)
			{
				bfN.GetValueBF(bfIndex(j), u2, v2, w2, Egp[j]);
				Egp[j] = InvJac * Egp[j];
			}
			for (int i = 0; i < 3; i++)
				for (int j = 0; j < 3; j++)
				{
					// MATLAB: sum(E(:,i).*cross(normal,cross(E(:,j),normal)))
					const Vector3d tj = crossNormal(normal, Egp[j]);
					Ae(i, j) += complex<double>(0, 1) * k0 * nn * integCoe * wg[k] * Egp[i].dot(tj) * 2.0;
				}
			if (isInc)
			{
				const Vector3d EincReal(_mEinc(0).real(), _mEinc(1).real(), _mEinc(2).real());
				const Vector3d tInc = crossNormal(normal, EincReal);
				for (int i = 0; i < 3; i++)
				{
					// MATLAB: sum(E(:,i).*cross(normal,cross(phy.Einc,normal)))
					Be(i) -= complex<double>(0, 1) * k0 * nn * 2.0 * integCoe * wg[k] * Egp[i].dot(tInc) * 2.0;
				}
			}
		}

		for (int i = 0; i < 3; i++)
			for (int j = 0; j < 3; j++)
				_mTripleA.emplace_back(mapIdx(i), mapIdx(j), Ae(i, j));
		if (isInc)
			for (int i = 0; i < 3; i++)
				_mB(mapIdx(i)) += Be(i);
	}
	delete[] u; delete[] v; delete[] w; delete[] wg;
}

void OpticsFEM_3D_SBC_Scatter::Assemble_SBC_Out()
{
	Assemble_SBC_Face(false);
}

void OpticsFEM_3D_SBC_Scatter::Assemble_SBC_Inc()
{
	Assemble_SBC_Face(true);
}

void OpticsFEM_3D_SBC_Scatter::DumpAssembly(const std::string& outDir) const
{
	MKDIR(outDir.c_str());
	using SpMat = SparseMatrix<std::complex<double>, RowMajor>;
	SpMat A(_mDof, _mDof);
	A.setFromTriplets(_mTripleA.begin(), _mTripleA.end());
	A.makeCompressed();
	std::ofstream coo(outDir + "/coo.txt");
	coo << A.nonZeros() << " " << _mDof << "\n";
	for (int r = 0; r < _mDof; r++)
		for (typename SpMat::InnerIterator it(A, r); it; ++it)
			coo << r << " " << it.col() << " " << it.value().real() << " " << it.value().imag() << "\n";
	std::ofstream br(outDir + "/b.txt");
	for (int i = 0; i < _mB.size(); i++)
		br << _mB(i).real() << " " << _mB(i).imag() << "\n";
}

void OpticsFEM_3D_SBC_Scatter::Assemble()
{
	_mDof = _mMesh->GetNbrEdge();
	_mTripleA.clear();
	_mB = VectorXcd::Zero(_mDof);
	Assemble_WaveEquation();
	Assemble_SBC_Out();
	Assemble_SBC_Inc();
	if (std::getenv("OPTICSFEM_DUMP"))
		DumpAssembly("cpp_dump");
}

bool OpticsFEM_3D_SBC_Scatter::Run()
{
	_mX = VectorXcd::Zero(_mDof);
	using SpMat = SparseMatrix<complex<double>, RowMajor>;
	SpMat A(_mDof, _mDof);
	A.setFromTriplets(_mTripleA.begin(), _mTripleA.end());
	A.makeCompressed();

	double bnorm = _mB.norm();
	std::cout << "3D SBC: DOF=" << _mDof << " nnz=" << A.nonZeros()
		<< " |b|=" << bnorm << std::endl;

	using ColMat = SparseMatrix<std::complex<double>>;
	ColMat Acol = A;
	SparseLU<ColMat> lu;
	lu.compute(Acol);
	if (lu.info() != Eigen::Success)
	{
		std::cerr << "3D SBC: SparseLU factorization failed." << std::endl;
		return false;
	}
	_mX = lu.solve(_mB);
	if (lu.info() != Eigen::Success)
	{
		std::cerr << "3D SBC: SparseLU solve failed." << std::endl;
		return false;
	}
	std::cout << "3D SBC: |x|=" << _mX.norm() << std::endl;
	return std::isfinite(_mX(0).real());
}

void OpticsFEM_3D_SBC_Scatter::Post(const std::string& outDir)
{
	_mEx = VectorXcd::Zero(_mMesh->GetNbrVertex());
	_mEy = VectorXcd::Zero(_mMesh->GetNbrVertex());
	_mEz = VectorXcd::Zero(_mMesh->GetNbrVertex());
	VectorXd nCnt = VectorXd::Zero(_mMesh->GetNbrVertex());

	BF bfN;
	bfN.GetNbrBF(THREEDIM, TETRAHEDRON, BF_NEDELEC, BF_LINEFUNC);
	Vector3d vertex[4];
	Matrix3d Jac, InvJac;

	for (int n = 0; n < _mMesh->GetNbrTet(); n++)
	{
		for (int i = 0; i < 4; i++)
			_mMesh->GetVertex(_mMesh->GetTet(n, i) - 1, vertex[i]);
		Jac(0, 0) = vertex[0](0) - vertex[3](0); Jac(0, 1) = vertex[0](1) - vertex[3](1); Jac(0, 2) = vertex[0](2) - vertex[3](2);
		Jac(1, 0) = vertex[1](0) - vertex[3](0); Jac(1, 1) = vertex[1](1) - vertex[3](1); Jac(1, 2) = vertex[1](2) - vertex[3](2);
		Jac(2, 0) = vertex[2](0) - vertex[3](0); Jac(2, 1) = vertex[2](1) - vertex[3](1); Jac(2, 2) = vertex[2](2) - vertex[3](2);
		const double detJ = Jac.determinant();
		if (std::abs(detJ) < 1e-30)
			continue;
		InvJac = Jac.inverse();

		const double u0[4] = { 1, 0, 0, 0 };
		const double v0[4] = { 0, 1, 0, 0 };
		const double w0[4] = { 0, 0, 1, 0 };

		for (int j = 0; j < 4; j++)
		{
			Vector3cd Esum = Vector3cd::Zero();
			for (int i = 0; i < 6; i++)
			{
				Vector3d bf;
				bfN.GetValueBF(i + 1, u0[j], v0[j], w0[j], bf);
				bf = InvJac * bf;
				const int eid = _mMesh->GetEdgeOfTet(n, i) - 1;
				const complex<double> xe = _mX(eid);
				Esum(0) += complex<double>(bf(0), 0.0) * xe;
				Esum(1) += complex<double>(bf(1), 0.0) * xe;
				Esum(2) += complex<double>(bf(2), 0.0) * xe;
			}
			const int vid = _mMesh->GetTet(n, j) - 1;
			_mEx(vid) += Esum(0);
			_mEy(vid) += Esum(1);
			_mEz(vid) += Esum(2);
			nCnt(vid) += 1.0;
		}
	}
	for (int i = 0; i < _mMesh->GetNbrVertex(); i++)
		if (nCnt(i) > 0)
		{
			_mEx(i) /= nCnt(i);
			_mEy(i) /= nCnt(i);
			_mEz(i) /= nCnt(i);
		}
	_mNormE = (_mEx.array().abs().square() + _mEy.array().abs().square() + _mEz.array().abs().square()).sqrt();

	MKDIR(outDir.c_str());
	auto writeField = [&](const std::string& name, const VectorXcd& f) {
		std::ofstream o(outDir + "/" + name);
		for (int i = 0; i < f.size(); i++)
			o << f(i).real() << " " << f(i).imag() << "\n";
	};
	writeField("Ex", _mEx);
	writeField("Ey", _mEy);
	writeField("Ez", _mEz);
	std::ofstream on(outDir + "/normE");
	for (int i = 0; i < _mNormE.size(); i++)
		on << _mNormE(i) << "\n";
	std::cout << "3D SBC: normE max=" << _mNormE.maxCoeff() << std::endl;
}