//#####################################################################
// 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_OLDROYD_FINITE_ELEMENTS.h>
#include <Core/Arrays/ARRAYS_ND.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_OLDROYD_FINITE_ELEMENTS<TV>::
MPM_OLDROYD_FINITE_ELEMENTS(MPM_FORCE_HELPER<TV>& force_helper,
    OLDROYD_CONSTITUTIVE_MODEL<TV>& constitutive_model,
    GATHER_SCATTER<TV>& gather_scatter_input,ARRAY<int>* affected_particles,
    const MPM_TRANSFER_OPTIONS<T>& opt)
    :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))),stored_dt(0),opt(opt)
{
    if(!affect_all){
        gather_scatter.weights=gather_scatter_input.weights;
        gather_scatter.threads=gather_scatter_input.threads;
        gather_scatter.Prepare_Scatter(particles);}
}
//#####################################################################
// Destructor
//#####################################################################
template<class TV> MPM_OLDROYD_FINITE_ELEMENTS<TV>:: 
~MPM_OLDROYD_FINITE_ELEMENTS()
{
    if(!affect_all){
        delete &gather_scatter.simulated_particles;
        delete &gather_scatter;}
    delete &constitutive_model;
}
//#####################################################################
// Function Precompute
//#####################################################################
template<class TV> void MPM_OLDROYD_FINITE_ELEMENTS<TV>:: 
Precompute(const T time,const T dt,bool want_dE,bool want_ddE)
{
    PHYSBAM_ASSERT(particles.store_S);
    stored_dt=dt;
    constitutive_model.Resize(particles.number);
#pragma omp parallel for
    for(int k=0;k<gather_scatter.simulated_particles.m;k++){
        int p=gather_scatter.simulated_particles(k);
        constitutive_model.Precompute(particles.F(p),particles.S(p),p);}
}
//#####################################################################
// Function Potential_Energy
//#####################################################################
template<class TV> typename TV::SCALAR MPM_OLDROYD_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(p);}
    return pe;
}
//#####################################################################
// Function Add_Forces
//#####################################################################
template<class TV> void MPM_OLDROYD_FINITE_ELEMENTS<TV>:: 
Add_Forces(REGION_ARRAY<TV,IV>& F,const T time) const
{
    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)
        {
            SYMMETRIC_MATRIX<T,TV::m> dQ,S=force_helper.Sn(p);
            MATRIX<T,TV::m> dP,Fn=force_helper.Fn(p),B=force_helper.dt_grad_u(p);
            constitutive_model.Gradient(dP,dQ,p);
            auto& f=F.u(particles.Region(p));
            auto K=Times_Diff_G2P_F(opt,Fn,-particles.volume(p)*dP,B);
            K+=Times_Diff_G2P_S(opt,S,-particles.volume(p)*dQ,B,stored_dt);
            Scatter(opt,it,f,K);
        });
}
//#####################################################################
// Function Add_Hessian_Times
//#####################################################################
template<class TV> void MPM_OLDROYD_FINITE_ELEMENTS<TV>:: 
Add_Hessian_Times(REGION_ARRAY<TV,IV>& F,const REGION_ARRAY<TV,IV>& V,const T time) const
{
    tmp.Resize(particles.X.m);

    //T zeta=1/(1+stored_dt*force_helper.opt.inv_Wi);
    //T c=force_helper.opt.use_quad_F*(T).5;
    auto [ng,np]=Transfer_G2P(force_helper.opt);
    gather_scatter.Gather(ng,np,0,0,
        [this,&V](int p,PARTICLE_GRID_ITERATOR<TV>& it,int data)
        {
            SYMMETRIC_MATRIX<T,TV::m> S=force_helper.Sn(p),G,dQ;
            MATRIX<T,TV::m> F=force_helper.Fn(p),B=force_helper.dt_grad_u(p),L,dP;
            constitutive_model.Gradient(dP,dQ,p);
            auto K=Gather(force_helper.opt,it,V.u(particles.Region(p)));
            auto D=Diff_Times_G2P_F(force_helper.opt,F,B,K);
            auto E=Diff_Times_G2P_S(force_helper.opt,S,B,K,stored_dt);
            constitutive_model.Hessian(D,E,L,G,p);
            auto Q=Times_Diff_G2P_F(force_helper.opt,F,L,B);
            Q+=Times_Hess_Times_G2P_F(force_helper.opt,F,dP,B,K);
            Q+=Times_Diff_G2P_S(force_helper.opt,S,G,B,stored_dt);
            Q+=Times_Hess_Times_G2P_S(force_helper.opt,S,dQ,B,K,stored_dt);
            tmp(p)=Q;
        });

    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_OLDROYD_FINITE_ELEMENTS<VECTOR<float,2> >;
template class MPM_OLDROYD_FINITE_ELEMENTS<VECTOR<float,3> >;
template class MPM_OLDROYD_FINITE_ELEMENTS<VECTOR<double,2> >;
template class MPM_OLDROYD_FINITE_ELEMENTS<VECTOR<double,3> >;
}
