#ifndef SOFA_BranchingImage_MassFromDensity_H
#define SOFA_BranchingImage_MassFromDensity_H

#include <BranchingImage/config.h>
#include <Flexible/mass/MassFromDensity.h>
#include "../../BranchingImage.h"


namespace sofa
{
namespace component
{
namespace engine
{


/// Specialization for branching Image
template <class DataTypes, class T>
struct MassFromDensitySpecialization<DataTypes, defaulttype::BranchingImage<T>>
{
    typedef MassFromDensity<DataTypes, defaulttype::BranchingImage<T>> MassFromDensityT;

    static void update(MassFromDensityT* This)
    {
        typedef typename MassFromDensityT::Real Real;
        typedef typename MassFromDensityT::VecCoord VecCoord;
        typedef typename MassFromDensityT::Coord Coord;
        typedef typename MassFromDensityT::ImageTypes ImageTypes;

        typename MassFromDensityT::raTransform inT(This->transform);
        typename MassFromDensityT::raImage rin(This->image);
        const ImageTypes& in = rin.ref();
        if(in.isEmpty())
        {
            This->serr<<This->getName()<<": input image is empty"<<This->sendl;
            return;
        }

        Data<helper::vector<int> >* d_cell = dynamic_cast< Data<helper::vector<int> >* >( This->deformationMapping->findData( "cell" ) );
        if(!d_cell)
        {
            This->serr<<This->getName()<<": branching cell offsets can't be provided to the mapping"<<This->sendl;
            return;
        }
        helper::WriteOnlyAccessor<Data<helper::vector<int> > > cell ( d_cell );

        // count non zero voxels
        unsigned int nb=0;
        bimg_foroff1D(in,off1D) nb+=in.imgList[This->time][off1D].size();

        // build mass and mapped dofs
        This->Me=typename MassFromDensityT::rmat(3*nb,3*nb);
        This->Me.reserve(3*nb);
        This->dofs->resize(nb);
        cell.resize(nb);
        helper::WriteOnlyAccessor<Data<VecCoord> > rpos ( This->dofs->writeOnlyRestPositions() );
        helper::WriteOnlyAccessor<Data<VecCoord> > pos ( This->dofs->writeOnlyPositions() );

        Real voxelVol = inT->getScale()[0]*inT->getScale()[1]*inT->getScale()[2];
        nb=0;
        Real totalMass=0;
        bimg_forXYZ(in,x,y,z)
        {
            unsigned off1D =in.index3Dto1D(x,y,z);
            for( unsigned int v=0 ; v<in.imgList[This->time][off1D].size() ; ++v )
            {
                rpos[nb]=pos[nb]=inT->fromImage(Coord(x,y,z));
                cell[nb]=v+1; // add 1 to be compatible with VoronoiShapeFunction<BranchingImage>
                Real m = voxelVol*(Real)in(off1D,v,0,This->time);
                for (int k=0; k<3; ++k)  This->Me.insert(3*nb+k,3*nb+k) = m;
                totalMass+=m;
                nb++;
            }
        }
        This->Me.makeCompressed();

        // output some information
        This->sout<<"Total Volume = "<<voxelVol*nb<<" ("<<pow((double)voxelVol*nb,1.0/3.0)<<")"<<This->sendl;
        This->sout<<"Total Mass = "<<totalMass<<This->sendl;
    }
};


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

#endif // SOFA_Flexible_MassFromDensity_H
