//#####################################################################
// Copyright 2017 Craig Schroeder
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
// Class Parallel_Partition
//#####################################################################
#ifndef __Parallel_Partition__
#define __Parallel_Partition__

#include <Core/Arrays/ARRAY.h>
#ifdef USE_OPENMP
#include <omp.h>
#endif
namespace PhysBAM{

// f(T) -> bin, where 0<=bin<out.m
// out.m need not be the number of threads.
template<class T,class F>
void Parallel_Partition(ARRAY<ARRAY<T>>& out,const ARRAY<T>& in,F f)
{
    if(out.m==1)
    {
        out(0)=in;
        return;
    }
#ifdef USE_OPENMP
    int threads=1;
#pragma omp parallel
#pragma omp single
    {
        threads=omp_get_num_threads();
    }

    struct PER_THREAD
    {
        ARRAY<ARRAY<T>> data;
        char spacer[128];
    };

    ARRAY<PER_THREAD> all(threads);
    for(auto&a:all) a.data.Resize(out.m);
#pragma omp parallel for
    for(int t=0;t<threads;t++)
    {
        int a=(in.m*t)/threads;
        int b=(in.m*(t+1))/threads;
        auto& data=all(t).data;
        for(int i=a;i<b;i++)
            data(f(in(i))).Append(in(i));
    }

#pragma omp parallel for
    for(int i=0;i<out.m;i++)
    {
        ARRAY<T> ar;
        for(auto& a:all) ar.Append(a.data(i));
        std::swap(ar,out(i));
    }
#else
    for(auto& a:in) out(f(a)).Append(a);
#endif
}
}
#endif
