#ifndef SOFA_BRANCHINGIMAGE_BranchingCellIndicesFromLabels_H
#define SOFA_BRANCHINGIMAGE_BranchingCellIndicesFromLabels_H

#include <BranchingImage/config.h>
#include <image/ImageTypes.h>
#include "BranchingImage.h"
#include <sofa/core/objectmodel/Event.h>
#include <sofa/simulation/AnimateEndEvent.h>
#include <sofa/defaulttype/Vec.h>

namespace sofa
{
namespace component
{
namespace engine
{

/**
 * Returns global index of branching image voxels given a list of labels
 */


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

    typedef SReal Real;

    typedef _BranchingImageTypes BranchingImageTypes;
    typedef typename BranchingImageTypes::T T;
    typedef helper::ReadAccessor<Data< BranchingImageTypes > > raBranchingImage;
    Data< BranchingImageTypes > branchingImage;

    typedef defaulttype::ImageLPTransform<Real> TransformType;
    typedef typename TransformType::Coord Coord;
    typedef helper::ReadAccessor<Data< TransformType > > raTransform;
    Data< TransformType > branchingImageTransform;

    typedef helper::vector<T> labelType;
    typedef helper::ReadAccessor<Data< labelType > > raLabels;
    Data< labelType > labels;

    typedef helper::vector<unsigned int> valuesType;
    typedef helper::WriteOnlyAccessor<Data< valuesType > > waValues;
    Data< valuesType > cell;  ///< output values

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

    BranchingCellIndicesFromLabels()    :   Inherited()
      , branchingImage(initData(&branchingImage,BranchingImageTypes(),"branchingImage",""))
      , branchingImageTransform(initData(&branchingImageTransform,TransformType(),"branchingImageTransform",""))
      , labels(initData(&labels,"labels","labels in input image used to select superimposed voxels"))
      , cell( initData ( &cell,"cell","cell indices." ) )
      , time((unsigned int)0)
    {
        this->addAlias(&branchingImageTransform, "branchingTransform");
        this->addAlias(&branchingImageTransform, "transform");
        this->addAlias(&branchingImage, "image");
        branchingImage.setReadOnly(true);
        branchingImageTransform.setReadOnly(true);
        f_listening.setValue(true);
    }

    virtual void init()
    {
        addInput(&branchingImage);
        addInput(&branchingImageTransform);
        addInput(&labels);
        addOutput(&cell);
        setDirtyValue();
    }

    virtual void reinit() { update(); }

protected:

    unsigned int time;

    virtual void update()
    {
        updateAllInputsIfDirty();
        cleanDirty();

        raBranchingImage bin(this->branchingImage);
        raLabels rlab(this->labels);

        // store labels in set to speed up search
        std::set<T> lab;
        for(unsigned int l = 0 ; l<rlab.size() ; l++) lab.insert(rlab[l]);

        // get image at time t
        const typename BranchingImageTypes::BranchingImage3D& img = bin->imgList[this->time];

        // output
        waValues val(this->cell);
        val.clear();

        unsigned int globalIndex = 0;
        bimg_foroff1D(bin.ref(),index1d)
                for(unsigned int offset = 0 ; offset<img[index1d].size() ; offset++)
        {
            T v= img[index1d][offset][0]; // assume that labels are stored in first channel
            if(lab.find(v)!=lab.end()) val.push_back(globalIndex);
            globalIndex++;
        }

        if(this->f_printLog.getValue())    std::cout<<this->name<<":"<<"updated"<<std::endl;
    }

    void handleEvent(sofa::core::objectmodel::Event *event)
    {
        if (simulation::AnimateEndEvent::checkEventType(event))
        {
            raBranchingImage in(this->branchingImage);
            raTransform inT(this->branchingImageTransform);

            // get current time modulo dimt
            const unsigned int dimt=in->getDimensions()[4];
            if(!dimt) return;
            Real t=inT->toImage(this->getContext()->getTime()) ;
            t-=(Real)((int)((int)t/dimt)*dimt);
            t=(t-floor(t)>0.5)?ceil(t):floor(t); // nearest
            if(t<0) t=0.0; else if(t>=(Real)dimt) t=(Real)dimt-1.0; // clamp

            if(this->time!=(unsigned int)t) { this->time=(unsigned int)t; update(); }
        }
    }

};


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

#endif // SOFA_BRANCHINGIMAGE_BranchingCellIndicesFromLabels_H
