/*---------------------------------------------------------------------------*\
  =========                 |
  \\      /  F ield         | OpenFOAM: The Open Source CFD Toolbox
   \\    /   O peration     |
    \\  /    A nd           | www.openfoam.com
     \\/     M anipulation  |
-------------------------------------------------------------------------------
    Copyright (C) 2016 OpenFOAM Foundation
-------------------------------------------------------------------------------
License
    This file is part of OpenFOAM.

    OpenFOAM is free software: you can redistribute it and/or modify it
    under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    OpenFOAM is distributed in the hope that it will be useful, but WITHOUT
    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
    FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
    for more details.

    You should have received a copy of the GNU General Public License
    along with OpenFOAM.  If not, see <http://www.gnu.org/licenses/>.

\*---------------------------------------------------------------------------*/

#include "rigidBodyModel.H"
#include "rigidBodyModelState.H"
#include "rigidBodyRestraint.H"

// * * * * * * * * * * * * * * Member Functions  * * * * * * * * * * * * * * //

void Foam::RBD::rigidBodyModel::applyRestraints
(
    scalarField& tau,
    Field<spatialVector>& fx,
    const rigidBodyModelState& state
) const
{
    if (restraints_.empty())
    {
        return;
    }

    forAll(restraints_, ri)
    {
        DebugInfo << "Restraint " << restraints_[ri].name();

        // Accumulate the restraint forces
        restraints_[ri].restrain(tau, fx, state);
    }
}

void Foam::RBD::rigidBodyModel::forwardDynamics
(
    rigidBodyModelState& state,
    const scalarField& tau,
    const Field<spatialVector>& fx
) const
{
    const scalarField& q = state.q();
    const scalarField& qDot = state.qDot();
    scalarField& qDdot = state.qDdot();

    DebugInFunction
        << "q = " << q << nl
        << "qDot = " << qDot << nl
        << "tau = " << tau << endl;

    // Joint state returned by jcalc
    joint::XSvc J;

    v_[0] = Zero;

    for (label i=1; i<nBodies(); i++)
    {
        const joint& jnt = joints()[i];
        jnt.jcalc(J, q, qDot);

        S_[i] = J.S;
        S1_[i] = J.S1;

        Xlambda_[i] = J.X & XT_[i];

        const label lambdai = lambda_[i];

        if (lambdai != 0)
        {
            X0_[i] = Xlambda_[i] & X0_[lambdai];
        }
        else
        {
            X0_[i] = Xlambda_[i];
        }

        v_[i] = (Xlambda_[i] & v_[lambdai]) + J.v;
        c_[i] = J.c + (v_[i] ^ J.v);
        IA_[i] = I(i);
        pA_[i] = v_[i] ^* (I(i) & v_[i]);

        if (fx.size())
        {
            pA_[i] -= *X0_[i] & fx[i];
        }
    }

    for (label i=nBodies()-1; i>0; i--)
    {
        const joint& jnt = joints()[i];
        const label qi = jnt.qIndex();

        if (jnt.nDoF() == 1)
        {
            U1_[i] = IA_[i] & S1_[i];
            Dinv_[i].xx() = 1/(S1_[i] && U1_[i]);
            u_[i].x() = tau[qi] - (S1_[i] && pA_[i]);

            const label lambdai = lambda_[i];

            if (lambdai != 0)
            {
                const spatialTensor Ia
                (
                    IA_[i] - (U1_[i]*(Dinv_[i].xx()*U1_[i]))
                );

                const spatialVector pa
                (
                    pA_[i] + (Ia & c_[i]) + U1_[i]*(Dinv_[i].xx()*u_[i].x())
                );

                IA_[lambdai] +=
                    spatialTensor(Xlambda_[i].T())
                  & Ia
                  & spatialTensor(Xlambda_[i]);

                pA_[lambdai] += Xlambda_[i].T() & pa;
            }
        }
        else
        {
            U_[i] = IA_[i] & S_[i];
            Dinv_[i] = (S_[i].T() & U_[i]).inv();

            u_[i] = tau.block<vector>(qi) - (S_[i].T() & pA_[i]);

            const label lambdai = lambda_[i];

            if (lambdai != 0)
            {
                spatialTensor Ia
                (
                    IA_[i]
                  - (U_[i] & Dinv_[i] & U_[i].T())
                );

                spatialVector pa
                (
                    pA_[i]
                  + (Ia & c_[i])
                  + (U_[i] & Dinv_[i] & u_[i])
                );

                IA_[lambdai] +=
                    spatialTensor(Xlambda_[i].T())
                  & Ia
                  & spatialTensor(Xlambda_[i]);

                pA_[lambdai] += Xlambda_[i].T() & pa;
            }
        }
    }

    a_[0] = spatialVector(Zero, -g_);

    for (label i=1; i<nBodies(); i++)
    {
        const joint& jnt = joints()[i];
        const label qi = jnt.qIndex();

        a_[i] = (Xlambda_[i] & a_[lambda_[i]]) + c_[i];

        if (jnt.nDoF() == 1)
        {
            qDdot[qi] = Dinv_[i].xx()*(u_[i].x() - (U1_[i] && a_[i]));
            a_[i] += S1_[i]*qDdot[qi];
        }
        else
        {
            vector qDdoti(Dinv_[i] & (u_[i] - (U_[i].T() & a_[i])));

            // Need to add mutable "block<vector>" to Field
            qDdot[qi] = qDdoti.x();
            qDdot[qi+1] = qDdoti.y();
            qDdot[qi+2] = qDdoti.z();

            a_[i] += (S_[i] & qDdoti);
        }
    }

    DebugInfo
        << "qDdot = " << qDdot << nl
        << "a = " << a_ << endl;
}

