#include "superdirectivity.h"
#include "sph_bessel.h"

#include <algorithm>

using namespace std;
using namespace sph_bessel;

Real get_directivity_sphml_dZx_epsRe(vector<Real> &param, void *cls) {
	int ni = (size(param)-1)/2; // 1, ni, ni
	sphere_ml *sph = reinterpret_cast<sphere_ml*>(cls);
	vector<Real> kR(param.begin()+1, param.begin()+ni+1);
	sort(kR.begin(),kR.end());
	vector<Complex> eps;
	for (int i=0; i<ni; ++i)
		eps.push_back(Complex(param[ni+1+i],Real(0)));
	eps.push_back(Complex(Real(1),Real(0)));
	return sph->get_directivity_edipole_x(param[0],kR,eps);
}

Real get_directivity_sphml_dZx_epsRe_parallel(vector<Real> &param, void *cls) {
	int ni = (size(param)-1)/2;
	sphere_ml_parallel *sph = reinterpret_cast<sphere_ml_parallel*>(cls);
	vector<Real> kR(param.begin()+1, param.begin()+ni+1);
	sort(kR.begin(),kR.end());
	vector<Complex> eps;
	for (int i=0; i<ni; ++i)
		eps.push_back(Complex(param[ni+1+i],Real(0)));
	eps.push_back(Complex(Real(1),Real(0)));
	return sph->get_directivity_edipole_x(param[0],kR,eps);
}

//###############################################################################
// class sphere_ml

// public members

void sphere_ml::set_directivity_angles(Real theta_, Real phi_) {
	theta = theta_; phi = phi_;
	get_pitau_m01(cos(theta),pi1,tau0,tau1,deg_max);
}

Real sphere_ml::get_directivity_edipole_x(Real kz, const std::vector<Real> &kR, const std::vector<Complex> &eps) {
	int ni = size(kR);
	if (size(eps)-1 != ni) {std::cout << "get_directivity_edipole_x error\n"; return Real(0);}
	Real D = 0.;
	Complex kzn;
	if (kz <= kR[0]) {
			// s-matrix of a the multilayer
		sphere_sm(kR[0], eps[0], eps[1], rtml_out);
		//cout << "sphere coeff:\n";
		//for (int i=0; i<deg_max; ++i) {cout << i << " " << rtml_out[8*i+2] << " " << rtml_out[8*i+3] << endl;}
		for (int k=1; k<ni; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_out, rt, rtml_out);
		}
			// dipole amplitudes
		kzn = kz*sqrt(eps[0]);
		if (arg(kzn) < -1e-14) kzn = -kzn;
		get_edipole_amp_x(kzn, ampl_out, 1);
			// output amplitudes
		for (int n=0; n<deg_max; ++n) {
			ampl_out[3*n] *= rtml_out[8*n+2]; // e, m = +-1
			ampl_out[3*n+1] *= rtml_out[8*n+3]; // h, m = +-1
			ampl_out[3*n+2] *= rtml_out[8*n+3]; // h, m = 0
		}
	}
	else if (kz <= kR[kR.size()-1]) { // the dipole is located in a layer
			// determine a dipole layer
		int kd = 0;
		while (kz > kR[++kd]); // determine a first interface with radius larger than the dipole position
			// calculate s-matrices of multilayers inside and outside the dipole position:
		sphere_sm(kR[0], eps[0], eps[1], rtml_in);
		for (int k=1; k<kd; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_in, rt, rtml_in);
		}
		sphere_sm(kR[kd], eps[kd], eps[kd+1], rtml_out);
		for (int k=kd+1; k<ni; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_out, rt, rtml_out);
		}
			// dipole amplitudes
		kzn = kz*sqrt(eps[kd]);
		if (arg(kzn) < -1e-14) kzn = -kzn;
		get_edipole_amp_x(kzn, ampl_in, 0);
		get_edipole_amp_x(kzn, ampl_out, 1);
			// output self-consistent amplitudes
		for (int n=0; n<deg_max; ++n) {
			ampl_out[3*n] = ( ampl_in[3*n]*rtml_in[8*n+6] + ampl_out[3*n] ) * rtml_out[8*n+2]
										/ (Real(1) - rtml_in[8*n+6]*rtml_out[8*n]); // e, m = 1
			ampl_out[3*n+1] = ( ampl_in[3*n+1]*rtml_in[8*n+7] + ampl_out[3*n+1] ) * rtml_out[8*n+3]
										/ (Real(1) - rtml_in[8*n+7]*rtml_out[8*n+1]); // h, m = 1
			ampl_out[3*n+2] = ( ampl_in[3*n+2]*rtml_in[8*n+7] + ampl_out[3*n+2] ) * rtml_out[8*n+3]
										/ (Real(1) - rtml_in[8*n+7]*rtml_out[8*n+1]); // h, m = 0
		}
	}
	else { // the dipole is located outside the multilayer
			// s-matrix of a the multilayer
		sphere_sm(kR[0], eps[0], eps[1], rtml_out);
		for (int k=1; k<ni; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_out, rt, rtml_out);
		}
			// dipole amplitudes
		kzn = kz*sqrt(eps[ni]);
		if (arg(kzn) < -1e-14) kzn = -kzn;
		get_edipole_amp_x(kzn, ampl_in, 0);
		get_edipole_amp_x(kzn, ampl_out, 1);
			// output amplitudes
		for (int n=0; n<deg_max; ++n) {
			ampl_out[3*n] += ampl_in[3*n] * rtml_out[8*n+6]; // e, m = +-1
			ampl_out[3*n+1] += ampl_in[3*n+1] * rtml_out[8*n+7]; // h, m = +-1
			ampl_out[3*n+2] += ampl_in[3*n+2] * rtml_out[8*n+7]; // h, m = 0
		}
	}

	return directivity_dipole_z(ampl_out);
}

