/* -*- mia-c++  -*-
 *
 * This file is part of pymia - python bindings for MIA
 * Copyright (c) Leipzig, Madrid 1999-2013 Gert Wollny
 *
 * pymia is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with pymia; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <Python.h>
#include "mia_conversions.hh"
#include <mia/2d/filter.hh>
#include <mia/2d/imageio.hh>
#include <mia/2d/nonrigidregister.hh>
#include <mia/3d/filter.hh>
#include <mia/3d/imageio.hh>
#include <mia/3d/nonrigidregister.hh>



#if PY_MAJOR_VERSION >= 3
#define IS_PYTHON3
#endif


using namespace std;
using namespace mia;

static PyObject *MiaError; 

static PyArrayObject *run_filters_2d(PyArrayObject *input, const vector<string>& filter_list) 
{
	TRACE_FUNCTION; 
	auto image = mia_image_from_pyarray<T2DImage>(input);
	auto result = C2DImageFilterChain(filter_list).run(image); 
	return mia_pyarray_from_image(*result);
}

static PyArrayObject *run_filters_3d(PyArrayObject *input, const vector<string>& filter_list) 
{
	TRACE_FUNCTION; 
	auto image = mia_image_from_pyarray<T3DImage>(input); 
	auto result = C3DImageFilterChain(filter_list).run(image); 
	return mia_pyarray_from_image(*result);
}

#ifdef IS_PYTHON3	
std::string as_string(PyObject *obj)
{
	auto h2 =  PyUnicode_AsUTF8String(obj); 
	if (!h2) {
		throw create_exception<invalid_argument>("mia.get_strings_in_list: non-string value in list");  
	}
	std::string s(PyBytes_AsString(h2)); 
	Py_DECREF(h2); 
        return s; 
}
#else 
std::string as_string(PyObject *obj)
{
        auto s = PyString_AsString(obj); 
        if (!s)
		throw create_exception<invalid_argument>("mia.get_strings_in_list: non-string value in list");  
        return std::string(s); 
}
#endif 


static vector<string> get_strings_in_list(PyObject *obj)
{
	vector<string> result_list; 
	if ( PyList_Check(obj)) {
		auto length = PyList_Size(obj);
		result_list.reserve(length); 
		for (Py_ssize_t i = 0; i < length; ++i) {
			auto help = PyList_GET_ITEM(obj,i); 
			result_list.push_back(as_string(help)); 
		}
	} else { 
		result_list.push_back(as_string(obj)); 
	}
	return result_list;
}

PYTHON_MIA_CALL_self_args(run_filters)
{
	PyArrayObject *py_resultarray = NULL;
	PyArrayObject *py_inputarray;
	PyObject *py_filterlist;
	

	/* parse input */
	if (!PyArg_ParseTuple(args,"O!O",&PyArray_Type, &py_inputarray, &py_filterlist))
		return NULL;

	auto filter_list = get_strings_in_list(py_filterlist); 
	
	switch (PyArray_NDIM(py_inputarray)) {
	case 2: py_resultarray = run_filters_2d(py_inputarray, filter_list); 
		break; 
	case 3: py_resultarray = run_filters_3d(py_inputarray, filter_list); 
		break; 
	default:
		throw create_exception<invalid_argument>("mia dosn't support images of ", 
							 PyArray_NDIM(py_inputarray), " dimensions");
	}
	
	// return result
	return PyArray_Return(py_resultarray);
}

PYTHON_MIA_CALL_self_args(set_verbose)
{
	const char *verbosity; 
	if (!PyArg_ParseTuple(args,"s", &verbosity))
		return NULL;
	
	cverb.set_verbosity(g_verbose_dict.get_value( verbosity) ); 
	Py_INCREF(Py_None);
	return Py_None;
}