void Foam::RBD::rigidBodyModel::forwardDynamicsCOM
(
    rigidBodyModelState& state,
    const scalarField& tau,
    const Field<spatialVector>& fx
) const
{
//    const scalarField& q = state.q();
    const scalarField& qDot = state.qDot();
    scalarField& qDdot = state.qDdot();

//Info << "********** forwardDynamics ***************" <<  endl;

//Info << "bodyIDs_ = " << bodyIDs_ << endl;	// 0: root, 1:massless, 2: boxp300body
//Info << "bodies_ = " << bodies_ << endl;	// masslessBody, jointBody, rigidBody(mass, centreOfMass, inertia )


    a_[0] = spatialVector(Zero, Zero);	//spatialVector(Zero, -g_);

    for (label i=0; i<((nBodies()-1)/2); i++){
	const label num = 2*i+2;		// Rxyz body ID
        const joint& jnt = joints()[num-1];	// Pxyz body ID
        const label qi = jnt.qIndex();		// 
	Tensor<scalar> Rb = X0(num).E().T();	// orientation matrix of Rxyz body
	Tensor<scalar> Inertia(I(num).Ic());		// inertia of Rxyz body for COM
	const vector vel(qDot[qi+0], qDot[qi+1], qDot[qi+2]);	// velocity of Pxyz body (= Rxyz body)
	const vector omega(qDot[qi+3], qDot[qi+4], qDot[qi+5]);	// angular velocity of Rxyz 
	
	Inertia = Rb & Inertia & Rb.T();	// Rxyz body inertia viewing from parent (=reference=global coordinate system)

	vector force( fx[num][3], fx[num][4], fx[num][5]);
	vector moment(fx[num][0], fx[num][1], fx[num][2]);
	vector	acceleration( force/I(num).m() );
	vector	angular_acceleration( Zero );

	acceleration += g_;	// += mg/m
	angular_acceleration = Inertia.inv() & ( moment - ( omega ^ (Inertia & omega) ) );

        qDdot[qi+0] = acceleration[0];	// fx / m;
        qDdot[qi+1] = acceleration[1];	// fy / m;
        qDdot[qi+2] = acceleration[2];	// fz / m;
        qDdot[qi+3] = angular_acceleration[0];	// qDdoti.roll();
        qDdot[qi+4] = angular_acceleration[1];	// qDdoti.pitch();
        qDdot[qi+5] = angular_acceleration[2];	// qDdoti.yaw();

        a_[num-1] = spatialVector( Zero, acceleration );			// Pxyz		
        a_[num]   = spatialVector( angular_acceleration, acceleration );	// Pxyz and Rxyz	
    }
}

