#include "interaction.h"
#include "math.h"


#define raise2(x) (x)*(x)

static float 
interaction_force_function
	( float r
	, float * coefficients)
{
	float result = 0;
	int i;
	for(i = 0; i < 7; i++)
	{
		result += coefficients[i] * powf(r, i);
	}
	result *= coefficients[7] * expf(coefficients[8] * (r - coefficients[9]));
	result += coefficients[10] * expf(coefficients[11] * (r - coefficients[12]));
	result += coefficients[13] * expf(coefficients[14] * raise2(r - coefficients[15]));
	result += coefficients[16] * expf(coefficients[17] * raise2(r - coefficients[18]));
	return result;
}

static void
interaction_ufunc_force
	( char ** args
	, npy_intp * dimensions
	, npy_intp * steps
	, void * data)
{
	char * in = args[0]
		, * out = args[1];
	npy_intp n = dimensions[0];
	npy_intp in_step = steps[0]
		, out_step = steps[1];

	npy_intp i;

	float * coefficients = (float *) data;

	for(i = 0; i < n; i++)
	{
		*(float *)out = interaction_force_function(*(float *)in, coefficients);
		out += out_step;
		in += in_step;
	}
}
	

static void
interaction_ufunc_float2D
	( char ** args
	, npy_intp * dimensions
	, npy_intp * steps
	, void * data)
{
	npy_intp i;
	npy_intp j;
	npy_intp n = dimensions[0];

	char * x_old = args[0]
		, * y_old = args[1]
		, * p_x_old = args[2]
		, * p_y_old = args[3]
		, * p_x_new = args[4]
		, * p_y_new = args[5];

	npy_intp x_old_steps = steps[0]
		, y_old_steps = steps[1]
		, p_x_old_steps = steps[2]
		, p_y_old_steps = steps[3]
		, p_x_new_steps = steps[4]
		, p_y_new_steps = steps[5];

	float *coefficients = (float *) data;

	// Compute the new momenta: 

	// Stuff we will need:
	float r;         // Distance between interacting particles.
	float delta_x_e; // x-component of unit direction vector.
	float delty_y_e; // y-component of unit direction vector.
	float delta_p_x;
	float delta_p_y;

	// The current x_old[i], y_old[i] coordinates.
	float this_x_i;
	float this_y_i;
	// The current x_old[j], y_old[j] coordinates.
	float this_x_j;
	float this_y_j;


	for(i = 0; i < n; i++)
	{
		this_x_i = *(float *)(x_old + i*x_old_steps);
		this_y_i = *(float *)(y_old + i*y_old_steps);
		
		// copy current momenta
		*(float *)(p_x_new + i*p_x_new_steps) = *(float *)(p_x_old + i*p_x_old_steps);
		*(float *)(p_y_new + i*p_y_new_steps) = *(float *)(p_y_old + i*p_y_old_steps);

		// compute and add the momentum offset
		for(j = 0; j < i; j++)
		{
			this_x_j = *(float *)(x_old + j*x_old_steps);
			this_y_j = *(float *)(y_old + j*y_old_steps);

			// Compute distance and direction between particles i,j.
			r = sqrtf(raise2(this_x_i - this_x_j) + raise2(this_y_i - this_y_j));
			delta_x_e = (this_x_i - this_x_j) / r;
			delty_y_e = (this_y_i - this_y_j) / r;

			// Update the momenta.
			delta_p_x = delta_x_e * interaction_force_function(r, coefficients);
			delta_p_y = delty_y_e * interaction_force_function(r, coefficients);
			*(float *)(p_x_new + i*p_x_new_steps) -= delta_p_x;
			*(float *)(p_y_new + i*p_y_new_steps) -= delta_p_y;
			*(float *)(p_x_new + j*p_x_new_steps) += delta_p_x;
			*(float *)(p_y_new + j*p_y_new_steps) += delta_p_y;
		}
	}
}

typedef struct
{
	PyObject_HEAD
	float coefficients[16];
	PyObject * ufunc;
} interaction_UFuncWrapper;

