/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis 
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA 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 MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */

#ifndef mia_internal_seededwatershed_hh
#define mia_internal_seededwatershed_hh

#include <queue>
#include <mia/template/dimtrait.hh>

NS_MIA_BEGIN

///  @cond DOC_PLUGINS 

template <int dim> 
class TSeededWS : public watershed_traits<dim>::Handler::Product {
public:
	typedef typename watershed_traits<dim>::PNeighbourhood  PNeighbourhood; 
	typedef typename PNeighbourhood::element_type::value_type MPosition; 
	typedef typename watershed_traits<dim>::Handler Handler; 
	typedef typename watershed_traits<dim>::FileHandler FileHandler; 
	typedef typename Handler::Product CFilter; 
	typedef typename FileHandler::Instance::DataKey  DataKey; 
	typedef typename CFilter::Pointer PFilter; 
	typedef typename CFilter::Image CImage; 
	typedef typename CImage::Pointer PImage; 
	typedef typename CImage::dimsize_type Position; 
	

	TSeededWS(const DataKey& mask_image, PNeighbourhood neighborhood, 
		  bool with_borders, bool input_is_gradient);

	template <template <typename>  class Image, typename T>
	typename TSeededWS<dim>::result_type operator () (const Image<T>& data) const;
private:
	virtual PImage do_filter(const CImage& image) const;

	DataKey m_label_image_key; 
	PNeighbourhood m_neighborhood; 
	PFilter m_togradnorm; 
	bool m_with_borders; 
};

template <int dim> 
class TSeededWSFilterPlugin: public watershed_traits<dim>::Handler::Interface {
public:
	typedef typename watershed_traits<dim>::PNeighbourhood  PNeighbourhood; 
	typedef typename watershed_traits<dim>::Handler Handler; 
	typedef typename watershed_traits<dim>::FileHandler FileHandler; 
	typedef typename Handler::Product CFilter; 

	TSeededWSFilterPlugin();
private:
	virtual CFilter *do_create()const;
	virtual const std::string do_get_descr()const;
	std::string m_seed_image_file;
	PNeighbourhood m_neighborhood; 
	bool m_with_borders; 
	bool m_input_is_gradient; 
};

template <template <typename>  class Image, typename  T, typename  S, typename N, typename R, int dim, bool supported>
struct seeded_ws {
	static R apply(const Image<T>& image, const Image<S>& seed, N n, bool with_borders); 
}; 

template <template <typename>  class Image, typename  T, typename  S, typename N, typename R, int dim>
struct seeded_ws<Image, T, S, N, R, dim, false> {
	static R apply(const Image<T>& /*image*/, const Image<S>& /*seed*/, N /*n*/, bool /*with_borders*/) {
		throw create_exception<std::invalid_argument>("C2DRunSeededWS: seed data type '", __type_descr<S>::value, "' not supported");
	}
}; 


// helper to dispatch based on the pixel type of the seed image 
template <template <typename>  class Image, typename  T, typename N, typename R, int dim>
struct dispatch_RunSeededWS : public TFilter<R> {
	
	dispatch_RunSeededWS(N neighborhood, const Image<T>& image, bool with_borders):
		m_neighborhood(neighborhood), 
		m_image(image), 
		m_with_borders(with_borders)
		{}

	template <typename  S>
	R operator () (const Image<S>& seed) const {
		const bool supported = std::is_integral<S>::value && !std::is_same<S, bool>::value; 
		return seeded_ws<Image, T, S, N, R, dim, supported>::apply(m_image, seed, m_neighborhood, m_with_borders); 
	}
	N m_neighborhood; 
	const Image<T>& m_image; 
	bool m_with_borders; 
}; 


template <typename L, typename Position>
struct PixelWithLocation {
	Position pos; 
	float value; 
	L label; 
}; 

template <typename L, typename Position>
bool operator < (const PixelWithLocation<L,Position>& lhs, const PixelWithLocation<L,Position>& rhs) 
{		
	mia::less_then<Position> l; 
	return lhs.value > rhs.value|| 
		( lhs.value ==  rhs.value && l(rhs.pos, lhs.pos)); 
}

