#include"Assemble_Base.h"
#include"../function/BF.h"
#include"../function/Gauss.h"
#include"../common/define.h"
#include"../Eigen/Sparse"
#include <iostream>

#include <vector>
#include <algorithm>
#include <string>

#include "../nlohmann/json.hpp"

using namespace std;
using namespace Eigen;
using json = nlohmann::json;

void OpticsFEM_2D_EigenFreq::Assemble_WaveEquation()
{
	//init of Gauss point
	Gauss gauss;
	int NbrGuassPoints;
	double* u, * v, * w, * wght;
	NbrGuassPoints = gauss.GetNbrGaussPoints(TWODIM, TRIANGLE, BF_LINEFUNC * 2);
	u = new double[NbrGuassPoints];
	v = new double[NbrGuassPoints];
	w = new double[NbrGuassPoints];
	wght = new double[NbrGuassPoints];
	gauss.GetGaussPoints(TWODIM, TRIANGLE, u, v, w, wght);

	//init of geo
	Vector3d* vertex = new Vector3d[3];
	Matrix3d Jac, InvJac, TJac, JacS;

	//init of basis function
	BF BF_Lagrange, BF_Curl_Lagrange, BF_Nedelec, BF_Curl_Nedelec;
	int sdof, vdof;
	Vector3d** Et, ** curlEt, ** Ez, ** curlEz;
	sdof = BF_Lagrange.GetNbrBF(TWODIM, TRIANGLE, BF_LAGRANGE, BF_LINEFUNC);
	BF_Curl_Lagrange.GetNbrBF(TWODIM, TRIANGLE, BF_CURL_LAGRANGE, BF_LINEFUNC);
	vdof = BF_Nedelec.GetNbrBF(TWODIM, TRIANGLE, BF_NEDELEC, BF_LINEFUNC);
	BF_Curl_Nedelec.GetNbrBF(TWODIM, TRIANGLE, BF_CURL_NEDELEC, BF_LINEFUNC);
	Et = new Vector3d * [NbrGuassPoints];
	curlEt = new Vector3d * [NbrGuassPoints];
	Ez = new Vector3d * [NbrGuassPoints];
	curlEz = new Vector3d * [NbrGuassPoints];
	for (int i = 0; i < NbrGuassPoints; i++)
	{
		Et[i] = new Vector3d[vdof];
		curlEt[i] = new Vector3d[vdof];
		Ez[i] = new Vector3d[sdof];
		curlEz[i] = new Vector3d[sdof];
	}

	//loop over tri
	int NbrTri = _mMesh->GetNbrTri();
	for (int n = 0; n < NbrTri; n++)
	{
		//coordinate of vertex
		for (int i = 0; i < 3; i++)
			_mMesh->GetVertex(_mMesh->GetTri(n, i), vertex[i]);

		//Jac
		Jac(0, 0) = vertex[1](0) - vertex[0](0); Jac(0, 1) = vertex[1](1) - vertex[0](1); Jac(0, 2) = 0.;
		Jac(1, 0) = vertex[2](0) - vertex[0](0); Jac(1, 1) = vertex[2](1) - vertex[0](1); Jac(1, 2) = 0.;
		Jac(2, 0) = 0.;			Jac(2, 1) = 0.;			Jac(2, 2) = 1.;
		InvJac = Jac.inverse();
		JacS(0, 0) = InvJac(1, 1);  JacS(0, 1) = -InvJac(1, 0); JacS(0, 2) = 0.;
		JacS(1, 0) = -InvJac(0, 1); JacS(1, 1) = InvJac(0, 0);  JacS(1, 2) = 0.;
		JacS(2, 0) = 0.;			JacS(2, 1) = 0.;		    JacS(2, 2) = 1.;
		double DetJac = fabs(Jac.determinant());
		TJac = Jac.transpose() / Jac.determinant();

		//basis function
		for (int i = 0; i < NbrGuassPoints; i++)
		{
			for (int j = 0; j < sdof; j++)
			{
				BF_Lagrange.GetValueBF(j + 1, u[i], v[i], w[i], Ez[i][j]);
				BF_Curl_Lagrange.GetValueBF(j + 1, u[i], v[i], w[i], curlEz[i][j]);
				curlEz[i][j] = JacS * curlEz[i][j];
			}
			for (int j = 0; j < vdof; j++)
			{
				BF_Nedelec.GetValueBF(j + 1, u[i], v[i], w[i], Et[i][j]);
				Et[i][j] = InvJac * Et[i][j];
				BF_Curl_Nedelec.GetValueBF(j + 1, u[i], v[i], w[i], curlEt[i][j]);
				curlEt[i][j] = TJac * curlEt[i][j];
			}
		}

		//material
		int domain = _mMesh->GetDomainOfTri(n);
		Matrix3cd epsr = _mMatLib->GetEpsr(domain);
		Matrix3cd Mur = _mMatLib->GetMur(domain);
		//PML
		Matrix3cd invMur;
		Eigen::VectorXd PMLData;
		double lambda;
		int PMLType;
		if (_mPhy->GetPML(domain, PMLType, PMLData, lambda))
		{
			if (PMLType == 0)
			{
				Eigen::Matrix3cd invT = Eigen::Matrix3cd::Zero();

				complex<double> diff_coordx_x = lambda / PMLData(1) * complex<double>(1.0, -1.0);
				complex<double> diff_coordx_y = complex<double>(0.0, 0.0);
				complex<double> diff_coordy_x = complex<double>(0.0, 0.0);
				complex<double> diff_coordy_y = lambda / PMLData(3) * complex<double>(1.0, -1.0);

				invT(0, 0) = (PMLData(1) == 0) ? complex<double>(1.0, 0.0) : diff_coordx_x;
				invT(1, 1) = (PMLData(3) == 0) ? complex<double>(1.0, 0.0) : diff_coordy_y;
				invT(2, 2) = std::complex<double>(1.0, 0.0);
				std::complex<double> detInvT = invT.determinant();
				Eigen::Matrix3cd T = invT.inverse();
				epsr = T * epsr * T.transpose() * detInvT;
				Mur = T * Mur * T.transpose() * detInvT;
				invMur = Mur.inverse();
			}
			else if (PMLType == 1)
			{
				int R0 = 10;
				double k0 = 2 * Pi / lambda;
				double averX = (vertex[0](0) + vertex[1](0) + vertex[2](0)) / 3.0;
				double averY = (vertex[0](1) + vertex[1](1) + vertex[2](1)) / 3.0;
				Matrix3cd Lambda = Matrix3cd::Zero();
				double rho = sqrt(averX * averX + averY * averY);
				double sigma = pow((rho - PMLData(1)) / PMLData(3), 2) * R0 / PMLData(3) / k0;
				complex<double> s1{ 1, -PMLData(3) / 2 / rho * sigma };
				complex<double> s2{ 1, -sigma };
				complex<double> aa = s1 / s2; complex<double> bb = s2 / s1; complex<double> cc = s1 * s2;
				Lambda(0, 0) = (aa * averX * averX + bb * averY * averY) / rho / rho;
				Lambda(0, 1) = (aa - bb) * averX * averY / rho / rho;
				Lambda(1, 0) = (aa - bb) * averX * averY / rho / rho;
				Lambda(1, 1) = (bb * averX * averX + aa * averY * averY) / rho / rho;
				Lambda(2, 2) = cc;
				epsr = epsr * Lambda;
				invMur = (Mur * Lambda).inverse();
			}			
		}
		else
		{
			invMur = Mur.inverse();
		}

		//mapping 
		VectorXi MappingIndexS = VectorXi::Zero(sdof);
		VectorXi MappingIndexV = VectorXi::Zero(vdof);
		for (int i = 0; i < sdof; i++)
			MappingIndexS(i) = _mMesh->GetTri(n, i);
		for (int i = 0; i < vdof; i++)
			MappingIndexV(i) = _mMesh->GetEdgeOfTri(n, i) + _mMesh->GetNbrVertex();

		//submatrix
		MatrixXcd St, Sz, Tt, Tz;
		St = MatrixXcd::Zero(vdof, vdof);
		Sz = MatrixXcd::Zero(sdof, sdof);
		Tt = MatrixXcd::Zero(vdof, vdof);
		Tz = MatrixXcd::Zero(sdof, sdof);
		for (int i = 0; i < sdof; i++)
			for (int j = 0; j < sdof; j++)
				for (int k = 0; k < NbrGuassPoints; k++)
				{
					Sz(i, j) = Sz(i, j) + wght[k] * DetJac * curlEz[k][i].dot(invMur * curlEz[k][j]);
					Tz(i, j) = Tz(i, j) + wght[k] * DetJac * Ez[k][i].dot(epsr * Ez[k][j]);
				}
		for (int i = 0; i < vdof; i++)
			for (int j = 0; j < vdof; j++)
				for (int k = 0; k < NbrGuassPoints; k++)
				{
					St(i, j) = St(i, j) + wght[k] * DetJac * curlEt[k][i].dot(invMur * curlEt[k][j]);
					Tt(i, j) = Tt(i, j) + wght[k] * DetJac * Et[k][i].dot(epsr * Et[k][j]);
				}


		//store in triplet
		if (_mIsReal)
		{
			for (int i = 0; i < sdof; i++)
			{
				for (int j = 0; j < sdof; j++)
				{
					_mTripleA_real.push_back(Eigen::Triplet<double>(MappingIndexS(i), MappingIndexS(j), Sz(i, j).real()));
					_mTripleB_real.push_back(Eigen::Triplet<double>(MappingIndexS(i), MappingIndexS(j), Tz(i, j).real()));
				}
			}
			for (int i = 0; i < vdof; i++)
			{
				for (int j = 0; j < vdof; j++)
				{
					_mTripleA_real.push_back(Eigen::Triplet<double>(MappingIndexV(i), MappingIndexV(j), St(i, j).real()));
					_mTripleB_real.push_back(Eigen::Triplet<double>(MappingIndexV(i), MappingIndexV(j), Tt(i, j).real()));
				}
			}
		}
		else
		{
			for (int i = 0; i < sdof; i++)
			{
				for (int j = 0; j < sdof; j++)
				{
					_mTripleA_complex.push_back(Triplet<complex<double>>(MappingIndexS(i), MappingIndexS(j), Sz(i, j)));
					_mTripleB_complex.push_back(Triplet<complex<double>>(MappingIndexS(i), MappingIndexS(j), Tz(i, j)));
				}
			}
			for (int i = 0; i < vdof; i++)
			{
				for (int j = 0; j < vdof; j++)
				{
					_mTripleA_complex.push_back(Triplet<complex<double>>(MappingIndexV(i), MappingIndexV(j), St(i, j)));
					_mTripleB_complex.push_back(Triplet<complex<double>>(MappingIndexV(i), MappingIndexV(j), Tt(i, j)));
				}
			}
		}
	}

	delete[] u, v, w, wght;

	delete[] vertex;

	for (int i = 0; i < NbrGuassPoints; i++)
	{
		delete[] Et[i], curlEt[i], Ez[i], curlEz[i];
	}
	delete[] Et, curlEt, Ez, curlEz;
}

