#include "PenetrationForceField.h"
#include <iostream>

using std::cerr;
using std::endl;

namespace sofa
{
namespace component
{
namespace forcefield
{

template<class DataTypes>
PenetrationForceField<DataTypes>::PenetrationForceField( core::behavior::MechanicalState<DataTypes> *mm )
    : Inherit(mm)
    , hardStiffness( initData(&hardStiffness, (SReal)1E10, "hardStiffness", "Stiffness when position <0."))
    , separatingStiffness( initData(&separatingStiffness, (SReal)0, "separatingStiffness", "used instead of hardStiffness when velocity >0."))
    , proximity( initData(&proximity, (SReal)0, "proximity", "Position value above which stiffness=0."))
    , damping( initData(&damping, (SReal)0, "damping", "uniform viscous damping."))
    , colors( initData(&colors, "colors", "output colors depending on stiffness."))
    , computeColors( initData(&computeColors, false, "computeColors", "compute output colors depending on stiffness."))
{
    colors.setReadOnly(true);
    this->isCompliance.setValue(false);
}

template<class DataTypes>
void PenetrationForceField<DataTypes>::init()
{
    Inherit::init();
    if( this->getMState()==NULL ) serr<<"init(), no mstate !" << sendl;
    reinit();
}

template<class DataTypes>
void PenetrationForceField<DataTypes>::reinit()
{
    core::behavior::BaseMechanicalState* state = this->getContext()->getMechanicalState();
    assert(state);

    unsigned int n = state->getMatrixSize();

    if( damping.getValue() > 0 )
    {
        matB.resize(n, n);
        for(unsigned i=0; i < n; i++)
        {
            matB.compressedMatrix.startVec(i);
            matB.compressedMatrix.insertBack(i, i) = -damping.getValue();
        }
        matB.compressedMatrix.finalize();
    }
    else matB.compressedMatrix.resize(0,0);

    reset();
}

template<class DataTypes>
void PenetrationForceField<DataTypes>::reset()
{
    colors.beginWriteOnly()->clear(); colors.endEdit();
    Inherit::reset();
}

template<class DataTypes>
const sofa::defaulttype::BaseMatrix* PenetrationForceField<DataTypes>::getComplianceMatrix(const core::MechanicalParams*)
{
    const DataVecCoord& x = *this->mstate->read(core::ConstVecCoordId::position());
    const DataVecDeriv& v = *this->mstate->read(core::ConstVecDerivId::velocity());
    unsigned int n = x.getValue().size();
    matC.resize(n, n);
    for(unsigned i=0; i < n; i++)
    {
        Real k = getStiffness(x.getValue()[i],v.getValue()[i]);
        if(k)
        {
            matC.compressedMatrix.startVec(i);
            matC.compressedMatrix.insertBack(i, i) = 1./k;
        }
        else
        {
            matC.compressedMatrix.startVec(i);
            matC.compressedMatrix.insertBack(i, i) = 1000.;
        }
    }
    matC.compressedMatrix.finalize();
    return &matC;
}

template<class DataTypes>
void PenetrationForceField<DataTypes>::addKToMatrix( sofa::defaulttype::BaseMatrix * matrix, SReal kFact, unsigned int &offset )
{
    matK.addToBaseMatrix( matrix, kFact, offset );
}

template<class DataTypes>
void PenetrationForceField<DataTypes>::addBToMatrix( sofa::defaulttype::BaseMatrix * matrix, SReal bFact, unsigned int &offset )
{
    matB.addToBaseMatrix( matrix, bFact, offset );
}


template<class DataTypes>
void PenetrationForceField<DataTypes>::addForce(const core::MechanicalParams *, DataVecDeriv& f, const DataVecCoord& xx, const DataVecDeriv& vv)
{
    VecDeriv& force = *f.beginEdit();
    const VecCoord& x = xx.getValue();
    const VecCoord& v = vv.getValue();

    size_t n = x.size();
    matK.resize(n, n);
    for(unsigned i=0; i < n; i++)
    {
        Real k = getStiffness(x[i],v[i]);
        if(k)
        {
            matK.compressedMatrix.startVec(i);
            matK.compressedMatrix.insertBack(i, i) = -k;
            force[i][0]-=(x[i][0]-this->proximity.getValue())*k;
        }
    }
    if( damping.getValue() > 0 )
        for(unsigned i=0; i < n; i++)
            force[i][0]-=v[i][0]*this->damping.getValue();

    matK.compressedMatrix.finalize();
    f.endEdit();


    // for external display (e.g. mapping)
    if( computeColors.getValue() )
    {
        helper::vector<defaulttype::Vec4f>&c = *colors.beginWriteOnly();
        c.resize(n);
        for(unsigned i=0; i < n; i++)
            c[i] = getColor(x[i],v[i]);
        colors.endEdit();
    }
}


template<class DataTypes>
typename PenetrationForceField<DataTypes>::Real PenetrationForceField<DataTypes>::getStiffness(const Coord& p,const Deriv& v) const
{
    Real k=0.;
    if(p[0]<0)
    {
        if(v[0]>0) k=this->separatingStiffness.getValue();
        else k=this->hardStiffness.getValue();
    }
    else if(p[0]<this->proximity.getValue())
    {
        if(v[0]>0) k=(1.-p[0]/this->proximity.getValue())*this->separatingStiffness.getValue();
        else k=(1.-p[0]/this->proximity.getValue())*this->hardStiffness.getValue();
    }
    return k;
}


template<class DataTypes>
const defaulttype::Vec4f& PenetrationForceField<DataTypes>::getColor(const Coord& p,const Deriv& v) const
{
    static const defaulttype::Vec4f separatingColor(0,1,0,1);
    static const defaulttype::Vec4f hardColor(1,0,0,1);
    static const defaulttype::Vec4f proximityColor(1,.7,0,1);
    static const defaulttype::Vec4f toofarColor(1,1,0,1);

    const Real&prox = this->proximity.getValue();

    if( p[0]>prox ) return toofarColor;

    if(p[0]<prox/5.0) return hardColor;

    if(v[0]>1e-5) return separatingColor;

    return proximityColor;
}

template<class DataTypes>
void PenetrationForceField<DataTypes>::addDForce(const core::MechanicalParams *mparams, DataVecDeriv& df,  const DataVecDeriv& dx)
{
    Real kfactor = (Real)mparams->kFactorIncludingRayleighDamping(this->rayleighStiffness.getValue());
    if( kfactor ) matK.addMult( df, dx, kfactor );
    if( damping.getValue() > 0 )  matB.addMult( df, dx, (Real)mparams->bFactor() );
}


}
}
}