void sphere_ml::get_edipole_amp_x(Complex kz, Complex *ampl, bool out) {
	if (!out) {
		Real tv = -0.25/sqrt(_pi), tvn, tv1, tv2;
		Complex zf, zfd;
		memset(ampl,0,3*deg_max*sizeof(Complex));
		for (int n=0; n<deg_max; ++n) {
			tvn = tv*sqrt(Real(2*n+3));
			zf = besh1(kz,n+1); zfd = besh1d(kz,n+1);
			ampl[3*n] = tvn*zf; // e-polarization, m = 1
			ampl[3*n+1] = -im1*tvn*(zfd + zf/kz); // h-polarization m = 1
			tv = -tv;
		}
	}
	else {
		Real tv = -0.25/sqrt(_pi), tvn, tv1, tv2;
		Complex zf, zfd;
		memset(ampl,0,3*deg_max*sizeof(Complex));
		for (int n=0; n<deg_max; ++n) {
			tvn = tv*sqrt(Real(2*n+3));
			zf = besj(kz,n+1); zfd = besjd(kz,n+1);
			ampl[3*n] = tvn*zf; // m = -1,1
			ampl[3*n+1] = -im1*tvn*(zfd + zf/kz);
			tv = -tv;
		}
	}
}

std::vector<Complex> sphere_ml::get_last_ampl() const {
	std::vector<Complex> res;
	res.resize(3*deg_max);
	for (int i=0; i<3*deg_max; ++i) res[i] = ampl_out[i];
	return res;
}

std::vector<Complex> sphere_ml::ampl_far(Real th_rad, Real ph_rad) {
	Complex tci = -im1, tc, ts = Real(0);
	std::vector<Complex> E_out = {Real(0),Real(0)};
	get_pitau_m01(cos(th_rad),pi1,tau0,tau1,deg_max);
	for (int n=0; n<deg_max; ++n) {
		tc = tci/sqrt(Real((n+1)*(n+2)));
		E_out[0] -= tc * ( ampl_out[3*n] * pi1[n] + ampl_out[3*n+1] * tau1[n] );
		E_out[1] += tc * ( ampl_out[3*n] * tau1[n] + ampl_out[3*n+1] * pi1[n] );
		ts += tc * ampl_out[3*n+2] * tau0[n];
		tci *= -im1;
	}
	E_out[0] *= sqrt(2/_pi) * cos(ph_rad);
	E_out[1] = sqrt(2/_pi) * ( sin(ph_rad)*E_out[1] + 0.5*ts );
	return E_out;
}

