//#####################################################################
// Copyright 2015, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Hybrid_Methods/Forces/MPM_FINITE_ELEMENTS.h>
#include <Core/Arrays/ARRAYS_ND.h>
#include <Core/Log/FINE_TIMER.h>
#include <Core/Matrices/DIAGONAL_MATRIX.h>
#include <Deformables/Constitutive_Models/DIAGONALIZED_ISOTROPIC_STRESS_DERIVATIVE.h>
#include <Deformables/Constitutive_Models/ISOTROPIC_CONSTITUTIVE_MODEL.h>
#include <Hybrid_Methods/Examples_And_Drivers/MPM_PARTICLES.h>
#include <Hybrid_Methods/Examples_And_Drivers/MPM_TRANSFERS.h>
#include <Hybrid_Methods/Forces/MPM_FORCE_HELPER.h>
#include <Hybrid_Methods/Iterators/GATHER_SCATTER.h>
namespace PhysBAM{
//#####################################################################
// Constructor
//#####################################################################
template<class TV> MPM_FINITE_ELEMENTS<TV>::
MPM_FINITE_ELEMENTS(const MPM_FORCE_HELPER<TV>& force_helper,
    ISOTROPIC_CONSTITUTIVE_MODEL<T,TV::m>& constitutive_model,
    GATHER_SCATTER<TV>& gather_scatter_input,ARRAY<int>* affected_particles)
    :BASE(force_helper),constitutive_model(constitutive_model),affect_all(!affected_particles),
    gather_scatter(affect_all?gather_scatter_input:*new GATHER_SCATTER<TV>(gather_scatter_input.grid,*new ARRAY<int>(*affected_particles)))
{
    PHYSBAM_ASSERT(affect_all);
    if(!affect_all){
        gather_scatter.weights=gather_scatter_input.weights;
        gather_scatter.threads=gather_scatter_input.threads;
        gather_scatter.Prepare_Scatter(particles);}
    constitutive_model.mu=&particles.mu;
    constitutive_model.lambda=&particles.lambda;
}
//#####################################################################
// Destructor
//#####################################################################
template<class TV> MPM_FINITE_ELEMENTS<TV>:: 
~MPM_FINITE_ELEMENTS()
{
    if(!affect_all){
        delete &gather_scatter.simulated_particles;
        delete &gather_scatter;}
    delete &constitutive_model;
}
//#####################################################################
// Function Precompute
//#####################################################################
template<class TV> void MPM_FINITE_ELEMENTS<TV>:: 
Precompute(const T time,const T dt,bool want_dE,bool want_ddE)
{
    TIMER_SCOPE_FUNC;
    PHYSBAM_ASSERT(!want_ddE||want_dE);
    U.Resize(particles.X.m);
    V.Resize(particles.X.m);
    sigma.Resize(particles.X.m);
    dPi_dF.Resize(particles.X.m);
    PP.Resize(particles.X.m);

    TIMER_SCOPE("Precompute-A");
#pragma omp parallel for
    for(int k=0;k<gather_scatter.simulated_particles.m;k++){
        int p=gather_scatter.simulated_particles(k);
        auto [U_,s_,V_]=particles.F(p).Singular_Value_Decomposition();
        U(p)=U_;
        sigma(p)=s_;
        V(p)=V_;
        PP(p)=U_*constitutive_model.P_From_Strain(s_,p).Times_Transpose(V_);
        constitutive_model.Isotropic_Stress_Derivative(s_,dPi_dF(p),p);}
}
//#####################################################################
// Function Potential_Energy
//#####################################################################
template<class TV> typename TV::SCALAR MPM_FINITE_ELEMENTS<TV>:: 
Potential_Energy(const T time) const
{
    T pe=0;
#pragma omp parallel for reduction(+:pe)
    for(int k=0;k<gather_scatter.simulated_particles.m;k++){
        int p=gather_scatter.simulated_particles(k);
        pe+=particles.volume(p)*constitutive_model.Energy_Density(sigma(p),p);}
    return pe;
}
//#####################################################################
// Function Add_Forces
//#####################################################################
template<class TV> void MPM_FINITE_ELEMENTS<TV>:: 
Add_Forces(REGION_ARRAY<TV,IV>& F,const T time) const
{
    TIMER_SCOPE_FUNC;
    auto [ng,np]=Transfer_G2P(force_helper.opt);
    gather_scatter.Scatter(ng,np,0,0,
        [this,&F](int p,PARTICLE_GRID_ITERATOR<TV>& it,int)
        {
            auto K=Times_Diff_G2P_F(force_helper.opt,force_helper.Fn(p),
                -particles.volume(p)*PP(p),force_helper.dt_grad_u(p));
            Scatter(force_helper.opt,it,F.u(particles.Region(p)),K);
        });
}
//#####################################################################
// Function Add_Hessian_Times
//#####################################################################
template<class TV> void MPM_FINITE_ELEMENTS<TV>:: 
Add_Hessian_Times(REGION_ARRAY<TV,IV>& F,const REGION_ARRAY<TV,IV>& Z,const T time) const
{
    tmp.Resize(particles.X.m);

    auto [ng,np]=Transfer_G2P(force_helper.opt);
    gather_scatter.Gather(ng,np,0,0,
        [this,&Z](int p,PARTICLE_GRID_ITERATOR<TV>& it,int data)
        {
            auto K=Gather(force_helper.opt,it,Z.u(particles.Region(p)));
            auto dF=Diff_Times_G2P_F(force_helper.opt,force_helper.Fn(p),
                force_helper.dt_grad_u(p),K);
            auto G=U(p).Transpose_Times(dF*V(p));
            auto P=U(p)*dPi_dF(p).Differential(G).Times_Transpose(V(p));
            auto L=Times_Diff_G2P_F(force_helper.opt,force_helper.Fn(p),
                P,force_helper.dt_grad_u(p));
            L+=Times_Hess_Times_G2P_F(force_helper.opt,force_helper.Fn(p),PP(p),
                force_helper.dt_grad_u(p),K);
            tmp(p)=particles.volume(p)*L;
        });

    gather_scatter.Scatter(ng,np,0,0,
        [this,&F](int p,PARTICLE_GRID_ITERATOR<TV>& it,int data)
        {
            Scatter(force_helper.opt,it,F.u(particles.Region(p)),tmp(p));
        });
}
template class MPM_FINITE_ELEMENTS<VECTOR<float,1> >;
template class MPM_FINITE_ELEMENTS<VECTOR<float,2> >;
template class MPM_FINITE_ELEMENTS<VECTOR<float,3> >;
template class MPM_FINITE_ELEMENTS<VECTOR<double,1> >;
template class MPM_FINITE_ELEMENTS<VECTOR<double,2> >;
template class MPM_FINITE_ELEMENTS<VECTOR<double,3> >;
}
