/* -*- 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/msgstream.hh>
#include <mia/core/disctmap.hh>
#include <mia/core/histogram.hh>
#include <mia/3d/filter.hh>

NS_MIA_USE
using namespace std;


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

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

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<FEdgeStopping>::Table edge_stop_map[] = {
       { "tuckey", psi_tuckey},
       { "pm1", psi_pm1},
       { "pm2", psi_pm2},
       { NULL, 0}
};

typedef struct {
       int xp, xm, yp, ym, zp, zp;
} SGradients;


typedef vector<SGradients> CGradientField;

typedef std::pair<float, float> HScale;
inline bool operator < (const HScale& a, const HScale& b)
{
       return a.first < b.first;
}


template <typename T>
static inline get_gradient(T3DImage<T>::const_iterator i, SGradients& g)
{
       int idd = *i;
       g.zm =  idd - i[-nxy];
       g.ym =  idd - i[-size.x];
       g.xm =  idd - i[-1];
       g.xm =  idd - i[ 1];
       g.ym =  idd - i[size.x];
       g.zm =  idd - i[nxy];
}

void test_get_gradient()
{
       unsigned char init_data[9] = { 0,  0, 0,  0, 30,  0, 0,  0, 0,
                                      0, 20, 0, 40, 30, 10, 0, 40, 0,
                                      0,  0, 0,  0, 25,  0, 0,  0, 0
                                    };
       C3DBounds size(3, 3, 3);
       C3DUBImage ubimage(size, init_data);
       SGradients g;
       get_gradient(ubimage.begin() + 13, g);
       assert(g.xm == -10);
       assert(g.xp == 20);
       assert(g.ym == 10);
       assert(g.yp == -10);
       assert(g.zm ==   0);
       assert(g.zp ==   5);
}

template <typename T>
float evaluate_MAD(const T3DImage<T>& image, CGradientField& grad)
{
       assert(grad.size() == image.size());
       typedef  CHistogram<CHistogramFeeder<T>> CHisto;
       CHisto h = CHisto(CHistogramFeeder<T>(0, 0, 0));
       C3DBounds size = image.get_size();
       size_t nxy = size.x * size.y;
       T3DImage<T>::const_iterator i = image.begin() + nxy;
       CGradientField::iterator    g = grad.begin()  + nxy;
       size_t nelements = 0;

       for (size_t z = 1; z < size.z - 1; ++z) {
              i + = size.x;
              g + = size.x;

              for (size_t y = 1; y < size.y - 1; ++y) {
                     ++id;
                     ++g;

                     for (size_t x = 1; x < size.x - 1; ++x, ++i, ++g) {
                            if (*id > 0) {
                                   get_gradient(i, *g);
                                   h.push(g->xm > 0 ? g->xm : -g->xm);
                                   h.push(g->xp > 0 ? g->xp : -g->xp);
                                   h.push(g->ym > 0 ? g->ym : -g->ym);
                                   h.push(g->yp > 0 ? g->yp : -g->yp);
                                   h.push(g->zm > 0 ? g->zm : -g->zm);
                                   h.push(g->zp > 0 ? g->zp : -g->zp);
                                   nelements += 6;
                            }
                     }

                     ++g;
                     ++i;
              }

              i += size.x;
              g += size.x;
       }

       T median = h.median();
       CHisto h2 = CHisto(CHisto(CHistogramFeeder<T>(0, 0, 0)));
       T n = 0;

       for (CHisto::const_iterator i = h.begin(); i != h.end(); ++i, n++) {
              int k = n - median;
              h2.push(k > 0 ? k : -k, *i);
       }

       return h2.median();
}




class C3DAnisoDiffFilter: public C3DFilter
{
public:
       C3DAnisoDiffFilter(FEdgeStopping es, int maxiter, float eps, float edge_stop);

       void self_test() const;
private:
       int   m_maxiter;
       float m_epsilon;
       float m_edge_stop;

};

class C3DAnisoFilterFactory: public C3DFilterPlugin
{
public:
       C3DAnisoFilterFactory();
private:
       virtual C3DFilterPlugin::ProductPtr do_create()const;
       virtual const string do_get_descr()const;

       TDictMap<EEdgeStop> m_edge_stop_idmap;

       int   m_maxiter;
       float m_epsilon;
       float m_edge_stop;
       string m_edge_stop_func;
};

C3DAnisoFilterFactory::C3DAnisoFilterFactory():
       C3DFilterPlugin("anisodiff"),
       m_edge_stop_idmap(edge_stop_map),
       m_maxiter(100),
       m_epsilon(1.0),
       m_edge_stop(5.0),
       m_edge_stop_func("tuckey")
{
       add_parameter("iter", CParamList::PParameter(new CIntParameter(m_maxiter, 1, numeric_limits<int>::max(),
                     false, "maximum number of iterations")));
       add_parameter("e", CParamList::PParameter(new CFloatParameter(m_epsilon, numeric_limits<float>::min(), numeric_limits<float>::max(),
                     false, "delta value to finish iteration")));
       add_parameter("k", CParamList::PParameter(new CFloatParameter(m_epsilon, -1, numeric_limits<float>::max(),
                     false, "edge stopping value (<=0.0 = adaptive) ")));
       add_parameter("edgestop",
                     CParamList::PParameter(new CStringParameter(m_edge_stop_func, false, "Edge stopping function (pm1|pm2|tuckey)")));
}

C3DFilter *C3DAnisoFilterFactory::do_create()const
{
       FEdgeStopping estop = m_edge_stop_idmap.get_value(m_edge_stop_func.c_str());
       return new C3DAnisoDiffFilter(estop, m_maxiter, m_epsilon, m_edge_stop);
}

const string C3DAnisoFilterFactory::do_get_descr()const
{
       return "anisotropic diffusion filter"
}

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