// private members

void sphere_ml::sphere_sm(Real kr, Complex eps_in, Complex eps_out, Complex *res) {
	Complex t_besj_in, t_besdzj_in, t_besj_out, t_besdzj_out;
	Complex t_besh_in, t_besdzh_in, t_besh_out, t_besdzh_out;
	Complex kr_in, kr_out, te, tc, tcc;

	kr_in = kr*sqrt(eps_in); if (arg(kr_in) < -1.e-8) kr_in = -kr_in;
	kr_out = kr*sqrt(eps_out); if (arg(kr_out) < -1.e-8) kr_out = -kr_out;
	te = eps_in/eps_out;

	for (int n=0; n<deg_max; ++n) {
		t_besj_in = besj(kr_in,n+1);
		t_besdzj_in = bes_dzj(kr_in,n+1);
		t_besj_out = besj(kr_out,n+1);
		t_besdzj_out = bes_dzj(kr_out,n+1);

		t_besh_in = besh1(kr_in,n+1);
		t_besdzh_in = bes_dzh1(kr_in,n+1);
		t_besh_out = besh1(kr_out,n+1);
		t_besdzh_out = bes_dzh1(kr_out,n+1);

		tc = Real(1) / (t_besj_in*t_besdzh_out - t_besh_out*t_besdzj_in);
		res[0+8*n] = tc * (t_besh_out*t_besdzh_in - t_besh_in*t_besdzh_out); // 00e
		res[2+8*n] = tc * im1 / kr_in; // 01e
		res[6+8*n] = tc * (t_besj_out*t_besdzj_in - t_besj_in*t_besdzj_out); // 11e
		res[4+8*n] = tc * im1 / kr_out; // 10e
		tc = Real(1) / (te*t_besj_in*t_besdzh_out - t_besh_out*t_besdzj_in);
		res[1+8*n] = tc * (t_besh_out*t_besdzh_in - te*t_besh_in*t_besdzh_out); // 00h
		res[3+8*n] = tc * im1 / kr_out; // 01h
		res[7+8*n] = tc * (t_besj_out*t_besdzj_in - te*t_besj_in*t_besdzj_out); // 11h
		res[5+8*n] = tc * im1 * te / kr_in; // 10h
	}
}

void sphere_ml::compose_sm(Complex *in, Complex *out, Complex *res) {
	Complex tmp, tmp1;
	for (int n=0; n<deg_max; ++n) {
		tmp = Real(1) / (Real(1) - in[8*n+6]*out[8*n+0]);
		tmp1 = out[8*n+2] * out[8*n+4];
		res[8*n+0] = in[8*n+0] + tmp * in[8*n+2] * in[8*n+4] * out[8*n+0]; // 00e
		res[8*n+2] = tmp * in[8*n+2] * out[8*n+2]; // 01e
		res[8*n+4] = tmp * in[8*n+4] * out[8*n+4]; // 10e
		res[8*n+6] = out[8*n+6] + tmp * tmp1 * in[8*n+6]; // 11e
		
		tmp = Real(1)/(Real(1) - in[8*n+7]*out[8*n+1]);
		tmp1 = out[8*n+3] * out[8*n+5];
		res[8*n+1] = in[8*n+1] + tmp * in[8*n+3] * in[8*n+5] * out[8*n+1]; // 00h
		res[8*n+3] = tmp * in[8*n+3] * out[8*n+3]; // 01h
		res[8*n+5] = tmp * in[8*n+5] * out[8*n+5]; // 10h
		res[8*n+7] = out[8*n+7] + tmp * tmp1 * in[8*n+7]; // 11h
	}
}

