#ifndef SOFA_BRANCHINGIMAGE_ImageMeshIntersection_H
#define SOFA_BRANCHINGIMAGE_ImageMeshIntersection_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 the index of the mesh cell containing each non empty voxel
 */


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

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

/// Specialization for regular Image
template <class T>
struct ImageMeshIntersectionSpecialization<defaulttype::Image<T>>
{
    typedef ImageMeshIntersection<defaulttype::Image<T>> ImageMeshIntersectionT;

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

    static void compute( ImageMeshIntersectionT* This)
    {
        typename ImageMeshIntersectionT::raImage rimage(This->d_image);
        const cimg_library::CImg<T>& image = rimage->getCImg(0); // warning: take time 0
        typename ImageMeshIntersectionT::waIndices voxelIndices(This->d_voxelIndices);
        voxelIndices.clear();

        if(!This->d_computeImage.getValue())
        {
            cimg_forXYZ(image,x,y,z)
                    if(image(x,y,z))
                    voxelIndices.push_back(This->getIntersectedCell(x,y,z));
        }
        else
        {
            //allocate
            typename ImageMeshIntersectionT::waImageIndices woutputImage(This->d_outputImage);
            typename ImageMeshIntersectionT::imCoord dim = rimage->getDimensions();
            dim[ImageMeshIntersectionT::ImageTypes::DIMENSION_S]=dim[ImageMeshIntersectionT::ImageTypes::DIMENSION_T]=1;
            woutputImage->setDimensions(dim);
            typename ImageMeshIntersectionT::IndTypes::CImgT& outputImage = woutputImage->getCImg();
            outputImage.fill(0);
            //fill
            cimg_forXYZ(image,x,y,z)
                    if(image(x,y,z))
            {
                int i=This->getIntersectedCell(x,y,z);
                voxelIndices.push_back(i);
                outputImage(x,y,z)=(unsigned int)(i+1);
            }
        }
    }


};


/// Specialization for branching Image
template <class T>
struct ImageMeshIntersectionSpecialization<defaulttype::BranchingImage<T>>
{
    typedef ImageMeshIntersection<defaulttype::BranchingImage<T>> ImageMeshIntersectionT;

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

    static void compute( ImageMeshIntersectionT* This)
    {
        typename ImageMeshIntersectionT::raImage image(This->d_image);
        typename ImageMeshIntersectionT::waIndices voxelIndices(This->d_voxelIndices);
        voxelIndices.clear();

        bimg_forVXYZT(image.ref(),v,x,y,z,t)
                if(t==0) // warning: take time 0
                voxelIndices.push_back(This->getIntersectedCell(x,y,z));
        if(This->d_computeImage.getValue())
        {
            //allocate
            typename ImageMeshIntersectionT::waImageIndices woutputImage(This->d_outputImage);
            typename ImageMeshIntersectionT::IndTypes& outputImage = woutputImage.wref();
            typename ImageMeshIntersectionT::imCoord dim = image->getDimensions();
            dim[ImageMeshIntersectionT::ImageTypes::DIMENSION_S]=dim[ImageMeshIntersectionT::ImageTypes::DIMENSION_T]=1;
            outputImage.setDimensions(dim);
            outputImage.cloneTopology (image.ref(),0);
            //fill
            unsigned int count=0;
            bimg_forVoffT(outputImage,v,off1D,t)
            {
                outputImage(off1D,v)=(unsigned int)(voxelIndices[count]+1);
                count++;
            }
        }
    }
};



