//Copyright (c) 2004-2005, Baris Sumengen
//All rights reserved.
//
// CIMPL Matrix Performance Library
//
//Redistribution and use in source and binary
//forms, with or without modification, are
//permitted provided that the following
//conditions are met:
//
//    * No commercial use is allowed. 
//    This software can only be used
//    for non-commercial purposes. This 
//    distribution is mainly intended for
//    academic research and teaching.
//    * Redistributions of source code must
//    retain the above copyright notice, this
//    list of conditions and the following
//    disclaimer.
//    * Redistributions of binary form must
//    mention the above copyright notice, this
//    list of conditions and the following
//    disclaimer in a clearly visible part 
//    in associated product manual, 
//    readme, and web site of the redistributed 
//    software.
//    * Redistributions in binary form must
//    reproduce the above copyright notice,
//    this list of conditions and the
//    following disclaimer in the
//    documentation and/or other materials
//    provided with the distribution.
//    * The name of Baris Sumengen may not be
//    used to endorse or promote products
//    derived from this software without
//    specific prior written permission.
//
//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT
//HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
//EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT
//NOT LIMITED TO, THE IMPLIED WARRANTIES OF
//MERCHANTABILITY AND FITNESS FOR A PARTICULAR
//PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
//CONTRIBUTORS BE LIABLE FOR ANY
//DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
//EXEMPLARY, OR CONSEQUENTIAL DAMAGES
//(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
//OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
//DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
//HOWEVER CAUSED AND ON ANY THEORY OF
//LIABILITY, WHETHER IN CONTRACT, STRICT
//LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
//OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
//OF THIS SOFTWARE, EVEN IF ADVISED OF THE
//POSSIBILITY OF SUCH DAMAGE.



#include "./PDE.h"





namespace PDE
{

	Vector<float> tridiagonal_solve(Vector<float>& a, Vector<float>& b, Vector<float>& c, Vector<float>& w)
	{
		// Finds the tridiagonal solution for Mu = w, M is NxN
		// a,b,c are the diagonals from left to right all length N
		// with a(0) and c(end) are irrelevant;
		
		int n = b.Numel();
		Vector<float> x(n);
		Vector<float> y(n);
		Vector<float> u(n);
		
		x[n-1] = -a[n-1]/b[n-1];
		y[n-1] = w[n-1]/b[n-1];
		
		for(int i=n-2;i>=1;i--)
		{
			x[i] = -a[i]/(b[i] + c[i]*x[i+1]);
			y[i] = (w[i] - c[i]*y[i+1])/(b[i] + c[i]*x[i+1]);
		}
		
		x[0] = 0;
		y[0] = (w[0] - c[0]*y[1])/(b[0] + c[0]*x[1]);
		
		u[0] = y[0];
		for(int i=1; i<=n-1; i++)
		{
			u[i] = x[i]*u[i-1]+y[i];
		}
		
		return u;
	}



	Vector<float> Poisson1D(Vector<float>& v, float dx, BoundaryCondition boundary)
	{

		//tridiagonals
		Vector<float> a;
		Vector<float> b;
		Vector<float> c;

		Vector<float> ou; 
		if(boundary == Neumann)
		{
			a = Vector<float>(v.Numel(), 1);
			b = Vector<float>(v.Numel(), -2);
			c = Vector<float>(v.Numel(), 1);

			b(0) = -1;
			b(b.Numel()-1) = -1;
			ou = tridiagonal_solve(a, b, c, v*(dx*dx));
		}
		else
		{
			a = Vector<float>(v.Numel()-2, 1);
			b = Vector<float>(v.Numel()-2, -2);
			c = Vector<float>(v.Numel()-2, 1);
			ou = Vector<float>(v.Numel(), 0);
			ou(1, v.Numel()-2) = tridiagonal_solve(a, b, c, v.Slice(1, v.Numel()-2)*(dx*dx));
		}

		//Vector<float> a(v.Numel()-2, 1);
		//Vector<float> b(v.Numel()-2, -2);
		//Vector<float> c(v.Numel()-2, 1);

		return ou;
	}

	Vector<float> Poisson1DFFT(Vector<float>& v, float dx, BoundaryCondition boundary)
	{
		Vector<float> V;
		if(boundary == Neumann)
		{
			V = FFTCos(v);
		}
		else
		{
			V = FFTSin(v);
		}

		V(0) = 0;
		for(int i=1; i<v.Numel(); i++)
		{
			V(i) *= float(4*dx*dx/(2*cos(2*PI*i/(v.Numel()))-2));
		}

		Vector<float> poi;
		if(boundary == Neumann)
		{
			poi = IFFTCos(V);
		}
		else
		{
			poi = IFFTSin(V);
		}
		return poi;
	}

	Matrix<float> Poisson2DFFT(Matrix<float>& v, float dx, BoundaryCondition boundary)
	{
		Matrix<float> V;
		if(boundary == Neumann)
		{
			V = FFT2Cos(v);
		}
		else
		{
			//V = FFTSin(v);
		}

		for(int i=0; i<v.Rows(); i++)
		{
			for(int j=0; j<v.Columns(); j++)
			{
				V(i,j) *= float(4*dx*dx/(2*(cos(2*PI*i/(v.Rows()))+cos(2*PI*j/(v.Columns()))-2)));
			}
		}
		V(0,0) = 0;

		Matrix<float> poi;
		if(boundary == Neumann)
		{
			poi = IFFT2Cos(V);
		}
		else
		{
			//poi = IFFTSin(V);
		}
		return poi;
	}

