/* -*- 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/probmap.hh>
#include <mia/core/sparse_histogram.hh>
#include <mia/core/factory.hh>

NS_MIA_BEGIN



class EXPORT_CORE CMeans
{
public:
       typedef std::vector<double> DVector;
       typedef CSparseHistogram::Compressed SparseHistogram;
       typedef std::vector<std::pair<double, double>> NormalizedHistogram;

       class EXPORT_CORE SparseProbmap
       {
       public:
              typedef std::pair<unsigned short, DVector> value_type;
              typedef std::vector<value_type> Map;

              SparseProbmap() = delete;
              SparseProbmap (size_t size);
              SparseProbmap (const std::string& filename);

              value_type& operator [](int i)
              {
                     return m_map[i];
              }

              const value_type& operator [](int i) const
              {
                     return m_map[i];
              }

              Map::const_iterator begin() const
              {
                     return m_map.begin();
              }
              Map::iterator begin()
              {
                     return m_map.begin();
              }

              Map::const_iterator end() const
              {
                     return m_map.end();
              }
              Map::iterator end()
              {
                     return m_map.end();
              }

              bool save(const std::string& filename) const;


              DVector get_fuzzy(double x) const;

              size_t size() const
              {
                     return m_map.size();
              }
       private:
              Map m_map;

       };


       class EXPORT_CORE Initializer : public CProductBase
       {
       public:
              typedef Initializer plugin_data;
              typedef Initializer plugin_type;
              static const char *data_descr;
              static const char *type_descr;
              virtual DVector run(const NormalizedHistogram& nh) const = 0;
       };
       typedef std::shared_ptr<Initializer> PInitializer;

       CMeans(double epsilon, PInitializer class_center_initializer);

       ~CMeans();

       SparseProbmap run(const SparseHistogram& histogram,  DVector& class_centers) const;

       SparseProbmap run(const SparseHistogram& histogram,  DVector& class_centers, bool de_normalize_results) const;

private:
       PInitializer m_cci;
       struct CMeansImpl *impl;

};

/**
   \brief evaluate the probabilities for a c-means classification with gain field

   This function evaluates the per-pixel class probabilities for a c-means
   classification with gain field correction.

   With n classes the evalaution is done aoocrding to

   \f[
   p_{k,i} := \left\{\begin{array}{lcl}
   I_k < c_0 & & p_{k,0} = 1, p_{k,i} = 0 \: \forall \: i \in [1, n-1]\\
   c_j < I_k < c_{j+1} & &
        p_{k,l} = \frac{(I_k - g_k * c_{m})^2}{(I_k - g_k * c_{m})^2 + (I_k - g_k * c_{l})^2}
       \: \forall \: (l,m) \in \{(j, j+1), (j+1, j)\}\\
	I_k > c_{n-1} & & p_{k,n} = 1,  p_{k,i} = 0 \: \forall \: i \in [0, n-2]
   \end{array} \right.
   \f]

   \tparam T input pixel type of the data to be classified
   \tparam Field template of the data field type
   \param[in] image image the classification is applied to
   \param[in] gain multiplicative gain field
   \param[in] class_centers
   \param[out] pv probability fields containing the evaluated probabilities

*/


template <typename T, template <class > class  Field>
void cmeans_evaluate_probabilities(const Field<T>& image, const Field<float>& gain,
                                   const std::vector<double>& class_centers,
                                   std::vector<Field<float>>& pv)
{
       assert(image.size() == gain.size());
       assert(class_centers.size() == pv.size());
#ifndef NDEBUG

       for (auto i : pv)
              assert(image.size() == i.size());

#endif
       auto ii = image.begin();
       auto ie = image.end();
       auto ig = gain.begin();
       typedef typename Field<float>::iterator prob_iterator;
       std::vector<prob_iterator> ipv(pv.size());
       transform(pv.begin(), pv.end(), ipv.begin(), [](Field<float>& p) {
              return p.begin();
       });
       std::vector<double> gain_class_centers(class_centers.size());

       while (ii != ie) {
              double x = *ii;

              for (auto iipv : ipv)
                     *iipv = 0.0;

              const double  vgain = *ig;
              transform(class_centers.begin(), class_centers.end(), gain_class_centers.begin(),
              [vgain](double x) {
                     return vgain * x;
              });

              if ( x < gain_class_centers[0]) {
                     *ipv[0] = 1.0;
              } else {
                     unsigned j = 1;
                     bool value_set = false;

                     while (!value_set && (j < class_centers.size()) ) {
                            // between two centers
                            if (x < gain_class_centers[j]) {
                                   double p0 = x - gain_class_centers[j - 1];
                                   double p1 = x - gain_class_centers[j];
                                   double p02 = p0 * p0;
                                   double p12 = p1 * p1;
                                   double normalizer = 1.0 / (p02 + p12);
                                   *ipv[j] = p02  * normalizer;
                                   *ipv[j - 1] = p12  * normalizer;
                                   value_set = true;
                            }

                            ++j;
                     }

                     if (!value_set)
                            *ipv[class_centers.size() - 1] = 1.0;
              }

              ++ii;
              ++ig;

              for (unsigned i = 0; i < class_centers.size(); ++i)
                     ++ipv[i];
       }
}