template <template <typename>  class Image, typename T, typename S, int dim>
class TRunSeededWatershed {
public: 
	typedef typename watershed_traits<dim>::PNeighbourhood  PNeighbourhood; 
	typedef typename PNeighbourhood::element_type::value_type MPosition; 
	typedef typename watershed_traits<dim>::Handler Handler; 
	typedef typename Handler::Product CFilter; 
	typedef typename CFilter::Pointer PFilter; 
	typedef typename CFilter::Image CImage; 
	typedef typename CImage::Pointer PImage; 
	typedef typename CImage::dimsize_type Position; 
	

	TRunSeededWatershed(const Image<T>& image, const Image<S>& seed, PNeighbourhood neighborhood, bool with_borders); 
	PImage run(); 
private: 
	void add_neighborhood(const PixelWithLocation<S, Position>& pixel); 	
	const Image<T>& m_image;
	const Image<S>& m_seed; 
	PNeighbourhood m_neighborhood; 
	Image<unsigned char> m_visited; 
	Image<unsigned char>  m_stored; 
	Image<S>       *m_result;
	PImage m_presult; 
	std::priority_queue<PixelWithLocation<S, Position> > m_seeds; 
	S m_watershed; 
	bool m_with_borders; 
}; 

template <template <typename>  class Image, typename T, typename S, int dim>
TRunSeededWatershed<Image, T, S, dim>::TRunSeededWatershed(const Image<T>& image, const Image<S>& seed, 
							   PNeighbourhood neighborhood, bool with_borders):
	m_image(image), 
	m_seed(seed), 
	m_neighborhood(neighborhood), 
	m_visited(seed.get_size()),
	m_stored(seed.get_size()),
	m_watershed(std::numeric_limits<S>::max()), 
	m_with_borders(with_borders)
{
	m_result = new Image<S>(seed.get_size(), image); 
	m_presult.reset(m_result); 
}

template <template <typename>  class Image, typename T, typename S, int dim>
void TRunSeededWatershed<Image, T, S, dim>::add_neighborhood(const PixelWithLocation<S, Position>& pixel)
{
	PixelWithLocation<S, Position> new_pixel; 
	new_pixel.label = pixel.label; 
	bool hit_boundary = false; 
	for (auto i = m_neighborhood->begin(); i != m_neighborhood->end(); ++i) {
		Position new_pos( pixel.pos + *i); 
		if (new_pos != pixel.pos && new_pos < m_seed.get_size()) {
			if (!m_visited(new_pos)) {
				if (!m_stored(new_pos) ) {
					new_pixel.value = m_image(new_pos); 
					new_pixel.pos = new_pos; 
					m_seeds.push(new_pixel); 
					m_stored(new_pos) = 1; 
				}
			}else{
				hit_boundary |= (*m_result)(new_pos) != pixel.label &&
					(*m_result)(new_pos) != m_watershed; 
			}
		}
	}
	// set pixel to new label 
	if (!m_visited(pixel.pos)) {
		m_visited(pixel.pos) = true; 
		(*m_result)(pixel.pos) = (m_with_borders && hit_boundary) ? m_watershed : pixel.label;
	}
}

template <template <typename>  class Image, typename T, typename S, int dim>
typename TRunSeededWatershed<Image,T,S,dim>::PImage 
TRunSeededWatershed<Image,T,S,dim>::run()
{
	// copy seed and read initial pixels 
	auto iv = m_visited.begin(); 
	auto is = m_seed.begin_range(Position::_0, m_seed.get_size());
	auto es = m_seed.end_range(Position::_0, m_seed.get_size());
	auto ir = m_result->begin();
	auto ii = m_image.begin(); 
	auto ist = m_stored.begin(); 
	
	PixelWithLocation<S, Position> pixel; 
	while (is != es) {
		*ir = *is; 
		if (*is) {
			*iv = 1; 
			*ist; 
			pixel.pos = is.pos(); 
			pixel.value = *ii; 
			pixel.label = *is; 
			m_seeds.push(pixel); 
		}
		++iv; ++is; ++ir; ++ist; ++ii; 
	}
	
	while (!m_seeds.empty()) {
		auto p = m_seeds.top(); 
		m_seeds.pop(); 
		add_neighborhood(p); 
	}
	return m_presult; 
}

