#ifndef SOFA_COMPONENT_MAPPING_ASSEMBLEDMULTIMAPPING_H
#define SOFA_COMPONENT_MAPPING_ASSEMBLEDMULTIMAPPING_H

#include <sofa/core/MultiMapping.h>
#include <SofaEigen2Solver/EigenSparseMatrix.h>

#include <Compliant/config.h>

namespace sofa
{

namespace component
{

namespace mapping
{


/**  
     Provides an implementation basis for assembled, sparse multi
     mappings.

     TODO: add .inl to minimize bloat / compilation times
     
     @author: maxime.tournier@inria.fr
*/

template <class TIn, class TOut >
class AssembledMultiMapping : public core::MultiMapping<TIn, TOut>
{
	typedef AssembledMultiMapping self;
  public:
	SOFA_CLASS(SOFA_TEMPLATE2(AssembledMultiMapping,TIn,TOut), SOFA_TEMPLATE2(core::MultiMapping,TIn,TOut));
	
	typedef core::MultiMapping<TIn, TOut> Inherit;
	typedef TIn In;
	typedef TOut Out;
	typedef typename Out::VecCoord OutVecCoord;
	typedef typename Out::VecDeriv OutVecDeriv;
	typedef typename Out::Coord OutCoord;
	typedef typename Out::Deriv OutDeriv;
	typedef typename Out::MatrixDeriv OutMatrixDeriv;
	typedef typename Out::Real Real;
    typedef Data<OutVecCoord> OutDataVecCoord;
    typedef Data<OutVecDeriv> OutDataVecDeriv;
	typedef typename In::Deriv InDeriv;
	typedef typename In::MatrixDeriv InMatrixDeriv;
	typedef typename In::Coord InCoord;
	typedef typename In::VecCoord InVecCoord;
	typedef typename In::VecDeriv InVecDeriv;
    typedef Data<InVecCoord> InDataVecCoord;
    typedef Data<InVecDeriv> InDataVecDeriv;
	typedef linearsolver::EigenSparseMatrix<TIn,TOut>  SparseMatrixEigen;

	typedef typename helper::vector <const InVecCoord*> vecConstInVecCoord;
	typedef typename helper::vector<OutVecCoord*> vecOutVecCoord;


	virtual void init() {
		assert( (this->getTo().size() == 1) && 
		        "sorry, multi mapping to multiple output dofs unimplemented" );
		
        Inherit::init();
	}

    void update() {
        this->reinit();
        Inherit::apply(core::MechanicalParams::defaultInstance() , core::VecCoordId::position(), core::ConstVecCoordId::position());
        Inherit::applyJ(core::MechanicalParams::defaultInstance() , core::VecDerivId::velocity(), core::ConstVecDerivId::velocity());
        if (this->f_applyRestPosition.getValue())
            Inherit::apply(core::MechanicalParams::defaultInstance(), core::VecCoordId::restPosition(), core::ConstVecCoordId::restPosition());
    }

    typedef linearsolver::EigenSparseMatrix<In, In> geometric_type;
    geometric_type geometric;
    
    virtual const defaulttype::BaseMatrix* getK() {
        if( geometric.compressedMatrix.nonZeros() ) return &geometric;
        else return NULL;
    }


    virtual void updateK( const core::MechanicalParams* /*mparams*/, core::ConstMultiVecDerivId force ) {

		const unsigned n = this->getFrom().size();

        helper::vector<const_in_coord_type> in_vec; in_vec.reserve(n);

        static const core::ConstMultiVecCoordId pos = core::ConstVecCoordId::position();
        
		for( unsigned i = 0; i < n; ++i ) {
            const core::State<TIn>* from = this->getFromModels()[i];
            const_in_coord_type in_pos( *pos[from].read() );
			in_vec.push_back( in_pos );
		}

        const core::State<TOut>* to = this->getToModels()[0];
        const_out_deriv_type out_force( *force[to].read() );

        this->assemble_geometric(in_vec, out_force);
    }

	
	virtual void apply(const core::MechanicalParams* , 
	                   const helper::vector<OutDataVecCoord*>& dataVecOutPos,
	                   const helper::vector<const InDataVecCoord*>& dataVecInPos) {
		alloc();
	
		const unsigned n = this->getFrom().size();

        helper::vector<in_pos_type> in_vec; in_vec.reserve(n);

		for( unsigned i = 0; i < n; ++i ) {
			in_vec.push_back( in_pos_type(dataVecInPos[i]) );
		}
		
		out_pos_type out(dataVecOutPos[0]);
		
		apply(out, in_vec);
		assemble(in_vec);
	}



