#ifndef SOFA_BranchingImageGaussPointSAMPLER_H
#define SOFA_BranchingImageGaussPointSAMPLER_H

#include <BranchingImage/config.h>
#include <Flexible/quadrature/ImageGaussPointSampler.h>
#include "../../BranchingImage.h"
#include "../../ImageAlgorithms.h"



namespace sofa
{
namespace component
{
namespace engine
{


/**
 * This class samples an object represented by an image with gauss points
 */


/// Specialization for branching Image
template <class ImageT, class MaskT>
struct ImageGaussPointSamplerSpecialization<defaulttype::BranchingImage<ImageT>,defaulttype::BranchingImage<MaskT>>
{
    typedef ImageGaussPointSampler<defaulttype::BranchingImage<ImageT>,defaulttype::BranchingImage<MaskT>> ImageGaussPointSamplerT;

    typedef unsigned int IndT;
    typedef defaulttype::BranchingImage<IndT> IndTypes;

    static void init(ImageGaussPointSamplerT* This)
    {
        typedef typename ImageGaussPointSamplerT::IndTypes IndTypes;
        typedef typename ImageGaussPointSamplerT::waInd waInd;
        typedef typename ImageGaussPointSamplerT::DistTypes DistTypes;
        typedef typename ImageGaussPointSamplerT::raDist raDist;
        typedef typename ImageGaussPointSamplerT::waDist waDist;
        typedef typename ImageGaussPointSamplerT::waPositions waPositions;
        typedef typename ImageGaussPointSamplerT::waVolume waVolume;

        // retrieve data
        raDist rweights(This->f_w);         const DistTypes& weights = rweights.ref();
        if(weights.isEmpty())  { This->serr<<"Weights not found"<<This->sendl; return; }

        // init pos, vol, reg data; voronoi (=region data) and distances (=error image)
        typename ImageGaussPointSamplerT::imCoord dim = weights.getDimensions();
        dim[DistTypes::DIMENSION_S]=dim[DistTypes::DIMENSION_T]=1;

        waPositions pos(This->f_position);          pos.clear();                // pos is cleared since it is always initialized with one point, so user placed points are not allowed for now..
        waVolume vol(This->f_volume);   vol.clear();

        waInd wreg(This->f_region);        IndTypes& regimg = wreg.wref();
        regimg.setDimensions(dim);
        regimg.cloneTopology(weights,0);

        waDist werr(This->f_error);       DistTypes& dist = werr.wref();
        dist.setDimensions(dim);
        dist.cloneTopology(weights,-1.0);

        This->Reg.clear();
    }


