#ifndef SOFA_BRANCHINGIMAGE_BranchingImageSubsampler_H
#define SOFA_BRANCHINGIMAGE_BranchingImageSubsampler_H

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

namespace sofa
{
namespace component
{
namespace engine
{

/**
 * Returns subsampled branching image and volumes of coarse cells.
 * Each coarse cell represents a connected region of fine cells with similar intensities.
 */


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

    typedef SReal Real;

    typedef _BranchingImageTypes BranchingImageTypes;
    typedef typename BranchingImageTypes::T T;
    typedef typename BranchingImageTypes::imCoord imCoord;
    typedef helper::ReadAccessor<Data< BranchingImageTypes > > raImage;
    typedef helper::WriteOnlyAccessor<Data< BranchingImageTypes > > waImage;
    Data< BranchingImageTypes > inputImage;
    Data< BranchingImageTypes > outputImage;
    
    typedef defaulttype::ImageLPTransform<Real> TransformType;
    typedef typename TransformType::Coord Coord;
    typedef helper::ReadAccessor<Data< TransformType > > raTransform;
    typedef helper::WriteOnlyAccessor<Data< TransformType > > waTransform;
    Data< TransformType > inputTransform;
    Data< TransformType > outputTransform;
    
    Data< unsigned int > coarseningFactor;
    Data< bool > splitLabels;
    
    typedef helper::vector<Real> valuesType;
    typedef helper::WriteOnlyAccessor<Data< valuesType > > waValues;
    Data< valuesType > volumes;
    
    typedef int indType;
    Data<helper::vector<indType> > cell;
    
    virtual std::string getTemplateName() const    { return templateName(this);    }
    static std::string templateName(const BranchingImageSubsampler<BranchingImageTypes>* = NULL) { return BranchingImageTypes::Name();    }
    
    BranchingImageSubsampler()    :   Inherited()
      , inputImage(initData(&inputImage,BranchingImageTypes(),"inputImage",""))
      , outputImage(initData(&outputImage,BranchingImageTypes(),"outputImage",""))
      , inputTransform(initData(&inputTransform,TransformType(),"inputTransform",""))
      , outputTransform(initData(&outputTransform,TransformType(),"outputTransform",""))
      , coarseningFactor(initData(&coarseningFactor,(unsigned)1,"coarseningFactor","Voxel size multiplication factor (=nb fine voxels in a coarse voxel in each direction x/y/z)"))
      , splitLabels(initData(&splitLabels,true,"splitLabels","duplicate coarse cells to contain fine cells with similar intensities, otherwise average intensities"))
      , volumes( initData ( &volumes,"volumes","cell volumes." ) )
      , cell( initData ( &cell,"voxelIndices","global ouput voxel index of each non empty input voxel." ) )
    {
        inputImage.setReadOnly(true);
        inputTransform.setReadOnly(true);
        this->addAlias(&outputImage, "branchingImage");
        this->addAlias(&outputTransform, "branchingTransform");
        this->addAlias(&outputImage, "image");
        this->addAlias(&outputTransform, "transform");
    }
    
    virtual void init()
    {
        addInput(&inputImage);
        addInput(&inputTransform);
        addInput(&coarseningFactor);
        addInput(&splitLabels);
        addOutput(&outputImage);
        addOutput(&outputTransform);
        addOutput(&volumes);
        addOutput(&cell);
        setDirtyValue();
    }
    
    virtual void reinit() { update(); }
    
protected:
    
    typedef typename BranchingImageTypes::ConnectionVoxel ConnectionVoxel;
    typedef typename BranchingImageTypes::VoxelIndex VoxelIndex;
    typedef std::pair<unsigned,unsigned> OffsetPair; // use pair since VoxelIndex does not work with maps
    typedef std::set<OffsetPair> OffsetPairSet;
    typedef std::map<OffsetPair,OffsetPairSet> coarseTofineMap;
    typedef std::map<OffsetPair,OffsetPair> fineToCoarseMap;

