/* -*- 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 <limits>
#include <algorithm>
#include <cmath>
#include <boost/algorithm/minmax_element.hpp>
#include <mia/core/dictmap.hh>
#include <mia/2d/filter/aniso.hh>


NS_BEGIN(aniso_2dimage_filter)
NS_MIA_USE;
using namespace std;

// scale factor for MAD to zero-mean normal distribution
static const float zmn_weight = 1.4826;
static const float cinv_sqrt2 = 1.0f / sqrt(2.0f);
static const char *plugin_name = "aniso";

float psi_guess(float x, float sigma)
{
       const float val = x / sigma;
       return x * expf(- val * val * 0.25f);
}

float psi_test(float /*x*/, float /*sigma*/)
{
       return 1.0;
}


float psi_tuckey(float x, float sigma)
{
       if (x > sigma || -x > sigma)
              return 0.0;

       const float val = x / sigma;
       const float val2 = 1.0 - val * val;
       return x * val2 * val2;
}

float psi_pm1(float x, float sigma)
{
       const float val = x / sigma;
       return 2.0f * x / (2.0f + val * val);
}

float psi_pm2(float x, float sigma)
{
       const float val = x / sigma;
       return x * expf(- val * val * 0.25f);
}

static const TDictMap<C2DAnisoDiff::FEdgeStopping>::Table edge_stop_table[] = {
       { "tuckey", psi_tuckey, "tukey stopping function"},
       { "pm1", psi_pm1, "stopping function 1"},
       { "pm2", psi_pm2, "stopping function 2"},
       { "guess", psi_guess, "test stopping function"},
       { NULL, NULL, ""}
};

static const TDictMap<C2DAnisoDiff::FEdgeStopping> edge_stop_dict(edge_stop_table);


C2DAnisoDiff::C2DAnisoDiff(int maxiter, float epsilon, float k, FEdgeStopping edge_stop, int neighbourhood):
       m_maxiter(maxiter),
       m_epsilon(epsilon),
       m_k(k),
       m_edge_stop(edge_stop),
       m_neighbourhood(neighbourhood),
       m_histogramfeeder(0, 256, 256),
       m_sigma_e(0.0),
       m_gamma(0.0),
       m_sigma(0.0)
{
       if (m_neighbourhood != 4 && m_neighbourhood != 8) {
              stringstream errmsg;
              errmsg << "neighbourhood " << m_neighbourhood << "not supported. Select 4 or 8";
              throw invalid_argument(errmsg.str());
       }
}

typedef pair<float, float> HScale;

inline bool operator < (const HScale& a, const HScale& b)
{
       return a.first < b.first;
}

void C2DAnisoDiff::create_histogramfeeder(const C2DFImage& data) const
{
       pair<C2DFImage::const_iterator, C2DFImage::const_iterator>
       range = ::boost::minmax_element(data.begin(), data.end());
       // this is to avoid some warnings
       float dist = *range.second - *range.first;
       cvdebug() << "Histogram spread = " << dist
                 << " in [" << *range.first
                 << ", " << *range.second  << "]\n";
       size_t bins;

       if (dist <= 4096.0f)
              bins = 4096;
       else if (dist >= 16384)
              bins = 16384;
       else
              bins = (size_t)(dist + 1);

       m_histogramfeeder = THistogramFeeder<float>(*range.first, *range.second, bins);
}

/* estimate the MAD */
float C2DAnisoDiff::estimate_MAD(const C2DFImage& data)const
{
       THistogram<THistogramFeeder<float>>  gradient_histogram(m_histogramfeeder);
       C2DFImage::const_iterator id = data.begin();

       for (size_t y = 0; y < data.get_size().y - 1; ++y) {
              for (size_t x = 0; x < data.get_size().x - 1; ++x, ++id) {
                     const float idd = *id;
                     const float d3 = id[+1] - idd;
                     const float d4 = id[data.get_size().x]  - idd;
                     gradient_histogram.push ( d3 > 0 ? d3 : -d3 );
                     gradient_histogram.push ( d4 > 0 ? d4 : -d4 );
              }

              const float d4 = id[data.get_size().x]  - *id;
              gradient_histogram.push ( d4 > 0 ? d4 : -d4 );
              ++id;
       }

       for (size_t x = 0; x < data.get_size().x - 1; ++x, ++id) {
              const float d3 = id[+1] - *id;
              gradient_histogram.push ( d3 > 0 ? d3 : -d3 );
       }

       return gradient_histogram.MAD();
}