template <typename Handler>
PyObject *load_image(const Handler& handler, PyObject *args) 
{
	const char *filename; 
	if (!PyArg_ParseTuple(args,"s", &filename))
		return NULL;
	
	auto images = handler.load(filename);
	if (images && !images->empty()) {
		if (images->size() == 1) 
			return (PyObject*)mia_pyarray_from_image(*(*images)[0]);
		else {
			PyObject* output = PyList_New(images->size()); 
			for(size_t i = 0; i < images->size(); ++i) {
				auto img = mia_pyarray_from_image(*(*images)[i]);
				PyList_SetItem(output, i, (PyObject*)img);
			}
			return output; 
		}
	}
	throw create_exception<runtime_error>("No images found in '", filename, "'"); 
}

PYTHON_MIA_CALL_self_args(load_image2d)
{
	return load_image(C2DImageIOPluginHandler::instance(), args); 
}

PYTHON_MIA_CALL_self_args(load_image3d)
{
	return load_image(C3DImageIOPluginHandler::instance(), args); 
}

PYTHON_MIA_CALL_self_args(set_filter_plugin_cache) 
{
	int enabled; 
	if (!PyArg_ParseTuple(args,"i", &enabled))
		return NULL;
	
	C2DFilterPluginHandler::instance().set_caching(enabled); 
	C3DFilterPluginHandler::instance().set_caching(enabled); 
	
	Py_INCREF(Py_None);
	return Py_None;
}

struct SRegistrationParameters {
	const char *transform; 
	PyObject *cost; 
	
	// predefined 
	int mg_levels; 
	const char *optimizer; 
	const char *refiner; 
}; 

template <int dim> 
struct register_images_d {
	typedef typename dimension_traits<dim>::PTransformationFactory PTransformationFactory; 
	typedef typename dimension_traits<dim>::PFullCost PFullCost; 
	typedef typename dimension_traits<dim>::FullCostList FullCostList; 
	typedef typename FactoryTrait<PTransformationFactory>::type TransformationFactoryHandler; 
	typedef typename FactoryTrait<PFullCost>::type FullCostHandler; 
	typedef typename dimension_traits<dim>::PImage PImage; 

	static PyObject *apply(PImage src, PImage ref, const SRegistrationParameters& p); 
}; 

template <int dim> 
PyObject *register_images_d<dim>::apply(PImage src, PImage ref, const SRegistrationParameters& p)
{
	if (src->get_size() != ref->get_size()) 
		cvwarn() << "mia.register_images: Images are of different size, this hasnÄt beed tested\n"; 
	
	auto transform_creator = TransformationFactoryHandler::instance().produce(p.transform); 
	auto optimizer = produce_minimizer(p.optimizer); 
	auto cost_list = get_strings_in_list(p.cost);
	
	if (cost_list.empty())
		throw invalid_argument("mia.register_images: Got empty cost function list"); 
	
	FullCostList costs; 
	for (auto i = cost_list.begin(); i != cost_list.end(); ++i)
		costs.push(FullCostHandler::instance().produce(*i)); 
	
	
	TNonrigidRegister<dim> nrr(costs, optimizer,  transform_creator, p.mg_levels);
	if (p.refiner) {
		auto refiner = produce_minimizer(p.refiner); 
		nrr.set_refinement_minimizer(refiner); 
	}
	
	auto transform = nrr.run(src, ref);
	auto result = (*transform)(*src);
	
	// currently the transformation type is not defined, return only the transformed image
	return (PyObject *)mia_pyarray_from_image(*result); 
}

