#include "Stabilization.h"

#include <sofa/core/ObjectFactory.h>
#include "../utils/map.h"

namespace sofa {
namespace component {
namespace odesolver {


SOFA_DECL_CLASS(Stabilization)
int StabilizationClass = core::RegisterObject("Kinematic constraint stabilization").add< Stabilization >();


using namespace utils;

Stabilization::Stabilization( mstate_type* mstate )
    : BaseConstraintValue( mstate )
    , mask(initData(&mask, "mask", "dofs to be stabilized"))
{}

void Stabilization::correction(SReal* dst, unsigned n, unsigned dim, const core::MultiVecCoordId& posId, const core::MultiVecDerivId&) const {
	assert( mstate );

    unsigned size = n*dim;

    assert( mask.getValue().empty() || mask.getValue().size() == n );
	
    mstate->copyToBuffer(dst, posId.getId(mstate.get()), size);
    map(dst, size) *= -1.0 / this->getContext()->getDt();

	const mask_type& mask = this->mask.getValue();

    if( mask.empty() ) return; // all stabilized
	
    // zero for non-stabilized
	unsigned i = 0;
    for(SReal *it=dst, *last = dst + size; it < last; it+=dim, ++i) {
        if( !mask[i] ) memset( it, 0, dim*sizeof(SReal) );
    }
}


void Stabilization::dynamics(SReal* dst, unsigned n, unsigned dim, bool stabilization, const core::MultiVecCoordId& posId, const core::MultiVecDerivId&) const {
    assert( mstate );

    const unsigned size = n*dim;

    // warning iff stabilization, only cancelling relative velocities of violated constraints (given by mask)


    assert( mask.getValue().empty() || mask.getValue().size() == n );
    const mask_type& mask = this->mask.getValue();

    if( !stabilization )
    {
        // elastic constraints
        mstate->copyToBuffer(dst, posId.getId(mstate.get()), size);
        map(dst, size) *= -1.0 / this->getContext()->getDt();
    }
    else
    {
        if( mask.empty() ) memset( dst, 0, size*sizeof(SReal) ); // all violated
        else
        {
            // for possible elastic constraint
            mstate->copyToBuffer(dst, posId.getId(mstate.get()), size);

            unsigned i = 0;
            for(SReal *it=dst, *last = dst + size; it < last; it+=dim, ++i) {
                if( mask[i] ) memset( it, 0, dim*sizeof(SReal) ); // already violated
                else map(it, dim) *= -1.0 / this->getContext()->getDt(); // not violated
            }
        }
    }
}



void Stabilization::filterConstraints( helper::vector<bool>*& /*activateMask*/,
                                       const core::MultiVecCoordId& posId, unsigned n, unsigned dim )
{
    // All the constraints remain active
    // but non-violated constraint must not be stabilized

    if( !mask.getValue().empty() ) return; // already be done (contact for example)

    unsigned size = n*dim;

    mask_type& mask = *this->mask.beginWriteOnly();
    mask.resize( n );

    std::vector<SReal> violation(size);
    mstate->copyToBuffer(violation.data(), posId.getId(mstate.get()), size);

    for(unsigned block = 0; block < n; ++block) {
        unsigned line = block*dim; // first constraint line
        mask[block] = ( violation[line]<0 ); // violated constraint
    }

    this->mask.endEdit();
    // activateMask = &mask;
}


}
}
}