    virtual void update()
    {
        
        // get inputs
        raImage in(this->inputImage);
        raTransform inT(this->inputTransform);
        unsigned nb = coarseningFactor.getValue();
        unsigned splitLab = splitLabels.getValue();
        
        cleanDirty();
        
        // get outputs
        waValues vol(this->volumes);
        waImage out(this->outputImage);
        waTransform outT(this->outputTransform);
        outT->operator=(inT);
        helper::WriteOnlyAccessor<Data< helper::vector<indType> > > cellInd(this->cell);
        
        // compute dimensions
        const imCoord inDim = in->getDimensions();
        imCoord outDim = inDim;
        outDim[BranchingImageTypes::DIMENSION_X] = ceil( outDim[BranchingImageTypes::DIMENSION_X]/(float)nb );
        outDim[BranchingImageTypes::DIMENSION_Y] = ceil( outDim[BranchingImageTypes::DIMENSION_Y]/(float)nb );
        outDim[BranchingImageTypes::DIMENSION_Z] = ceil( outDim[BranchingImageTypes::DIMENSION_Z]/(float)nb );
        out->setDimensions( outDim );
        unsigned int spectrum = inDim[BranchingImageTypes::DIMENSION_S];

        // compute transform
        outT->getScale() *= nb;
        outT->getTranslation() += inT->fromImage(Coord(-0.5,-0.5,-0.5)) - outT->fromImage(Coord (-0.5,-0.5,-0.5)); // same corner
        Real voxelVol=inT->getScale()[0]*inT->getScale()[1]*inT->getScale()[2];
        
        // copy and quit in case of no subsampling
        if(nb<=1)
        {
            out->operator=(in);
            unsigned int num_vox=0; bimg_foroff1D(out.ref(),i) num_vox+=out->imgList[0][i].size(); // number of voxels taken a t=0
            vol.resize(num_vox); for(unsigned int i=0;i<num_vox;i++) vol[i]=voxelVol;
            cellInd.resize(num_vox); for(unsigned int i=0;i<num_vox;i++) cellInd[i]=(indType)i;
            return;
        }
        
        // map from coarse voxels to fine ones
        std::vector< coarseTofineMap > CFMap((size_t)inDim[BranchingImageTypes::DIMENSION_T]);
        
        bimg_forXYZ(in.ref(),x,y,z)
        {
            OffsetPair fine(in->index3Dto1D(x,y,z),0), coarse(out->index3Dto1D(x/nb,y/nb,z/nb),0);
            bimg_forT(in.ref(),t) for(fine.second=0 ; fine.second<in->imgList[t][fine.first].size() ; ++fine.second )        CFMap[t][coarse].insert(fine);
        }
        
        // duplicate coarse cells when necessary
        std::vector< coarseTofineMap > CFMap2((size_t)inDim[BranchingImageTypes::DIMENSION_T]); // the extended coarse to fine map
        
        bimg_forT(in.ref(),t) // time
                for(typename coarseTofineMap::iterator it=CFMap[t].begin(); it!=CFMap[t].end(); ++it) // coarse cell
        {
            std::vector<OffsetPairSet> vecFines; //groups of fine cells to be computed
            vecFines.push_back(it->second);
            
            if(splitLab) // split fines voxels in groups that have unique values
            {
                // map values to a set of fine cells
                const OffsetPairSet& fines=it->second;
                typedef std::map<std::vector<T>,OffsetPairSet> valueToFineMap;
                valueToFineMap valToFine;
                std::vector<T> vals((size_t)spectrum);
                for(typename OffsetPairSet::iterator fine=fines.begin(); fine!=fines.end(); ++fine)
                {
                    bimg_forC(in.ref(),c)     vals[c]=in->imgList[t][fine->first][fine->second][c];
                    valToFine[vals].insert(*fine);
                }
                // copy sets
                vecFines.clear();
                vecFines.reserve(valToFine.size());
                for(typename valueToFineMap::iterator VF=valToFine.begin(); VF!=valToFine.end(); ++VF) vecFines.push_back(VF->second);
            }
            
            // split by connex components
            std::vector<OffsetPairSet> newVecFines;
            for(size_t v=0;v<vecFines.size();++v)
            {
                const OffsetPairSet& fines=vecFines[v];
                SplitInConnexComponents(newVecFines,fines,in->imgList[t]);
            }
            newVecFines.swap(vecFines);
            newVecFines.clear();
            
            // add voxels to coarse image
            unsigned int index1d = it->first.first;
            out->imgList[t][index1d].resize(vecFines.size(),spectrum);
            if(splitLab) // use value of first fine voxel
            {
                for(size_t v=0;v<vecFines.size();++v)
                {
                    const OffsetPair& firstfine=*(vecFines[v].begin());
                    out->imgList[t][index1d][v].equal(in->imgList[t][firstfine.first][firstfine.second],spectrum);
                }
            }
            else  // average values of all fine voxels
            {
                for(size_t v=0;v<vecFines.size();++v)
                {
                    const OffsetPairSet& fines=vecFines[v];
                    bimg_forC(in.ref(),c) // channels
                    {
                        Real val=0;
                        for(typename OffsetPairSet::iterator fine=fines.begin(); fine!=fines.end(); ++fine) val+=(Real)in->imgList[t][fine->first][fine->second][c];
                        out->imgList[t][index1d][v][c]=(T)( val/(Real)fines.size() );
                    }
                }
            }
            
            // insert new voxels in map
            for(size_t v=0;v<vecFines.size();++v) CFMap2[t][OffsetPair(index1d,v)]=vecFines[v];
        }
        CFMap.clear();
        
        // compute volume
        vol.clear();
        bimg_foroff1D(out.ref(),off1D)
                for( unsigned v=0 ; v<out->imgList[0][off1D].size() ; ++v )
                vol.push_back(CFMap2[0][OffsetPair(off1D,v)].size() * voxelVol);
        
        // build fine to coarse map to speed up connections retrieval
        std::vector< fineToCoarseMap > FCMap((size_t)inDim[BranchingImageTypes::DIMENSION_T]);
        bimg_forT(in.ref(),t)
                for(typename coarseTofineMap::iterator it=CFMap2[t].begin(); it!=CFMap2[t].end(); ++it)
        {
            OffsetPair coarse(it->first.first,it->first.second);
            for(typename OffsetPairSet::iterator it2=it->second.begin(); it2!=it->second.end(); ++it2) FCMap[t][OffsetPair(it2->first,it2->second)]=coarse;
        }
        CFMap2.clear();
        
        // output global voxel indices for each non empty input voxel.
        {
            std::map<OffsetPair,unsigned int> offsetToGobalIndexMap_C;
            unsigned int count=0; bimg_foroff1D(out.ref(),i) for(size_t v=0;v<out->imgList[0][i].size();v++) { offsetToGobalIndexMap_C[OffsetPair(i,v)]=count; count++; }
            std::map<OffsetPair,unsigned int> offsetToGobalIndexMap_F;
            count=0;bimg_foroff1D(in.ref(),i) for(size_t v=0;v<in->imgList[0][i].size();v++) { offsetToGobalIndexMap_F[OffsetPair(i,v)]=count; count++; }
            
            cellInd.resize(FCMap[0].size());
            for(typename fineToCoarseMap::iterator it=FCMap[0].begin(); it!=FCMap[0].end(); ++it)
                cellInd[offsetToGobalIndexMap_F[OffsetPair(it->first.first,it->first.second)] ] = (indType)offsetToGobalIndexMap_C[OffsetPair(it->second.first,it->second.second)];
            
            offsetToGobalIndexMap_F.clear();
            offsetToGobalIndexMap_C.clear();
        }
        
        // make connections
        typedef typename BranchingImageTypes::VoxelIndex VoxelIndex;
        bimg_forT(in.ref(),t)
                for(typename fineToCoarseMap::iterator it=FCMap[t].begin(); it!=FCMap[t].end(); ++it)
        {
            VoxelIndex coarse1(it->second.first,it->second.second);
            const ConnectionVoxel& v = in->imgList[t][it->first.first][it->first.second];
            for(size_t n=0;n<v.neighbours.size();n++)
            {
                const OffsetPair& c2=FCMap[t][OffsetPair(v.neighbours[n].index1d,v.neighbours[n].offset)];
                VoxelIndex coarse2(c2.first,c2.second);
                if(coarse1.index1d!=coarse2.index1d || coarse1.offset!=coarse2.offset)
                {
                    out->imgList[t][coarse1.index1d][coarse1.offset].addNeighbour(coarse2);
                    out->imgList[t][coarse2.index1d][coarse2.offset].addNeighbour(coarse1);
                }
            }
        }
    }
    