    /// midpoint integration : put samples uniformly and weight them by their volume
    static void midpoint(ImageGaussPointSamplerT* This)
    {
        typedef typename ImageGaussPointSamplerT::Real Real;
        typedef typename ImageGaussPointSamplerT::IndTypes IndTypes;
        typedef typename ImageGaussPointSamplerT::raInd raInd;
        typedef typename ImageGaussPointSamplerT::waInd waInd;
        typedef typename ImageGaussPointSamplerT::DistTypes DistTypes;
        typedef typename ImageGaussPointSamplerT::DistT DistT;
        typedef typename ImageGaussPointSamplerT::waDist waDist;
        typedef typename ImageGaussPointSamplerT::Coord Coord;
        typedef typename ImageGaussPointSamplerT::waPositions waPositions;
        typedef typename ImageGaussPointSamplerT::indList indList;
        typedef typename ImageGaussPointSamplerT::raTransform raTransform;
        typedef typename ImageGaussPointSamplerT::factType factType;

        typedef typename DistTypes::VoxelIndex VoxelIndex;
        typedef std::pair<DistT,VoxelIndex > DistanceToPoint;
        typedef helper::vector<VoxelIndex> SeqPositions;

        // retrieve data
        raInd rindices(This->f_index);          const IndTypes& indices = rindices.ref();
        if(indices.isEmpty())  { This->serr<<"Indices not found"<<This->sendl; return; }
        raTransform transform(This->f_transform);
        const Coord voxelsize(transform->getScale());

        waPositions pos(This->f_position);
        waInd wreg(This->f_region);        IndTypes& regimg = wreg.wref();
        waDist werr(This->f_error);        DistTypes& dist = werr.wref();   // use error image as a container for distances

        // init soft regions (=more than one parent) where uniform sampling will be done
        // rigid regions (one parent) are kept in the list of region (and dist remains=-1 so they will not be sampled)
        for(unsigned int i=0; i<This->Reg.size();i++)
        {
            if(This->Reg[i].parentsToNodeIndex.size()>1 || This->sampleRigidParts.getValue())
            {
                bimg_forCVoffT(regimg,c,v,off1D,t) if(regimg(off1D,v,c,t)==*(This->Reg[i].voronoiIndices.begin()) )
                {
                    dist(off1D,v,c,t)=cimg_library::cimg::type<DistT>::max();
                    regimg(off1D,v,c,t)=0;
                }
                This->Reg.erase (This->Reg.begin()+i); i--;  //  erase region (soft regions will be generated after uniform sampling)
            }
        }
        unsigned int nbrigid = This->Reg.size();

        // fixed points = points set by the user in soft regions.
        // Disabled for now since pos is cleared
        SeqPositions fpos_voxelIndex;
        helper::vector<unsigned int> fpos_voronoiIndex;
        for(unsigned int i=0; i<pos.size(); i++)
        {
            Coord p = transform->toImageInt(pos[i]);
            if(indices.isInside(p[0],p[1],p[2]))
            {
                VoxelIndex vi(indices.index3Dto1D(p[0],p[1],p[2]),0);
                indList l;
                bimg_forC(indices,v) if(indices(vi,v)) l.insert(indices(vi,v)-1);
                if(l.size()>1) { fpos_voxelIndex.push_back(vi); fpos_voronoiIndex.push_back(i+1+nbrigid); }
            }
        }

        // target nb of points
        unsigned int nb = (fpos_voxelIndex.size()+nbrigid>This->targetNumber.getValue())?fpos_voxelIndex.size()+nbrigid:This->targetNumber.getValue();
        unsigned int nbsoft = nb-nbrigid;
        if(This->f_printLog.getValue()) std::cout<<This->getName()<<": Number of rigid/soft regions : "<<nbrigid<<"/"<<nbsoft<< std::endl;

        // init seeds for uniform sampling
        std::set<DistanceToPoint> trial;

        // farthest point sampling using geodesic distances
        SeqPositions newpos_voxelIndex;
        helper::vector<unsigned int> newpos_voronoiIndex;

        for(unsigned int i=0; i<fpos_voxelIndex.size(); i++) AddSeedPoint<DistT>(trial,dist,regimg, fpos_voxelIndex[i],fpos_voronoiIndex[i]);
        while(newpos_voxelIndex.size()+fpos_voxelIndex.size()<nbsoft)
        {
            DistT dmax=0;  VoxelIndex indMax;
            bimg_forCVoffT(dist,c,v,off1D,t) if(dist(off1D,v,c,t)>dmax) { dmax=dist(off1D,v,c,t); indMax = VoxelIndex(off1D,v); }
            if(dmax)
            {
                newpos_voxelIndex.push_back(indMax);
                newpos_voronoiIndex.push_back(fpos_voxelIndex.size()+nbrigid+newpos_voxelIndex.size());
                AddSeedPoint<DistT>(trial,dist,regimg, newpos_voxelIndex.back(),newpos_voronoiIndex.back());
                if(This->useDijkstra.getValue()) dijkstra<DistT,DistT>(trial,dist, regimg,voxelsize);
                else fastMarching<DistT,DistT>(trial,dist, regimg,voxelsize);
            }
            else break;
        }

        // Loyd
        unsigned int it=0;
        bool converged =(it>=This->iterations.getValue())?true:false;
        while(!converged)
        {
            converged=!(Lloyd<DistT>(newpos_voxelIndex,newpos_voronoiIndex,regimg));
            // recompute voronoi
            bimg_forCVoffT(dist,c,v,off1D,t) if(dist(off1D,v,c,t)!=-1) dist(off1D,v,c,t)=cimg_library::cimg::type<Real>::max();
            for(unsigned int i=0; i<fpos_voxelIndex.size(); i++) AddSeedPoint<DistT>(trial,dist,regimg, fpos_voxelIndex[i],fpos_voronoiIndex[i]);
            for(unsigned int i=0; i<newpos_voxelIndex.size(); i++) AddSeedPoint<DistT>(trial,dist,regimg, newpos_voxelIndex[i],newpos_voronoiIndex[i]);
            if(This->useDijkstra.getValue()) dijkstra<DistT,DistT>(trial,dist, regimg,voxelsize); else fastMarching<DistT,DistT>(trial,dist, regimg,voxelsize);
            it++; if(it>=This->iterations.getValue()) converged=true;
        }
        if(This->f_printLog.getValue()) std::cout<<This->getName()<<": Completed in "<< it <<" Lloyd iterations"<<std::endl;

        // create soft regions and update teir data
        for(unsigned int i=0; i<fpos_voxelIndex.size(); i++)           // Disabled for now since fpos is empty
        {
            indList l;
            bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==fpos_voronoiIndex[i])  {    bimg_forC(indices,c) if(indices(off1D,v,c,t0)) l.insert(indices(off1D,v,c,t0)-1); }   // collect indices over the region
            if(l.size())
            {
                factType reg(l,fpos_voronoiIndex[i]);
                unsigned x,y,z; regimg.index1Dto3D(fpos_voxelIndex[i].index1d, x,y,z);
                reg.center=transform->fromImage(Coord(x,y,z));
                // todo: memorize fpos_voxelIndex[i].offset in cell
                This->Reg.push_back(reg);
            }
        }
        for(unsigned int i=0; i<newpos_voxelIndex.size(); i++)
        {
            indList l;
            bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==newpos_voronoiIndex[i])  {    bimg_forC(indices,c) if(indices(off1D,v,c,t0)) l.insert(indices(off1D,v,c,t0)-1); }   // collect indices over the region
            if(l.size())
            {
                factType reg(l,newpos_voronoiIndex[i]);
                unsigned x,y,z; regimg.index1Dto3D(newpos_voxelIndex[i].index1d, x,y,z);
                reg.center=transform->fromImage(Coord(x,y,z));
                // todo: memorize newpos_voxelIndex[i].offset in cell
                This->Reg.push_back(reg);
            }
        }
        // update rigid regions (might contain soft material due to voronoi proximity)
        for(unsigned int i=0; i<nbrigid; i++)
        {
            indList l;
            bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==*(This->Reg[i].voronoiIndices.begin()))  {    bimg_forC(indices,c) if(indices(off1D,v,c,t0)) l.insert(indices(off1D,v,c,t0)-1); }   // collect indices over the region
            This->Reg[i].setParents(l);
        }