template <template <typename>  class Image, typename  T, typename  S, typename N, typename R, int dim, bool supported>
R seeded_ws<Image,T,S,N,R,dim,supported>::apply(const Image<T>& image, 
						const Image<S>& seed, N neighborhood, bool with_borders) 
{
	TRunSeededWatershed<Image, T, S, dim> ws(image, seed, neighborhood, with_borders);
	return ws.run(); 
}

template <int dim> 
TSeededWS<dim>::TSeededWS(const DataKey& label_image_key, PNeighbourhood neighborhood, bool with_borders, bool input_is_gradient):
	m_label_image_key(label_image_key), 
	m_neighborhood(neighborhood), 
	m_with_borders(with_borders)
	
{
	if (!input_is_gradient) 
		m_togradnorm = Handler::instance().produce("gradnorm"); 
}

template <int dim> 
template <template <typename>  class Image, typename T>
typename TSeededWS<dim>::result_type TSeededWS<dim>::operator () (const Image<T>& data) const 
{
	// read start label image
	auto in_image_list = m_label_image_key.get();
	if (!in_image_list || in_image_list->empty())
		throw create_exception<std::runtime_error>( "C2DSeededWS: no seed image could be loaded");
	
	if (in_image_list->size() > 1) 
		cvwarn() << "C2DSeededWS:got more than one seed image. Ignoring all but first"; 
	
	auto seed = (*in_image_list)[0]; 
	if (seed->get_size() != data.get_size()) {
		throw create_exception<std::invalid_argument>( "C2DSeededWS: seed and input differ in size: seed " , seed->get_size() 
		      , ", input " , data.get_size()); 
	}
	dispatch_RunSeededWS<Image, T, PNeighbourhood, PImage, dim> ws(m_neighborhood, data, m_with_borders); 
	return mia::filter(ws, *seed); 
}

template <int dim> 
typename TSeededWS<dim>::PImage TSeededWS<dim>::do_filter(const CImage& image) const
{
	PImage result;  
	if (m_togradnorm) {
		auto grad = m_togradnorm->filter(image); 
		result = mia::filter(*this, *grad);
	}
	else
		result = mia::filter(*this, image);
	
	return result;		
}

template <int dim> 
TSeededWSFilterPlugin<dim>::TSeededWSFilterPlugin(): 
	Handler::Interface("sws"), 
	m_with_borders(false), 
	m_input_is_gradient(false)
{
	this->add_parameter("seed", new CStringParameter(m_seed_image_file, CCmdOptionFlags::required_input, 
						   "seed input image containing the lables for the initial regions"));
	this->add_parameter("n", make_param(m_neighborhood, "sphere:r=1", false, "Neighborhood for watershead region growing")); 
	this->add_parameter("mark", new CBoolParameter(m_with_borders, false, "Mark the segmented watersheds with a special gray scale value")); 
	this->add_parameter("grad", new CBoolParameter(m_input_is_gradient, false, "Interpret the input image as gradient. ")); 
}

template <int dim> 
typename TSeededWSFilterPlugin<dim>::CFilter *TSeededWSFilterPlugin<dim>::do_create()const
{
	auto seed = FileHandler::instance().load_to_pool(m_seed_image_file);
	return new TSeededWS<dim>(seed, m_neighborhood, m_with_borders, m_input_is_gradient);
}

template <int dim> 
const std::string TSeededWSFilterPlugin<dim>::do_get_descr()const
{
	return "seeded watershead. The algorithm extracts exactly so "
		"many reagions as initial labels are given in the seed image.";
}

NS_MIA_END

///  @endcond
#endif
