//#####################################################################
// Copyright 2025 Craig Schroeder
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Parallel_Computation/Partial_Sum.h>
#ifdef USE_OPENMP
#include <omp.h>
#endif
namespace PhysBAM{

template<class T>
void ps(T* x,int m)
{
    T z{};
    for(int i=0;i<m;i++)
        x[i]=(z+=x[i]);
}

template<class T>
void ps2(T* x,int m,int n,int s)
{
    for(int i=1;i<m;i++)
    {
        int l=s*i;
        T z=x[l]+x[l-s];
        x[l]=z;
        for(int j=1;j<n;j++)
        {
            int l=s*i+j;
            x[l]=(z+=x[l]+x[l-s]-x[l-s-1]);
        }
    }
}

template<class T>
void ps(T* x,int m,int n)
{
    ps(x,n);
    ps2(x,m,n,n);
}

template<class T>
void ps(T* x,int m,int n,int p)
{
    ps(x,p);
    ps2(x,n,p,p);
    int np=n*p;
    ps2(x,m,p,np);
    for(int i=1;i<m;i++)
        for(int j=1;j<n;j++)
        {
            int l=np*i+p*j;
            T z=x[l]+x[l-p]+x[l-np]-x[l-np-p];
            x[l]=z;
            for(int k=1;k<p;k++)
            {
                int l=np*i+p*j+k;
                x[l]=(z+=x[l]+x[l-p]+x[l-np]-x[l-np-p]-x[l-np-1]-x[l-p-1]+x[l-np-p-1]);
            }
        }
}

template<class T>
void Partial_Sum(ARRAY_VIEW<T> x)
{
    ps(x.Get_Array_Pointer(),x.m);
}

template<class T,int d>
void Partial_Sum(ARRAYS_ND_BASE<T,VECTOR<int,d>>& x)
{
    static_assert(d>0 && d<=3);
    T* p=x.Get_Array_Pointer();
    auto s=x.domain.Edge_Lengths();
    if constexpr(d==1) ps(p,s.x);
    else if constexpr(d==2) ps(p,s.x,s.y);
    else ps(p,s.x,s.y,s.z);
}

#ifdef USE_OPENMP

template<class T>
void Partial_Sum_Parallel(ARRAY_VIEW<T> x)
{
    static constexpr int min_per_thread=8;
#pragma omp parallel
    {
        int threads=omp_get_num_threads();
        if(x.m<min_per_thread*threads)
        {
#pragma omp single
            Partial_Sum(x);
        }
        else
        {
            int tid=omp_get_thread_num();
            int a=tid*x.m/threads;
            int b=(tid+1)*x.m/threads;
            T z=x(a+1);
            for(int i=a+2;i<b;i++)
                x(i)=(z+=x(i));
#pragma omp barrier
#pragma omp single
            {
                T z=x(0);
                for(int i=1;i<threads;i++)
                {
                    int a=i*x.m/threads;
                    x(a)=(z+=x(a-1)+x(a));
                }
            }
#pragma omp barrier
            z=x(a);
            for(int i=a+1;i<b;i++)
                x(i)+=z;
        }
    }    
}

template<class T>
void ps_helper(T* x,int m,int n,int threads,int s)
{
    for(int j=0;j<s;j++)
    {
        T z=x[j];
        for(int i=1;i<threads;i++)
        {
            int a=i*m/threads,l=a*n+j;
            x[l]=(z+=x[l-n]+x[l]);
        }
    }
}

template<class T>
void Partial_Sum_Parallel_(ARRAYS_ND_BASE<T,VECTOR<int,2>>& x)
{
    static constexpr int m_per_thread=4;
    static constexpr int n_per_thread=16;
    T* p=x.Get_Array_Pointer();
    auto s=x.domain.Edge_Lengths();
    int m=s.x,n=s.y;

#pragma omp parallel
    {
        int threads=omp_get_num_threads();
        int mt=std::min(threads,m/m_per_thread);
        if(mt<3)
        {
#pragma omp single
            Partial_Sum(x);
        }
        else
        {
            int tid=omp_get_thread_num();
            int a=tid*m/mt;
            int b=(tid+1)*m/mt;
            if(tid<mt)
            {
                ps(p+a*n,n);
                ps(p+n*(a+1),b-a-1,n);
            }
#pragma omp barrier
            int nt=std::clamp(n/n_per_thread,1,threads);
            if(tid<nt)
            {
                int c=tid*n/nt;
                int d=(tid+1)*n/nt;
                ps_helper(p+c,m,n,mt,d-c);
            }
#pragma omp barrier
            if(tid<mt)
                for(int i=a+1;i<b;i++)
                    for(int j=0;j<n;j++)
                        p[i*n+j]+=p[a*n+j];
        }
    }    
}

template<class T>
void Partial_Sum_Parallel_(ARRAYS_ND_BASE<T,VECTOR<int,3>>& x)
{
    static constexpr int m_per_thread=4;
    static constexpr int n_per_thread=16;
    T* q=x.Get_Array_Pointer();
    auto s=x.domain.Edge_Lengths();
    int m=s.x,n=s.y,p=s.z,np=n*p;

#pragma omp parallel
    {
        int threads=omp_get_num_threads();
        int mt=std::min(threads,m/m_per_thread);
        if(mt<3)
        {
#pragma omp single
            Partial_Sum(x);
        }
        else
        {
            int tid=omp_get_thread_num();
            int a=tid*m/mt;
            int b=(tid+1)*m/mt;
            if(tid<mt)
            {
                ps(q+a*np,n,p);
                ps(q+np*(a+1),b-a-1,n,p);
            }
#pragma omp barrier
            int nt=std::clamp(np/n_per_thread,1,threads);
            if(tid<nt)
            {
                int c=tid*np/nt;
                int d=(tid+1)*np/nt;
                ps_helper(q+c,m,np,mt,d-c);
            }
#pragma omp barrier
            if(tid<mt)
                for(int i=a+1;i<b;i++)
                    for(int j=0;j<np;j++)
                        q[i*np+j]+=q[a*np+j];
        }
    }    
}

template<class T,int d>
void Partial_Sum_Parallel(ARRAYS_ND_BASE<T,VECTOR<int,d>>& x)
{
    if constexpr(d==1) return Partial_Sum_Parallel(x.array);
    else return Partial_Sum_Parallel_(x);
}

#else

template<class T>
void Partial_Sum_Parallel(ARRAY_VIEW<T> x)
{
    return Partial_Sum(x);
}
template<class T,int d>
void Partial_Sum_Parallel(ARRAYS_ND_BASE<T,VECTOR<int,d>>& x)
{
    return Partial_Sum(x);
}

#endif
template void Partial_Sum<double,1>(ARRAYS_ND_BASE<double,VECTOR<int,1> >&);
template void Partial_Sum<double,2>(ARRAYS_ND_BASE<double,VECTOR<int,2> >&);
template void Partial_Sum<double,3>(ARRAYS_ND_BASE<double,VECTOR<int,3> >&);
template void Partial_Sum<double>(ARRAY_VIEW<double,int>);
template void Partial_Sum<float,1>(ARRAYS_ND_BASE<float,VECTOR<int,1> >&);
template void Partial_Sum<float,2>(ARRAYS_ND_BASE<float,VECTOR<int,2> >&);
template void Partial_Sum<float,3>(ARRAYS_ND_BASE<float,VECTOR<int,3> >&);
template void Partial_Sum<float>(ARRAY_VIEW<float,int>);
template void Partial_Sum<int,1>(ARRAYS_ND_BASE<int,VECTOR<int,1> >&);
template void Partial_Sum<int,2>(ARRAYS_ND_BASE<int,VECTOR<int,2> >&);
template void Partial_Sum<int,3>(ARRAYS_ND_BASE<int,VECTOR<int,3> >&);
template void Partial_Sum<int>(ARRAY_VIEW<int,int>);
template void Partial_Sum_Parallel<double,1>(ARRAYS_ND_BASE<double,VECTOR<int,1> >&);
template void Partial_Sum_Parallel<double,2>(ARRAYS_ND_BASE<double,VECTOR<int,2> >&);
template void Partial_Sum_Parallel<double,3>(ARRAYS_ND_BASE<double,VECTOR<int,3> >&);
template void Partial_Sum_Parallel<double>(ARRAY_VIEW<double,int>);
template void Partial_Sum_Parallel<float,1>(ARRAYS_ND_BASE<float,VECTOR<int,1> >&);
template void Partial_Sum_Parallel<float,2>(ARRAYS_ND_BASE<float,VECTOR<int,2> >&);
template void Partial_Sum_Parallel<float,3>(ARRAYS_ND_BASE<float,VECTOR<int,3> >&);
template void Partial_Sum_Parallel<float>(ARRAY_VIEW<float,int>);
template void Partial_Sum_Parallel<int,1>(ARRAYS_ND_BASE<int,VECTOR<int,1> >&);
template void Partial_Sum_Parallel<int,2>(ARRAYS_ND_BASE<int,VECTOR<int,2> >&);
template void Partial_Sum_Parallel<int,3>(ARRAYS_ND_BASE<int,VECTOR<int,3> >&);
template void Partial_Sum_Parallel<int>(ARRAY_VIEW<int,int>);
}