Real sphere_ml::directivity_dipole_z(Complex *ampl) const {
	Real w_sca = 0;
	Complex ti = -im1, tc;
	Complex sum1, sum2, sum3;
	
	sum1 = sum2 = sum3 = Real(0);
	
	for (int n=0; n<deg_max; ++n) {
		tc = ti/sqrt(Real((n+1)*(n+2)));
		sum1 += tc * ( ampl[3*n]*pi1[n] + ampl[3*n+1]*tau1[n] );
		sum2 += tc * ( ampl[3*n]*tau1[n] + ampl[3*n+1]*pi1[n] );
		sum3 += tc * ampl[3*n+2]*tau0[n];
		ti *= -im1;
		w_sca += ( norm(ampl[3*n]) + norm(ampl[3*n+1]) ) + 0.5*norm(ampl[3*n+2]);
	}
	sum1 *= Real(2)*cos(phi);
	sum2 *= Real(2)*im1*sin(phi);
	
//	return Real(2)*((sum1+sum3)*conj(sum1+sum3) + sum2*conj(sum2)).real() / w_sca;
	return (norm(sum1+sum3) + norm(sum2)) / w_sca;
}

//############################################################################################

void sphere_ml_parallel::set_directivity_angles(Real theta_, Real phi_) {
	theta = theta_; phi = phi_;
	get_pitau_m01(cos(theta),pi1,tau0,tau1,deg_max);
}

Real sphere_ml_parallel::get_directivity_edipole_x(Real kz, const std::vector<Real> &kR, const std::vector<Complex> &eps) {
	int ni = size(kR);
	if (size(eps)-1 != ni) {std::cout << "get_directivity_edipole_x error\n"; return Real(0);}
	
	Real D = 0.;
	Complex kzn;
	Complex rtml_in[8*_degree], rtml_out[8*_degree], rt[8*_degree];
	Complex ampl_in[3*_degree], ampl_out[3*_degree];
	
	if (kz <= kR[0]) {
			// s-matrix of a the multilayer
		sphere_sm(kR[0], eps[0], eps[1], rtml_out);
		//cout << "sphere coeff:\n";
		//for (int i=0; i<deg_max; ++i) {cout << i << " " << rtml_out[8*i+2] << " " << rtml_out[8*i+3] << endl;}
		for (int k=1; k<ni; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_out, rt, rtml_out);
		}
			// dipole amplitudes
		kzn = kz*sqrt(eps[0]);
		if (arg(kzn) < -1e-14) kzn = -kzn;
		get_edipole_amp_x(kzn, ampl_out, 1);
			// output amplitudes
		for (int n=0; n<deg_max; ++n) {
			ampl_out[3*n] *= rtml_out[8*n+2]; // e, m = +-1
			ampl_out[3*n+1] *= rtml_out[8*n+3]; // h, m = +-1
			ampl_out[3*n+2] *= rtml_out[8*n+3]; // h, m = 0
		}
	}
	else if (kz <= kR[kR.size()-1]) { // the dipole is located in a layer
			// determine a dipole layer
		int kd = 0;
		while (kz > kR[++kd]); // determine a first interface with radius larger than the dipole position
			// calculate s-matrices of multilayers inside and outside the dipole position:
		sphere_sm(kR[0], eps[0], eps[1], rtml_in);
		for (int k=1; k<kd; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_in, rt, rtml_in);
		}
		sphere_sm(kR[kd], eps[kd], eps[kd+1], rtml_out);
		for (int k=kd+1; k<ni; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_out, rt, rtml_out);
		}
			// dipole amplitudes
		kzn = kz*sqrt(eps[kd]);
		if (arg(kzn) < -1e-14) kzn = -kzn;
		get_edipole_amp_x(kzn, ampl_in, 0);
		get_edipole_amp_x(kzn, ampl_out, 1);
			// output self-consistent amplitudes
		for (int n=0; n<deg_max; ++n) {
			ampl_out[3*n] = ( ampl_in[3*n]*rtml_in[8*n+6] + ampl_out[3*n] ) * rtml_out[8*n+2]
										/ (Real(1) - rtml_in[8*n+6]*rtml_out[8*n]); // e, m = 1
			ampl_out[3*n+1] = ( ampl_in[3*n+1]*rtml_in[8*n+7] + ampl_out[3*n+1] ) * rtml_out[8*n+3]
										/ (Real(1) - rtml_in[8*n+7]*rtml_out[8*n+1]); // h, m = 1
			ampl_out[3*n+2] = ( ampl_in[3*n+2]*rtml_in[8*n+7] + ampl_out[3*n+2] ) * rtml_out[8*n+3]
										/ (Real(1) - rtml_in[8*n+7]*rtml_out[8*n+1]); // h, m = 0
		}
	}
	else { // the dipole is located outside the multilayer
			// s-matrix of a the multilayer
		sphere_sm(kR[0], eps[0], eps[1], rtml_out);
		for (int k=1; k<ni; ++k) {
			sphere_sm(kR[k], eps[k], eps[k+1], rt);
			compose_sm(rtml_out, rt, rtml_out);
		}
			// dipole amplitudes
		kzn = kz*sqrt(eps[ni]);
		if (arg(kzn) < -1e-14) kzn = -kzn;
		get_edipole_amp_x(kzn, ampl_in, 0);
		get_edipole_amp_x(kzn, ampl_out, 1);
			// output amplitudes
		for (int n=0; n<deg_max; ++n) {
			ampl_out[3*n] += ampl_in[3*n] * rtml_out[8*n+6]; // e, m = +-1
			ampl_out[3*n+1] += ampl_in[3*n+1] * rtml_out[8*n+7]; // h, m = +-1
			ampl_out[3*n+2] += ampl_in[3*n+2] * rtml_out[8*n+7]; // h, m = 0
		}
	}

	return directivity_dipole_z(ampl_out);
}

