#include "task_group_split.h"

void task_group_split_cpu(
        int* group, bool* value, bool* output,
        const int batch_size, const int task_num, const int group_num)
{
    auto temp = torch::make_unique<bool[]>(group_num);
    for(int b=0; b<batch_size; b++)
    {
        for(int i=0; i<group_num; i++){
            temp[i] = false;
        }
        
        for(int i=0; i<task_num; i++){
            if(value[i]){
                int g = group[i];
                temp[g] = true;
            }
        }
        
        output[b] = false;
        for(int i=0; i<task_num; i++){
            int g = group[i];
            if(temp[g] && !value[i]){
                output[b] = true;
                break;
            }
        }
        
        group += task_num;
        value += task_num;
    }
};


auto task_group_split(
    const Tensor& group, const Tensor& value) -> Tensor
{
    auto device = group.device();
    const int batch_size = group.size(0);
    const int task_num = group.size(1);
    const int group_num = group.max().item<int>() + 1;
    const int _group_num = group.min().item<int>();

    GRL_CHECK(group_num <= task_num && _group_num >= 0, "group value error");

    GRL_CHECK_TENSOR(group, device, false, false, batch_size, task_num);
    GRL_CHECK_TENSOR(value, device, false, false, batch_size, task_num);

    auto output = torch::zeros({batch_size}, torch::dtype(torch::kBool).device(device));

    switch(device.type())
    {
        case torch::kCPU:
            task_group_split_cpu(group.data_ptr<int>(), value.data_ptr<bool>(),
                                 output.data_ptr<bool>(), batch_size, task_num, group_num); 
            break;
#ifdef CUDA_FOUND
        case torch::kCUDA:
            task_group_split_cuda(group.data_ptr<int>(), value.data_ptr<bool>(),
                                  output.data_ptr<bool>(), batch_size, task_num, group_num, device.index());
            break;
#endif
        default:
            GRL_ERROR("unsupported device: %s", device.str().c_str());
    }
    
    return output;
};