#ifndef STIM_PLANEWAVE_H
#define STIM_PLANEWAVE_H

#include <string>
#include <sstream>
#include <cmath>

#include "../math/vector.h"
#include "../math/quaternion.h"
#include "../math/constants.h"
#include "../math/plane.h"
#include "../math/complex.h"

namespace stim{
	namespace optics{

		/// evaluate the scalar field produced by a plane wave at a point (x, y, z)

		/// @param x is the x-coordinate of the point
		/// @param y is the y-coordinate of the point
		/// @param z is the z-coordinate of the point
		/// @param A is the amplitude of the plane wave, specifically the field at (0, 0, 0)
		/// @param kx is the k-vector component in the x direction
		/// @param ky is the k-vector component in the y direction
		/// @param kz is the k-vector component in the z direction
		template<typename T>
		stim::complex<T> planewave_scalar(T x, T y, T z, stim::complex<T> A, T kx, T ky, T kz){
			T d = x * kx + y * ky + z * kz;						//calculate the dot product between k and p = (x, y, z) to find the distance p is along the propagation direction
			stim::complex<T> di = stim::complex<T>(0, d);		//calculate the phase shift that will have to be applied to propagate the wave distance d
			return A * exp(di);									//multiply the phase term by the amplitude at (0, 0, 0) to propagate the wave to p
		}

		/// evaluate the scalar field produced by a plane wave at several positions

		/// @param field is a pre-allocated block of memory that will store the complex field at all points
		/// @param N is the number of field values to be evaluated
		/// @param x is a set of x coordinates defining positions within the field (NULL implies that all values are zero)
		/// @param y is a set of y coordinates defining positions within the field (NULL implies that all values are zero)
		/// @param z is a set of z coordinates defining positions within the field (NULL implies that all values are zero)
		/// @param A is the amplitude of the plane wave, specifically the field at (0, 0, 0)
		/// @param kx is the k-vector component in the x direction
		/// @param ky is the k-vector component in the y direction
		/// @param kz is the k-vector component in the z direction
		template<typename T>
		void cpu_planewave_scalar(stim::complex<T>* field, size_t N, T* x, T* y = NULL, T* z = NULL, stim::complex<T> A = 1.0, T kx = 0.0, T ky = 0.0, T kz = 0.0){
			T px, py, pz;
			for(size_t i = 0; i < N; i++){										// for each element in the array
				(x == NULL) ? px = 0 : px = x[i];								// test for NULL values
				(y == NULL) ? py = 0 : py = y[i];
				(z == NULL) ? pz = 0 : pz = z[i];

				field[i] = planewave_scalar(px, py, pz, A, kx, ky, kz);			// call the single-value plane wave function
			}
		}

template<typename T>
class planewave{

protected:

	stim::vec<T> k;							//k-vector, pointed in propagation direction with magnitude |k| = tau / lambda = 2pi / lambda
	stim::vec< stim::complex<T> > E0;		//amplitude (for a scalar plane wave, only E0[0] is used)

	/// Bend a plane wave via refraction, given that the new propagation direction is known
	CUDA_CALLABLE planewave<T> bend(stim::vec<T> kn) const{

		stim::vec<T> kn_hat = kn.norm();				//normalize the new k
		stim::vec<T> k_hat = k.norm();					//normalize the current k

		planewave<T> new_p;								//create a new plane wave

		T k_dot_kn = k_hat.dot(kn_hat);					//if kn is equal to k or -k, handle the degenerate case

		//if k . n < 0, then the bend is a reflection
		if(k_dot_kn < 0) k_hat = -k_hat;				//flip k_hat

		if(k_dot_kn == -1){
			new_p.k = -k;
			new_p.E0 = E0;
			return new_p;
		}
		else if(k_dot_kn == 1){
			new_p.k = k;
			new_p.E0 = E0;
			return new_p;
		}

		vec<T> r = k_hat.cross(kn_hat);					//compute the rotation vector
		T theta = asin(r.len());						//compute the angle of the rotation about r
		quaternion<T> q;								//create a quaternion to capture the rotation
		q.CreateRotation(theta, r.norm());		
		vec< complex<T> > E0n = q.toMatrix3() * E0;		//apply the rotation to E0
		new_p.k = kn_hat * kmag();
		new_p.E0 = E0n;

		return new_p;
	}

public:

	///constructor: create a plane wave propagating along k
	CUDA_CALLABLE planewave(vec<T> kvec = stim::vec<T>(0, 0, stim::TAU), 
							vec< complex<T> > E = stim::vec<T>(1, 0, 0))
	{
		//phi = phase;

		k = kvec;
		vec< complex<T> > k_hat = k.norm();

		if(E.len() == 0)			//if the plane wave has an amplitude of 0
			E0 = vec<T>(0);			//just return it
		else{
			vec< complex<T> > s = (k_hat.cross(E)).norm();		//compute an orthogonal side vector
			vec< complex<T> > E_hat = (s.cross(k)).norm();	//compute a normalized E0 direction vector
			E0 = E_hat;// * E_hat.dot(E);					//compute the projection of _E0 onto E0_hat
		}

		E0 = E0 * exp( complex<T>(0, phase) );
	}

	///multiplication operator: scale E0
    CUDA_CALLABLE planewave<T> & operator* (const T & rhs){		
		E0 = E0 * rhs;
		return *this;
	}

	CUDA_CALLABLE T lambda() const{
		return stim::TAU / k.len();
	}

	CUDA_CALLABLE T kmag() const{
		return k.len();
	}

	CUDA_CALLABLE vec< complex<T> > E(){
		return E0;
	}

	CUDA_CALLABLE vec<T> kvec(){
		return k;
	}

	/// calculate the value of the field produced by the plane wave given a three-dimensional position
	CUDA_CALLABLE vec< complex<T> > pos(T x, T y, T z){
		return pos( stim::vec<T>(x, y, z) );
	}

	CUDA_CALLABLE vec< complex<T> > pos(vec<T> p = vec<T>(0, 0, 0)){
		vec< complex<T> > result;

		T kdp = k.dot(p);
		complex<T> x = complex<T>(0, kdp);
		complex<T> expx = exp(x);

		result[0] = E0[0] * expx;
		result[1] = E0[1] * expx;
		result[2] = E0[2] * expx;

		return result;
	}

	//scales k based on a transition from material ni to material nt
	CUDA_CALLABLE planewave<T> n(T ni, T nt){
		return planewave<T>(k * (nt / ni), E0);
	}

	CUDA_CALLABLE planewave<T> refract(stim::vec<T> kn) const{
		return bend(kn);
	}

	/// Calculate the result of a plane wave hitting an interface between two refractive indices

	/// @param P is a plane representing the position and orientation of the surface
	/// @param n0 is the refractive index outside of the surface (in the direction of the normal)
	/// @param n1 is the refractive index inside the surface (in the direction away from the normal)
	/// @param r is the reflected component of the plane wave
	/// @param t is the transmitted component of the plane wave
	void scatter(stim::plane<T> P, T n0, T n1, planewave<T> &r, planewave<T> &t){
		scatter(P, n1/n0, r, t);
	}

	/// Calculate the scattering result when nr = n1/n0