static PyTypeObject interaction_UFuncWrapperType
{
	PyTypeObject_HEAD_INIT(NULL, 0)
	.tp_name = "brown.interaction.UFuncWrapper",
	.tp_doc = "A wrapper that wraps the ufuncs for interaction and force, storing the coefficients",
	.tp_basicsize = sizeof(interaction_UFuncWrapper),
	.tp_itemsize = 0,
	.tp_flags = Py_TPFLAGS_DEFAULT,
	.tp_new = PyType_GenericNew,

};


static char interaction_types[] = 
	{ NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT, NPY_FLOAT};
static char force_types[] =
	{ NPY_FLOAT, NPY_FLOAT};

static int
interaction_UFuncWrapper_init
	( interaction_UFuncWrapper * self
	, PyObject * args
	, PyObject * kwds)
{
	// 0: ufunc_force
	// 1: interaction2D
	char type;
	PyObject * coefficients;
	int i;
	PyObject * this_coefficient;

	if(!PyArgs_ParseTupleAndKeywords(args, kwds, "BO!", &type, &PyList_Type, &coefficients))
	{
		return -1;
	}

	if(PyList_Size(coefficients) != 19)
	{
		PyErr_SetString("coefficients must have length 19", PyExc_ValueError);
		return -1;
	}

	// copy the coefficients.
	for(i = 0; i < 19; i++)
	{
		this_coefficient = PyList_GetItem(coefficients, i);
		// XXX: PyFloat_AsDouble might call python code,
		// so make sure that nothing bad can happen.
		Py_INCREF(this_coefficient);
		self->coefficients[i] = PyFloat_AsDouble(this_coefficient);
		Py_DECREF(this_coefficient);
		if(PyErr_Occurred())
		{
			return -1;
		}
	}

	switch(type)
	{
		case 0:
		{
			self->ufunc = PyUFunc_FromFuncAndData(
				force_funcs
				, self->coefficients
				, force_types
				, 1
				, 2
				, 1
				, PyUFunc_None
				, "force_function"
				, "computes the scalar force between two particles with given coefficients"
				, 0);
			break;
		}
		case 1:
		{
			self->ufunc = PyUFunc_FromFuncAndData(
				interaction_funcs
				, interaction_data
				, interaction_types
				, 1
				, 5
				, 2
				, PyUFunc_None
				, "interaction2D"
				, "Update the momenta according to the given coefficients and positions"
				, 0);
			break;
		}
		default:
		{
			PyErr_SetString(PyExc_ValueError, "unknown ufunc type, must be 0 or 1");
			return -1
		}

	}
}


static PyObject *
interaction_UFuncWrapper_call
	(interaction_UFuncWrapper * self
	 , PyObject * args
	 , PyObject * kwargs)
{
	return PyObject_Call(self->ufunc, args, kwargs);
}

static PyMethodDef InteractionMethods[] = {
	{NULL, NULL, 0, NULL}
};


static struct PyModuleDef moduledef = {
	PyModuleDef_HEAD_INIT
	, "brown.interaction"
	, NULL
	, -1
	, InteractionMethods
	, NULL
	, NULL
	, NULL
	, NULL
};

PyMODINIT_FUNC 
PyInit_interaction(void)
{
	PyObject * module
		, * ufunc_interaction
		, * ufunc_force
		, * dct;

	module = PyModule_Create(&moduledef);
	if(!module)
	{
		return NULL;
	}
	import_array();
	import_ufunc();

	ufunc_interaction = PyUFunc_FromFuncAndDataAndSignature(
	ufunc_force = PyUFunc_FromFuncAndDataAndSignature(

	dct = PyModule_GetDict(module);
	PyDict_SetItemString(dct, "interaction2D", ufunc_interaction);
	PyDict_SetItemString(dct, "force_function", ufunc_force);
	Py_DECREF(ufunc_interaction);
	Py_DECREF(ufunc_force);

	return module;
}