//#####################################################################
// Copyright 2004-2008, Ron Fedkiw, Geoffrey Irving, Andrew Selle.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Log/LOG.h>
#include <Core/Log/DEBUG_SUBSTEPS.h>
#include "MPM_BC.h"
using namespace PhysBAM;

//#####################################################################
// Constructor
//#####################################################################
template<class TV> BOUNDARY_CONDITION<TV>::
BOUNDARY_CONDITION(REGION_ARRAY<T,IV>& reflect_mass,REGION_ARRAY<TV,IV>& reflect_momentum,VECTOR<MPM::BC_TYPE,TV::m*2>& side_bc_type,
    bool test_bc,T& friction,GRID<TV> &grid, int& ghost,VECTOR<bool,TV::m>& is_periodic)
    :reflect_mass(reflect_mass),reflect_momentum(reflect_momentum),side_bc_type(side_bc_type),test_bc(test_bc),friction(friction),grid(grid),
    ghost(ghost),is_periodic(is_periodic)
{
}
//#####################################################################
// Constructor
//#####################################################################
template<class TV> BOUNDARY_CONDITION<TV>::
~BOUNDARY_CONDITION()
{}
//#####################################################################
// Function Reflect_Helper
//#####################################################################
template<class TV,class T,class TYPE>
pair<TV,TV> Reflect_Helper(T m_i,T m_o,TV v_i,TV v_o,TYPE type,T mu,int a,int s)
{
    T m_t=m_i+m_o, a_i=m_i/m_t,a_o=1-a_i;
    TV v_in=a_i*v_i+a_o*v_o,v_on=v_in;
    if(type==MPM::BC_FREE) return {v_in,v_on};
    TV v_id=a_i*v_i+a_o*(-v_o),v_od=a_o*v_o+a_i*(-v_i);
    if(type==MPM::BC_NOSLIP) {
        return {v_id,v_od};

    }

    if(type==MPM::BC_SLIP){
        v_in(a)=v_id(a);
        v_on(a)=v_od(a);
        return {v_in,v_on};}

    // LOG::cout<<"mu="<<mu<<std::endl;
    TV v_id_n,v_od_n;
    v_id_n(a)=v_id(a);
    v_od_n(a)=v_od(a);
    v_in(a)=0;

    //Mehrnaz
    T d_v=2*a_o*a_i*(-v_i(a)-v_o(a));
    T d_v_n=s?-d_v:d_v;
    if(d_v_n<0) {//separating
        //slip: tangential unchanged, normal zero out(on boundary)
        //tangential v_in
        TV vi_f=v_in;
        TV vo_f=v_on;
        vi_f(a)=v_i(a);
        vo_f(a)=v_o(a);
        return {vi_f,vo_f};
    }

    T d_v_t_mag=v_in.Normalize();
    T r=d_v_t_mag-mu*d_v_n;
    if(r<0) return {v_id_n,v_od_n};

    return {v_id_n+r*v_in,v_od_n+r*v_in};
}

template<class TV,class T,class TYPE>
TV Reflect_Helper(TV v,TYPE type,T mu,int a,int s)
{
    TV bc;
    if(type==MPM::BC_FREE) return v;
    if(type==MPM::BC_NOSLIP) return bc;
    if(type==MPM::BC_SLIP){
        v(a)=bc(a);
        return v;}

    T d_v=-v(a);
    T d_v_n=s?-d_v:d_v;
    if(d_v_n<0) return v;

    v(a)=0;

    T d_v_t_mag=v.Normalize();
    T r=d_v_t_mag-mu*d_v_n;
    if(r<0) return bc;
    return r*v;
}
//#####################################################################
// Function Reflect_Momentum_Helper
//#####################################################################
template<class TV,class T,class TYPE>
pair<TV,TV> Reflect_Momentum_Helper(TV mv_i,TV mv_o,TYPE type,T mu,int a,int s)
{
    TV mv_in=mv_i+mv_o,mv_on=mv_in;
    TV mv_id=mv_i+(-mv_o),mv_od=mv_o+(-mv_i);
    if(type==MPM::BC_NOSLIP) {
        return {mv_id,mv_od};
    }
    // LOG::cout<<"reflecting momentum"<<std::endl;
    mv_in(a)=mv_id(a);
    mv_on(a)=mv_od(a);
    return {mv_in,mv_on};
}