// private members

void sphere_ml_parallel::get_edipole_amp_x(Complex kz, Complex *ampl, bool out) {
	if (!out) {
		Real tv = -0.25/sqrt(_pi), tvn, tv1, tv2;
		Complex zf, zfd;
		memset(ampl,0,3*deg_max*sizeof(Complex));
		for (int n=0; n<deg_max; ++n) {
			tvn = tv*sqrt(Real(2*n+3));
			zf = besh1(kz,n+1); zfd = besh1d(kz,n+1);
			ampl[3*n] = tvn*zf; // e-polarization, m = 1
			ampl[3*n+1] = -im1*tvn*(zfd + zf/kz); // h-polarization m = 1
			tv = -tv;
		}
	}
	else {
		Real tv = -0.25/sqrt(_pi), tvn, tv1, tv2;
		Complex zf, zfd;
		memset(ampl,0,3*deg_max*sizeof(Complex));
		for (int n=0; n<deg_max; ++n) {
			tvn = tv*sqrt(Real(2*n+3));
			zf = besj(kz,n+1); zfd = besjd(kz,n+1);
			ampl[3*n] = tvn*zf; // m = -1,1
			ampl[3*n+1] = -im1*tvn*(zfd + zf/kz);
			tv = -tv;
		}
	}
}

