/*
Vnd:  n-dimensional vector class with conjugate gradient Ax=b solver
Author:  Aaron Dwyer, 2003
Note:  This class hasn't been rigorously tested outside of my physics simulations and therefore its correctness is in no way gauranteed.  If anyone does determine it has problems (or doesn't) I'd be glad to hear about it.
*/

#include "Vnd.h"

Vnd::Vnd()
{
	n = 0;
	vec = 0;
}

Vnd::Vnd(int nParam)
{
	n = nParam;
	vec = new double[n];
}

Vnd::~Vnd()
{
	delete [] vec;
}

void Vnd::SetSize(int size)
{
	if (n != size)
	{
		double *newStorage = new double[size];
		for (int i = 0; i < n; i++)
			newStorage[i] = vec[i];
		
		if (vec)
			delete [] vec;
		
		vec = newStorage;
		n = size;
	}
}

void Vnd::SetZero()
{
	for (int index = 0; index < n; index++)
	{
		vec[index] = 0;
	}
}

void Vnd::SetAll(double s)
{
	for (int index = 0; index < n; index++)
	{
		vec[index] = s;
	}
}

void Vnd::Add(Vnd &dest, Vnd &lhs, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		dest.vec[index] = lhs.vec[index] + rhs.vec[index];
	}
}

void Vnd::AddTo(Vnd &lhs, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		lhs.vec[index] += rhs.vec[index];
	}
}

void Vnd::AddScaled(Vnd &dest, Vnd &lhs, double s, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		dest.vec[index] = lhs.vec[index] + rhs.vec[index] * s;
	}
}

void Vnd::AddScaledTo(Vnd &lhs, double s, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		lhs.vec[index] += rhs.vec[index] * s;
	}
}

void Vnd::Subtract(Vnd &dest, const Vnd &lhs, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		dest.vec[index] = lhs.vec[index] - rhs.vec[index];
	}
}

void Vnd::Scale(Vnd &dest, double lhs, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		dest.vec[index] = lhs * rhs.vec[index];
	}
}

void Vnd::Scale(double lhs, Vnd &rhs)
{
	for (int index = 0; index < rhs.n; index++)
	{
		rhs.vec[index] *= lhs;
	}
}

double Vnd::Dot(Vnd &lhs, Vnd &rhs)
{
	double sum = 0;

	for (int index = 0; index < lhs.n; index++)
	{
		sum += lhs.vec[index] * rhs.vec[index];
	}
	
	return sum;
}

double Vnd::DotSelf(Vnd &v)
{
	double sum = 0;

	for (int index = 0; index < v.n; index++)
	{
		sum += v.vec[index] * v.vec[index];
	}
	
	return sum;
}

void Vnd::Copy(Vnd &dest, Vnd &source)
{
	for (int index = 0; index < dest.n; index++)
	{
		dest.vec[index] = source.vec[index];
	}
}

void Vnd::ConjugateGradient(MatrixTimesVector matrixTimesVector, void *userData, Vnd &x, const Vnd &b, int maxIterations, double epsilon)
{
	Vnd temp(x.n), r(x.n), d(x.n), q(x.n);
	double rDotrZero, rDotrOld, rDotrNew, alpha, beta;
	
	//r = b - Ax
	matrixTimesVector(userData, temp, x);
	Subtract(r, b, temp);
	
	//d = r
	Copy(d, r);
	
	//rDotrNew = rT * r
	rDotrNew = DotSelf(r);
	
	//rDotrZero = rDotrNew
	rDotrZero = rDotrNew;

	for (int iteration = 0;
		iteration < maxIterations && rDotrNew > epsilon*epsilon*rDotrZero;
		iteration++)
	{
		//q = Ad
		matrixTimesVector(userData, q, d);
		
		//alpha = rDotrNew / (d dot q)
		alpha = rDotrNew / Dot(d, q);
		
		//x = x + alpha * d
		AddScaledTo(x, alpha, d);
		
		//r = r - alpha * q
		AddScaledTo(r, -alpha, q);
		
		//rDotrOld = rDotrNew
		rDotrOld = rDotrNew;
		
		//rDotrNew = r dot r
		rDotrNew = DotSelf(r);
		
		//beta = rDotrNew / rDotrOld
		beta = rDotrNew / rDotrOld;
		
		//d = r + beta * d
		AddScaled(d, r, beta, d);
	}
}