	/// @param P is a plane representing the position and orientation of the surface
	/// @param r is the ration n1/n0
	/// @param n1 is the refractive index inside the surface (in the direction away from the normal)
	/// @param r is the reflected component of the plane wave
	/// @param t is the transmitted component of the plane wave
	void scatter(stim::plane<T> P, T nr, planewave<T> &r, planewave<T> &t){

		int facing = P.face(k);		//determine which direction the plane wave is coming in

		if(facing == -1){		//if the wave hits the back of the plane, invert the plane and nr
			P = P.flip();			//flip the plane
			nr = 1/nr;				//invert the refractive index (now nr = n0/n1)
		}

		//use Snell's Law to calculate the transmitted angle
		T cos_theta_i = k.norm().dot(-P.norm());				//compute the cosine of theta_i
		T theta_i = acos(cos_theta_i);							//compute theta_i
		T sin_theta_t = (1/nr) * sin(theta_i);						//compute the sine of theta_t using Snell's law
		T theta_t = asin(sin_theta_t);							//compute the cosine of theta_t

		bool tir = false;						//flag for total internal reflection
		if(theta_t != theta_t){
			tir = true;
			theta_t = stim::PI / (T)2;
		}

		//handle the degenerate case where theta_i is 0 (the plane wave hits head-on)
		if(theta_i == 0){
			T rp = (1 - nr) / (1 + nr);		//compute the Fresnel coefficients
			T tp = 2 / (1 + nr);
			vec<T> kr = -k;
			vec<T> kt = k * nr;			//set the k vectors for theta_i = 0
			vec< complex<T> > Er = E0 * rp;		//compute the E vectors
			vec< complex<T> > Et = E0 * tp;
			T phase_t = P.p().dot(k - kt);	//compute the phase offset
			T phase_r = P.p().dot(k - kr);

			//create the plane waves
			r = planewave<T>(kr, Er, phase_r);
			t = planewave<T>(kt, Et, phase_t);
			return;
		}


		//compute the Fresnel coefficients
		T rp, rs, tp, ts;
		rp = tan(theta_t - theta_i) / tan(theta_t + theta_i);
		rs = sin(theta_t - theta_i) / sin(theta_t + theta_i);
		
		if(tir){
			tp = ts = 0;
		}
		else{
			tp = ( 2 * sin(theta_t) * cos(theta_i) ) / ( sin(theta_t + theta_i) * cos(theta_t - theta_i) );
			ts = ( 2 * sin(theta_t) * cos(theta_i) ) / sin(theta_t + theta_i);
		}

		//compute the coordinate space for the plane of incidence
		vec<T> z_hat = -P.norm();
		vec<T> y_hat = P.parallel(k).norm();
		vec<T> x_hat = y_hat.cross(z_hat).norm();

		//compute the k vectors for r and t
		vec<T> kr, kt;
		kr = ( y_hat * sin(theta_i) - z_hat * cos(theta_i) ) * kmag();
		kt = ( y_hat * sin(theta_t) + z_hat * cos(theta_t) ) * kmag() * nr;

		//compute the magnitude of the p- and s-polarized components of the incident E vector
		complex<T> Ei_s = E0.dot(x_hat);
		int sgn = E0.dot(y_hat).sgn();
		vec< complex<T> > cx_hat = x_hat;
		complex<T> Ei_p = ( E0 - cx_hat * Ei_s ).len() * sgn;
		//compute the magnitude of the p- and s-polarized components of the reflected E vector
		complex<T> Er_s = Ei_s * rs;
		complex<T> Er_p = Ei_p * rp;
		//compute the magnitude of the p- and s-polarized components of the transmitted E vector
		complex<T> Et_s = Ei_s * ts;
		complex<T> Et_p = Ei_p * tp;

		//compute the reflected E vector
		vec< complex<T> > Er = vec< complex<T> >(y_hat * cos(theta_i) + z_hat * sin(theta_i)) * Er_p + cx_hat * Er_s;
		//compute the transmitted E vector
		vec< complex<T> > Et = vec< complex<T> >(y_hat * cos(theta_t) - z_hat * sin(theta_t)) * Et_p + cx_hat * Et_s;

		T phase_t = P.p().dot(k - kt);
		T phase_r = P.p().dot(k - kr);

		//create the plane waves
		r.k = kr;
		r.E0 = Er * exp( complex<T>(0, phase_r) );

		t.k = kt;
		t.E0 = Et * exp( complex<T>(0, phase_t) );
	}

	std::string str()
	{
		std::stringstream ss;
		ss<<"Plane Wave:"<<std::endl;
		ss<<"	"<<E0<<" e^i ( "<<k<<" . r )";
		return ss.str();
	}
};					//end planewave class
}					//end namespace optics
}					//end namespace stim

template <typename T>
std::ostream& operator<<(std::ostream& os, stim::optics::planewave<T> p)
{
    os<<p.str();
    return os;
}

#endif