template <class _ImageTypes>
class ImageMeshIntersection : public core::DataEngine
{
public:
    typedef core::DataEngine Inherited;
    SOFA_CLASS(SOFA_TEMPLATE(ImageMeshIntersection,_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;
    typedef defaulttype::ImageLPTransform<Real> TransformType;
    typedef helper::ReadAccessor<Data< TransformType > > raTransform;
    Data< TransformType > d_transform;
    /**@}*/

    //@name parameters
    /**@{*/
    Data< Real > d_tolerance;
    Data< bool > d_computeImage;
    /**@}*/

    //@name mesh data
    /**@{*/
    typedef defaulttype::Vec<3,Real> Coord;
    typedef helper::vector<Coord> SeqPositions;
    typedef helper::ReadAccessor<Data< SeqPositions > > raPositions;
    Data< SeqPositions > d_position;
    typedef typename core::topology::BaseMeshTopology::SeqEdges SeqEdges;
    Data< SeqEdges > d_edges;
    typedef typename core::topology::BaseMeshTopology::SeqTriangles SeqTriangles;
    Data< SeqTriangles > d_triangles;
    typedef typename core::topology::BaseMeshTopology::SeqQuads SeqQuads;
    Data< SeqQuads > d_quads;
    typedef typename core::topology::BaseMeshTopology::SeqTetrahedra SeqTetrahedra;
    Data< SeqTetrahedra > d_tetrahedra;
    typedef typename core::topology::BaseMeshTopology::SeqHexahedra SeqHexahedra;
    Data< SeqHexahedra > d_hexahedra;
    /**@}*/

    //@name output = cell map (voxel index->cell index) and optionnaly index image
    /**@{*/
    typedef helper::vector<int> VecIndices;
    typedef helper::WriteOnlyAccessor<Data< VecIndices > > waIndices;
    Data<VecIndices > d_voxelIndices;

    typedef typename ImageMeshIntersectionSpecialization<ImageTypes>::IndT IndT;
    typedef typename ImageMeshIntersectionSpecialization<ImageTypes>::IndTypes IndTypes;
    typedef helper::WriteOnlyAccessor<Data< IndTypes > > waImageIndices;
    Data< IndTypes > d_outputImage;
    /**@}*/

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

    ImageMeshIntersection()    :   Inherited()
      , d_image(initData(&d_image,ImageTypes(),"image",""))
      , d_transform(initData(&d_transform,TransformType(),"transform",""))
      , d_tolerance(initData(&d_tolerance,(Real)-1.0,"tolerance","minimum barycentric weight (allows for fuzzy intersection test)"))
      , d_computeImage(initData(&d_computeImage,false,"computeImage","compute a 'label' image containing cell indices"))
      , d_position(initData(&d_position,SeqPositions(),"position","position"))
      , d_edges(initData(&d_edges,SeqEdges(),"edges","edges"))
      , d_triangles(initData(&d_triangles,SeqTriangles(),"triangles","triangles"))
      , d_quads(initData(&d_quads,SeqQuads(),"quads","quads"))
      , d_tetrahedra(initData(&d_tetrahedra,SeqTetrahedra(),"tetrahedra","tetrahedra"))
      , d_hexahedra(initData(&d_hexahedra,SeqHexahedra(),"hexahedra","hexahedra"))
      , d_voxelIndices( initData ( &d_voxelIndices,"voxelIndices","cell containing each voxel (-1 when outside all cells)" ) )
      , d_outputImage(initData(&d_outputImage,IndTypes(),"outputImage",""))
    {
        d_image.setReadOnly(true);
        d_transform.setReadOnly(true);
    }

    virtual void init()
    {
        addInput(&d_image);
        addInput(&d_transform);
        addInput(&d_position);
        addInput(&d_tolerance);
        addInput(&d_computeImage);
        addInput(&d_edges);
        addInput(&d_triangles);
        addInput(&d_quads);
        addInput(&d_tetrahedra);
        addInput(&d_hexahedra);
        addOutput(&d_voxelIndices);
        addOutput(&d_outputImage);
        setDirtyValue();
    }

    virtual void reinit() { update(); }

protected:

    static const unsigned int spatial_dimensions=3;
    typedef defaulttype::Mat<spatial_dimensions,spatial_dimensions,Real> Basis;
    sofa::helper::vector<Basis> bases;

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

        cleanDirty();

        computeBasis();

        ImageMeshIntersectionSpecialization<ImageTypes>::compute( this);
    }



    // from barycentricShapeFunction

    void computeBasis()
    {
        raPositions position(d_position);
        const SeqTetrahedra& tetrahedra = d_tetrahedra.getValue();
        const SeqHexahedra& cubes = d_hexahedra.getValue();
        const SeqTriangles& triangles = d_triangles.getValue();
        const SeqQuads& quads = d_quads.getValue();
        const SeqEdges& edges = d_edges.getValue();

        if ( tetrahedra.empty() && cubes.empty() )
        {
            if ( triangles.empty() && quads.empty() )
            {
                if ( edges.empty() ) return;

                //no 3D elements, nor 2D elements -> map on 1D elements
                bases.resize ( edges.size() );
                for (unsigned int e=0; e<edges.size(); e++ )
                {
                    Coord V12 = ( position[edges[e][1]]-position[edges[e][0]] );
                    bases[e][0] = V12/V12.norm2();
                }
            }
            else
            {
                // no 3D elements -> map on 2D elements
                bases.resize ( triangles.size() +quads.size() );
                for ( unsigned int t = 0; t < triangles.size(); t++ )
                {
                    Basis m,mt;
                    getBasisFrom2DElements( m, position[triangles[t][0]], position[triangles[t][1]], position[triangles[t][2]] );
                    mt.transpose ( m );
                    bases[t].invert ( mt );
                }
                int c0 = triangles.size();
                for ( unsigned int c = 0; c < quads.size(); c++ )
                {
                    Basis m,mt;
                    getBasisFrom2DElements( m, position[quads[c][0]], position[quads[c][1]], position[quads[c][3]] );
                    mt.transpose ( m );
                    bases[c0+c].invert ( mt );
                }
            }
        }
        else
        {
            // map on 3D elements
            bases.resize ( tetrahedra.size() +cubes.size() );
            for ( unsigned int t = 0; t < tetrahedra.size(); t++ )
            {
                Basis m,mt;
                m[0] = position[tetrahedra[t][1]]-position[tetrahedra[t][0]];
                m[1] = position[tetrahedra[t][2]]-position[tetrahedra[t][0]];
                m[2] = position[tetrahedra[t][3]]-position[tetrahedra[t][0]];
                mt.transpose ( m );
                bases[t].invert ( mt );
            }
            int c0 = tetrahedra.size();
            for ( unsigned int c = 0; c < cubes.size(); c++ )
            {
                Basis m,mt;
                m[0] = position[cubes[c][1]]-position[cubes[c][0]];
                m[1] = position[cubes[c][3]]-position[cubes[c][0]];
                m[2] = position[cubes[c][4]]-position[cubes[c][0]];
                mt.transpose ( m );
                bases[c0+c].invert ( mt );
            }
        }
    }

public:
    int getIntersectedCell(const unsigned int x,const unsigned int y,const unsigned int z)
    {
        raPositions position(d_position);
        const SeqTetrahedra& tetrahedra = d_tetrahedra.getValue();
        const SeqHexahedra& cubes = d_hexahedra.getValue();
        const SeqTriangles& triangles = d_triangles.getValue();
        const SeqQuads& quads = d_quads.getValue();
        const SeqEdges& edges = d_edges.getValue();
        raTransform transform(d_transform);
        Coord childPosition = transform->fromImage(Coord(x,y,z));

        // compute barycentric weights by projection in cell basis
        int index = -1;
        double distance = -d_tolerance.getValue();

        if ( tetrahedra.empty() && cubes.empty() )
        {
            if ( triangles.empty() && quads.empty() )
            {
                if ( edges.empty() ) return -1;
                //no 3D elements, nor 2D elements -> map on 1D elements
                for ( unsigned int i = 0; i < edges.size(); i++ )
                {
                    Coord v = bases[i] * ( childPosition - position[edges[i][0]] );
                    double d = std::max ( -v[0], v[0]-(Real)1. );
                    if ( d<=distance ) { distance = d; index = i; }
                }
            }
            else
            {
                // no 3D elements -> map on 2D elements
                for ( unsigned int i = 0; i < triangles.size(); i++ )
                {
                    Coord v = bases[i] * ( childPosition - position[triangles[i][0]] );
                    double d = getDistanceTriangle( v );
                    if ( d<=distance ) {distance = d; index = i; }
                }
                int c0 = triangles.size();
                for ( unsigned int i = 0; i < quads.size(); i++ )
                {
                    Coord v = bases[c0+i] * ( childPosition - position[quads[i][0]] );
                    double d = getDistanceQuad( v );
                    if ( d<=distance ) { distance = d; index = c0+i; }
                }
            }
        }
        else
        {
            // map on 3D elements
            for ( unsigned int i = 0; i < tetrahedra.size(); i++ )
            {
                Coord v = bases[i] * ( childPosition - position[tetrahedra[i][0]] );
                double d = getDistanceTetra( v );
                if ( d<=distance ) { distance = d; index = i; }
            }
            int c0 = tetrahedra.size();
            for ( unsigned int i = 0; i < cubes.size(); i++ )
            {
                Coord v = bases[c0+i] * ( childPosition - position[cubes[i][0]] );  // for cuboid hexahedra
                // Coord v; Coord ph[8];  for ( unsigned int j = 0; j < 8; j++ ) ph[j]=position[cubes[i][j]]; computeHexaTrilinearWeights(v,ph,childPosition,1E-10); // for arbitrary hexahedra
                double d = getDistanceHexa( v );
                if ( d<=distance ) { distance = d; index = c0+i; }
            }
        }

        return index;
    }

protected:
    static void getBasisFrom2DElements( Basis& b, const Coord& p0, const Coord& p1, const Coord& p2 )
    {
        b[0] = p1 - p0;
        b[1] = p2 - p0;
        b[2] = cross ( b[0], b[1] );
    }


    static double getDistanceTriangle( const Coord& v )
    {
        return std::max ( std::max ( -v[0],-v[1] ),std::max ( ( v[2]<0?-v[2]:v[2] )-(Real)0.01,v[0]+v[1]-(Real)1. ) );
    }
    static double getDistanceQuad( const Coord& v )
    {
        return std::max ( std::max ( -v[0],-v[1] ),std::max ( std::max ( v[1]-(Real)1.,v[0]-(Real)1. ),std::max ( v[2]-(Real)0.01,-v[2]-(Real)0.01 ) ) );
    }
    static double getDistanceTetra( const Coord& v )
    {
        return std::max ( std::max ( -v[0],-v[1] ),std::max ( -v[2],v[0]+v[1]+v[2]-(Real)1. ) );
    }
    static double getDistanceHexa( const Coord& v )
    {
        return std::max ( std::max ( -v[0],-v[1] ),std::max ( std::max ( -v[2],v[0]-(Real)1. ),std::max ( v[1]-1,v[2]-(Real)1. ) ) );
    }
};


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

#endif // SOFA_BRANCHINGIMAGE_ImageMeshIntersection_H