        // update nb voxels in each region (used later in weight fitting)
        for(unsigned int i=0; i<This->Reg.size(); i++)
        {
            This->Reg[i].nb=0;
            bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==*(This->Reg[i].voronoiIndices.begin()))  This->Reg[i].nb++;
        }
    }

    /// returns true if (x,y,z) in the region of interest
    static bool isInMask(ImageGaussPointSamplerT* This,unsigned off1D,unsigned v)
    {
        typename ImageGaussPointSamplerT::raMask rmask(This->f_mask);
        if(rmask->isEmpty()) return true;
        typename ImageGaussPointSamplerT::raMaskLabels labels(This->f_maskLabels);
        typename ImageGaussPointSamplerT::MaskT val = rmask.ref()(off1D,v);
        for(unsigned int i=0;i<labels.size();i++) if(labels[i]==val) return true;
        return false;
    }

    /// Identify regions sharing similar parents
    /// returns a list of region containing the parents, the number of voxels and center; and fill the voronoi image
    static void Cluster_SimilarIndices(ImageGaussPointSamplerT* This)
    {
        typedef typename ImageGaussPointSamplerT::Real Real;
        typedef typename ImageGaussPointSamplerT::IndTypes IndTypes;
        typedef typename ImageGaussPointSamplerT::raInd raInd;
        typedef typename ImageGaussPointSamplerT::waInd waInd;
        typedef typename ImageGaussPointSamplerT::indList indList;
        typedef typename ImageGaussPointSamplerT::raTransform raTransform;
        typedef typename ImageGaussPointSamplerT::Coord Coord;
        typedef typename ImageGaussPointSamplerT::raPositions raPositions;
        typedef typename ImageGaussPointSamplerT::factType factType;

        typedef typename IndTypes::VoxelIndex VoxelIndex;

        // retrieve data
        raInd rindices(This->f_index);     if(rindices->isEmpty())  { This->serr<<"Indices not found"<<This->sendl; return; }        const IndTypes& indices = rindices.ref();
        waInd wreg(This->f_region);        IndTypes& regimg = wreg.wref();
        raTransform transform(This->f_transform);

        // map to find repartitions-> region index
        typedef std::map<indList, unsigned int> indMap;
        indMap List;

        // allows user to fix points. Currently disabled since pos is cleared
        raPositions pos(This->f_position);
        const unsigned int initialPosSize=pos.size();
        for(unsigned int i=0; i<initialPosSize; i++)
        {
            Coord p = transform->toImageInt(pos[i]);
            if(indices.isInside(p[0],p[1],p[2]))
            {
                VoxelIndex vi(indices.index3Dto1D(p[0],p[1],p[2]),0);
                indList l;
                bimg_forC(indices,v) if(indices(vi,v)) l.insert(indices(vi,v)-1);
                List[l]=i;
                This->Reg.push_back(factType(l,i+1));
                regimg(vi)=i+1;
            }
        }

        // traverse index image to identify regions with unique indices
        bimg_forVoffT(indices,v,off1D,t)
                if(indices(off1D,v))
                if(isInMask(This,off1D,v))
        {
            indList l;
            bimg_forC(indices,c) if(indices(off1D,v,c)) l.insert(indices(off1D,v,c)-1);
            typename indMap::iterator it=List.find(l);
            unsigned int index;
            if(it==List.end()) { index=List.size(); List[l]=index;  This->Reg.push_back(factType(l,index+1)); This->Reg.back().nb=1; }
            else { index=it->second; This->Reg[index].nb++;}
            unsigned x,y,z; regimg.index1Dto3D(off1D, x,y,z);
            This->Reg[index].center+=transform->fromImage(Coord(x,y,z));
            // todo: memorize v in cell
            regimg(off1D,v)=*(This->Reg[index].voronoiIndices.begin());
        }

        // average to get centroid (may not be inside the region if not convex)
        for(unsigned int i=0; i<This->Reg.size(); i++) This->Reg[i].center/=(Real)This->Reg[i].nb;
    }

    /// subdivide region[index] in two regions
    static void subdivideRegion(ImageGaussPointSamplerT* This,const unsigned int index)
    {
        typedef typename ImageGaussPointSamplerT::Real Real;
        typedef typename ImageGaussPointSamplerT::IndTypes IndTypes;
        typedef typename ImageGaussPointSamplerT::waInd waInd;
        typedef typename ImageGaussPointSamplerT::DistTypes DistTypes;
        typedef typename ImageGaussPointSamplerT::DistT DistT;
        typedef typename ImageGaussPointSamplerT::waDist waDist;
        typedef typename ImageGaussPointSamplerT::raTransform raTransform;
        typedef typename ImageGaussPointSamplerT::Coord Coord;
        typedef typename ImageGaussPointSamplerT::factType factType;

        typedef typename DistTypes::VoxelIndex VoxelIndex;
        typedef std::pair<DistT,VoxelIndex > DistanceToPoint;

        // retrieve data
        raTransform transform(This->f_transform);
        const Coord voxelsize(transform->getScale());

        waInd wreg(This->f_region);        IndTypes& regimg = wreg.wref();
        waDist werr(This->f_error);        DistTypes& dist = werr.wref();

        // compute
        helper::vector<VoxelIndex> pos(2);
        helper::vector<unsigned int> vorindex;
        vorindex.push_back(*(This->Reg[index].voronoiIndices.begin()));
        vorindex.push_back(This->Reg.size()+1);
        for(unsigned int i=0; i<This->Reg.size(); i++) if(vorindex[1]==*(This->Reg[i].voronoiIndices.begin())) vorindex[1]++; // check that the voronoi index is unique. not necessary in principle

        // get closest/farthest point from c and init distance image
        Real dmin=cimg_library::cimg::type<Real>::max(),dmax=0;
        bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==vorindex[0])
        {
            dist(off1D,v)=cimg_library::cimg::type<DistT>::max();
            unsigned x,y,z; regimg.index1Dto3D(off1D, x,y,z);
            Coord p = Coord(x,y,z);
            Real d = (transform->fromImage(p)-This->Reg[index].center).norm2();
            if(dmin>d) {dmin=d; pos[0]=VoxelIndex(off1D,v);}
            if(dmax<d) {dmax=d; pos[1]=VoxelIndex(off1D,v);}
        }
        else dist(off1D,v)=(DistT)(-1);

        // Loyd relaxation
        std::set<DistanceToPoint> trial;
        unsigned int it=0;
        bool converged =(it>=This->iterations.getValue())?true:false;

        for(unsigned int i=0; i<2; i++) AddSeedPoint<DistT>(trial,dist,regimg, pos[i],vorindex[i]);
        if(This->useDijkstra.getValue()) dijkstra<DistT,DistT>(trial,dist, regimg,voxelsize); else fastMarching<DistT,DistT>(trial,dist, regimg,voxelsize);
        //dist.display();
        //regimg.display();
        while(!converged)
        {
            converged=!(Lloyd<DistT>(pos,vorindex,regimg));
            // recompute voronoi
            bimg_forVoffT(dist,v,off1D,t) if(dist(off1D,v)!=-1) dist(off1D,v)=cimg_library::cimg::type<DistT>::max();
            for(unsigned int i=0; i<2; i++) AddSeedPoint<DistT>(trial,dist,regimg, pos[i],vorindex[i]);
            if(This->useDijkstra.getValue()) dijkstra<DistT,DistT>(trial,dist, regimg,voxelsize); else fastMarching<DistT,DistT>(trial,dist, regimg,voxelsize);
            it++; if(it>=This->iterations.getValue()) converged=true;
        }

        // add region
        factType reg;
        reg.parentsToNodeIndex=This->Reg[index].parentsToNodeIndex;
        reg.voronoiIndices.insert(vorindex[1]);
        unsigned x,y,z; regimg.index1Dto3D(pos[1].index1d, x,y,z);
        reg.center=transform->fromImage(Coord(x,y,z));
        // todo: memorize pos[1].offset in cell
        reg.nb=0;
        bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==vorindex[1]) reg.nb++;
        This->Reg.push_back(reg);

        // update old region data
        regimg.index1Dto3D(pos[0].index1d, x,y,z);
        This->Reg[index].center=transform->fromImage(Coord(x,y,z));
        // todo: memorize pos[0].offset in cell
        This->Reg[index].nb=0;
        bimg_forVoffT(regimg,v,off1D,t0) if(regimg(off1D,v)==vorindex[0])  This->Reg[index].nb++;
    }

    /// update Polynomial Factors from the voxel map
    static void fillPolynomialFactors(ImageGaussPointSamplerT* This,const unsigned int factIndex, const bool writeErrorImg=false)
    {
        typedef typename ImageGaussPointSamplerT::Real Real;
        typedef typename ImageGaussPointSamplerT::IndTypes IndTypes;
        typedef typename ImageGaussPointSamplerT::raInd raInd;
        typedef typename ImageGaussPointSamplerT::DistTypes DistTypes;
        typedef typename ImageGaussPointSamplerT::raDist raDist;
        typedef typename ImageGaussPointSamplerT::waDist waDist;
        typedef typename ImageGaussPointSamplerT::Coord Coord;
        typedef typename ImageGaussPointSamplerT::indListIt indListIt;
        typedef typename ImageGaussPointSamplerT::raTransform raTransform;
        typedef typename ImageGaussPointSamplerT::factType factType;

        // retrieve data
        raDist rweights(This->f_w);             if(rweights->isEmpty())  { This->serr<<"Weights not found"<<This->sendl; return; }  const DistTypes& weights = rweights.ref();
        raInd rindices(This->f_index);          if(rindices->isEmpty())  { This->serr<<"Indices not found"<<This->sendl; return; }  const IndTypes& indices = rindices.ref();
        raInd rreg(This->f_region);        const IndTypes& regimg = rreg.ref();
        raTransform transform(This->f_transform);
        const Coord voxelsize(transform->getScale());

        // list of absolute coords
        factType &fact = This->Reg[factIndex];
        helper::vector<Coord> pi(fact.nb);

        // weights (one line for each parent)
        typename ImageGaussPointSamplerT::Matrix wi(fact.parentsToNodeIndex.size(),fact.nb); wi.setZero();

        // get them from images
        unsigned int count=0;
        bimg_forVoffT(regimg,v,off1D,t0)
        {
            indListIt it=fact.voronoiIndices.find(regimg(off1D,v));
            if(it!=fact.voronoiIndices.end())
            {
                bimg_forC(indices,c) if(indices(off1D,v,c))
                {
                    std::map<unsigned int,unsigned int>::iterator pit=fact.parentsToNodeIndex.find(indices(off1D,v,c)-1);
                    if(pit!=fact.parentsToNodeIndex.end())  wi(pit->second,count)= (Real)weights(off1D,v,c);
                }
                unsigned x,y,z; regimg.index1Dto3D(off1D, x,y,z);
                pi[count]= transform->fromImage(Coord(x,y,z));
                count++;
            }
        }

        fact.fill(wi,pi,This->fillOrder(),voxelsize,This->volOrder());

        //  std::cout<<"pt "<<*(fact.voronoiIndices.begin())-1<<" : "<<fact.center<<std::endl<<std::endl<<std::endl<<pi<<std::endl<<std::endl<<wi<<std::endl;
        //test: fact.directSolve(wi,pi); std::cout<<"Jacobi err="<<fact.getError()<<std::endl;

        // write error into output image
        if(writeErrorImg)
        {
            waDist werr(This->f_error); DistTypes& outimg = werr.wref();
            count=0;
            bimg_forVoffT(regimg,v,off1D,t0)
            {
                indListIt it=fact.voronoiIndices.find(regimg(off1D,v));
                if(it!=fact.voronoiIndices.end()) { outimg(off1D,v)=fact.getError(pi[count],wi.col(count)); count++; }
            }
        }
    }

};


}
}
}

#endif