void sphere_ml_parallel::sphere_sm(Real kr, Complex eps_in, Complex eps_out, Complex *res) {
	Complex t_besj_in, t_besdzj_in, t_besj_out, t_besdzj_out;
	Complex t_besh_in, t_besdzh_in, t_besh_out, t_besdzh_out;
	Complex kr_in, kr_out, te, tc, tcc;

	kr_in = kr*sqrt(eps_in); if (arg(kr_in) < -1.e-8) kr_in = -kr_in;
	kr_out = kr*sqrt(eps_out); if (arg(kr_out) < -1.e-8) kr_out = -kr_out;
	te = eps_in/eps_out;

	for (int n=0; n<deg_max; ++n) {
		t_besj_in = besj(kr_in,n+1);
		t_besdzj_in = bes_dzj(kr_in,n+1);
		t_besj_out = besj(kr_out,n+1);
		t_besdzj_out = bes_dzj(kr_out,n+1);

		t_besh_in = besh1(kr_in,n+1);
		t_besdzh_in = bes_dzh1(kr_in,n+1);
		t_besh_out = besh1(kr_out,n+1);
		t_besdzh_out = bes_dzh1(kr_out,n+1);

		tc = Real(1) / (t_besj_in*t_besdzh_out - t_besh_out*t_besdzj_in);
		res[0+8*n] = tc * (t_besh_out*t_besdzh_in - t_besh_in*t_besdzh_out); // 00e
		res[2+8*n] = tc * im1 / kr_in; // 01e
		res[6+8*n] = tc * (t_besj_out*t_besdzj_in - t_besj_in*t_besdzj_out); // 11e
		res[4+8*n] = tc * im1 / kr_out; // 10e
		tc = Real(1) / (te*t_besj_in*t_besdzh_out - t_besh_out*t_besdzj_in);
		res[1+8*n] = tc * (t_besh_out*t_besdzh_in - te*t_besh_in*t_besdzh_out); // 00h
		res[3+8*n] = tc * im1 / kr_out; // 01h
		res[7+8*n] = tc * (t_besj_out*t_besdzj_in - te*t_besj_in*t_besdzj_out); // 11h
		res[5+8*n] = tc * im1 * te / kr_in; // 10h
	}
}

void sphere_ml_parallel::compose_sm(Complex *in, Complex *out, Complex *res) {
	Complex tmp, tmp1;
	for (int n=0; n<deg_max; ++n) {
		tmp = Real(1) / (Real(1) - in[8*n+6]*out[8*n+0]);
		tmp1 = out[8*n+2] * out[8*n+4];
		res[8*n+0] = in[8*n+0] + tmp * in[8*n+2] * in[8*n+4] * out[8*n+0]; // 00e
		res[8*n+2] = tmp * in[8*n+2] * out[8*n+2]; // 01e
		res[8*n+4] = tmp * in[8*n+4] * out[8*n+4]; // 10e
		res[8*n+6] = out[8*n+6] + tmp * tmp1 * in[8*n+6]; // 11e

		tmp = Real(1)/(Real(1) - in[8*n+7]*out[8*n+1]);
		tmp1 = out[8*n+3] * out[8*n+5];
		res[8*n+1] = in[8*n+1] + tmp * in[8*n+3] * in[8*n+5] * out[8*n+1]; // 00h
		res[8*n+3] = tmp * in[8*n+3] * out[8*n+3]; // 01h
		res[8*n+5] = tmp * in[8*n+5] * out[8*n+5]; // 10h
		res[8*n+7] = out[8*n+7] + tmp * tmp1 * in[8*n+7]; // 11h
	}
}

Real sphere_ml_parallel::directivity_dipole_z(Complex *ampl) const {
	Real w_sca = 0;
	Complex ti = -im1, tc;
	Complex sum1, sum2, sum3;

	sum1 = sum2 = sum3 = Real(0);

	for (int n=0; n<deg_max; ++n) {
		tc = ti/sqrt(Real((n+1)*(n+2)));
		sum1 += tc * ( ampl[3*n]*pi1[n] + ampl[3*n+1]*tau1[n] );
		sum2 += tc * ( ampl[3*n]*tau1[n] + ampl[3*n+1]*pi1[n] );
		sum3 += tc * ampl[3*n+2]*tau0[n];
		ti *= -im1;
		w_sca += ( norm(ampl[3*n]) + norm(ampl[3*n+1]) ) + 0.5*norm(ampl[3*n+2]);
	}
	sum1 *= Real(2)*cos(phi);
	sum2 *= Real(2)*im1*sin(phi);

	return (norm(sum1+sum3) + norm(sum2)) / w_sca;
}