    virtual void applyJ(const core::MechanicalParams*, const helper::vector<OutDataVecDeriv*>& outDeriv, const helper::vector<const InDataVecDeriv*>& inDeriv)
    {
        bool first = true;

        for(unsigned i = 0, n = js.size(); i < n ; ++i ) {
            if( jacobian(i).rowSize() > 0 ) {
                if(first) {
                    jacobian(i).mult(*outDeriv[0], *inDeriv[i]);
                    first = false;
                } else {
                    jacobian(i).addMult(*outDeriv[0], *inDeriv[i]);
                }
                    
            }
        }

        if( first ) {
            
            // nothing written: jacobian is zero and we still have to clear
            helper::WriteOnlyAccessor<Data<OutVecDeriv> > out (*outDeriv[0]);

            for(auto& x : out) {
                x = typename TOut::Deriv();
            }

        }
       


    }

    void debug() {
		std::cerr << this->getClassName() << std::endl;
		for( unsigned i = 0, n = js.size(); i < n; ++i) {
			std::cerr << "from: " <<  this->getFrom()[i]->getContext()->getName() 
					  << "/" << this->getFrom()[i]->getName() << std::endl;
		}
		std::cerr << "to: " << this->getTo()[0]->getContext()->getName() << "/" << this->getTo()[0]->getName() << std::endl;
		std::cerr << std::endl;
	}

	virtual void applyJT(const core::MechanicalParams*,
						 const helper::vector< InDataVecDeriv*>& outDeriv, 
	                     const helper::vector<const OutDataVecDeriv*>& inDeriv) {
		for( unsigned i = 0, n = js.size(); i < n; ++i) {
			if( jacobian(i).rowSize() > 0 ) {
				jacobian(i).addMultTranspose(*outDeriv[i], *inDeriv[0]);
			}
		}
	}

    virtual void applyDJT(const core::MechanicalParams* mparams,
                          core::MultiVecDerivId inForce,
                          core::ConstMultiVecDerivId /*outForce*/)
    {
        if( geometric.compressedMatrix.nonZeros() )
        {
            const SReal& kfactor = mparams->kFactor();
            unsigned size = this->getFromModels().size();

            // TODO not optimized but at least it is implemented

            // merge all parent displacements
            InVecDeriv parentDisplacements;

            for( unsigned i=0; i< size ; i++ )
            {
                const core::State<In>* fromModel = this->getFromModels()[i];
                const InVecDeriv& parentDisplacement = mparams->readDx(fromModel)->getValue();
                parentDisplacements.insert(parentDisplacements.end(), parentDisplacement.begin(), parentDisplacement.end());
            }

            // merged parent forces
            InVecDeriv parentForces(parentDisplacements.size());
            geometric.addMult( parentForces, parentDisplacements, kfactor );

            // un-merge parent forces
            size_t offset=0;
            for( unsigned i=0; i< size ; i++ )
            {
                core::State<In>* fromModel = this->getFromModels()[i];
                InVecDeriv& parentForce = *inForce[fromModel].write()->beginEdit();
                for( size_t j=0;j<parentForce.size();++j)
                    parentForce[j] += parentForces[offset+j];
                offset += parentForce.size();
                inForce[fromModel].write()->endEdit();
            }
        }
    }

