/* -*- 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/>.
 *
 */

#include <mia/core/fastica_nonlinearity.hh>
#include <mia/core/gsl_matrix_vector_ops.hh>
#include <mia/core/export_handler.hh>
#include <mia/core/plugin_base.cxx>
#include <mia/core/handler.cxx>
#include <algorithm>

namespace mia {

using std::transform; 
using std::vector; 
using namespace gsl; 
using std::ostream_iterator; 

CFastICANonlinearityBase::CFastICANonlinearityBase():
	m_mu(1.0),
	m_sample_scale(1.0), 
        m_signal(nullptr)
{
}

void CFastICANonlinearityBase::set_signal(const Matrix *signal)
{
	m_signal = signal; 
	assert(m_signal); 
	m_sample_scale = 1.0 / m_signal->cols(); 
	
	post_set_signal(); 
}

void CFastICANonlinearityBase::set_mu(double mu)
{
	m_mu = mu; 
}



const Matrix& CFastICANonlinearityBase::get_signal() const
{
	assert(m_signal); 
	return *m_signal; 
}

void CFastICADeflNonlinearity::post_set_signal()
{
	m_workspace = Vector(get_signal().rows(), false); 
	m_XTw = Vector(get_signal().cols(), false); 
}

const char *CFastICANonlinearityBase::data_descr = "fastica"; 
const char *CFastICADeflNonlinearity::type_descr = "deflation"; 


void CFastICADeflNonlinearity::apply(gsl::Vector& w)
{
	multiply_v_m(m_XTw, w, get_signal());
	const double scale = get_correction_and_scale(m_XTw, m_workspace); 

	if (get_mu() >= 1.0) 
		sum_final(w, scale);  
	else 
		sum_final_stabelized(w, scale); 
}

void CFastICADeflNonlinearity::apply(gsl::Matrix& W)
{
	for (unsigned c = 0; c < W.cols(); ++c) {
		auto wc = gsl_matrix_column(W, c); 
		gsl::Vector w(&wc.vector); 
		
		multiply_v_m(m_XTw, w, get_signal());
		const double scale = get_correction_and_scale(m_XTw, m_workspace); 

		if (get_mu() >= 1.0) 
			sum_final(w, scale);  
		else 
			sum_final_stabelized(w, scale); 
	}
}

vector<double> CFastICADeflNonlinearity::get_saddle_test_table(const gsl::Matrix& ics) const
{
	vector<double> result(ics.rows()); 
	for (unsigned i = 0; i < ics.rows(); ++i) {
		auto row = ics.get_row(i); 
		result[i] = do_get_saddle_test_value(row);
	}
	return result; 
}

double CFastICADeflNonlinearity::get_saddle_test_value(const gsl::Vector& ic) const
{
	return do_get_saddle_test_value(ic); 
}


void CFastICADeflNonlinearity::sum_final(gsl::Vector& w, double scale)
{
	const double inv_m = get_sample_scale(); 

	transform(m_workspace.begin(), m_workspace.end(), w.begin(), w.begin(), 
		  [scale, inv_m](double x, double y) { return (x - scale * y) * inv_m;}); 

}

void CFastICADeflNonlinearity::sum_final_stabelized(gsl::Vector& w, double scale)
{
	const double beta = dot(w, m_workspace); 
	const double a2 = get_mu() / (scale - beta); 
	const double a1 = 1 + beta * a2; 

	transform(m_workspace.begin(), m_workspace.end(), w.begin(), w.begin(), 
		  [a1, a2](double x, double y){return a1 * y - a2 * x;}); 

}

template<>  const char * const 
TPluginHandler<TFactory<CFastICADeflNonlinearity>>::m_help = 
	"These plug-ins provide various non-linearity models for the FastICA algorithm "
        "that uses deflation.";

EXPORT_CORE PFastICADeflNonlinearity produce_fastica_nonlinearity(const std::string& descr)
{
	return CFastICADeflNonlinearityPluginHandler::instance().produce(descr);
}

EXPLICIT_INSTANCE_HANDLER(CFastICADeflNonlinearity); 


}// end namespace 
