//#####################################################################
// Copyright 2004-2007, Kevin Der, Craig Schroeder, Andrew Selle, Tamar Shinar, Eftychios Sifakis, Rachel Weinstein.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Rigids/Articulated_Rigid_Bodies/ARTICULATED_RIGID_BODY_3D.h>
#include <Core/Matrices/MATRIX.h>
#include <Core/Matrices/MATRIX_MXN.h>
#include <Core/Read_Write/FILE_UTILITIES.h>
#include <Core/Utilities/VIEWER_DIR.h>
#include <Rigids/Joints/JOINT.h>
#include <Rigids/Joints/JOINT_FUNCTION.h>
#include <Rigids/Joints/JOINT_MESH.h>
#include <Rigids/Rigid_Bodies/RIGID_BODY.h>
#include <Rigids/Rigid_Bodies/RIGID_BODY_COLLECTION.h>
#include <Rigids/Rigid_Bodies/RIGID_BODY_MASS.h>
using namespace PhysBAM;
//#####################################################################
// Constructor
//#####################################################################
template<class T> ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
ARTICULATED_RIGID_BODY(RIGID_BODY_COLLECTION<TV>& rigid_body_collection_input)
    :ARTICULATED_RIGID_BODY_BASE<TV>(rigid_body_collection_input)
{}
//#####################################################################
// Destructor
//#####################################################################
template<class T> ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
~ARTICULATED_RIGID_BODY()
{}
//####################################################################################
// Function Compute_Desired_PD_Velocity
//####################################################################################
template<class T> void ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
Compute_Desired_PD_Velocity(const T dt,const T time)
{
    for(JOINT_ID i(0);i<joint_mesh.Size();i++)
        if(joint_mesh.Is_Active(i) && joint_mesh(i)->joint_function && joint_mesh(i)->joint_function->active && (!Parent(i)->Has_Infinite_Inertia() || !Child(i)->Has_Infinite_Inertia()))
            joint_mesh(i)->joint_function->Compute_Desired_PD_Velocity(dt,time);
}
//####################################################################################
// Function Compute_Position_Based_State
//####################################################################################
template<class T> void ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
Compute_Position_Based_State(const T dt,const T time)
{
    const ARRAY<int>& dynamic_rigid_body_particles=rigid_body_collection.dynamic_rigid_body_particles;
    Compute_Desired_PD_Velocity(dt,time);

    if(!global_post_stabilization) return;

    // build global poststabilization matrix
    ARRAY<ARRAY<pair<int,bool> >,int> joints_on_rigid_body(rigid_body_collection.rigid_body_particles.Size()); // stores joint id and whether it's parent for joint.
    ARRAY<TV> joint_location(joint_mesh.Num_Joints());
    ARRAY<MATRIX_MXN<T> > joint_constraint_matrix(joint_mesh.Num_Joints());
    joint_constrained_dimensions.Resize(joint_mesh.Num_Joints());joint_constrained_dimensions.Fill(0);
    joint_angular_constraint_matrix.Resize(joint_mesh.Num_Joints()); // angular directions constrained by joint or PD

    for(int i=0;i<joint_mesh.Num_Joints();i++){
        JOINT<TV>& joint=*joint_mesh.Joints(i);RIGID_BODY<TV>& parent=*Parent(joint.id_number),&child=*Child(joint.id_number);
        if(parent.Has_Infinite_Inertia() && child.Has_Infinite_Inertia()) continue;
        joints_on_rigid_body(Parent_Id(joint.id_number)).Append(pair<int,bool>(i,true));
        joints_on_rigid_body(Child_Id(joint.id_number)).Append(pair<int,bool>(i,false));

        if(!joint.global_post_stabilization) continue;

        joint_location(i)=joint.Location(parent,child);

        // assumes full positional constraint and allows variable angular constraints
        // note: PD control counts as a constrained direction (it simply has a nonzero desired relative velocity for that direction)
        assert(joint.Has_Prismatic_Constraint()); // assume this for now
        joint.Angular_Constraint_Matrix(parent.Frame(),child.Frame(),joint_angular_constraint_matrix(i));
        if(joint.joint_function && joint.joint_function->active && use_pd_actuators)
            joint.joint_function->Angular_Constraint_Matrix(parent.Frame(),joint_angular_constraint_matrix(i));

        // full constraint matrix (linear and angular)
        joint_constrained_dimensions(i)=3+joint_angular_constraint_matrix(i).Columns(); // includes PD dimensions
        joint_constraint_matrix(i)=MATRIX_MXN<T>(6,joint_constrained_dimensions(i));
        joint_constraint_matrix(i).Set_Submatrix(0,0,MATRIX<T,3>::Identity_Matrix());
        joint_constraint_matrix(i).Set_Submatrix(3,3,joint_angular_constraint_matrix(i));}

    joint_offset_in_post_stabilization_matrix.Resize(joint_mesh.Num_Joints());
    for(int i=0;i<joint_mesh.Num_Joints();i++) joint_offset_in_post_stabilization_matrix(i)=(i==0)?0:joint_offset_in_post_stabilization_matrix(i-1)+joint_constrained_dimensions(i-1);
    int total_constrained_dimensions=joint_constrained_dimensions.Sum();
    global_post_stabilization_matrix_11=MATRIX_MXN<T>(total_constrained_dimensions,total_constrained_dimensions);
    LOG::cout<<"*** Total constrained dimensions "<<total_constrained_dimensions<<std::endl;

    for(int i=0;i<dynamic_rigid_body_particles.m;i++){int p=dynamic_rigid_body_particles(i);
        RIGID_BODY<TV>& rigid_body=rigid_body_collection.Rigid_Body(p);SYMMETRIC_MATRIX<T,3> I_inverse=rigid_body.World_Space_Inertia_Tensor_Inverse();
        int id=rigid_body.particle_index;
        for(int i=0;i<joints_on_rigid_body(id).m;i++) for(int j=i;j<joints_on_rigid_body(id).m;j++){
            int joint_index_1=joints_on_rigid_body(id)(i).first,joint_index_2=joints_on_rigid_body(id)(j).first;

            // joints that won't take part in global post stabilization are skipped
            // we currently assume that a connected component must either be all in or all out of global post stabilization but not mixed
            if(!joint_mesh.Joints(joint_index_1)->global_post_stabilization || !joint_mesh.Joints(joint_index_2)->global_post_stabilization){
                //assert(!joint_mesh.Joints(joint_index_1)->global_post_stabilization && !joint_mesh.Joints(joint_index_2)->global_post_stabilization);
                continue;}

            // compute 6x6 matrix assuming impulse applied at r2 and measured at r1
            TV r1=joint_location(joint_index_1)-rigid_body.Frame().t,r2=joint_location(joint_index_2)-rigid_body.Frame().t;
            MATRIX<T,3> c10=I_inverse*Cross_Product_Matrix(r2),c01=Cross_Product_Matrix(-r1)*I_inverse;
            MATRIX<T,3> c00=Cross_Product_Matrix(-r1)*c10+1/rigid_body.Mass();
            SYMMETRIC_MATRIX<T,3> c11=I_inverse;
            MATRIX_MXN<T> C(6,6);C.Add_To_Submatrix(0,0,c00);C.Add_To_Submatrix(0,3,c01);C.Add_To_Submatrix(3,0,c10);C.Add_To_Submatrix(3,3,c11);
            // important: sign of block is negative if the joints at which impulse is applied and measured are opposite types (parent/child)
            if(joints_on_rigid_body(id)(i).second!=joints_on_rigid_body(id)(j).second) C*=-1;

            // compute projected matrix
            MATRIX_MXN<T> C_times_joint_constraint_matrix_2=C*joint_constraint_matrix(joint_index_2);
            MATRIX_MXN<T> projected_constrained_C=joint_constraint_matrix(joint_index_1).Transpose_Times(C_times_joint_constraint_matrix_2);

            // build global_post_stabilization_matrix_11 which relates constraint impulses (post stabilization impulses) to constrained velocity directions
            int istart=joint_offset_in_post_stabilization_matrix(joint_index_1),jstart=joint_offset_in_post_stabilization_matrix(joint_index_2);
            global_post_stabilization_matrix_11.Add_To_Submatrix(istart,jstart,projected_constrained_C);
            if(joint_index_1!=joint_index_2) global_post_stabilization_matrix_11.Add_To_Submatrix(jstart,istart,projected_constrained_C.Transposed());}}
}
//####################################################################################
// Function Solve_Minimum_Norm_Solution_For_Linear_Constraints
//####################################################################################
template<class T> void ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
Solve_Minimum_Norm_Solution_For_Linear_Constraints(MATRIX_MXN<T>& A,const ARRAY<T>& b,ARRAY<T>& x,const T zero_row_tolerance,const bool verbose)
{
    assert(A.m==A.n);
    if(verbose){
        LOG::cout<<"*** SYSTEM BEFORE TRANSFORM: "<<std::endl;
        LOG::cout<<A<<std::endl<<b<<std::endl;}

    // get initial factorization
    MATRIX_MXN<T> A_after_QR(A);ARRAY<T> b_after_QR(b);ARRAY<int> permute(A.n);
    A_after_QR.In_Place_Robust_Householder_QR_Solve(b_after_QR,permute);
    T threshold_squared=sqr(zero_row_tolerance);
    int equations_to_keep=0;
    for(int i=0;i<A_after_QR.Rows();i++){
        T row_norm_squared=0;
        for(int j=0;j<A_after_QR.Columns();j++) row_norm_squared+=sqr(A_after_QR(i,j));
        if(row_norm_squared>threshold_squared) equations_to_keep++;}

    if(verbose){
        LOG::cout<<"*** SYSTEM AFTER TRANSFORM: "<<std::endl;
        LOG::cout<<A_after_QR<<std::endl<<b_after_QR<<std::endl<<permute<<std::endl;}
    LOG::cout<<"Global poststabilization: got "<<equations_to_keep<<" dof"<<std::endl;

    MATRIX_MXN<T> B(equations_to_keep);A_after_QR.Get_Submatrix(0,0,B);

    if(equations_to_keep==A.n) x.Subset(permute)=B.Upper_Triangular_Solve(b_after_QR); // unique solution
    else{
        // extract submatrices and subvectors
        int epsilon_size=A.n-equations_to_keep;
        MATRIX_MXN<T> epsilon(epsilon_size);A_after_QR.Get_Submatrix(equations_to_keep,equations_to_keep,epsilon);
        MATRIX_MXN<T> S(equations_to_keep,epsilon_size);A_after_QR.Get_Submatrix(0,equations_to_keep,S);
        ARRAY<T> b_B(equations_to_keep),b_epsilon(epsilon_size);
        b_after_QR.Extract(b_B,b_epsilon);
        // construct the matrix
        MATRIX_MXN<T> B_inverse_S=B.Upper_Triangular_Solve(S);
        MATRIX_MXN<T> C(epsilon_size);
        MATRIX_MXN<T>::Add_Transpose_Times(B_inverse_S,B_inverse_S,C);
        MATRIX_MXN<T>::Add_Transpose_Times(epsilon,epsilon,C);
        C+=1;
        // construct the rhs
        ARRAY<T> x_B(B.Upper_Triangular_Solve(b_B));
        ARRAY<T> rhs(epsilon_size);MATRIX_MXN<T>::Add_Transpose_Times(epsilon,b_epsilon,rhs);MATRIX_MXN<T>::Add_Transpose_Times(B_inverse_S,x_B,rhs);
        // compute solution vector 
        ARRAY<T> x_S=C.Cholesky_Solve(rhs);
        MATRIX_MXN<T>::Subtract_Times(S,x_S,b_B);
        x_B=B.Upper_Triangular_Solve(b_B);
        x.Resize(A.n);
        x.Subset(permute).Combine(x_B,x_S);}

    if(verbose) LOG::cout<<"Result of linearly constrained optimization:\n"<<x<<std::endl;
}
//####################################################################################
// Function Solve_Velocities_for_PD
//####################################################################################
template<class T> void ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
Solve_Velocities_for_PD(const T time,const T dt,bool test_system,bool print_matrix)
{
    assert(global_post_stabilization || !use_pd_actuators); // must use global post stabilization for pd actuation

    if(global_post_stabilization && global_post_stabilization_matrix_11.n+global_post_stabilization_matrix_21.m){
        // build constrained delta relative joint velocities (also get joint locations)
        ARRAY<T> constrained_delta_relative_joint_velocities(global_post_stabilization_matrix_11.n),joint_constraint_impulses;
        ARRAY<TV> joint_locations(joint_mesh.Num_Joints());
        for(int i=0;i<joint_mesh.Num_Joints();i++) if(joint_constrained_dimensions(i)){
            TWIST<TV> delta_relative_twist;Delta_Relative_Twist(joint_mesh.Joints(i)->id_number,true,joint_locations(i),delta_relative_twist);
            if(joint_constrained_dimensions(i)){
                assert(joint_constrained_dimensions(i)>=3); // assume unconstrained linear directions, constrained angular directions
                int istart=joint_offset_in_post_stabilization_matrix(i);
                constrained_delta_relative_joint_velocities.Array_View(istart,(int)TV::m)=delta_relative_twist.linear;
                for(int ii=0;ii<joint_constrained_dimensions(i)-3;ii++){
                    TV u(joint_angular_constraint_matrix(i).Get_Column(ii));
                    constrained_delta_relative_joint_velocities(istart+ii+2)=u.Dot(delta_relative_twist.angular);}}}

        Solve_Minimum_Norm_Solution_For_Linear_Constraints(global_post_stabilization_matrix_11,constrained_delta_relative_joint_velocities,
            joint_constraint_impulses,zero_row_tolerance,verbose);

        // apply poststabilization impulses
        for(int i=0;i<joint_mesh.Num_Joints();i++) if(joint_constrained_dimensions(i)){
            assert(joint_constrained_dimensions(i)>=3); // unconstrained linear directions, constrained angular directions
            int istart=joint_offset_in_post_stabilization_matrix(i);
            TV impulse(joint_constraint_impulses.Array_View(istart,(int)TV::m));
            TV angular_impulse;for(int ii=0;ii<joint_constrained_dimensions(i)-3;ii++){
                TV u(joint_angular_constraint_matrix(i).Get_Column(ii));
                angular_impulse+=joint_constraint_impulses(istart+ii+2)*u;}
            RIGID_BODY<TV>::Apply_Impulse(*Parent(joint_mesh.Joints(i)->id_number),*Child(joint_mesh.Joints(i)->id_number),joint_locations(i),impulse,angular_impulse);}}

    // the second flag will cause poststabilization to be applied to joints not processed by global post stabilization
    Apply_Poststabilization(use_pd_actuators,global_post_stabilization);

    if(use_angular_damping){
        // precompute damped stuff
        precomputed_desired_damped_angular_velocities.Resize(joint_mesh.Num_Joints());
        int njoints=0;
        for(int i=0;i<joint_mesh.Num_Joints();i++) if(joint_mesh.Joints(i)->angular_damping){
            njoints++;JOINT<TV>& joint=*joint_mesh.Joints(i);RIGID_BODY<TV> &parent=*Parent(joint.id_number),&child=*Child(joint.id_number);
            parent.Update_Angular_Velocity();child.Update_Angular_Velocity();
            precomputed_desired_damped_angular_velocities(i)=(exp(-joint.angular_damping*dt)-1)*RIGID_BODY<TV>::Relative_Angular_Velocity(parent,child);}
            LOG::cout<<"Angular damping of "<<njoints<<" joints"<<std::endl;
        Apply_Poststabilization(use_pd_actuators,false,true);}
}
//####################################################################################
// Function Create_Joint_Function
//####################################################################################
template<class T> JOINT_FUNCTION<VECTOR<T,3> >* ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
Create_Joint_Function(const JOINT_ID joint_id)
{
    joint_mesh(joint_id)->Set_Joint_Function(new JOINT_FUNCTION<TV>(*joint_mesh(joint_id),*Parent(joint_id),*Child(joint_id)));
    return joint_mesh(joint_id)->joint_function;
}
//####################################################################################
// Function Output_Articulation_Points
//####################################################################################
template<class T> void ARTICULATED_RIGID_BODY<VECTOR<T,3> >::
Output_Articulation_Points(const STREAM_TYPE stream_type,const VIEWER_DIR& viewer_dir) const
{
    if(joint_mesh.Num_Joints()==0) return;
    FILE_OSTREAM output(stream_type,viewer_dir.Frame_File("arb_info"));
    Write_Binary(output,joint_mesh.Num_Joints()*2);
    for(int i=0;i<joint_mesh.Num_Joints();i++){
        JOINT<TV>& joint=*joint_mesh.Joints(i);
        const RIGID_BODY<TV> &parent=*Parent(joint.id_number),&child=*Child(joint.id_number);
        TV ap1=parent.World_Space_Point(joint.F_pj().t),ap2=child.World_Space_Point(joint.F_cj().t);
        Write_Binary(output,ap1,parent.Frame()*joint.F_pj(),ap2,child.Frame()*joint.F_pj());}
}
//####################################################################################
namespace PhysBAM{
template class ARTICULATED_RIGID_BODY<VECTOR<float,3> >;
template class ARTICULATED_RIGID_BODY<VECTOR<double,3> >;
}