    virtual void applyJT( const core::ConstraintParams*,
						  const helper::vector< typename self::InDataMatrixDeriv* >& , 
						  const helper::vector< const typename self::OutDataMatrixDeriv* >&  ) {
		// throw std::logic_error("not implemented");
	}

	
    virtual const helper::vector<sofa::defaulttype::BaseMatrix*>* getJs() {

        if( js.empty() )
        {
            serr << "empty js for " << this->getPathName() << sendl;
            assert(false);
        }
		
		return &js;
    }
 
    
  protected:
	typedef linearsolver::EigenSparseMatrix<In, Out> jacobian_type;

	
  // returns the i-th jacobian matrix
	jacobian_type& jacobian(unsigned i) {
		assert( i < js.size() );
		assert( dynamic_cast<jacobian_type*>(js[i]) );
	
		return *static_cast<jacobian_type*>(js[i]);
	}


	enum {Nin = In::deriv_total_size, Nout = Out::deriv_total_size };


    // let's do this once and for all
    typedef helper::ReadAccessor< Data< typename self::InVecCoord > > const_in_coord_type;
    typedef helper::WriteAccessor< Data< typename self::InVecCoord > > in_coord_type;

    typedef helper::ReadAccessor< Data< typename self::OutVecCoord > > const_out_coord_type;
    typedef helper::WriteAccessor< Data< typename self::OutVecCoord > > out_coord_type;

    typedef helper::ReadAccessor< Data< typename self::InVecDeriv > > const_in_deriv_type;
    typedef helper::WriteAccessor< Data< typename self::InVecDeriv > > in_deriv_type;
    
    typedef helper::ReadAccessor< Data< typename self::OutVecDeriv > > const_out_deriv_type;
    typedef helper::WriteAccessor< Data< typename self::OutVecDeriv > > out_deriv_type;


	// TODO rename in_coord_type/out_coord_type
    typedef helper::ReadAccessor< Data< typename self::InVecCoord > > in_pos_type;
    typedef helper::WriteOnlyAccessor< Data< typename self::OutVecCoord > > out_pos_type;

    typedef helper::ReadAccessor< Data< typename self::InVecDeriv > > in_vel_type;
    typedef helper::WriteAccessor< Data< typename self::OutVecDeriv > > out_vel_type;

    typedef helper::ReadAccessor< Data< typename self::OutVecDeriv > > in_force_type;
    typedef helper::WriteAccessor< Data< typename self::InVecDeriv > > out_force_type;
	
	
	// perform a jacobian blocs assembly
	// TODO pass out value as well ?
    virtual void assemble( const helper::vector<in_pos_type>& in ) = 0;

    virtual void assemble_geometric( const helper::vector<const_in_coord_type>& /*in*/,
                                     const const_out_deriv_type& /*out*/) { }
    
    using Inherit::apply;
	// perform mapping operation on positions
    virtual void apply(out_pos_type& out, 
                       const helper::vector<in_pos_type>& in ) = 0;

  protected:

	// allocate jacobians
	virtual void alloc() {
		
		const unsigned n = this->getFrom().size();

        // add needed jacobians
        for(unsigned i = js.size(); i < n; ++i) {
            js.push_back( new SparseMatrixEigen );
        }


        // release unneeded jacobians
        for(unsigned i = n; i < js.size(); ++i) {
            delete js[i];
        }

        js.resize(n);
	}

	// delete jacobians
    void clear() {
        for( unsigned i = 0, n = js.size(); i < n; ++i) {
			delete js[i];
		}

        js.clear();
	}

	
    typedef helper::vector< sofa::defaulttype::BaseMatrix* > js_type;
	js_type js;


  protected:
	
	core::behavior::BaseMechanicalState* from(unsigned i) {
		// TODO assert
        return this->getFrom()[i]->toBaseMechanicalState();
	}

	core::behavior::BaseMechanicalState* to(unsigned i = 0) {
		// TODO assert
        return this->getTo()[i]->toBaseMechanicalState();
	}

	
  public:
	
    virtual ~AssembledMultiMapping() {
		clear();
	}


};


}
}
}



#endif