float C2DAnisoDiff::diffuse(C2DFImage& dest, const C2DFImage& src)const
{
       float sum = 0.0;
       C2DFImage::iterator idest = dest.begin() + src.get_size().x;
       C2DFImage::const_iterator id = src.begin() + src.get_size().x;
       int sxm1 = src.get_size().x - 1;
       int sxp1 = src.get_size().x + 1;
       int sx   = src.get_size().x;
       // lazy way out to get the boundaries copied
       // just let the test pass for now
       copy(src.begin(), src.end(), dest.begin());

       for (size_t y = 1; y < src.get_size().y - 1; ++y) {
              ++id;
              ++idest;

              for (size_t x = 1; x < src.get_size().x - 1; ++x, ++id, ++idest) {
                     float val  = 0.0;
                     float idd = *id;

                     if (m_neighbourhood == 8) {
                            val += m_edge_stop(cinv_sqrt2 * (id[-sxp1]  - idd), m_sigma);
                            val += m_edge_stop(cinv_sqrt2 * (id[-sxm1]  - idd), m_sigma);
                            val += m_edge_stop(cinv_sqrt2 * (id[ sxm1]  - idd), m_sigma);
                            val += m_edge_stop(cinv_sqrt2 * (id[ sxp1]  - idd), m_sigma);
                     }

                     val += m_edge_stop(id[-sx] - idd, m_sigma);
                     val += m_edge_stop(id[ -1] - idd, m_sigma);
                     val += m_edge_stop(id[ +1] - idd, m_sigma);
                     val += m_edge_stop(id[ sx] - idd, m_sigma);
                     val *= m_gamma;
                     sum += val * val;
                     *idest = idd + val;
              }

              ++id;
              ++idest;
       }

       return sum;
}

void C2DAnisoDiff::update_gamma_sigma(const mia::C2DFImage& src)const
{
       if (m_k <= 0) {
              m_sigma_e = zmn_weight *  estimate_MAD(src);
              m_sigma = sqrt(5.0f) *  m_sigma_e;
       } else {
              m_sigma_e = m_k / sqrt(5.0f);
              m_sigma = m_k;
       }

       m_gamma = 1.0 / (m_edge_stop(m_sigma_e, m_sigma) * m_neighbourhood);
}

template <class T>
typename C2DAnisoDiff::result_type C2DAnisoDiff::operator () (const T2DImage<T>& image) const
{
       cvdebug() << "begin C2DAnisoDiff::operator (), k=" << m_k << "\n";
       int iter = 0;
       C2DFImage *dest = new C2DFImage(image.get_size());;
       C2DFImage *src = new C2DFImage(image.get_size());
       std::copy(image.begin(), image.end(), src->begin());
       create_histogramfeeder(*src);
       char endline = cverb.show_debug() ? '\n' : '\r';
       float delta;

       do {
              ++iter;
              update_gamma_sigma(*src);

              if (m_sigma_e == 0.0)  // image contains only one colour
                     break;

              delta = diffuse(*dest, *src);
              cvmsg() << iter << ": " << " m_sigma_e = " << m_sigma_e << " m_gamma = " << m_gamma << " ";
              cvmsg() << " delta " << delta << "       " << endline;
              swap(src, dest);
       } while (delta > m_epsilon && iter < m_maxiter);

       cvmsg() << '\n';
       const T max_val = numeric_limits<T>::max();
       const T min_val = numeric_limits<T>::min();
       T2DImage<T> *result = new T2DImage<T>(image.get_size(), image);
       typename T2DImage<T>::iterator di = result->begin();

       for (C2DFImage::const_iterator wi = dest->begin();
            wi != dest->end(); ++wi, ++di) {
              *di = (*wi > min_val) ?
                    ( (*wi < max_val) ? static_cast<T>(*wi) : max_val )
                    : min_val;
       }

       delete src;
       delete dest;
       return P2DImage(result);
}

P2DImage C2DAnisoDiff::do_filter(const C2DImage& image) const
{
       return mia::filter(*this, image);
}

CAnisoDiff2DImageFilterFactory::CAnisoDiff2DImageFilterFactory():
       C2DFilterPlugin(plugin_name),
       m_maxiter(100),
       m_epsilon(1.0),
       m_k ( -1.0),
       m_edge_stop(psi_tuckey),
       m_neighbourhood(8)
{
       TRACE("CAnisoDiff2DImageFilterFactory::CAnisoDiff2DImageFilterFactory()");
       cvdebug() << "CAnisoDiff2DImageFilterFactory::CAnisoDiff2DImageFilterFactory()\n";
       add_parameter("iter", make_ci_param(m_maxiter, 1, 10000, false,  "number of iterations"));
       add_parameter("epsilon", make_positive_param(m_epsilon, false,  "iteration change threshold"));
       add_parameter("k", make_ci_param(m_k, 0.0f, 100.0f, false, "k the noise threshold (<=0 -> adaptive)"));
       add_parameter("psi", new CDictParameter<C2DAnisoDiff::FEdgeStopping>(m_edge_stop, edge_stop_table,
                     "edge stopping function"));
       set<int> nset;
       nset.insert(4);
       nset.insert(8);
       add_parameter("n", new CSetParameter<int>(m_neighbourhood, nset, "neighbourhood"));
}

C2DFilter *CAnisoDiff2DImageFilterFactory::do_create() const
{
       return new C2DAnisoDiff(m_maxiter, m_epsilon, m_k, m_edge_stop, m_neighbourhood);
}

const string CAnisoDiff2DImageFilterFactory::do_get_descr()const
{
       return string("2D Anisotropic image filter");
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{
       return new CAnisoDiff2DImageFilterFactory();
}

NS_END