    // split voxels in V in N connex groups, and insert them in vecV
    void SplitInConnexComponents(std::vector<OffsetPairSet>& vecV,const OffsetPairSet& V,const typename BranchingImageTypes::BranchingImage3D& im)
    {
        OffsetPairSet trial = V; // queue
        while( !trial.empty() )
        {
            // take a seed (first non inserted voxel)
            OffsetPair seed = *trial.begin();
            trial.erase(trial.begin());
            // move neighbors recursively from trial to newV
            OffsetPairSet newV;
            newV.insert(seed);
            insertNeighbors(newV,trial,seed,im);
            vecV.push_back(newV);
        }
    }
    
    void insertNeighbors(OffsetPairSet& V,OffsetPairSet& trial,const OffsetPair& seed,const typename BranchingImageTypes::BranchingImage3D& im)
    {
        const typename BranchingImageTypes::ConnectionVoxel& v = im[seed.first][seed.second];
        for(size_t n=0;n<v.neighbours.size();n++)
        {
            OffsetPair nV(v.neighbours[n].index1d,v.neighbours[n].offset);
            typename OffsetPairSet::iterator it=trial.find(nV);
            if(it!=trial.end())  // test if in queue
            {
                trial.erase(it);
                V.insert(nV);
                insertNeighbors(V,trial,nV,im);
            }
        }
    }
};


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

#endif // SOFA_BRANCHINGIMAGE_BranchingImageSubsampler_H