	//Matrix<float> Poisson2D(Matrix<float>& v, float dx, BoundaryCondition boundary)
	//{

	//}

	Vector<float> FFTSin(Vector<float>& m)
	{
		int end = m.Numel()-1;
		Vector<float> mm(2*end);
		mm(0) = m(0);
		mm(end) = m(0);
		for(int i=1; i<end;i++)
		{
			mm(i) = m(i);
			mm(2*end-i) = -m(i);
		}
		Vector<ComplexFloat> MM = FFT(mm);
		
		Vector<float> M(end+1);
		M(0) = 0;
		M(end) = 0;
		for(int i=1; i<end;i++)
		{
			M(i) = -2*imag(MM(i));
		}
		return M;
	}
	
	Vector<float> IFFTSin(Vector<float>& M)
	{
		int end = M.Numel()-1;
		Vector<ComplexFloat> MM(2*end);
		MM(0) = 0;
		MM(end) = 0;
		for(int i=1; i<end;i++)
		{
			MM(i) = ComplexFloat(0, -M(i)/2);
			MM(2*end-i) = ComplexFloat(0, M(i)/2);
		}
		Vector<ComplexFloat> mm = IFFT(MM);
		
		Vector<float> m(end+1);
		m(0) = 0;
		m(end) = 0;
		for(int i=1; i<end;i++)
		{
			m(i) = real(mm(i));
		}
		return m;
	}

	Vector<float> FFTCos(Vector<float>& m)
	{
		int end = m.Numel()-1;
		Vector<float> mm(2*end);
		mm(0) = m(0);
		mm(end) = m(end);
		for(int i=1; i<end;i++)
		{
			mm(i) = m(i);
			mm(2*end-i) = m(i);
		}
		Vector<ComplexFloat> MM = FFT(mm);
		
		Vector<float> M(end+1);
		M(0) = real(MM(0));
		M(end) = real(MM(end));
		for(int i=1; i<end;i++)
		{
			M(i) = 2*real(MM(i));
		}
		return M;
	}
	
	Vector<float> IFFTCos(Vector<float>& M)
	{
		int end = M.Numel()-1;
		Vector<ComplexFloat> MM(2*end);
		MM(0) = M(0);
		MM(end) = M(end);
		for(int i=1; i<end;i++)
		{
			MM(i) = M(i)/2;
			MM(2*end-i) = M(i)/2;
		}
		Vector<ComplexFloat> mm = IFFT(MM);
		
		Vector<float> m(end+1);
		m(0) = real(mm(0));
		m(end) = real(mm(end));
		for(int i=1; i<end;i++)
		{
			m(i) = real(mm(i));
		}
		return m;
	}


	Matrix<float> FFT2Cos(Matrix<float>& m)
	{
		int endR = m.Rows()-1;
		int endC = m.Columns()-1;
		Matrix<float> mm = Matrix<float>::Cat(1, (m , FlipLRI(m.Slice(0,endR,1,endC-1))), (FlipUDI(m.Slice(1,endR-1,0,endC)), FlipUDI(FlipLRI(m.Slice(1,endR-1,1,endC-1)))) );

		Matrix<ComplexFloat> MM = FFT2(ToComplexFloat(mm));
		
		Matrix<float> M(m.Rows(), m.Columns());
		for(int i=1; i<endR; i++)
		{
			for(int j=1; j<endC; j++)
			{
				M(i,j) = 4*real(MM(i,j));
			}
		}
		for(int i=1; i<endR; i++)
		{
			M(i,0) = 2*real(MM(i,0));
			M(i,endC) = 2*real(MM(i,endC));
		}
		for(int i=1; i<endC; i++)
		{
			M(0,i) = 2*real(MM(0,i));
			M(endR,i) = 2*real(MM(endR,i));
		}
		M(0,0) = real(MM(0,0));
		M(0,endC) = real(MM(0,endC));
		M(endR,0) = real(MM(endR,0));
		M(endR,endC) = real(MM(endR,endC));

		return M;
	}
	
	Matrix<float> IFFT2Cos(Matrix<float>& M)
	{
		int endR = M.Rows()-1;
		int endC = M.Columns()-1;

		Matrix<float> MM(M.Rows(),M.Columns());
		for(int i=1; i<endR;i++)
		{
			for(int j=1; j<endC;j++)
			{
				MM(i,j) = M(i,j)/4;
			}
		}
		for(int i=1; i<endR; i++)
		{
			MM(i,0) = M(i,0)/2;
			MM(i,endC) = M(i,endC)/2;
		}
		for(int i=1; i<endC; i++)
		{
			MM(0,i) = M(0,i)/2;
			MM(endR,i) = M(endR,i)/2;
		}
		MM(0,0) = M(0,0);
		MM(0,endC) = M(0,endC);
		MM(endR,0) = M(endR,0);
		MM(endR,endC) = M(endR,endC);

		MM = Matrix<float>::Cat(1, (MM , FlipLRI(MM.Slice(0,endR,1,endC-1))), (FlipUDI(MM.Slice(1,endR-1,0,endC)), FlipUDI(FlipLRI(MM.Slice(1,endR-1,1,endC-1)))) );

		Matrix<float> mm = Real(IFFT2(ToComplexFloat(MM)));
		return mm.Slice(0,endR,0,endC);
	}



};