#ifndef SOFA_BRANCHINGIMAGE_MeshLabeler_H
#define SOFA_BRANCHINGIMAGE_MeshLabeler_H

#include <BranchingImage/config.h>
#include "BranchingImage.h"
#include <image/ImageTypes.h>
#include <sofa/core/topology/BaseMeshTopology.h>

#include <sofa/core/objectmodel/Event.h>
#include <sofa/simulation/AnimateEndEvent.h>

#include <sofa/defaulttype/Vec.h>
#include <sofa/helper/SVector.h>

namespace sofa
{
namespace component
{
namespace engine
{

/**
 * Returns image labels associated to each part of a mesh
 */


/// Default implementation does not compile
template <class ImageType>
struct MeshLabelerSpecialization
{
};

/// forward declaration
template <class ImageType> class MeshLabeler;

/// Specialization for regular Image
template <class T>
struct MeshLabelerSpecialization<defaulttype::Image<T>>
{
    typedef MeshLabeler<defaulttype::Image<T>> MeshLabelerT;

    static bool label( MeshLabelerT* This, std::vector< std::set<typename MeshLabelerT::ValueType> >& hexaSets )
    {
        typename MeshLabelerT::raImage rimage(This->d_image);
        typename MeshLabelerT::raIndices voxelToHexahedra(This->d_voxelToHexahedra);

        const cimg_library::CImg<typename MeshLabelerT::T>& image = rimage->getCImg(0); // warning: take time 0
        unsigned int count=0;
        cimg_forXYZ(image,x,y,z)
                if(image(x,y,z))
        {
            int h= count<voxelToHexahedra.size() ? voxelToHexahedra[count]:count; // is not provided: suppose one to one mapping
            if(h>=(int)hexaSets.size())
            {
                This->serr<<"not enough hexahedra"<<This->sendl;
                return false;
            }
            if(h>=0) cimg_forC(image,c)   hexaSets[h].insert((typename MeshLabelerT::ValueType)image(x,y,z,c));
            count++;
        }
        return true;
    }
};


/// Specialization for branching Image
template <class T>
struct MeshLabelerSpecialization<defaulttype::BranchingImage<T>>
{
    typedef MeshLabeler<defaulttype::BranchingImage<T>> MeshLabelerT;

    static bool label( MeshLabelerT* This , std::vector< std::set<typename MeshLabelerT::ValueType> >& hexaSets)
    {
        typename MeshLabelerT::raImage image(This->d_image);
        typename MeshLabelerT::raIndices voxelToHexahedra(This->d_voxelToHexahedra);

        unsigned int count=0;
        bimg_forVoffT(image.ref(),v,off1D,t)
                if(t==0) // warning: take time 0
        {
            int h= count<voxelToHexahedra.size() ? voxelToHexahedra[count]:count; // is not provided: suppose one to one mapping
            if(h>=(int)hexaSets.size())
            {
                This->serr<<"not enough hexahedra"<<This->sendl;
                return false;
            }
            if(h>=0) bimg_forC(image.ref(),c)   hexaSets[h].insert((typename MeshLabelerT::ValueType)image.ref()(off1D,v,c,t));
            count++;
        }
        return true;
    }
};



template <class _ImageTypes>
class MeshLabeler : public core::DataEngine
{
public:
    typedef core::DataEngine Inherited;
    SOFA_CLASS(SOFA_TEMPLATE(MeshLabeler,_ImageTypes),Inherited);

    typedef SReal Real;

    //@name Image data
    /**@{*/
    typedef _ImageTypes ImageTypes;
    typedef typename ImageTypes::T T;
    typedef typename ImageTypes::imCoord imCoord;
    typedef helper::ReadAccessor<Data< ImageTypes > > raImage;
    Data< ImageTypes > d_image;
    /**@}*/

    //@name sample data (points+connectivity)
    /**@{*/
    typedef typename core::topology::BaseMeshTopology::Hexa Hexa;
    typedef typename core::topology::BaseMeshTopology::SeqHexahedra SeqHexahedra;
    typedef helper::ReadAccessor<Data< SeqHexahedra > > raHexa;
    Data< SeqHexahedra > d_hexahedra;
    /**@}*/