//#####################################################################
// Function Reflect_Boundary_Friction
//#####################################################################
template <class TV> void BOUNDARY_CONDITION<TV>::
Reflect_Boundary_Friction(const REGION_ARRAY<T,IV>& mass,REGION_ARRAY<TV,IV>& velocity)
{
    for(REGION_FLAT_ID i=REGION_FLAT_ID(0);i<mass.Flat_Size();i++)
    {
        reflect_mass(i)=mass(i);
        reflect_momentum(i)=velocity(i)*mass(i);
    }
    PHYSBAM_DEBUG_WRITE_SUBSTEP("updating reflect mass",4);

    // TIMER_SCOPE_FUNC;
        Visit_Nodes(RI::duplicate_corners,mass,
        [this,&mass,&velocity](int side,const auto& out,const auto& in)
        {
            auto bc_type=side_bc_type(side);
            T m_i=mass(in),m_o=mass(out);
            if(!m_i && !m_o) return;
            TV v_i,v_o;
            if(m_i) {
                v_i=velocity(in);
                if(test_bc)
                    bc_test_velocity(in)=v_i;
            }
            if(m_o) {
                v_o=velocity(out);
                if(test_bc)
                    bc_test_velocity(out)=v_o;
            }
            if(bc_type!=MPM::BC_SLIP && bc_type!=MPM::BC_NOSLIP){
                auto pr=Reflect_Helper(m_i,m_o,v_i,v_o,bc_type,friction,side/2,side%2);
                if(m_i) {
                velocity(in)=pr.first;
                }
                if(m_o) {
                    velocity(out)=pr.second;
                }
                return;
            }
            auto pr=Reflect_Momentum_Helper(reflect_momentum(in),reflect_momentum(out),bc_type,friction,side/2,side%2);
            T tmp=reflect_mass(in)+reflect_mass(out);
            if(m_i || m_o) {
                reflect_momentum(in)=pr.first;
                reflect_mass(in)=tmp;
                reflect_momentum(out)=pr.second;
                reflect_mass(out)=tmp;
            }
        },
        [this,&mass,&velocity](int side,const auto& out)
        {
            auto bc_type=side_bc_type(side);
            T m=mass(out);
            if(!m) return;
            TV &v_o=velocity(out);
            if(test_bc)
                bc_test_velocity(out)=v_o;
            v_o=Reflect_Helper(v_o,bc_type,friction,side/2,side%2);
    });

    Visit_Nodes(RI::duplicate_corners,mass,
        [this,&mass,&velocity](int side,const auto& out,const auto& in)
        {
            auto bc_type=side_bc_type(side);
            if(bc_type==MPM::BC_FREE || bc_type==MPM::BC_SEP){
                return;
            }
            // LOG::cout<<"going over index in= "<<in.i<<" ,out= "<<out.i<<std::endl;
            T m_i=mass(in),m_o=mass(out);
            if(!m_i && !m_o) return;
            TV v_i,v_o;
            if(m_i) {
                velocity(in)=reflect_momentum(in)/reflect_mass(in);
            }
            if(m_o) {
                velocity(out)=reflect_momentum(out)/reflect_mass(out);;
            }
        },
        [this,&mass,&velocity](int side,const auto& out)
        {
            auto bc_type=side_bc_type(side);
            T m=mass(out);
            if(!m) return;
            // TV &v_o=velocity(out);
            TV &v_o=reflect_momentum(out);
            if(test_bc)
                bc_test_velocity(out)=v_o;
            v_o=Reflect_Helper(v_o,bc_type,friction,side/2,side%2);
    });
    PHYSBAM_DEBUG_WRITE_SUBSTEP("velocities after boundary collisions",3);
}
//#####################################################################
// Function Test_Boundary_Friction
//#####################################################################
template <class TV> void BOUNDARY_CONDITION<TV>::
Test_Reflect_Boundary(const REGION_ARRAY<T,IV>& mass,REGION_ARRAY<TV,IV>& v0,REGION_ARRAY<TV,IV>& v1)
{
    bool passed=true;
    Visit_Nodes(RI::duplicate_corners,mass,
        [this,&mass,&v0,&v1,&passed](int side,const auto& out,const auto& in)
        {
            auto bc_type=side_bc_type(side);
            T m_i=mass(in),m_o=mass(out);
            if(!m_i && !m_o) return;
            TV v1_i,v1_o,v0_i,v0_o;
            if(m_i) {v1_i=v1(in); v0_i=v0(in);}
            if(m_o) {v1_o=v1(out); v0_o=v0(out);}
            TV bc;
            if(m_i && m_o){
                if(bc_type==MPM::BC_FREE){
                    if((v1_i-v0_i).Magnitude()>1e-7){
                        passed=false;
                        LOG::cout<<"[bc test NOT passed FREE]v_i!=v_in_past"<<std::endl;
                        LOG::cout<<"[bc test NOT passed FREE]velocity of v_i(velocity( "<<in.i<<")="<<v1_i <<"v_i_past ="<<v0_i<<std::endl;
                    }
                    if((v1_o-v0_o).Magnitude()>1e-7){
                        passed=false;
                        LOG::cout<<"[bc test NOT passed FREE]v_o!=v_out_past"<<std::endl;
                        LOG::cout<<"[bc test NOT passed FREE]velocity of v_o(velocity( "<<out.i<<")="<<v1_o <<"v_o_past ="<<v0_o<<std::endl;
                    }
                }
            }
        },
         [this,&mass,&v1,&v0,&passed](int side,const auto& out)
        {
            auto type=side_bc_type(side);
            T m=mass(out);
            if(!m) return;
            TV &v1_o=v1(out);
            TV &v0_o=v0(out);
            TV bc;
            // int a=side/2;
            if(type==MPM::BC_FREE){
                if((v1_o-v0_o).Magnitude()>1e-7){
                    passed=false;
                    LOG::cout<<"[bc test FREE]v_o_boudary!=v_out_past"<<std::endl;
                    LOG::cout<<"[bc test FREE]velocity of v_o_boudanry(velocity( "<<out.i<<")="<<v1_o <<"v_o_past ="<<v0_o<<std::endl;
                }
            } 

        });
}
//#####################################################################
// Function Test_Boundary_Friction
//#####################################################################
template <class TV> void BOUNDARY_CONDITION<TV>::
Test_Boundary_Friction(const REGION_ARRAY<T,IV>& mass,REGION_ARRAY<TV,IV>& velocity)
{
    test_bc_visit_nodes=true;
    LOG::cout<<"running boundary condition tests"<<std::endl;
    bool passed=true;
    bool printed=false;
    static const char* bc_type_name[6] = {
        "BC_FREE",
        "BC_SLIP",
        "BC_NOSLIP",
        "BC_PERIODIC",
        "BC_SEP",
        "BC_MIXED"
    };
    Visit_Nodes(RI::duplicate_corners,mass,
        [this,&mass,&velocity,&passed,&printed](int side,const auto& out,const auto& in)
        {
            auto bc_type=side_bc_type(side);
            if(!printed){
                LOG::cout<<"[bc test] testing side "<<side<<" with bc type "<<bc_type_name[bc_type]<<std::endl;
                printed=true;
            }
            T m_i=mass(in),m_o=mass(out);
            if(!m_i && !m_o) return;
            TV v_i,v_o;
            if(m_i) v_i=velocity(in);
            if(m_o) v_o=velocity(out);
            TV bc;
            if(test_bc && m_i && m_o){
                if(bc_type==MPM::BC_FREE){
                    if((velocity(in)-bc_test_velocity(in)).Magnitude()>1e-7){
                        passed=false;
                        LOG::cout<<"[bc test NOT passed FREE]v_i!=v_in_past"<<std::endl;
                        LOG::cout<<"[bc test NOT passed FREE]velocity of v_i(velocity( "<<in.i<<")="<<v_i <<"v_i_past ="<<bc_test_velocity(in)<<std::endl;
                    }
                    if((velocity(out)-bc_test_velocity(out)).Magnitude()>1e-7){
                        passed=false;
                        LOG::cout<<"[bc test NOT passed FREE]v_o!=v_out_past"<<std::endl;
                        LOG::cout<<"[bc test NOT passed FREE]velocity of v_o(velocity( "<<out.i<<")="<<v_o <<"v_o_past ="<<bc_test_velocity(out)<<std::endl;
                    }
                }
                else if(bc_type==MPM::BC_SLIP){
                    if(velocity(in).Magnitude()-velocity(out).Magnitude()>1e-5){//magni
                        passed=false;
                        LOG::cout<<"[bc test NOT passed SLIP]v_i!=v_out diff is "<<velocity(in).Magnitude()-velocity(out).Magnitude()<<std::endl;
                        LOG::cout<<"[bc test NOT passed SLIP]velocity of v_i(velocity( "<<in.i<<")="<<v_i <<"v_o (velocity("<<out.i<<"))="<<v_o<<std::endl;
                        LOG::cout<<"[bc test NOT passed SLIP]side = "<<side<<std::endl;
                    }
                    if((velocity(in)(int(side/2))+velocity(out)(int(side/2)))>1e-3){
                        passed=false;
                        LOG::cout<<"[bc test NOT passed SLIP] normal elemets are wrong -- sum = "<<(velocity(in)(side/2)+velocity(out)(side/2)-2*bc).Magnitude()<<std::endl;
                        LOG::cout<<"[bc test NOT passed SLIP] normal elemets are wrong -- velocity out normal"<<out.i<< "= "<<velocity(out)(side/2)<<std::endl;
                        LOG::cout<<"[bc test NOT passed SLIP] normal elemets are wrong -- velocity in normal" <<in.i<<" = "<<velocity(in)(side/2)<<std::endl;
                    }
                }else if(bc_type==MPM::BC_NOSLIP){
                    if(velocity(in).Magnitude()-velocity(out).Magnitude()>1e-3){//magni
                        passed=false;
                        LOG::cout<<"[bc test NOT passed STICK]v_i!=v_out -- diff is "<<velocity(in).Magnitude()-velocity(out).Magnitude()<<std::endl;
                        LOG::cout<<"[bc test NOT passed STICK]velocity of v_i(velocity( "<<in.i<<")="<<v_i <<"v_o (velocity("<<out.i<<"))="<<v_o<<std::endl;
                        LOG::cout<<"[bc test NOT passed STICK]side = "<<side<<std::endl;
                    }
                    if((velocity(in)+velocity(out)).Magnitude()>1e-3){
                        passed=false;
                        LOG::cout<<"[bc test NOT passed STICK] sums do not match "<<(velocity(in)+velocity(out)).Magnitude()<<std::endl;
                        LOG::cout<<"[bc test NOT passed STICK] velocity of v_i(velocity( "<<in.i<<")="<<v_i <<"v_o (velocity("<<out.i<<"))="<<v_o<<std::endl;
                    }
                }else if(bc_type==MPM::BC_SEP){
                    
                }

                if (bc_type==MPM::BC_PERIODIC) {
                    REGION_INDEX<IV::m> expected_index = out;
                    int axis = side / 2;
                    int shift = grid.numbers_of_cells(axis);
                    int cells = grid.numbers_of_cells(axis); // total number of cells in this axis
                    expected_index.i(axis) = (side % 2 == 0) 
                        ? (out.i(axis) + shift) 
                        : (out.i(axis)-shift);
                    if (!grid.Valid_Node(expected_index.i)){
                        LOG::cout<<"[bc test NOT passed PERIODIC] Wrapped index out of bounds: "<<expected_index.i<<std::endl;
                        LOG::cout<<"[bc test NOT passed PERIODIC] out index: "<<out.i<<std::endl;
                        LOG::cout<<"[bc test NOT passed PERIODIC] in index: "<<in.i<<std::endl;
                        LOG::cout<<"[bc test NOT passed PERIODIC] cells: "<<cells<<std::endl;
                        return;
                    }

                    if(velocity.Valid_Index(expected_index)){
                        TV wrapped_value=velocity(expected_index);
                        TV ghost_value=velocity(out);
                        if ((wrapped_value-ghost_value).Magnitude()>1e-7){
                            passed=false;
                            LOG::cout<<"[bc test NOT passed PERIODIC] Velocity mismatch at periodic boundary (side = "<<side<<")"<<std::endl;
                            LOG::cout<<"[bc test NOT passed PERIODIC] velocity(ghost = "<<out.i<<") = "<<ghost_value 
                                    <<", expected = velocity(wrapped = "<<expected_index.i<<") = "<<wrapped_value
                                    <<" cells= "<<cells<<std::endl;
                        }
                    }
                       
                }

            }
        },
    //     //TODO check to see if it's pointing to the right direction??
        [this,&mass,&velocity,&passed](int side,const auto& out)
        {
            auto type=side_bc_type(side);
            T m=mass(out);
            if(!m || !test_bc) return;
            TV &v_o=velocity(out);
            TV bc;
            int a=side/2;
            // if(type==MPM::BC_FREE){
    //             if((velocity(out)-example.bc_test_velocity(out)).Magnitude()>1e-7){
    //                 passed=false;
    //                 LOG::cout<<"[bc test FREE]v_o_boudary!=v_out_past"<<std::endl;
    //                 LOG::cout<<"[bc test FREE]velocity of v_o_boudanry(velocity( "<<out.i<<")="<<v_o <<"v_o_past ="<<example.bc_test_velocity(out)<<std::endl;
    //             }
    //         } 
            if(type==MPM::BC_NOSLIP){
                if((v_o-bc).Magnitude()>1e-7) {
                    passed=false;
                    LOG::cout<<"[bc test NOT passed STICK] v boundary!=bc -- diff is "<<(v_o-bc).Magnitude()<<std::endl;
                    LOG::cout<<"[bc test NOT passed STICK] v boundary!=bc -- vo= "<<out.i<<" value= "<<v_o<<std::endl;
                    LOG::cout<<"[bc test NOT passed STICK] v boundary!=bc -- bc= "<<bc<<std::endl;
                }
            } 
            if(type==MPM::BC_SLIP){
                if(v_o(a)-bc(0)>1e-7) {
                    passed=false;
                    LOG::cout<<"[bc test NOT passed SLIP] v boundary!=bc -- diff is "<<v_o(a)-bc(0)<<std::endl;
                    LOG::cout<<"[bc test NOT passed SLIP] v boundary!=bc -- v_o"<<out.i<<"is "<<v_o(a)<<std::endl;
                    LOG::cout<<"[bc test NOT passed SLIP] v boundary!=bc -- bc is "<<bc(a)<<std::endl;
                }
            }
            if (type==MPM::BC_PERIODIC) {
                    REGION_INDEX<IV::m> expected_index = out;
                    int axis = side / 2;
                    int shift = grid.numbers_of_cells(axis);
                    int cells = grid.numbers_of_cells(axis); // total number of cells in this axis
                    expected_index.i(axis) = (side % 2 == 0) 
                        ? (out.i(axis) + shift) 
                        : (out.i(axis)-shift);
                    if (!grid.Valid_Node(expected_index.i)){
                        LOG::cout<<"[bc test NOT passed PERIODIC] boundary velocity Wrapped index out of bounds: "<<expected_index.i<<std::endl;
                        LOG::cout<<"[bc test NOT passed PERIODIC] boundary velocity out index: "<<out.i<<std::endl;
                        LOG::cout<<"[bc test NOT passed PERIODIC] boundary velocity cells: "<<cells<<std::endl;
                        return;
                    }

                    if(velocity.Valid_Index(expected_index)){
                        TV wrapped_value=velocity(expected_index);
                        TV ghost_value=velocity(out);
                        if ((wrapped_value-ghost_value).Magnitude()>1e-7){
                            passed=false;
                            LOG::cout<<"[bc test NOT passed PERIODIC] boundary Velocity mismatch at periodic boundary (side = "<<side<<")"<<std::endl;
                            LOG::cout<<"[bc test NOT passed PERIODIC] boundary velocity(ghost = "<<out.i<<") = "<<ghost_value 
                                    <<", expected = velocity(wrapped = "<<expected_index.i<<") = "<<wrapped_value
                                    <<" cells= "<<cells<<std::endl;
                        }
                    }
                }
    //         //TODO MEHRNAZ I don't know how check sep
        });
    test_bc_visit_nodes=false;
}
//#####################################################################
// Function Visit_Nodes
//#####################################################################
template<class TV> void BOUNDARY_CONDITION<TV>::
Visit_Nodes(RI flags,const REGION_ARRAY<T,IV>& a,
    const std::function<void(int side,const REGION_INDEX<IV::m>& in,
        const REGION_INDEX<IV::m>& out)>& f,
    const std::function<void(int side,const REGION_INDEX<IV::m>& in)>& bf) const
{
    IV ranges[2]={IV(),grid.numbers_of_cells};
    int side_mask=0;
    if(test_bc_visit_nodes){
        for(int i=0;i<IV::m;i++)
            side_mask|=3<<(2*i);
    }else{
        for(int i=0;i<IV::m;i++)
            if(side_bc_type(2*i)!=MPM::BC_PERIODIC)
                side_mask|=3<<(2*i);
    }
   
    int prev_side=0;
    for(RANGE_ITERATOR<IV::m> it(grid.Node_Indices(),ghost,0,
            RI::ghost|RI::side_mask|flags,side_mask);it.Valid();it.Next_Side())
    {
        if(prev_side!=it.side){
            prev_side=it.side;
            PHYSBAM_DEBUG_WRITE_SUBSTEP("changing side to %i",4,it.side);
        }
        for(REGION_ID r(0);r<a.u.m;r++)
        {
            RANGE<IV> range=a.u(r).domain.Intersect(it.Current_Region());
            if(range.Empty()) continue;
            for(RANGE_ITERATOR<IV::m> ri(range);ri.Valid();ri.Next())
            {
                // LOG::cout<<"goin over index "<<ri.index<<std::endl;
                REGION_INDEX<IV::m> in{r,ri.index},out=in;
                int axis=it.side/2,side=it.side%2;
                out.i(axis)=2*ranges[side](axis)-in.i(axis);
                if(!a.Valid_Index(out)) continue;
                for(int a=0;a<IV::m;a++){
                    if(a!=axis){
                        if(out.i(a)==0){
                            bf(2*a,out);
                        }else if(out.i(a)==ranges[1](a)){
                            bf(2*a+1,out);
                        }
                        if(in.i(a)==ranges[0](a)){
                            bf(2*a,in);
                        }else if(in.i(a)==ranges[1](a)){
                            bf(2*a+1,in);
                        }
                        
                    }
                }
                f(it.side,in,out);
            }
        }
    }
    
    for(RANGE_ITERATOR<IV::m> it(grid.Node_Indices(),0,-1,
        RI::ghost|RI::side_mask|flags,side_mask);it.Valid();it.Next_Side())
    {
        for(REGION_ID r(0);r<a.u.m;r++)
        {
            RANGE<IV> range=a.u(r).domain.Intersect(it.Current_Region());
            if(range.Empty()) continue;
            for(RANGE_ITERATOR<IV::m> ri(range);ri.Valid();ri.Next()){
                // LOG::cout<<"going over border index"<<ri.index<<std::endl;
                // if(test_bc_visit_nodes) LOG::cout<<"[bc test] going over index "<<ri.index<<std::endl;
                bf(it.side,REGION_INDEX<IV::m>{r,ri.index});
            }
        }
    }
    PHYSBAM_DEBUG_WRITE_SUBSTEP("done with the borders",4);

}
//#####################################################################
// Function Fix_Periodic_Accum
//#####################################################################
template<class TV> template<class T2,class OP> void BOUNDARY_CONDITION<TV>::
Fix_Periodic_Accum(ARRAY_VIEW<T2,IV> u,const OP& op) const
{
    if(!side_bc_type.Contains(MPM::BC_PERIODIC)) return;
    Fill_Ghost_Nodes_Periodic_Accum(grid,u,is_periodic,op);
}
//#####################################################################
// Function Fix_Periodic
//#####################################################################
template<class TV> template<class T2> void BOUNDARY_CONDITION<TV>::
Fix_Periodic(ARRAY_VIEW<T2,IV> u) const
{
    if(!side_bc_type.Contains(MPM::BC_PERIODIC)) return;
    Fill_Ghost_Nodes_Periodic(grid,u,is_periodic);
}
//#####################################################################
namespace PhysBAM{
template class BOUNDARY_CONDITION<VECTOR<float,1> >;
template class BOUNDARY_CONDITION<VECTOR<float,2> >;
template class BOUNDARY_CONDITION<VECTOR<float,3> >;
template class BOUNDARY_CONDITION<VECTOR<double,1> >;
template class BOUNDARY_CONDITION<VECTOR<double,2> >;
template class BOUNDARY_CONDITION<VECTOR<double,3> >;
}