void OpticsFEM_3D_EigenFreq::Assemble_WaveEquation()
{
	//init of Gauss point
	Gauss gauss;
	int NbrGuassPoints;
	double* u, * v, * w, * wght;
	NbrGuassPoints = gauss.GetNbrGaussPoints(THREEDIM, TETRAHEDRON, BF_LINEFUNC * 2);
	u = new double[NbrGuassPoints];
	v = new double[NbrGuassPoints];
	w = new double[NbrGuassPoints];
	wght = new double[NbrGuassPoints];
	gauss.GetGaussPoints(THREEDIM, TETRAHEDRON, u, v, w, wght);

	//init of geo
	Vector3d* vertex = new Vector3d[4];
	Matrix3d Jac, TJac, InvJac;

	//init of basis function
	BF BF_Nedelec, BF_Curl_Nedelec;
	int dof;
	Vector3d** E, ** curlE;
	dof = BF_Nedelec.GetNbrBF(THREEDIM, TETRAHEDRON, BF_NEDELEC, BF_LINEFUNC);
	BF_Curl_Nedelec.GetNbrBF(THREEDIM, TETRAHEDRON, BF_CURL_NEDELEC, BF_LINEFUNC);
	E = new Vector3d * [NbrGuassPoints];
	curlE = new Vector3d * [NbrGuassPoints];
	for (int i = 0; i < NbrGuassPoints; i++)
	{
		E[i] = new Vector3d[dof];
		curlE[i] = new Vector3d[dof];
	}

	//loop over tri
	int NbrTet = _mMesh->GetNbrTet();
	for (int n = 0; n < NbrTet; n++)
	{
		//coordinate of vertex
		for (int i = 0; i < 4; i++)
			_mMesh->GetVertex(_mMesh->GetTet(n, i), vertex[i]);

		//Jac
		Jac(0, 0) = vertex[0](0) - vertex[3](0); Jac(0, 1) = vertex[0](1) - vertex[3](1); Jac(0, 2) = 0.;
		Jac(1, 0) = vertex[1](0) - vertex[3](0); Jac(1, 1) = vertex[1](1) - vertex[3](1); Jac(1, 2) = 0.;
		Jac(2, 0) = vertex[2](0) - vertex[3](0); Jac(2, 1) = vertex[2](1) - vertex[3](1); Jac(2, 2) = 1.;
		double DetJac = fabs(Jac.determinant());
		TJac = Jac.transpose() / Jac.determinant();
		InvJac = Jac.inverse();

		//basis function
		for (int i = 0; i < NbrGuassPoints; i++)
		{
			for (int j = 0; j < dof; j++)
			{
				BF_Nedelec.GetValueBF(j + 1, u[i], v[i], w[i], E[i][j]);
				E[i][j] = InvJac * E[i][j];
				BF_Curl_Nedelec.GetValueBF(j + 1, u[i], v[i], w[i], curlE[i][j]);
				curlE[i][j] = TJac * curlE[i][j];
			}
		}

		//material
		int domain = _mMesh->GetDomainOfTri(n);
		Matrix3cd epsr = _mMatLib->GetEpsr(domain);
		Matrix3cd invMur = _mMatLib->GetEpsr(domain).inverse();

		//mapping 
		VectorXi MappingIndex = VectorXi::Zero(dof);
		for (int i = 0; i < dof; i++)
			MappingIndex(i) = _mMesh->GetEdgeOfTet(n, i);

		//submatrix
		MatrixXcd Se, Te;
		Se = MatrixXcd::Zero(dof, dof);
		Te = MatrixXcd::Zero(dof, dof);
		for (int i = 0; i < dof; i++)
			for (int j = 0; j < dof; j++)
				for (int k = 0; k < NbrGuassPoints; k++)
				{
					Se(i, j) = Se(i, j) + wght[k] * DetJac * curlE[k][i].dot(invMur * curlE[k][j]);
					Te(i, j) = Te(i, j) + wght[k] * DetJac * E[k][i].dot(epsr * E[k][j]);
				}


		//store in triplet
		if (_mIsReal)
		{
			for (int i = 0; i < dof; i++)
			{
				for (int j = 0; j < dof; j++)
				{
					_mTripleA_real.push_back(Eigen::Triplet<double>(MappingIndex(i), MappingIndex(j), Se(i, j).real()));
					_mTripleB_real.push_back(Eigen::Triplet<double>(MappingIndex(i), MappingIndex(j), Te(i, j).real()));
				}
			}
		}
		else
		{
			for (int i = 0; i < dof; i++)
			{
				for (int j = 0; j < dof; j++)
				{
					_mTripleA_complex.push_back(Eigen::Triplet<complex<double>>(MappingIndex(i), MappingIndex(j), Se(i, j)));
					_mTripleB_complex.push_back(Eigen::Triplet<complex<double>>(MappingIndex(i), MappingIndex(j), Te(i, j)));
				}
			}
		}
	}

	delete[] u, v, w, wght;

	delete[] vertex;

	for (int i = 0; i < NbrGuassPoints; i++)
	{
		delete[] E[i], curlE[i];
	}
	delete[] E, curlE;
}