    //@name cell map (voxel index->hexa index)
    /**@{*/
    typedef helper::vector<int> VecIndices;
    typedef helper::ReadAccessor<Data< VecIndices > > raIndices;
    Data<VecIndices > d_voxelToHexahedra;
    /**@}*/

    //@name output = list of labels for each part of a mesh
    /**@{*/
    typedef T ValueType;
    typedef helper::SVector<ValueType> Vvalues;
    typedef helper::vector< Vvalues > VecVvalues;
    typedef helper::WriteOnlyAccessor<Data< VecVvalues > > waLabels;
    Data< VecVvalues > d_pointLabels;
    Data< VecVvalues > d_hexahedraLabels;
    /**@}*/

    virtual std::string getTemplateName() const    { return templateName(this);    }
    static std::string templateName(const MeshLabeler<ImageTypes>* = NULL) { return ImageTypes::Name();    }

    MeshLabeler()    :   Inherited()
      , d_image(initData(&d_image,ImageTypes(),"image",""))
      , d_hexahedra(initData(&d_hexahedra,SeqHexahedra(),"hexahedra","hexahedra"))
      , d_voxelToHexahedra( initData ( &d_voxelToHexahedra,"voxelToHexahedra","hexahedra associated to each voxel (if not provided, one to one mapping is assumed)" ) )
      , d_pointLabels ( initData ( &d_pointLabels,"pointLabels","lists of labels for each vertex" ) )
      , d_hexahedraLabels ( initData ( &d_hexahedraLabels,"hexahedraLabels","lists of labels for each hexahedra" ) )
    {
        d_image.setReadOnly(true);
    }

    virtual void init()
    {
        addInput(&d_image);
        addInput(&d_hexahedra);
        addInput(&d_voxelToHexahedra);
        addOutput(&d_pointLabels);
        addOutput(&d_hexahedraLabels);
        setDirtyValue();
    }

    virtual void reinit() { update(); }

protected:

    virtual void update()
    {
        updateAllInputsIfDirty(); // easy to ensure that all inputs are up-to-date

        cleanDirty();

        // allocate sets for unique insertion
        raHexa hexa(d_hexahedra);
        std::vector< std::set<ValueType> > hexaSets(hexa.size());

        // read image and insert hexa labels
        if(!MeshLabelerSpecialization<ImageTypes>::label( this , hexaSets)) return;

        // count nb points as the max index in hexa (suppose that all mesh points are in an hexa)
        size_t nb=0;
        for(size_t i=0;i<hexa.size();++i) for(unsigned int j=0;j<8;++j) if(hexa[i][j]>=nb) nb=hexa[i][j]+1;

        // insert in points
        std::vector< std::set<ValueType> > pointSets(nb);
        for(size_t i=0;i<hexa.size();++i) for(unsigned int j=0;j<8;++j) pointSets[hexa[i][j]].insert(hexaSets[i].begin(),hexaSets[i].end());

        // copy to data
        waLabels hexahedraLabels(d_hexahedraLabels);
        hexahedraLabels.resize(hexa.size());
        for(size_t i=0;i<hexa.size();++i)
        {
            hexahedraLabels[i].clear();
            for (typename std::set<ValueType>::iterator it=hexaSets[i].begin(); it!=hexaSets[i].end(); ++it) hexahedraLabels[i].push_back(*it);
        }

        waLabels pointLabels(d_pointLabels);
        pointLabels.resize(nb);
        for(size_t i=0;i<nb;++i)
        {
            pointLabels[i].clear();
            for (typename std::set<ValueType>::iterator it=pointSets[i].begin(); it!=pointSets[i].end(); ++it) pointLabels[i].push_back(*it);
        }

    }

};


} // namespace engine
} // namespace component
} // namespace sofa

#endif // SOFA_BRANCHINGIMAGE_MeshLabeler_H