void Foam::RBD::rigidBodyModel::forwardDynamicsCorrection
(
    const rigidBodyModelState& state
) const
{
    DebugInFunction << endl;

    const scalarField& q = state.q();
    const scalarField& qDot = state.qDot();
    const scalarField& qDdot = state.qDdot();

    // Joint state returned by jcalc
    joint::XSvc J;

    v_[0] = Zero;
    a_[0] = spatialVector(Zero, -g_);

    for (label i=1; i<nBodies(); i++)
    {
        const joint& jnt = joints()[i];
        const label qi = jnt.qIndex();

        jnt.jcalc(J, q, qDot);

        S_[i] = J.S;
        S1_[i] = J.S1;

        Xlambda_[i] = J.X & XT_[i];

        const label lambdai = lambda_[i];

        if (lambdai != 0)
        {
            X0_[i] = Xlambda_[i] & X0_[lambdai];
        }
        else
        {
            X0_[i] = Xlambda_[i];
        }

        v_[i] = (Xlambda_[i] & v_[lambdai]) + J.v;
        c_[i] = J.c + (v_[i] ^ J.v);
        a_[i] = (Xlambda_[i] & a_[lambdai]) + c_[i];

        if (jnt.nDoF() == 1)
        {
            a_[i] += S1_[i]*qDdot[qi];
        }
        else
        {
            a_[i] += S_[i] & qDdot.block<vector>(qi);
        }
    }

    DebugInfo<< "a = " << a_ << endl;
}

void Foam::RBD::rigidBodyModel::forwardDynamicsCorrectionCOM
(
    const rigidBodyModelState& state
) const
{
//Info << " forwardDynamicsCorrectionCOM: " << endl;

    Tensor<scalar> R;

    v_[0] = Zero;
    a_[0] = Zero;

    for (label i=1; i<nBodies(); i++)	// 0, first rigid body:1, 2, second rigid body: 3, 4, .....(if there are n rigid bodies, nBodies()=2*n+1 )
    {
	if( (i%2) == 1 ){		// for Pxyz ???
	   Xlambda_[i] = XT_[i];	// XT means initial attitude and position ?
	   X0_[i] = XT_[i];
    	   v_[i] = Zero;
    	   c_[i] = Zero;	
    	   a_[i] = Zero;
	}else{			// for Pxyz and Rxyz
	  label	ID = i/2 - 1;	// ridig body ID

	  X0_[i].r() = X0_[i-1].r() + vector( state.q()[6*ID+0], state.q()[6*ID+1], state.q()[6*ID+2] );	// r() means origin = rigid body center

	  vector n(state.q()[6*ID+3], state.q()[6*ID+4], state.q()[6*ID+5]);
	  scalar theta = mag(n);
	  n /= (theta+VSMALL);

	  /*** Rodrigues matrix ***/
	  R.xx() = cos(theta) + n[0]*n[0] * ( 1.0 - cos(theta) );		R.xy() = n[0]*n[1] * ( 1.0 - cos(theta) ) - n[2] * sin(theta);	R.xz() = n[2]*n[0] * ( 1.0 - cos(theta) ) + n[1] * sin(theta);
	  R.yx() = n[0]*n[1] * ( 1.0 - cos(theta) ) + n[2] * sin(theta);	R.yy() = cos(theta) + n[1]*n[1] * ( 1.0 - cos(theta) );		R.yz() = n[1]*n[2] * ( 1.0 - cos(theta) ) - n[0] * sin(theta);
	  R.zx() = n[2]*n[0] * ( 1.0 - cos(theta) ) - n[1] * sin(theta);	R.zy() = n[1]*n[2] * ( 1.0 - cos(theta) ) + n[0] * sin(theta);	R.zz() = cos(theta) + n[2]*n[2] * ( 1.0 - cos(theta) );

	  R = R & X0_[i-1].E().T();	// E() means attitude matrix, R and X0[i-1] must be unit matrix, I.
	  X0_[i].E() = R.T();
	  
	  const vector vel(  state.qDot()[6*ID+0], state.qDot()[6*ID+1], state.qDot()[6*ID+2]);	// velocity of Pxyz body (= Rxyz body)
	  const vector omega(state.qDot()[6*ID+3], state.qDot()[6*ID+4], state.qDot()[6*ID+5]);	// angular velocity of Rxyz 
	  v_[i] = spatialVector(omega, vel);

	  const vector velDot(  state.qDdot()[6*ID+0], state.qDdot()[6*ID+1], state.qDdot()[6*ID+2]);	// velocity of Pxyz body (= Rxyz body)
	  const vector omegaDot(state.qDdot()[6*ID+3], state.qDdot()[6*ID+4], state.qDdot()[6*ID+5]);	// angular velocity of Rxyz 
	  c_[i] = spatialVector(omegaDot, velDot);
	  a_[i] = spatialVector(omegaDot, velDot);
	}
    }
/*
for (label i=0; i<nBodies(); i++){
Info << "i = " << i << endl;
Info << "X0 " << X0_[i] << endl;
Info << "v  " << v_[i] << endl;
Info << "c  " << c_[i] << endl;
Info << "a  " << a_[i] << endl << endl;

}*/
}

// ************************************************************************* //