PYTHON_MIA_CALL_self_args_kwdict(register_images)
{
	SRegistrationParameters p = {
		NULL, NULL, 3, 
		"nlopt:opt=ld-var1,xtola=0.001,ftolr=0.001,maxiter=300", 
		NULL
	}; 

	// keywords:  
	// study - source or floating image 
	// reference - reference image 
	// transform = transformation 
	// cost - list of cost functions

	// multigrid - number of multigrid levels 
	// optimizer - optimizer string 
	// refiner - refinement 
	static const char *kwlist[] = {"src", "ref", "transform", "cost",
				 "mglevels", "optimizer", "refiner", NULL};

	PyArrayObject *src = NULL; 
	PyArrayObject *ref = NULL; 

	if (!PyArg_ParseTupleAndKeywords(args, kwdict, "O!O!sO|iss", 
					 const_cast<char**>(kwlist), 
					 &PyArray_Type, &src, &PyArray_Type, &ref, 
					 &p.transform, &p.cost, &p.mg_levels, &p.optimizer, &p.refiner))
		return NULL;
	
	if (PyArray_NDIM(src) != PyArray_NDIM(ref)) {
		throw create_exception<invalid_argument>("Images must be of the same dimensions, but the source "
							 "is of dimension ",  PyArray_NDIM(src), " and the reference of dimension ", PyArray_NDIM(ref)); 
	}
	switch (PyArray_NDIM(src)) {
	case 2: {
		auto miasrc = mia_image_from_pyarray<T2DImage>(src);
		auto miaref = mia_image_from_pyarray<T2DImage>(ref);
		return register_images_d<2>::apply(miasrc, miaref, p); 
	}
	case 3: {
		auto miasrc = mia_image_from_pyarray<T3DImage>(src);
		auto miaref = mia_image_from_pyarray<T3DImage>(ref);
		return register_images_d<3>::apply(miasrc, miaref, p); 
	}
	default:
		throw create_exception<invalid_argument>("mia dosn't support images of ",  
							 PyArray_NDIM(src), " dimensions");
	}
}


static struct PyMethodDef mia_methods[]={
	{ "load_image2d", load_image2d, METH_VARARGS, "loads one or more images from a file. "
	  "If the file contains only one images, it is returned directly as numpy.array, "
	  "otherwise a list of numpy.array objects is returned." },
	{ "load_image3d", load_image3d, METH_VARARGS, "loads one or more images from a file. "
	  "If the file contains only one images, it is returned directly as numpy.array, "
	  "otherwise a list of numpy.array objects is returned." },
	{ "filter",  run_filters, METH_VARARGS, "run the given filter(s) on an input image. "
	  "The image must be given as a numpy.array of scalar values. Either you give one filter "
	  "as string, or a series of filters as list. If no filter as given then you must give at "
	  "least an empty list."},
	{ "register_images", (PyCFunction)register_images, METH_VARARGS | METH_KEYWORDS, 
	  "run the registration of the image 'src' to the image 'ref' allowing for the transformation "
	  "'transform' and optimizing cost functions 'cost'. Use 'mglevels' multi-resolution levels "
	  " for registration and optimize by using 'optimizer'. A refinement optimizer 'refiner' can be given " 
	  "for post-iteration."},
	{ "set_filter_plugin_cache", set_filter_plugin_cache, METH_VARARGS, "Enable or disable the "
	  "2D and 3D filter plugin caches"}, 
	{ "set_verbose",  set_verbose, METH_VARARGS, "set the verbosity of the MIA output" },
	{ NULL, NULL, 0, NULL} /* stop mark */
};

/* ----------- module initialization -------------------------- */

#ifdef IS_PYTHON3	



static struct PyModuleDef moduledef = {
        PyModuleDef_HEAD_INIT,
        "mia",
        NULL,
        0,
        mia_methods,
        NULL,
        NULL,
        NULL,
        NULL
};

#endif

#ifdef IS_PYTHON3	
extern "C" 
PyObject *
PyInit_mia(void){
	PyObject *m,*d ;

	m = PyModule_Create(&moduledef);


	/* initialize exception object */
	d=PyModule_GetDict(m) ; /* get module dictionary */
	MiaError=PyErr_NewException(const_cast<char*>("mia.error"),NULL,NULL) ;
	PyDict_SetItemString(d,"error",MiaError) ;
	
	import_array();
	
	if (PyErr_Occurred()) /* something went wrong ?*/
		Py_FatalError("can't initialize module mia") ;
        return m; 

}



#else 

extern "C" 
PyMODINIT_FUNC
initmia() {
	PyObject *m,*d ;

	m=Py_InitModule("mia", mia_methods);

	/* initialize exception object */
	d=PyModule_GetDict(m) ; /* get module dictionary */
	MiaError=PyErr_NewException(const_cast<char*>("mia.error"),NULL,NULL) ;
	PyDict_SetItemString(d,"error",MiaError) ;
	
	import_array();
	
	if (PyErr_Occurred()) /* something went wrong ?*/
		Py_FatalError("can't initialize module mia") ;
}
#endif 