/**
   Evaluate the new clas centers from
   \f[
   \sum_{k,i} (p_{i,k} I_k - g_k c_i)^2 \rightarrow \min
   \f]
   In order to avoid a ping-pong effect, the actual class center update is evaluated
   according to
   \f[
   c_i^{(t+1)} = \frac{1}{2} \left( \sum_{k} \frac{p_{i,k}^2 g_k I_k}{ p_{i,k}^2 g_k^2 } - c_i^{(t)} \right)
   \f]

   \tparam T input pixel type of the data to be classified
   \tparam Field template of the data field type
   \param[in] image image the classification is applied to
   \param[in] gain multiplicative gain field
   \param[in] pv probability fields
   \param[in,out] class_centers
   \returns sum of absolute change applied to the class centers

*/

template <typename T, template <class> class Field>
double cmeans_update_class_centers(const Field<T>& image, const Field<float>& gain,
                                   const std::vector<Field<float>>& pv,
                                   std::vector<double>& class_centers)
{
       double residuum = 0.0;

       for (size_t i = 0; i < class_centers.size(); ++i) {
              double cc = class_centers[i];
              double sum_prob = 0.0;
              double sum_weight = 0.0;
              auto ie = image.end();
              auto ii = image.begin();
              auto ig = gain.begin();
              auto ip = pv[i].begin();

              while (ii != ie)  {
                     if (*ip > 0.0) {
                            auto v = *ip * *ip * *ig;
                            sum_prob += v * *ig;
                            sum_weight += v * *ii;
                     }

                     ++ii;
                     ++ig;
                     ++ip;
              }

              if (sum_prob  != 0.0) // move slowly in the direction of new center
                     cc = sum_weight / sum_prob;
              else {
                     cvwarn() << "class[" << i << "] has no probable members, keeping old value:" <<
                              sum_prob << ":" << sum_weight << "\n";
              }

              double delta = (cc - class_centers[i]) * 0.5;
              residuum += delta * delta;
              class_centers[i] +=  delta;
       }// end update class centers

       return sqrt(residuum);
}


typedef TFactory<CMeans::Initializer> CMeansInitializerPlugin;

// the class that has only the size as a parameter
class EXPORT_CORE CMeansInitializerSizedPlugin : public CMeansInitializerPlugin
{
public:
       CMeansInitializerSizedPlugin(const char *name);
protected:
       size_t get_size_param() const;
private:
       size_t m_size;

};

/// @cond NEVER
#ifdef __GNUC__
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wattributes"
#endif
extern template class EXPORT_CORE TPlugin<CMeans::Initializer, CMeans::Initializer>;
extern template class EXPORT_CORE TFactory<CMeans::Initializer>;
#ifdef __GNUC__
#pragma GCC diagnostic pop
#endif
extern template class EXPORT_CORE TFactoryPluginHandler<TFactory<CMeans::Initializer>>;
extern template class EXPORT_CORE THandlerSingleton<TFactoryPluginHandler<TFactory<CMeans::Initializer>>>;
/// @endcond

              template <>  const char *const TPluginHandler<TFactory<CMeans::Initializer>>::m_help;


              typedef THandlerSingleton<TFactoryPluginHandler<CMeansInitializerPlugin>>CMeansInitializerPluginHandler;

/// @cond NEVER
              FACTORY_TRAIT(CMeansInitializerPluginHandler);
/// @endcond


              NS_MIA_END

