/* skeleton.cpp */

#include <math.h>
#include "afmm/include/mfmm.h"
#include "afmm/include/skeldt.h"
#include "afmm/include/genrl.h"
#include "afmm/include/flags.h"
#include "skeleton.hpp"
#include "include/connected.hpp"
#include "include/messages.h"

unsigned int SKELETON_SALIENCY_THRESHOLD =10;
unsigned int SKELETON_ISLAND_THRESHOLD   =10;
float        SKELETON_DT_THRESHOLD       =10;


#define EPSILON 0.00000001

float skeletonThreshold = 15;

FIELD<float>* computeDT(FIELD<float> *im) {
    FLAGS *flags;
    FIELD<float> *cnt1;


    FIELD<std::multimap<float, int> >* origs = 0;
    int N = 1000000; /* Max afmm iterations */
    float length;

    /* Default to AFMM-Star, see afmm/README.txt */
    ModifiedFastMarchingMethod::METHOD meth = ModifiedFastMarchingMethod::AFMM_STAR;
    flags = new FLAGS(*im, -1);
    
    cnt1 = do_one_pass(im, flags, origs, N, 1, meth, length);
 
    delete cnt1;

    delete flags;
    delete origs;
    return im;

}

FIELD<float>* saliencyThreshold(FIELD<float>* distMap, FIELD<float> *grad) {
    ConnectedComponents *CC = new ConnectedComponents(255);
    int                 *ccaOut = new int[grad->dimX() * grad->dimY()];
    int                 hLabel;
    unsigned int        *hist;
    FIELD<float>        *salskel = grad->dupe();
    int                 i,
                        nPix = salskel->dimX() * salskel->dimY();

    float               *dmd = distMap->data(),
                        *grd = grad->data(),
                        *ssd = salskel->data();

    /* Compute the saliency metric and threshold */
    i=0;
   
    while (i < nPix) {
        ssd[i] = (dmd[i] < EPSILON) ? 0 : grd[i] / dmd[i];
        ssd[i] = ssd[i] < SKELETON_SALIENCY_THRESHOLD ? 0 : 255;
        ++i;
    }
    
    /* Perform connected component analysis */
    hLabel = CC->connected(ssd, ccaOut, grad->dimX(), grad->dimY(), std::equal_to<float>(), true);
    hist = (unsigned int *) SAFE_CALLOC((hLabel + 1) * sizeof (unsigned int));
    for (i = 0; i < nPix; i++) hist[ccaOut[i]]++;

    /* Remove small islands */
    for (i = 0; i < nPix; i++) {
        ssd[i] = (hist[ccaOut[i]] >= SKELETON_ISLAND_THRESHOLD) ? ssd[i] : 0;
    }

    delete CC;
    delete [] ccaOut;
    free(hist);
    return salskel;
}

FIELD<float>* do_one_pass(FIELD<float>* fi, FLAGS* flagsi, FIELD<std::multimap<float, int> >* origs, int N, int dir, ModifiedFastMarchingMethod::METHOD meth, float& length) {
    int iter;
    FIELD<float>* f = (!dir) ? new FIELD<float>(*fi) : fi; //Copy input field
    FLAGS* flags = new FLAGS(*flagsi); //Copy flags field
    FIELD<float>* count = new FIELD<float>;

    ModifiedFastMarchingMethod *fmm = new ModifiedFastMarchingMethod(f, flags, count, origs, N); //Create fast marching method engine

    fmm->setScanDir(dir); //Set its various parameters
    fmm->setMethod(meth);

    int nfail, nextr;

    iter = fmm->execute(nfail, nextr); //...and execute the skeletonization
    length = fmm->getLength();

    delete flags;
    if (!dir) delete f;
    delete fmm;
    return count;
}

void postprocess(FIELD<float>* grad, FIELD<char>* skel, float level)
//Simple reduction of gradient to binary skeleton via thresholding
{
    for (int j = 0; j < grad->dimY(); j++) //Threshold 'grad' to get real skeleton
        for (int i = 0; i < grad->dimX(); i++) {
            float g = grad->value(i, j);
            if (g > INFINITY / 2) g = 1.0;
            else if (g > level) g = 255;
            else g = 0;
            skel->value(i, j) = (char) g;
        }
}

void comp_grad2(FLAGS* forig, FIELD<float>* cnt1, FIELD<float>* cnt2, FIELD<float>* grad)
//Gradient computation using 2-pass method (i.e. grad computed out of 2 fields)
{
    int i, j;
    float MYINFINITY_2 = MYINFINITY / 2;

    for (j = 0; j < grad->dimY(); j++) //Compute grad in a special way, i.e. on a 2-pixel
        for (i = 0; i < grad->dimX(); i++) //neighbourhood - this ensures pixel-size skeletons!
        {
            float ux1 = cnt1->value(i + 1, j) - cnt1->value(i, j);
            float uy1 = cnt1->value(i, j + 1) - cnt1->value(i, j);
            float g1 = max(fabs(ux1), fabs(uy1));
            float ux2 = cnt2->value(i + 1, j) - cnt2->value(i, j);
            float uy2 = cnt2->value(i, j + 1) - cnt2->value(i, j);
            float g2 = max(fabs(ux2), fabs(uy2));
            grad->value(i, j) = (g1 > MYINFINITY_2 || g2 > MYINFINITY_2) ? 0 : min(g1, g2);
        }
}
